API

JAX surface operators matching BIEST conventions.

virtual_casing_jax.surface_ops.complete_vec_field(Y, is_surf: bool, half_period: bool, nfp: int, nt: int, npol: int, dtheta: float)

Match BIEST SurfaceOp::CompleteVecField.

Y shape: (dof, nt, npol) Returns X shape: (dof, nfp*nt, npol)

virtual_casing_jax.surface_ops.cross_prod(A, B)

SoA cross product: A,B shape (3, nt, npol) -> (3, nt, npol).

virtual_casing_jax.surface_ops.dot_prod(A, B)

SoA dot product: A,B shape (3, nt, npol) -> (nt, npol).

virtual_casing_jax.surface_ops.fft_c2r(y, nt: int, npol: int)

Unitary c2r FFT over (nt, npol) axes.

virtual_casing_jax.surface_ops.fft_r2c(x, nt: int, npol: int)

Unitary r2c FFT over (nt, npol) axes.

virtual_casing_jax.surface_ops.grad2d(X, nt: int, npol: int)

Spectral surface derivatives (BIEST SurfaceOp::Grad2D).

Returns dX with shape (dof * 2, nt, npol) where entries are ordered as [dX_t, dX_p] per component.

virtual_casing_jax.surface_ops.normal_orientation(X, normal)

Return +1 or -1 orientation used by BIEST for normals.

virtual_casing_jax.surface_ops.resample(X0, nt0: int, np0: int, nt1: int, np1: int)

Resample using upsample + decimation (BIEST SurfaceOp::Resample).

virtual_casing_jax.surface_ops.rotate_toroidal(X, nt: int, npol: int, dtheta)

Rotate field in toroidal angle by dtheta.

X shape: (dof, nt, npol)

virtual_casing_jax.surface_ops.surf_normal_area_elem(dX, X=None, *, return_orientation: bool = False)

Compute unit normal and area element (BIEST SurfNormalAreaElem).

dX: (6, nt, npol) for 3D surfaces (dX_t, dX_p per component). X: optional (3, nt, npol) coordinates for orientation. Returns (normal, area_elem) or (normal, area_elem, orient) when return_orientation=True.

virtual_casing_jax.surface_ops.upsample(X0, nt0: int, np0: int, nt1: int, np1: int)

Upsample using Fourier zero-padding (BIEST SurfaceOp::Upsample).

Kernel functions matching BIEST scaling and conventions.

virtual_casing_jax.kernels.biotsavart_fx_u(dx, fvec)

Biot-Savart kernel: (f x dx) / (4*pi*r^3).

dx: (…, 3) fvec: (…, 3) returns (…, 3)

virtual_casing_jax.kernels.biotsavart_fxd_u(dx, fvec)

Derivative of Biot-Savart kernel (matches BIEST FxdU).

Returns (…, 3, 3) tensor. The explicit formula matches the uker_FxdU in BIEST (see kernel.hpp).

virtual_casing_jax.kernels.laplace_dx_u(dx, n, f)

Laplace double-layer kernel: (n·dx) f / (4 pi r^3).

This is the outward-normal convention used by SCTL and BIEST since BIEST commit 0bf9ddc. For a closed surface, D[1] is -1 inside, 0 outside, and -1/2 on the surface.

virtual_casing_jax.kernels.laplace_fx_u(dx, f)

Laplace single-layer potential: f / (4*pi*r).

dx: (…, 3) f: (…,) returns (…,)

virtual_casing_jax.kernels.laplace_fxd2_u(dx, f)

Second derivatives of Laplace single-layer.

Returns a (…, 3, 3) tensor with entries: (-delta_ij * r^-3 + 3 r_i r_j r^-5) * f / (4*pi)

virtual_casing_jax.kernels.laplace_fxd_u(dx, f)

Gradient of Laplace single-layer: -(dx * f) / (4*pi*r^3).

dx: (…, 3) f: (…,) returns (…, 3)

Boundary integral evaluation (baseline direct-sum).

exception virtual_casing_jax.integrals.OffsurfaceConvergenceError(*, tolerance: float, achieved_error: float, nt: int, npol: int, levels: int)

Bases: RuntimeError

Raised when eager off-surface refinement exhausts its finite budget.

virtual_casing_jax.integrals.biotsavart_fx_u_eval(X_src, X_trg, density_vec, area_elem, *, chunk_size: int = 1024, target_chunk_size: int | None = None)

Evaluate Biot-Savart FxU by direct quadrature.

density_vec: (3, nt, np) or (3, nsrc) Returns: (3, ntrg) or (3, nt, np) matching X_trg layout.

virtual_casing_jax.integrals.biotsavart_fxd_u_eval(X_src, X_trg, density_vec, area_elem, *, chunk_size: int = 512, target_chunk_size: int | None = None)

Evaluate Biot-Savart FxdU by direct quadrature.

density_vec: (3, nt, np) or (3, nsrc) Returns: (3, 3, ntrg) or (3, 3, nt, np) matching X_trg layout.

virtual_casing_jax.integrals.computeB_offsurface_adaptive(X_src, BdotN, J, Xt, digits: int = 5, max_Nt: int = -1, max_Np: int = -1, max_levels: int = 6, ext: bool = True, chunk_size: int = 1024, target_chunk_size: int | None = None)

Adaptive off-surface evaluation matching ExtVacuumField logic.

virtual_casing_jax.integrals.computeB_offsurface_adaptive_schedule(X_src, BdotN, J, Xt, *, levels: tuple[tuple[int, int], ...], digits: int = 5, ext: bool = True, chunk_size: int = 1024, target_chunk_size: int | None = None, return_estimate: bool = False)

