Implementation

Structure

The code is organized into the following modules:

surface_ops: FFT-based surface operators.
kernels: Laplace and Biot-Savart kernels.
integrals: baseline direct-sum surface quadrature.
singular_quadrature: POU + polar quadrature.
boundary_integral: surface integrals with singular correction.
virtual_casing: high-level API for on-surface/off-surface fields and gradients.

Design Decisions

Static Shapes

JAX compiles for fixed shapes. We select quadrature sizes on the Python side, then use JIT-compiled kernels for the actual evaluation.

Mixed Precision

Mixed precision is allowed. The default policy is:

Inputs can be float32.
Accumulations are promoted to float64 when required for accuracy.

Custom VJPs

Boundary integrals can be memory-intensive. We plan to implement custom_vjp rules that re-evaluate kernels during the backward pass instead of storing large intermediates.

Compatibility with Reference Code

The goal is bitwise-identical results for small test cases, and numerical parity for larger grids. All critical numerics and normalizations follow the BIEST implementation [BIEST] and the formulation in [MCO2019].

Boundary Integrals (Baseline)

The first JAX implementation uses a direct-sum quadrature:

\[\begin{split}u(x_i) = \\sum_{j} K(x_i, x_j) \\, f_j \\, w_j\end{split}\]

where w_j is the area element returned by SurfNormalAreaElem. This reproduces the far-field part of BIEST’s Nyström evaluation.

The singular correction used by BIEST (partition-of-unity + polar quadrature) is not yet replicated. The baseline is therefore expected to deviate near the diagonal, and is used to build the parity harness and profiling instrumentation before introducing the full correction.

The direct-sum implementation is chunked to limit memory use and is JIT-compatible using jax.lax.scan.

Singular Correction (Laplace FxdU / Fxd2U)

We implement the BIEST partition-of-unity correction for both LaplaceFxdU (grad G) and LaplaceFxd2U (hyper-singular second derivatives):

A local patch is extracted around each target point.
A grid POU term subtracts the singular contribution from the trapezoidal rule.
A polar quadrature term adds back the singular part using Lagrange interpolation from the patch to polar nodes.
Hedgehog quadrature (order 8) is used for the Fxd2U corrections.

This supports parity for both ComputeB and ComputeGradB on on-surface targets.

Internal Field Variants

Internal fields are derived by sign flips relative to the external formulation:

\[\mathbf{B}_{\mathrm{int}} = \frac{1}{2}\mathbf{B} - \nabla G[\sigma] - \nabla \times G[\mathbf{K}],\]

\[\nabla \mathbf{B}_{\mathrm{int}} = - \nabla \mathbf{B}_{\mathrm{ext}}.\]

The JAX implementation exposes:

compute_internal_B (on-surface)
compute_internal_gradB (on-surface)
compute_internal_B_offsurf (off-surface)
compute_internal_gradB_offsurf (off-surface)

These mirror the C++ API names and are validated in parity tests.

Off-Surface Baseline

For off-surface targets the kernels are non-singular. The implementation evaluates:

\[\begin{split}\mathbf{B}_{\mathrm{ext}}(x) = \\nabla G[B\\cdot n](x) - \\text{BiotSavart}[J](x)\end{split}\]

and the corresponding field gradient:

\[(\nabla \mathbf{B}_{\mathrm{ext}})_{k i} = \varepsilon_{k \ell m}\,\partial_i\partial_\ell G[K_m] + \partial_i\partial_k G[\sigma].\]

Optional Fourier upsampling improves accuracy. For off-surface GradB, the default parity path uses the base resampled grid (matching the reference C++ implementation), with an optional adaptive mode enabled for higher accuracy.

Off-Surface GradB

Off-surface GradB is evaluated with the same second-derivative kernel used on-surface, but without singular correction. The JAX path uses laplace_fxd2_u_eval_vec for the vector density and assembles the curl term to match the reference implementation. An adaptive option is available to reuse the refined grid selected by the LaplaceDxU self-test.