Improvements for helicity, q(r) discritization, and d_phi advance

synthwave/magnetic_geometry/equilibrium_field.py

@@ -1,7 +1,7 @@
 import numpy as np
 
 from scipy.constants import mu_0
-from scipy.interpolate import RectBivariateSpline, make_interp_spline
+from scipy.interpolate import RectBivariateSpline, make_smoothing_spline
 from scipy.optimize import newton
 
 from synthwave.magnetic_geometry.utils import (
@@ -42,7 +42,7 @@ def biot_savart_cylindrical(
 
 
 class EquilibriumField:
-    def __init__(self, eqdsk):
+    def __init__(self, eqdsk, lam=1e-7):
         self.eqdsk = eqdsk
         self.psi = RectBivariateSpline(
             eqdsk.r_grid[:, 0], eqdsk.z_grid[0, :], eqdsk.psi, kx=3, ky=3, s=0
@@ -51,11 +51,14 @@ def __init__(self, eqdsk):
         # Linear grid of psi for 1D profiles
         # https://freeqdsk.readthedocs.io/en/stable/geqdsk.html
         self.psi_grid = np.linspace(eqdsk.simagx, eqdsk.sibdry, eqdsk.nx)
-
         # q(psi)
-        self.qpsi = make_interp_spline(self.psi_grid, eqdsk.qpsi, k=3, axis=0)
+        # RNC EDIT: Switching to smoothing spline to avoid strange "discretization" jumps
+        # Note: lam value (lower = less smoothing) should be reasonably consistent across q
+        # profiles, but this is not certain. Lam=1e-7 works for q(psi) and F(psi) so far.
+        self.qpsi = make_smoothing_spline(self.psi_grid, eqdsk.qpsi, lam=lam, axis=0)
+
         # F(psi)
-        self.F = make_interp_spline(self.psi_grid, eqdsk.fpol, k=3, axis=0)
+        self.F = make_smoothing_spline(self.psi_grid, eqdsk.fpol, lam=lam, axis=0)
 
     def get_field_at_point(self, R, Z) -> np.ndarray:
         # Bp = Br + Bz = (d(psi)/dZ - d(psi)/dR) / R
@@ -68,7 +71,7 @@ def get_field_at_point(self, R, Z) -> np.ndarray:
 
         return np.array([Br, Bt, Bz])
 
-    def get_psi_of_q(self, q, tol=1e-3):
+    def get_psi_of_q(self, q):
         """Get psi corresponding to a given q"""
         qpsi_grid = self.qpsi(self.psi_grid)
         psi_guess = self.psi_grid[np.argmin(np.abs(qpsi_grid - q))]
@@ -77,21 +80,22 @@ def get_psi_of_q(self, q, tol=1e-3):
             x0=psi_guess,
             fprime=lambda psi: self.qpsi.derivative(1)(psi),
             maxiter=400,
-            tol=tol,
+            tol=1e-3,
         )
 
         # RNC EDIT:
-        # CHECK: For q~<=1, depening on the resolution of the gEQDSK file,
+        # CHECK: For q~<=1, depending on the resolution of the gEQDSK file,
         # the interpolation function can request a psi value at or less
         # than the minimum in the psi_grid vector
         # Check to see if the value we want plausibly exists in the final set
         # of q values from the eqdsk
         # The issue appears to be that although psi is continuous and monotonic,
         # q values are "grouped" in an odd, stepwise fashion
         if psi <= self.psi_grid[0]:
-            # check if there's a range of possible q-values (e.g.) if this is plausably a resolution issue
+            # check if there's a range of possible q-values (e.g.) if this is plausibly a resolution issue
+
             if (
-                np.argwhere(qpsi_grid > (qpsi_grid[0] + tol)).squeeze()[0] > 1
+                np.argwhere(qpsi_grid > (qpsi_grid[0] + 1e-3)).squeeze()[0] > 1
             ):  # multiple identical q values in a row
                 lin_interp_q = np.polyfit(self.psi_grid[:30], qpsi_grid[:30], 1)
                 psi = self.psi_grid[

synthwave/magnetic_geometry/filaments.py

@@ -284,8 +284,30 @@ def _trace(
         self,
         num_filament_points: Optional[int] = None,
         method: TraceType = TraceType.SINGLE,
+        helicity_sign: int = +1,
     ) -> tuple[np.ndarray, np.ndarray]:
-        """Trace a filament"""
+        """Trace a filament along the equilibrium magnetic field.
+
+        Parameters
+        ----------
+        num_filament_points : int, optional
+            Number of points to use for tracing a single poloidal turn of the filament.
+            If ``None``, the value stored in ``self.num_points`` is used.
+        method : EquilibriumFilamentTracer.TraceType, optional
+            Tracing strategy to use. This controls how the field is followed when
+            computing the filament shape (e.g., cylindrical approximation, naive,
+            single-point, or averaged tracing).
+        helicity_sign : int, optional
+            Sign of the helicity used when following the field lines. A value of
+            ``+1`` traces in the default direction, while ``-1`` reverses the
+            direction.
+
+        Returns
+        -------
+        tuple of numpy.ndarray
+            Tuple containing arrays describing the traced filament coordinates.
+
+        """
         if num_filament_points is None:
             num_filament_points = self.num_points
 
@@ -322,17 +344,18 @@ def _R_a(eta, a):
             R = self.eq_field.eqdsk.rmagx + (a * np.cos(eta))
             return R
 
+        # Currently: +m/+n helicity matches empirical C-Mod pickup
         def _Z_a(eta, a):
-            Z = self.eq_field.eqdsk.zmagx - (a * np.sin(eta))
+            Z = self.eq_field.eqdsk.zmagx + (helicity_sign) * (a * np.sin(eta))
             return Z
 
         def psi_prime_a(eta, a):
             # Derivative of psi with respect to a at a given eta
             R = _R_a(eta, a)
             Z = _Z_a(eta, a)
-            return self.eq_field.psi.ev(R, Z, dx=1, dy=0) * np.cos(
-                eta
-            ) - self.eq_field.psi.ev(R, Z, dx=0, dy=1) * np.sin(eta)
+            return self.eq_field.psi.ev(R, Z, dx=1, dy=0) * np.cos(eta) + (
+                helicity_sign
+            ) * self.eq_field.psi.ev(R, Z, dx=0, dy=1) * np.sin(eta)
 
         for i, eta in enumerate(filament_etas):
             if i == 0:
@@ -349,13 +372,20 @@ def psi_prime_a(eta, a):
                 func=lambda a: self.eq_field.psi.ev(_R_a(eta, a), _Z_a(eta, a)) - psi_q,
                 x0=a_guess,
                 fprime=lambda a: psi_prime_a(eta, a),
+                maxiter=800,
+                tol=1e-3,
             )
             poloidal_points[i, :] = [_R_a(eta, a_next), _Z_a(eta, a_next), a_next]
 
         def _d_phi(r, R, Bp, Bt, d_eta):
             # https://wiki.fusion.ciemat.es/wiki/Rotational_transform
             return (Bt * r * d_eta) / (R * Bp)
 
+        # alternative form removing the assumption that dl = r d_eta (that assumption holds only for circular cross-sections)
+        def _d_phi_dl(dl, R, Bp, Bt):
+            # https://youjunhu.github.io/research_notes/tokamak_equilibrium_htlatex/tokamak_equilibrium.html
+            return (Bt * dl) / (R * Bp)
+
         # Finalize filament trace based on method
         if method == EquilibriumFilamentTracer.TraceType.CYLINDRICAL:
             # Circular cross section around the magnetic axis
@@ -377,10 +407,19 @@ def _d_phi(r, R, Bp, Bt, d_eta):
             # determine d(phi)/d(eta) from magnetic field
             R = poloidal_points[:, 0]
             Z = poloidal_points[:, 1]
-            r = poloidal_points[:, 2]
             B = self.eq_field.get_field_at_point(R, Z)
-            d_eta = np.mean(np.diff(filament_etas))
-            d_phi = _d_phi(r, R, np.sqrt(B[0] ** 2 + B[2] ** 2), B[1], d_eta)
+
+            # Switching to improved d_phi formula from the below:
+            # d_eta = np.mean(np.diff(filament_etas))
+            # d_phi = _d_phi(r, R, np.sqrt(B[0] ** 2 + B[2] ** 2), B[1], d_eta)
+
+            # Compute segment lengths with wraparound so the last segment goes from
+            # the final point back to the first
+            dR = np.roll(R, -1) - R
+            dZ = np.roll(Z, -1) - Z
+            dl = np.sqrt(dR**2 + dZ**2)
+            d_phi = _d_phi_dl(dl, R, np.sqrt(B[0] ** 2 + B[2] ** 2), B[1])
+
             if method == EquilibriumFilamentTracer.TraceType.SINGLE:
                 phi = np.cumsum(d_phi) - d_phi[0]
             else:
@@ -391,12 +430,16 @@ def _d_phi(r, R, Bp, Bt, d_eta):
             # Numerical correction to ensure final point is at the proper angle
             # We can do this multiple times, whenever we know for sure that the phi is a multiple of pi * m / n
 
-            # RNC EDIT: Sign flip necessary to account for helicity direction
+            # Sign flip necessary to account for helicity direction
+            # Note: The improved d_phi method should ensure this is always correct now [e.g. the correction factor is no longer strictly necessary]
+
             known_phis = np.linspace(
                 0, 2 * np.pi * self.m_local, (2 * self.n_local) + 1
             ) * np.sign(phi[-1])
+
             for i, known_phi_start in enumerate(known_phis[:-1]):
                 known_phi_end = known_phis[i + 1]
+
                 if np.sign(phi[-1]) == 1:
                     phi_indices = np.squeeze(
                         np.where((phi >= known_phi_start) & (phi <= known_phi_end))
@@ -411,6 +454,7 @@ def _d_phi(r, R, Bp, Bt, d_eta):
                 correction_factor = (known_phi_end - known_phi_start) / (
                     actual_phi_end - actual_phi_start
                 )
+
                 phi[phi_indices] = (
                     known_phi_start
                     + (phi[phi_indices] - actual_phi_start) * correction_factor

synthwave/mirnov/run_thincurr_model.py

@@ -377,57 +377,74 @@ def correct_frequency_response(
                 )
             )
 
+
 ################################################################################################
 ################################################################################################
 def plot_sensor_output(
-        working_files_directory,
-        probe_details,
-        mode,
-        freq,
-        save_Ext,
-        doSave,
+    working_files_directory,
+    probe_details,
+    mode,
+    freq,
+    save_Ext,
+    doSave,
 ):
-
     # Load in saved sensor signals from netCDF format and plot
 
-    fName = working_files_directory + \
-        "probe_signals_%s_m%02d_n%02d_f%1.1ekHz%s.nc" % \
-        (
-            probe_details.attrs["probe_set_name"], 
-            mode["m"],
-            mode["n"],
-            freq / 1e3,
-            save_Ext,
-        )
+    fName = working_files_directory + "probe_signals_%s_m%02d_n%02d_f%1.1ekHz%s.nc" % (
+        probe_details.attrs["probe_set_name"],
+        mode["m"],
+        mode["n"],
+        freq / 1e3,
+        save_Ext,
+    )
 
     # Load data
     sensors_bode = xr.load_dataset(fName)
 
     # Generate signal magnitudes
-    mags = [ np.linalg.norm(sensors_bode.sel(sensor=sig).signal.values.tolist()) \
-            for sig in sensors_bode.sensor.values]
+    mags = [
+        np.linalg.norm(sensors_bode.sel(sensor=sig).signal.values.tolist())
+        for sig in sensors_bode.sensor.values
+    ]
     sensor_names = sensors_bode.sensor.values.tolist()
 
     # Plot
-    plt.close('Sensor_Signals_m%02d_n%02d_f%1.1ekHz'%( mode["m"],mode["n"],freq/1e3))
-    fig,ax = plt.subplots(1,1,figsize=(4,3),tight_layout=True,num='Sensor_Signals_m%02d_n%02d_f%1.1ekHz'%(
-        mode["m"],mode["n"],freq/1e3))
-    sc = ax.plot(sensor_names, mags, '-*',label=r'$m/n=%d/%d,\,f=%1.1f$ kHz'%(mode["m"],mode["n"],freq/1e3))
-
+    plt.close(
+        "Sensor_Signals_m%02d_n%02d_f%1.1ekHz" % (mode["m"], mode["n"], freq / 1e3)
+    )
+
+    fig, ax = plt.subplots(
+        1,
+        1,
+        figsize=(4, 3),
+        tight_layout=True,
+        num="Sensor_Signals_m%02d_n%02d_f%1.1ekHz" % (mode["m"], mode["n"], freq / 1e3),
+    )
+
+    ax.plot(
+        sensor_names,
+        mags,
+        "-*",
+        label=r"$m/n=%d/%d,\,f=%1.1f$ kHz" % (mode["m"], mode["n"], freq / 1e3),
+    )
+
     ax.set_xticks(range(0, len(sensor_names), 3))
     ax.set_xticklabels([sensor_names[i] for i in range(0, len(sensor_names), 3)])
-    ax.tick_params(axis='x', rotation=90)
-    ax.legend(fontsize=8,handlelength=1)
-    ax.set_ylabel('Signal Magnitude [T/s]')
+    ax.tick_params(axis="x", rotation=90)
+    ax.legend(fontsize=8, handlelength=1)
+    ax.set_ylabel("Signal Magnitude [T/s]")
 
     ax.grid()
 
     if doSave:
-        fig.savefig(working_files_directory + 
-            "Sensor_Signals_m%02d_n%02d_f%1.1ekHz%s.svg" % 
-            (
+        fig.savefig(
+            working_files_directory
+            + "Sensor_Signals_m%02d_n%02d_f%1.1ekHz%s.svg"
+            % (
                 mode["m"],
                 mode["n"],
                 freq / 1e3,
                 save_Ext,
-            ), transparent=True)
+            ),
+            transparent=True,
+        )