JIT-friendly adaptive off-surface evaluation with fixed refinement schedule.

The refinement schedule is provided as a static tuple of (Nt, Np) pairs. Shapes are static per-level, so this function can be JIT-compiled with levels marked static. The level is chosen by the calibrated estimate of _refine_by_estimate(); return_estimate also returns that estimate, per target and dimensionless, to compare with 10**-digits.

virtual_casing_jax.integrals.computeB_offsurface_baseline(X_src, BdotN, J, Xt, upsample_factor: int = 1, chunk_size: int = 1024, target_chunk_size: int | None = None, ext: bool = True)

Baseline off-surface evaluation using direct quadrature.

This mirrors ExtVacuumField behavior (no singular correction) with optional upsampling for improved accuracy.

virtual_casing_jax.integrals.computeGradB_offsurface_adaptive_schedule(X_src, BdotN, J, Xt, *, levels: tuple[tuple[int, int], ...], digits: int = 5, ext: bool = True, chunk_size: int = 1024, target_chunk_size: int | None = None, return_estimate: bool = False)

JIT-friendly adaptive off-surface GradB evaluation with fixed schedule.

The level is chosen by the calibrated estimate of _refine_by_estimate(); return_estimate also returns that estimate, per target and dimensionless, to compare with 10**-digits.

virtual_casing_jax.integrals.curl_single_layer_gradient(grad)

Assemble curl(G[J]) from grad[J_component, derivative].

virtual_casing_jax.integrals.field_period_target_coords(X_quad, trg_nt: int, trg_np: int, nfp: int)

Select target coordinates used by FieldPeriodBIOp.

X_quad: (3, quad_nt, quad_np) for the full NFP surface. Returns X_trg: (3, trg_nt, trg_np) for the first field period.

virtual_casing_jax.integrals.laplace_dx_u_eval(X_src, n_src, X_trg, density, area_elem, *, chunk_size: int = 1024, target_chunk_size: int | None = None)

Evaluate Laplace DxU (double-layer) by direct quadrature.

virtual_casing_jax.integrals.laplace_dx_u_eval_singular(X_src, dX_src, density, trg_nt: int, trg_np: int, nfp: int, X_trg=None, digits: int = 5, patch_dim0: int | None = None, rad_dim: int | None = None, chunk_size: int = 1024, target_chunk_size: int | None = None, patch_idx=None, orient: float | None = None, pou_dtype=None, patch_dtype=None, interp_block_size: int | str | None = 'auto', remat: bool = False)

Evaluate Laplace DxU with singular correction on surface targets.

virtual_casing_jax.integrals.laplace_fxd2_u_eval(X_src, X_trg, density, area_elem, *, chunk_size: int = 1024, target_chunk_size: int | None = None)

Evaluate Laplace Fxd2U (second derivatives) by direct quadrature.

virtual_casing_jax.integrals.laplace_fxd2_u_eval_singular(X_src, dX_src, density, trg_nt: int, trg_np: int, nfp: int, X_trg=None, digits: int = 5, patch_dim0: int | None = None, rad_dim: int | None = None, hedgehog_order: int = 8, chunk_size: int = 1024, target_chunk_size: int | None = None, patch_idx=None, orient: float | None = None, pou_dtype=None, patch_dtype=None, interp_block_size: int | str | None = 'auto', remat: bool = False, scan_targets: bool = False, hedgehog_side: Literal['interior', 'exterior'] = 'interior')

Evaluate Laplace Fxd2U with singular correction (Hedgehog).

Args:

scan_targets: If True, use a lax.scan loop over targets instead

of vmap in the singular correction. This can reduce peak memory by avoiding large broadcasted temporaries, at the cost of lower parallelism per chunk.

hedgehog_side: Side from which the on-surface limit is taken.

virtual_casing_jax.integrals.laplace_fxd2_u_eval_vec(X_src, X_trg, density_vec, area_elem, *, chunk_size: int = 1024, target_chunk_size: int | None = None)

Vector-density wrapper for Laplace Fxd2U.

virtual_casing_jax.integrals.laplace_fxd2_u_eval_vec_singular(X_src, dX_src, density_vec, trg_nt: int, trg_np: int, nfp: int, X_trg=None, digits: int = 5, patch_dim0: int | None = None, rad_dim: int | None = None, hedgehog_order: int = 8, chunk_size: int = 1024, target_chunk_size: int | None = None, patch_idx=None, orient: float | None = None, pou_dtype=None, patch_dtype=None, interp_block_size: int | str | None = 'auto', remat: bool = False, scan_targets: bool = False, hedgehog_side: Literal['interior', 'exterior'] = 'interior')

Vector-density wrapper for Laplace Fxd2U with singular correction.

Args:

scan_targets: forwarded to laplace_fxd2_u_eval_singular.

hedgehog_side: forwarded to laplace_fxd2_u_eval_singular.

virtual_casing_jax.integrals.laplace_fxd_u_eval(X_src, X_trg, density, area_elem, *, chunk_size: int = 1024, target_chunk_size: int | None = None)

Evaluate Laplace FxdU (grad single-layer) by direct quadrature.

X_src: (3, nt, np) or (3, nsrc) X_trg: (3, nt, np) or (3, ntrg) density: (nt, np) or (nsrc,) area_elem: (nt, np) or (nsrc,) Returns: (3, ntrg) or (3, nt, np) matching X_trg layout.

