Test Data

This module mirrors the C++ VirtualCasingTestData utility to generate small, deterministic datasets for validation. The JAX port keeps the same mathematical constructions and discretization choices so parity checks can be performed without relying on the C++ backend.

Overview

The test data builder provides:

  • Prebuilt analytic surfaces (axisymmetric, rotating ellipse, and W7X)

  • Synthetic magnetic fields from internal/external current loops

  • On-surface target grids and off-surface evaluation points

  • GradB computed from the Biot–Savart gradient kernel

These utilities are intended for validation and regression tests. They are not performance-optimized and intentionally favor determinism over speed.

Surface Coordinates

The surface generator follows the BIEST definitions exactly. For the axisymmetric and rotating-ellipse families, the toroidal coordinate t and poloidal coordinate p are used to construct a twisted ellipse cross-section. Let theta = 2π t / Nt and phi = 2π p / Np. The base cross-section is

\[\mathbf{x}_\perp = (\cos\phi, \sin\phi),\]

optionally deformed by

\[x_\perp \leftarrow \frac{\cos\phi + \exp(\cos\phi) - 1.5}{2.4}\]

for the Axisym* variants without a strict ellipse. The cross-section is rotated by a toroidal-dependent angle and embedded in 3D:

\[R = x_\perp' + R_0, \quad X = R\cos\theta, \; Y = R\sin\theta, \; Z = y_\perp'.\]

For Stell and W7X, the surface is defined by Fourier series:

\[R(\theta,\phi) = \sum_{i,j} R_{bc}(i,j)\,\cos(j\phi - N_{fp} i \theta), \qquad Z(\theta,\phi) = \sum_{i,j} Z_{bs}(i,j)\,\sin(j\phi - N_{fp} i \theta),\]

with Nfp=5 for the built-in W7X coefficients. The coefficients are embedded in virtual_casing_jax.w7x_coeffs and match the BIEST source.

Bundled Geometry Assets

The C++ code ships additional surfaces (Quas3, LHD, W7X) as SCTL .mat files. For the JAX port these are converted to .npz and bundled under virtual_casing_jax/geom. The conversion is lossless and uses the same scaling as the reference implementation:

  • Quas3: scale 0.45 and drop the final row/column (periodic seam).

  • LHD: scale 0.25.

  • W7X: scale 0.45.

The conversion script lives in tools/convert_geom_mat.py. JAX loads the assets via importlib.resources and applies the same spectral upsampling as the C++ SurfaceOp::Upsample call.

Synthetic Loop Fields

Test fields are generated from a pair of loop families:

  • An internal loop constructed by shifting the surface inward by a fixed normal offset and averaging over poloidal angle.

  • A set of external loops placed at radii chosen from the internal loop, with random orientations generated by a deterministic drand48 sequence.

The magnetic field is computed using the Biot–Savart integral

\[\mathbf{B}(\mathbf{x}) = \frac{1}{4\pi} \int \frac{\mathbf{J}(\mathbf{x}') \times (\mathbf{x}-\mathbf{x}')} {|\mathbf{x}-\mathbf{x}'|^3}\,d\ell',\]

where the loop density uses the parametric derivative d\mathbf{X}/dt from the same spectral derivative operator used in the C++ implementation.

GradB is generated with the analytic derivative of the Biot–Savart kernel, matching the BIEST FxdU operator. This enables parity checks against the C++ reference without finite differencing.

Randomization and Determinism

The loop orientation uses a drand48-compatible linear congruential sequence with the standard default seed 0x1234abcd330e. The JAX port exposes an optional rng argument to allow reuse of the same generator across calls when exact sequencing is required. If no generator is provided, a fresh deterministic generator is used for each call.

API

The test data builder is available both as functions and as a lightweight class mirroring the C++ static API:

from virtual_casing_jax import SurfType, VirtualCasingTestData

X = VirtualCasingTestData.surface_coordinates(
    nfp=1, half_period=False, nt=6, npol=5, surf_type=SurfType.AxisymNarrow
)
Bext, Bint = VirtualCasingTestData.magnetic_field_data(
    1, False, 6, 5, X, trg_nt=4, trg_np=4
)
GradBext, GradBint = VirtualCasingTestData.magnetic_field_grad_data(
    1, False, 6, 5, X, trg_nt=4, trg_np=4
)

Validation

The parity suite compares JAX outputs against C++ dumps generated from VirtualCasingTestData. The datasets are stored under tests/data with prefix case_testdata_axisym and are exercised in tests/test_testdata_parity.py.