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

dx: (…, 3) n: (…, 3) source normals f: (…,) returns (…,)

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

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

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 method updates the result only while the double-layer self-test error exceeds the tolerance.

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)

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

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)

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.

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)

Vector-density wrapper for Laplace Fxd2U with singular correction.

Args:

scan_targets: 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)

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, 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, 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.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)

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, 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')

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, 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')

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)

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, 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')

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, 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')

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.

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.

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

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