virtual_casing_jax.integrals.laplace_fxd_u_eval_singular(X_src, dX_src, density, trg_nt: int, trg_np: int, nfp: int, X_trg=None, digits: int = 5, patch_dim0: int | None = None, rad_dim: int | None = None, chunk_size: int = 1024, target_chunk_size: int | None = None, patch_idx=None, orient: float | None = None, pou_dtype=None, patch_dtype=None, interp_block_size: int | str | None = 'auto', remat: bool = False)

Evaluate Laplace FxdU with singular correction on surface targets.

virtual_casing_jax.integrals.laplace_fxd_u_eval_vec(X_src, X_trg, density_vec, area_elem, *, chunk_size: int = 1024, target_chunk_size: int | None = None)

Vector-density wrapper for Laplace FxdU.

density_vec: (3, nt, np) or (3, nsrc) Returns: (3, 3, ntrg) with first index over density component.

virtual_casing_jax.integrals.laplace_fxd_u_eval_vec_singular(X_src, dX_src, density_vec, trg_nt: int, trg_np: int, nfp: int, X_trg=None, digits: int = 5, patch_dim0: int | None = None, rad_dim: int | None = None, chunk_size: int = 1024, target_chunk_size: int | None = None, patch_idx=None, orient: float | None = None, pou_dtype=None, patch_dtype=None, interp_block_size: int | str | None = 'auto', remat: bool = False)

virtual_casing_jax.integrals.source_field_scale(BdotN, J)

RMS field magnitude on the source surface.

J = B x n and BdotN = B . n with a unit normal, so |J|^2 + BdotN^2 = |B|^2. Off-surface field differences are divided by this to give a dimensionless error estimate comparable with 10**-digits.

Functional Virtual Casing API with differentiable geometry inputs.

class virtual_casing_jax.functional.FunctionalSetup(nfp: int, nfp_eff: int, half_period: bool, surf_nt: int, surf_np: int, src_nt: int, src_np: int, trg_nt: int, trg_np: int, quad_nt: int, quad_np: int, patch_dim0: int, patch_idx: Array, orient: float)

Bases: object

Static quadrature setup for functional API.

half_period: bool

nfp: int

nfp_eff: int

orient: float

patch_dim0: int

patch_idx: Array

quad_np: int

quad_nt: int

src_np: int

src_nt: int

surf_np: int

surf_nt: int

trg_np: int

trg_nt: int

virtual_casing_jax.functional.build_patch_idx(quad_nt: int, quad_np: int, trg_nt: int, trg_np: int, nfp_eff: int, patch_dim0: int)

Build patch indices for singular quadrature.

virtual_casing_jax.functional.build_quad_setup(surface_coord, quad_nt: int, quad_np: int, *, orient: float | None = None)

Compute quadrature coordinates, derivatives, and normals.

virtual_casing_jax.functional.build_surface_coord(X, nfp: int, half_period: bool, surf_nt: int, surf_np: int, trg_nt: int)

Build full-field-period surface coordinates from base grid.

virtual_casing_jax.functional.compute_external_B_functional(X, B0, *, digits: int, nfp: int, half_period: bool, surf_nt: int, surf_np: int, src_nt: int, src_np: int, trg_nt: int, trg_np: int, quad_nt: int, quad_np: int, patch_dim0: int | None = None, patch_idx=None, orient: float | None = None, X_trg=None, chunk_size: int | str | None = 'auto', target_chunk_size: int | str | None = 'auto', pou_dtype=None, patch_dtype=None, interp_block_size: int | str | None = 'auto', remat: bool | None = None)

Compute Bext with surface coordinates as differentiable inputs.

virtual_casing_jax.functional.compute_external_B_jvp_columns_functional(X, B0, X_tangents, B0_tangents=None, *, digits: int, nfp: int, half_period: bool, surf_nt: int, surf_np: int, src_nt: int, src_np: int, trg_nt: int, trg_np: int, quad_nt: int, quad_np: int, patch_dim0: int | None = None, patch_idx=None, orient: float | None = None, chunk_size: int | str | None = 'auto', target_chunk_size: int | str | None = 'auto', pou_dtype=None, patch_dtype=None, interp_block_size: int | str | None = 'auto', remat: bool | None = None)

Return Bext and multiple forward-mode tangent columns.

X_tangents has shape (ncols, 3, surf_nt, surf_np). B0_tangents has shape (ncols, 3, src_nt, src_np); when omitted the magnetic-field input is held fixed. The returned tangent array has shape (ncols, 3, trg_nt, trg_np).

virtual_casing_jax.functional.compute_external_B_normal_functional(X, B0, *, digits: int, nfp: int, half_period: bool, surf_nt: int, surf_np: int, src_nt: int, src_np: int, trg_nt: int, trg_np: int, quad_nt: int, quad_np: int, patch_dim0: int | None = None, patch_idx=None, orient: float | None = None, chunk_size: int | str | None = 'auto', target_chunk_size: int | str | None = 'auto', pou_dtype=None, patch_dtype=None, interp_block_size: int | str | None = 'auto', remat: bool | None = None)

Compute on-surface Bext dot n with differentiable geometry inputs.

virtual_casing_jax.functional.compute_external_B_normal_jvp_columns_functional(X, B0, X_tangents, B0_tangents=None, *, digits: int, nfp: int, half_period: bool, surf_nt: int, surf_np: int, src_nt: int, src_np: int, trg_nt: int, trg_np: int, quad_nt: int, quad_np: int, patch_dim0: int | None = None, patch_idx=None, orient: float | None = None, chunk_size: int | str | None = 'auto', target_chunk_size: int | str | None = 'auto', pou_dtype=None, patch_dtype=None, interp_block_size: int | str | None = 'auto', remat: bool | None = None)

Return Bext dot n and multiple forward-mode tangent columns.

