diff --git a/python/plot_expeq.py b/python/plot_expeq.py index 3dbe0c9b511cc571ce78f8756357948bd60ab450..bf605122f0f50039aceaf315fa6449193f92a9b9 100755 --- a/python/plot_expeq.py +++ b/python/plot_expeq.py @@ -5,6 +5,7 @@ DEFAULT_INTERACTIVE = matplotlib_interactive.DEFAULT_INTERACTIVE import matplotlib.pyplot as plt import argparse import os.path +import sys try: from scipy.optimize import curve_fit, minimize from scipy.constants import mu_0 @@ -59,11 +60,11 @@ def read_expeq(filename="EXPEQ"): expeq_str["pedge"] = read_scalar(expeq, float) expeq_str["nbps"] = read_scalar(expeq, int) expeq_str["boundary"] = read_field(expeq, expeq_str["nbps"]) - expeq_str["npsi"], expeq_str["nppfun"] = [int(x) for x in expeq.readline().split()] + expeq_str["nrho"], expeq_str["nppfun"] = [int(x) for x in expeq.readline().split()] expeq_str["nsttp"], expeq_str["nrhotype"] = [int(x) for x in expeq.readline().split()] - expeq_str["psi"] = read_field(expeq, expeq_str["npsi"]) - expeq_str["pressure"] = read_field(expeq, expeq_str["npsi"]) - expeq_str["q"] = read_field(expeq, expeq_str["npsi"]) + expeq_str["rho"] = read_field(expeq, expeq_str["nrho"]) + expeq_str["pressure"] = read_field(expeq, expeq_str["nrho"]) + expeq_str["q"] = read_field(expeq, expeq_str["nrho"]) footer = [] while True: line = expeq.readline() @@ -85,10 +86,10 @@ def write_expeq(expeq_str, filename="EXPEQ_out"): expeq.write(str(expeq_str["nbps"]) + "\n") for x, y in expeq_str["boundary"]: expeq.write(f"{x} {y}\n") - expeq.write(f"{expeq_str['npsi']} {expeq_str['nppfun']}\n") + expeq.write(f"{expeq_str['nrho']} {expeq_str['nppfun']}\n") expeq.write(f"{expeq_str['nsttp']} {expeq_str['nrhotype']}\n") - for psival in expeq_str["psi"]: - expeq.write(f"{psival}\n") + for rhoval in expeq_str["rho"]: + expeq.write(f"{rhoval}\n") for pressureval in expeq_str["pressure"]: expeq.write(f"{pressureval}\n") for qval in expeq_str["q"]: @@ -101,9 +102,9 @@ def test_plot(expeq, show=True): ax.plot(expeq["boundary"][:,0], expeq["boundary"][:,1]) ax.set_aspect(1.) fig, ax = plt.subplots() - ax.plot(np.sqrt(expeq["psi"]), expeq["q"]) + ax.plot(expeq["rho"], expeq["q"]) fig, ax = plt.subplots() - ax.plot(np.sqrt(expeq["psi"]), expeq["pressure"]) + ax.plot(expeq["rho"], expeq["pressure"]) print(expeq["pressure"][0]) if show: plt.show(block=True) @@ -291,39 +292,40 @@ def replace_profiles(expeq, args, B0=1.0, plot=True): With whatever is returned by the qfunc function in qfunc.py (if defined). Or equivalently pfunc in pfunc.py. """ - psi = expeq["psi"] - psi_norm = psi/psi[-1] - rhopol = np.sqrt(psi_norm) - qvals = replace_q(psi_norm, expeq["q"]) - expeq["q"] = qvals - print(f'{expeq["nrhotype"]=}') + rho = expeq["rho"] + rho_norm = rho/rho[-1] assert(expeq["nrhotype"] == 0) + print(f'{expeq["nrhotype"]=}') + assert(expeq["nsttp"] == 5) + qvals = replace_q(rho_norm, expeq["q"]) + expeq["q"] = qvals + assert(expeq["nppfun"] == 8) - print(f'{expeq["pressure"][-1]=}') - pvals = replace_p(psi_norm, expeq["pressure"], B0=B0) - print(f'{expeq["pedge"]=}') - print(f'{pvals[-1]=}') - expeq["pedge"] = pvals[-1] + pvals = replace_p(rho_norm, expeq["pressure"], B0=B0) expeq["pressure"] = pvals + expeq["pedge"] = pvals[-1] if plot: fig, ax = plt.subplots() - ax.plot(rhopol, qvals, label='q') + ax.plot(rho_norm, qvals, label='q') ax2 = ax.twinx() line, = ax.plot([], [], label='p') - ax2.plot(rhopol, pvals, label='p', color=line.get_color()) + ax2.plot(rho_norm, pvals, label='p', color=line.get_color()) ax.legend() fig.tight_layout() fig.savefig("new_profiles.pdf") def replace_q(rhopol, old_q): qvals = old_q + old_path = sys.path try: + sys.path.insert(0, ".") from qfunc import qfunc qvals = qfunc(rhopol) except ImportError as e: - pass + print("skipping q") + sys.path = old_path return qvals def replace_p(rhopol, old_p, B0=1.0): @@ -331,13 +333,16 @@ def replace_p(rhopol, old_p, B0=1.0): Since SI pressure is normalized by mu_0 B_0^{-2}, we need to pass in B0. """ pvals = old_p + old_path = sys.path try: + sys.path.insert(0, ".") from pfunc import pfunc pvals = pfunc(rhopol) # chease normalization from SI pvals *= (mu_0 / B0**2) except ImportError as e: - pass + print("skipping p") + sys.path = old_path return pvals def plot_expeq(expeq=None, filename='EXPEQ', directory='.', interactivePlots=False): @@ -357,13 +362,13 @@ def plot_expeq(expeq=None, filename='EXPEQ', directory='.', interactivePlots=Fal fig, ax = plt.subplots() label = NPPFUN_LABEL_DICT.get(expeq["nppfun"]) - ax.plot(expeq["psi"], expeq["pressure"], label=label) + ax.plot(expeq["rho"], expeq["pressure"], label=label) ax2 = ax.twinx() label = NSTTP_LABEL_DICT.get(expeq["nsttp"]) line, = ax.plot([], [], label=label) c = line.get_color() - ax2.plot(expeq["psi"], expeq["q"], color=c) + ax2.plot(expeq["rho"], expeq["q"], color=c) ax.legend()