plot_expeq.py: Rename psi to rho to avoid additional ambiguity; fix imports

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()