X_tangents has shape (ncols, 3, surf_nt, surf_np). B0_tangents has shape (ncols, 3, src_nt, src_np); when omitted the magnetic-field input is held fixed. The returned tangent array has shape (ncols, trg_nt, trg_np).

virtual_casing_jax.functional.compute_external_B_offsurf_functional(X, B0, *, X_trg, digits: int, nfp: int, half_period: bool, surf_nt: int, surf_np: int, src_nt: int, src_np: int, trg_nt: int, trg_np: int, max_Nt: int = -1, max_Np: int = -1, max_levels: int = 6, chunk_size: int | str | None = 'auto', target_chunk_size: int | str | None = 'auto', adaptive: bool = True)

Compute off-surface Bext with differentiable geometry inputs.

virtual_casing_jax.functional.compute_external_gradB_functional(X, B0, *, digits: int, nfp: int, half_period: bool, surf_nt: int, surf_np: int, src_nt: int, src_np: int, trg_nt: int, trg_np: int, quad_nt: int, quad_np: int, patch_dim0: int | None = None, patch_idx=None, orient: float | None = None, hedgehog_order: int = 8, chunk_size: int | str | None = 'auto', target_chunk_size: int | str | None = 'auto', pou_dtype=None, patch_dtype=None, interp_block_size: int | str | None = 'auto', remat: bool | None = None)

Compute GradBext with surface coordinates as differentiable inputs.

virtual_casing_jax.functional.compute_external_gradB_offsurf_functional(X, B0, *, X_trg, digits: int, nfp: int, half_period: bool, surf_nt: int, surf_np: int, src_nt: int, src_np: int, trg_nt: int, trg_np: int, max_Nt: int = -1, max_Np: int = -1, max_levels: int = 6, chunk_size: int | str | None = 'auto', target_chunk_size: int | str | None = 'auto', adaptive: bool = False)

Compute off-surface GradBext with differentiable geometry inputs.

virtual_casing_jax.functional.compute_internal_B_functional(X, B0, *, digits: int, nfp: int, half_period: bool, surf_nt: int, surf_np: int, src_nt: int, src_np: int, trg_nt: int, trg_np: int, quad_nt: int, quad_np: int, patch_dim0: int | None = None, patch_idx=None, orient: float | None = None, X_trg=None, chunk_size: int | str | None = 'auto', target_chunk_size: int | str | None = 'auto', pou_dtype=None, patch_dtype=None, interp_block_size: int | str | None = 'auto', remat: bool | None = None)

Compute Bint with surface coordinates as differentiable inputs.

virtual_casing_jax.functional.compute_internal_gradB_functional(X, B0, *, digits: int, nfp: int, half_period: bool, surf_nt: int, surf_np: int, src_nt: int, src_np: int, trg_nt: int, trg_np: int, quad_nt: int, quad_np: int, patch_dim0: int | None = None, patch_idx=None, orient: float | None = None, hedgehog_order: int = 8, chunk_size: int | str | None = 'auto', target_chunk_size: int | str | None = 'auto', pou_dtype=None, patch_dtype=None, interp_block_size: int | str | None = 'auto', remat: bool | None = None)

Compute GradBint with surface coordinates as differentiable inputs.

virtual_casing_jax.functional.prepare_functional_setup(X, *, digits: int, nfp: int, half_period: bool, surf_nt: int, surf_np: int, src_nt: int, src_np: int, trg_nt: int, trg_np: int, quad_nt: int, quad_np: int, patch_dim0: int | None = None, orient: float | None = None)

Prepare static quadrature setup for functional API.

This helper is intended to be called outside autodiff; it uses non-differentiable logic to choose patch sizes if not provided.

virtual_casing_jax.functional.select_patch_dim_from_geom(dX, quad_nt: int, quad_np: int, digits: int)

Select patch_dim0 using surface condition (non-differentiable).

virtual_casing_jax.functional.target_surface_normal(X, *, nfp: int, half_period: bool, surf_nt: int, surf_np: int, trg_nt: int, trg_np: int, orient: float | None = None)

Return unit normals on the virtual-casing target grid.

High-level Virtual Casing routines in JAX.

class virtual_casing_jax.virtual_casing.PrecisionPlan(quad_nt: int, quad_np: int, patch_dim0: int, patch_idx: Array)

Bases: object

Frozen adaptive-precision choices for a virtual-casing B evaluation.

Automatic quadrature and singular-patch selection depend on concrete surface geometry. Select them once with plan_precision(), then pass the resulting plan to compute_external_B() or compute_internal_B() when differentiating with respect to the surface coordinates.

patch_dim0: int

patch_idx: Array

quad_np: int

quad_nt: int

class virtual_casing_jax.virtual_casing.QuadSetup(quad_nt: 'int', quad_np: 'int', quad_coord: 'jnp.ndarray', dX: 'jnp.ndarray', normal: 'jnp.ndarray', orient: 'float', patch_idx_cache: 'dict[int, jnp.ndarray]'=<factory>)

Bases: object

dX: Array

normal: Array

orient: float

patch_idx_cache: dict[int, Array]

quad_coord: Array

quad_np: int

quad_nt: int

class virtual_casing_jax.virtual_casing.VirtualCasingJAX

Bases: object

JAX mirror of VirtualCasing for external field and GradB.

compute_external_B(B0, *, X_trg=None, quad_nt: int | None = None, quad_np: int | None = None, digits: int | None = None, chunk_size: int | str | None = 'auto', target_chunk_size: int | str | None = 'auto', pou_dtype=None, patch_dtype=None, interp_block_size: int | str | None = 'auto', remat: bool | None = None, patch_dim0: int | None = None, patch_idx=None, precision: PrecisionPlan | None = None)

Compute Bext from total B on the source grid.

compute_external_B_autodiff(B0, *, X_trg, quad_nt: int | None = None, quad_np: int | None = None, digits: int | None = None, chunk_size: int | str | None = 'auto', target_chunk_size: int | str | None = 'auto', pou_dtype=None, patch_dtype=None, interp_block_size: int | str | None = 'auto', remat: bool | None = None, hedgehog_order: int = 8)

Compute Bext with a custom JVP that matches ComputeGradB on-surface.

compute_external_B_batch(B0_batch, *, X_trg=None, **kwargs)

Vectorized compute_external_B over a batch dimension.

compute_external_B_jit(B0, **kwargs)

JIT-compiled version of compute_external_B.

compute_external_B_offsurf(B0, *, X_trg, digits: int | None = None, max_Nt: int = -1, max_Np: int = -1, max_levels: int = 6, chunk_size: int | str | None = 'auto', target_chunk_size: int | str | None = 'auto')

Compute Bext at off-surface targets using adaptive quadrature.

compute_external_B_offsurf_schedule(B0, *, X_trg, levels: tuple[tuple[int, int], ...] | str | None, digits: int | None = None, max_Nt: int = -1, max_Np: int = -1, max_levels: int = 6, chunk_size: int | str | None = 'auto', target_chunk_size: int | str | None = 'auto', return_estimate: bool = False)

Compute Bext off-surface using a fixed adaptive refinement schedule.

compute_external_B_offsurf_schedule_jit(B0, *, X_trg, levels: tuple[tuple[int, int], ...] | str | None = 'auto', digits: int | None = None, max_Nt: int = -1, max_Np: int = -1, max_levels: int = 6, chunk_size: int | str | None = 'auto', target_chunk_size: int | str | None = 'auto', donate: bool = False)

JIT-compiled schedule-based off-surface Bext.

compute_external_gradB(B0, *, quad_nt: int | None = None, quad_np: int | None = None, digits: int | None = None, hedgehog_order: int = 8, chunk_size: int | str | None = 'auto', target_chunk_size: int | str | None = 'auto', pou_dtype=None, patch_dtype=None, interp_block_size: int | str | None = 'auto', remat: bool | None = None, scan_targets: bool = False, patch_dim0: int | None = None, patch_idx=None)

Compute GradBext from total B on the source grid.

Args:

scan_targets: If True, use a lax.scan over target points

inside the singular correction to reduce peak memory.

compute_external_gradB_batch(B0_batch, **kwargs)

Vectorized compute_external_gradB over a batch dimension.

compute_external_gradB_jit(B0, **kwargs)

JIT-compiled version of compute_external_gradB.

compute_external_gradB_offsurf(B0, *, X_trg, digits: int | None = None, max_Nt: int = -1, max_Np: int = -1, max_levels: int = 6, adaptive: bool = False, chunk_size: int | str | None = 'auto', target_chunk_size: int | str | None = 'auto')

Compute GradBext at off-surface targets using direct quadrature.

The off-surface GradB path mirrors the reference implementation and currently uses the base resampled grid (no adaptive refinement).

compute_external_gradB_offsurf_schedule(B0, *, X_trg, levels: tuple[tuple[int, int], ...] | str | None, digits: int | None = None, max_Nt: int = -1, max_Np: int = -1, max_levels: int = 6, chunk_size: int | str | None = 'auto', target_chunk_size: int | str | None = 'auto', return_estimate: bool = False)

Compute GradBext off-surface using a fixed adaptive refinement schedule.

compute_external_gradB_offsurf_schedule_jit(B0, *, X_trg, levels: tuple[tuple[int, int], ...] | str | None = 'auto', digits: int | None = None, max_Nt: int = -1, max_Np: int = -1, max_levels: int = 6, chunk_size: int | str | None = 'auto', target_chunk_size: int | str | None = 'auto', donate: bool = False)

JIT-compiled schedule-based off-surface GradBext.

compute_internal_B(B0, *, X_trg=None, quad_nt: int | None = None, quad_np: int | None = None, digits: int | None = None, chunk_size: int | str | None = 'auto', target_chunk_size: int | str | None = 'auto', pou_dtype=None, patch_dtype=None, interp_block_size: int | str | None = 'auto', remat: bool | None = None, patch_dim0: int | None = None, patch_idx=None, precision: PrecisionPlan | None = None)

Compute Bint from total B on the source grid.

compute_internal_B_batch(B0_batch, *, X_trg=None, **kwargs)

Vectorized compute_internal_B over a batch dimension.

compute_internal_B_jit(B0, **kwargs)

JIT-compiled version of compute_internal_B.

compute_internal_B_offsurf(B0, *, X_trg, digits: int | None = None, max_Nt: int = -1, max_Np: int = -1, max_levels: int = 6, chunk_size: int | str | None = 'auto', target_chunk_size: int | str | None = 'auto')

Compute Bint at off-surface targets using adaptive quadrature.

compute_internal_B_offsurf_schedule(B0, *, X_trg, levels: tuple[tuple[int, int], ...] | str | None, digits: int | None = None, max_Nt: int = -1, max_Np: int = -1, max_levels: int = 6, chunk_size: int | str | None = 'auto', target_chunk_size: int | str | None = 'auto', return_estimate: bool = False)

Compute Bint off-surface using a fixed adaptive refinement schedule.

compute_internal_B_offsurf_schedule_jit(B0, *, X_trg, levels: tuple[tuple[int, int], ...] | str | None = 'auto', digits: int | None = None, max_Nt: int = -1, max_Np: int = -1, max_levels: int = 6, chunk_size: int | str | None = 'auto', target_chunk_size: int | str | None = 'auto', donate: bool = False)

JIT-compiled schedule-based off-surface Bint.

compute_internal_gradB(B0, *, quad_nt: int | None = None, quad_np: int | None = None, digits: int | None = None, hedgehog_order: int = 8, chunk_size: int | str | None = 'auto', target_chunk_size: int | str | None = 'auto', pou_dtype=None, patch_dtype=None, interp_block_size: int | str | None = 'auto', remat: bool | None = None, scan_targets: bool = False, patch_dim0: int | None = None, patch_idx=None)

Compute GradBint from total B on the source grid.

Args:

scan_targets: If True, use a lax.scan over target points

inside the singular correction to reduce peak memory.

compute_internal_gradB_batch(B0_batch, **kwargs)

Vectorized compute_internal_gradB over a batch dimension.

compute_internal_gradB_jit(B0, **kwargs)

JIT-compiled version of compute_internal_gradB.

compute_internal_gradB_offsurf(B0, *, X_trg, digits: int | None = None, max_Nt: int = -1, max_Np: int = -1, max_levels: int = 6, adaptive: bool = False, chunk_size: int | str | None = 'auto', target_chunk_size: int | str | None = 'auto')

Compute GradBint at off-surface targets using direct quadrature.

compute_internal_gradB_offsurf_schedule(B0, *, X_trg, levels: tuple[tuple[int, int], ...] | str | None, digits: int | None = None, max_Nt: int = -1, max_Np: int = -1, max_levels: int = 6, chunk_size: int | str | None = 'auto', target_chunk_size: int | str | None = 'auto', return_estimate: bool = False)

Compute GradBint off-surface using a fixed adaptive refinement schedule.

compute_internal_gradB_offsurf_schedule_jit(B0, *, X_trg, levels: tuple[tuple[int, int], ...] | str | None = 'auto', digits: int | None = None, max_Nt: int = -1, max_Np: int = -1, max_levels: int = 6, chunk_size: int | str | None = 'auto', target_chunk_size: int | str | None = 'auto', donate: bool = False)

JIT-compiled schedule-based off-surface GradBint.

plan_precision(*, digits: int | None = None, quad_nt: int | None = None, quad_np: int | None = None) → PrecisionPlan

Freeze geometry-dependent precision choices for traced evaluations.

Call this on a concrete surface, then reuse the returned plan for small differentiable perturbations of that surface. The plan does not mutate the object’s stored B quadrature setup.

setup(digits: int, nfp: int, half_period: bool, surf_nt: int, surf_np: int, X, src_nt: int, src_np: int, trg_nt: int, trg_np: int)

SIMSOPT-compatible VirtualCasing class backed by virtual_casing_jax.

class virtual_casing_jax.simsopt_virtual_casing.VirtualCasing

Bases: object

SIMSOPT-compatible VirtualCasing class backed by JAX.

This class mirrors simsopt.mhd.virtual_casing.VirtualCasing so it can be imported as:

from virtual_casing_jax import VirtualCasing

classmethod from_vmec(vmec, src_nphi, src_ntheta=None, trgt_nphi=None, trgt_ntheta=None, use_stellsym=True, digits=6, filename='auto')

Create a VirtualCasing object from a VMEC equilibrium.

This routine uses simsopt’s VMEC utilities and computes the external field using VirtualCasingJAX.

classmethod load(filename)

Load a virtual casing solution from a NetCDF file.

plot(ax=None, show=True)

Plot B_external_normal and B_external_normal_extended.

save(filename='vcasing.nc')

Save the results of a virtual casing calculation in a NetCDF file.

High-level exterior magnetic-field wrapper for virtual casing.

virtual_casing_jax.exterior_field.B_cyl_from_B_xyz(field_fn: Callable, R_phi_Z)

Evaluate a Cartesian field callback and return cylindrical components.

class virtual_casing_jax.exterior_field.ExteriorFieldConfig(digits: int = 8, src_nphi: int = 64, src_ntheta: int = 64, levels: tuple[tuple[int, int], ...] | None = None, chunk_size: int | str = 'auto', target_chunk_size: int | str = 'auto', branch: Literal['internal', 'external'] = 'internal', use_jit_schedule: bool = True, validate_orientation: bool = True, dtype: str = 'float64')

Bases: object

Configuration for fixed-schedule exterior-field evaluation.

branch: Literal['internal', 'external'] = 'internal'

chunk_size: int | str = 'auto'

digits: int = 8

dtype: str = 'float64'

levels: tuple[tuple[int, int], ...] | None = None

src_nphi: int = 64

src_ntheta: int = 64

target_chunk_size: int | str = 'auto'

use_jit_schedule: bool = True

validate_orientation: bool = True

class virtual_casing_jax.exterior_field.NearSurfaceTaylorPlan(center_rz: Array, surface_rz: Array, B_surface: Array, gradB_surface: Array, nfp: int)

Bases: object

Reusable interpolation data for first-order near-surface continuation.

B_surface: Array

center_rz: Array

gradB_surface: Array

nfp: int

surface_rz: Array

class virtual_casing_jax.exterior_field.VirtualCasingExteriorField(surface_data: VmecSurfaceFieldData, config: ExteriorFieldConfig | None = None, external_B_fn: Callable | None = None, external_gradB_fn: Callable | None = None)

Bases: object

JAX-native exterior field from VMEC-surface virtual-casing data.

B_plasma_xyz defaults to the internal off-surface branch because VMEC plasma currents are inside the LCFS. The external branch means currents outside the VMEC surface, not targets outside the VMEC surface.

B_and_derivatives_xyz(xyz, *, order=0, level=None, source_chunk=4096)

B and its spatial derivatives at xyz, from the closed forms.

One traversal of the source sum returns every order up to order, where nesting jacfwd over B_plasma_xyz() costs one traversal per order and differentiates the quadrature rather than the field.

Returns a tuple of order + 1 arrays, the k-th with shape xyz.shape[:-1] + (3,) * (k + 1).

B_cyl(R_phi_Z)

Evaluate the total field and return (B_R, B_phi, B_Z).

B_external_xyz(xyz)

Return the diagnostic external-branch virtual-casing field.

B_plasma_error_estimate(xyz, *, branch: Literal['internal', 'external'] | None = None)

Achieved relative error of B_plasma_xyz(), per target.

Compare with 10**-config.digits. Valid on either side of the source surface. A target closer to it than about two source-grid spacings reports a large error: the periodic trapezoid rule is near-singular there and the direct schedule is not the right tool, whatever level it reaches.

B_plasma_near_surface_xyz(xyz, plan: NearSurfaceTaylorPlan)

Evaluate a first-order continuation of the plasma field.

This path is intended for targets close enough to the source surface that direct periodic trapezoidal quadrature would require prohibitive refinement. Returned vectors are Cartesian and batched over xyz.

B_plasma_on_surface(*, digits=None, chunk_size=None, quad_nt=None, quad_np=None, precision=None)

Return the internal-current plasma field on the source surface.

This reuses the geometry prepared by the exterior field, avoiding a second quadrature setup when boundary diagnostics and off-surface evaluations are needed together.

B_plasma_xyz(xyz, *, branch: Literal['internal', 'external'] | None = None, return_estimate: bool = False)

Return the virtual-casing plasma-current field in Cartesian components.

With return_estimate also return the schedule’s achieved relative error per target, dimensionless and comparable with 10**-digits, so an unconverged schedule reports what it achieved instead of silently handing back its last level. See B_plasma_error_estimate().

B_xyz(xyz)

Return coil/external callback plus the plasma virtual-casing field.

B_xyz_sharded(xyz, *, devices=None)

Evaluate a large (n, 3) target batch across JAX devices.

Surface data are replicated while the independent target axis is partitioned. A single-device process follows B_xyz() directly.

export_rphiz_grid(R, phi, Z, *, chunk_size: int | str = 'auto')

Evaluate this field on a tensor-product cylindrical grid.

gradB_plasma_xyz(xyz, *, branch: Literal['internal', 'external'] | None = None)

Return dB_i/dx_j for the virtual-casing plasma field.

gradB_xyz(xyz)

Return dB_i/dx_j for the total field.

JAX differentiates external_B_fn point by point when an analytic external_gradB_fn was not supplied.

level_sources(level=None)

Quadrature nodes and layer densities of one schedule level.

Returns (nodes, sheet_current, charge) with shapes (m, 3), (m, 3) and (m,), the densities already carrying the quadrature weight, ready for layer_derivatives().

level defaults to the finest of the schedule – the one the two-level default returns unconditionally. The densities are built the way the schedule builds them: the field is recovered pointwise on the base grid and interpolated, and the products are formed on the level with the level’s own normal, so these nodes reproduce B_plasma_xyz() rather than approximating it.

plan_near_surface(*, digits=None, precision=None, B_surface=None)

Prepare a smooth first-order field continuation near the surface.

The singular on-surface field and gradient are computed once. Reusing the returned plan avoids increasingly fine direct quadrature for every nearby target in a field-line or grid evaluation.

plan_surface_precision(*, digits=None, quad_nt=None, quad_np=None)

Plan on-surface singular quadrature using this prepared geometry.

class virtual_casing_jax.exterior_field.VmecSurfaceFieldData(gamma: Array, B_total: Array, normal: Array, area_vector: Array, theta: Array, phi: Array, nfp: int, stellsym: bool, signgs: int, source_convention: str = 'vmex')

Bases: object

VMEC boundary data on a virtual-casing source grid.

Arrays use structure-of-arrays layout (3, nphi, ntheta). phi is the physical toroidal angle, not the VMEC field-period coordinate zeta.

B_total: Array

area_vector: Array

gamma: Array

nfp: int

normal: Array

phi: Array

signgs: int

source_convention: str = 'vmex'

stellsym: bool

theta: Array

virtual_casing_jax.exterior_field.build_near_surface_taylor_plan(surface_data, B_surface, gradB_surface)

Build a smooth first-order near-surface continuation plan.

Vector and gradient components are stored in the local cylindrical basis, making their Fourier representation periodic over one field period.

virtual_casing_jax.exterior_field.cyl_to_xyz(R_phi_Z)

Convert cylindrical point coordinates (R, phi, Z) to Cartesian.

virtual_casing_jax.exterior_field.default_schedule_levels(nphi: int, ntheta: int, nfp: int)

Two-level schedule sized from the source grid, in full-torus counts.

surface_data is sampled on ONE field period, while a schedule level counts the WHOLE torus, so the toroidal entries carry the nfp factor. Without it an nfp = 5 boundary sampled at 32 points per period was resolved by a finest level of 64 over the torus – 13 points per period – and the achieved error one minor radius off the boundary was 2.5e-02 against a requested 1e-6.

Two levels are enough: the coarse level exists only to supply the Richardson term of the error estimate, and it costs a quarter of the fine one. A schedule of one level is accepted and reports the double-layer quadrature error alone.

virtual_casing_jax.exterior_field.xyz_vec_to_cyl_vec(R_phi_Z, B_xyz)

Convert Cartesian vector components to cylindrical components.

Cylindrical grid evaluation and lightweight NetCDF export.

virtual_casing_jax.grid_export.evaluate_on_rphiz_grid(field, R, phi, Z, *, chunk_size: int | str = 'auto') → dict

Return BR, Bphi, BZ, and absB on an R, phi, Z grid.

virtual_casing_jax.grid_export.write_extended_field_netcdf(path, grid_data: Mapping, metadata: Mapping | None = None)

Write an extended-field cylindrical grid to a NetCDF file.

virtual_casing_jax.grid_export.write_mgrid_like(path, grid_data: Mapping, metadata: Mapping | None = None)

Write a lightweight MGRID-like NetCDF file with explicit convention metadata.

Smooth objective helpers for exterior-field workflows.

virtual_casing_jax.objectives.boundary_normal_residual(surface_data, field)

Return B_total . n on the VMEC boundary source grid.

virtual_casing_jax.objectives.exterior_B_magnitude_squared(field, target_xyz, weights, target_absB)

Weighted mean-square magnetic-field magnitude objective.

virtual_casing_jax.objectives.exterior_Bn_squared(field, target_xyz, target_normals, weights, target_Bn=0.0)

Weighted mean-square normal-field objective on target points.

virtual_casing_jax.objectives.grid_divergence_penalty(BR, Bphi, BZ, R, phi, Z, weights=None)

Finite-difference cylindrical div B mean-square diagnostic.

Test data utilities mirroring virtual-casing VirtualCasingTestData.

class virtual_casing_jax.testdata.Drand48(state: int = 20017429951246)

Bases: object

Deterministic drand48-compatible generator.

rand() → float

state: int = 20017429951246

class virtual_casing_jax.testdata.SurfType(value)

Bases: Enum

AxisymCircleNarrow = 1

AxisymCircleWide = 0

AxisymNarrow = 3

AxisymWide = 2

LHD = 7

NoneType = 11

Quas3 = 6

RotatingEllipseNarrow = 5

RotatingEllipseWide = 4

Stell = 9

W7X = 8

W7X_ = 10

class virtual_casing_jax.testdata.VirtualCasingTestData

Bases: object

JAX mirror of virtual-casing VirtualCasingTestData.

static magnetic_field_data(nfp: int, half_period: bool, nt: int, npol: int, X, trg_nt: int, trg_np: int)

static magnetic_field_data_offsurf(nfp: int, half_period: bool, nt: int, npol: int, X, X_trg)

static magnetic_field_grad_data(nfp: int, half_period: bool, nt: int, npol: int, X, trg_nt: int, trg_np: int)

static magnetic_field_grad_data_offsurf(nfp: int, half_period: bool, nt: int, npol: int, X, X_trg)

static surface_coordinates(nfp: int, half_period: bool, nt: int, npol: int, surf_type: SurfType = SurfType.AxisymNarrow)

virtual_casing_jax.testdata.magnetic_field_data(nfp: int, half_period: bool, surf_nt: int, surf_np: int, X, trg_nt: int, trg_np: int, *, rng: Drand48 | None = None)

Generate B field data for testing VirtualCasingJAX.

virtual_casing_jax.testdata.magnetic_field_data_offsurf(nfp: int, half_period: bool, surf_nt: int, surf_np: int, X, X_trg, *, rng: Drand48 | None = None, chunk_size: int = 1024)

Generate B field at arbitrary target points from synthetic loops.

virtual_casing_jax.testdata.magnetic_field_grad_data(nfp: int, half_period: bool, surf_nt: int, surf_np: int, X, trg_nt: int, trg_np: int, *, rng: Drand48 | None = None)

Generate GradB data for testing VirtualCasingJAX.

virtual_casing_jax.testdata.magnetic_field_grad_data_offsurf(nfp: int, half_period: bool, surf_nt: int, surf_np: int, X, X_trg, *, rng: Drand48 | None = None, chunk_size: int = 512)

Generate analytic GradB at arbitrary targets from synthetic loops.

virtual_casing_jax.testdata.surface_coordinates(nfp: int, half_period: bool, nt: int, npol: int, surf_type: SurfType = SurfType.AxisymNarrow)

Singular quadrature utilities (POU + polar correction).

class virtual_casing_jax.singular_quadrature.SingularPrecomp(patch_dim0: 'int', hedgehog_order: 'int', rad_dim_base: 'int', rad_dim: 'int', ang_dim: 'int', patch_dim: 'int', ngrid: 'int', npolar: 'int', qx: 'jnp.ndarray', qw: 'jnp.ndarray', Gpou: 'jnp.ndarray', Ppou: 'jnp.ndarray', I_G2P: 'jnp.ndarray', M_G2P: 'jnp.ndarray', interp_idx: 'jnp.ndarray', hedgehog_wts: 'jnp.ndarray')

Bases: object

Gpou: Array

I_G2P: Array

M_G2P: Array

Ppou: Array

ang_dim: int

hedgehog_order: int

hedgehog_wts: Array

interp_idx: Array

ngrid: int

npolar: int

patch_dim: int

patch_dim0: int

qw: Array

qx: Array

rad_dim: int

rad_dim_base: int

virtual_casing_jax.singular_quadrature.precompute_singular(patch_dim0: int, rad_dim: int, hedgehog_order: int = 1, pou_dtype=None, patch_dtype=None, index_dtype=None)

Build JAX constants from a bounded cache of host NumPy arrays.

JAX conversion is intentionally outside the cached function. Otherwise a JIT-first invocation can cache tracers globally and leak them into later eager calls.

virtual_casing_jax.singular_quadrature.select_patch_dim(digits: int, cond: float)