API — nefi.physics¶
Generated by scripts/gen_api_docs.py — do not edit by hand. Index
nefi.physics¶
Physics zoo: reusable differentiable forward models for many physical systems.
nefi.physics.elliptic¶
Variable-coefficient elliptic PDEs −∇·(σ∇u) + κu = f with an implicit-function adjoint.
class AxisConductance(interior: torch.Tensor, lo: torch.Tensor | None = None, hi: torch.Tensor | None = None)¶
Face conductances of one axis.
class CGInfo(iterations: int, residual: float, converged: bool, tol: float)¶
Convergence report of :func:conjugate_gradient.
class EllipticOperator(domain: Domain, rhs: torch.Tensor | Callable[[Domain], torch.Tensor], bc: BCSpec = 'dirichlet', *, field: str = 'sigma', sigma_transform: str | Callable[[to...¶
Steady-state elliptic forward model {σ} ↦ u solving −∇·(σ∇u) + κu = f.
at_resolution(self, shape: Sequence[int]) -> EllipticOperatorcoefficient(self, fields: Fields) -> torch.Tensor— σ (after the transform) from a field dict.extra_repr(self) -> strforward(self, fields: Fields) -> torch.Tensorkappa_value(self, fields: Fields, ref: torch.Tensor) -> torch.Tensor | Nonelast_adjoint_info(property) — Convergence report of the latest auxiliary solve (adjoint or forward-mode tangent).n_drives(property) — Number of right-hand sides (Nonefor a single unbatched source).output_shape(self, shape: Sequence[int]) -> tuple[int, ...]post(self, u: torch.Tensor, sigma: torch.Tensor, fields: Fields) -> torch.Tensor— Map the state to the prediction (default: optional weighted-mean referencing).reference_weights(self, device=None, dtype=None) -> torch.Tensor | Nonerequired_fields(self) -> tuple[str, ...]reset_warm_start(self)— Forget cached forward/adjoint solutions (e.g. before a fresh restart).rhs(self, device=None, dtype=None) -> torch.Tensor— The source density at this operator's resolution.solve(self, sigma: torch.Tensor, kappa: torch.Tensor | None = None) -> torch.Tensor— Solve for the stateugiven the (transformed) coefficientσ.spacing(property)transform(self, x: torch.Tensor) -> torch.Tensor— Applysigma_transformto a raw coefficient field.
class EllipticSolveConfig(spacing: tuple[float, ...], bc: tuple[tuple[str, str], ...], face_mode: str = 'harmonic', tol: float = 1e-08, atol: float = 0.0, max_iter: int | None = None...¶
Static settings of an elliptic solve (non-tensor argument of :class:ImplicitSolveFunction).
derived(self) -> EllipticSolveConfig— Copy for an auxiliary solve (adjoint / tangent) with its own diagnostics, no cache.has_dirichlet(property)iteration_cap(self, grid: Sequence[int]) -> intndim(property)singular(self, kappa: torch.Tensor | float | None) -> bool
class ImplicitSolveFunction(*args, **kwargs)¶
u = A(σ, κ)⁻¹ b with exact implicit-function derivative rules (O(N) memory).
backward(ctx, grad_u)forward(sigma, rhs, kappa, cfg: EllipticSolveConfig, x0=None)jvp(ctx, sigma_dot, rhs_dot, kappa_dot, _cfg_dot, _x0_dot)setup_context(ctx, inputs, output)vmap(info, in_dims, sigma, rhs, kappa, cfg, x0)
class LogBounded(lo: float, hi: float, init_value: float | None = None)¶
Positive coefficient in [lo, hi] with a log-uniform parameterization.
init_bias(self)inverse(self, value)transform(self, raw, others)
def apply_conductances(u: torch.Tensor, conds: Sequence[AxisConductance], bc: BCSpec, kappa: torch.Tensor | float | None = None) -> torch.Tensor¶
A u for precomputed face conductances (used inside the CG loop).
def apply_divgrad(u: torch.Tensor, sigma: torch.Tensor, spacing: SpacingSpec, bc: BCSpec = 'dirichlet', face_mode: str = 'harmonic', kappa: torch.Tensor | float | None = None...¶
Matrix-free −∇·(σ∇u) + κu on a uniform 1-/2-/3-D Cartesian grid.
def conjugate_gradient(matvec: Callable[[torch.Tensor], torch.Tensor], b: torch.Tensor, ndim: int, *, x0: torch.Tensor | None = None, precond: Callable[[torch.Tensor], torch.Tenso...¶
Batched preconditioned conjugate gradients for SPD A (Hestenes & Stiefel 1952).
def divgrad_diagonal(sigma: torch.Tensor, spacing: SpacingSpec, bc: BCSpec = 'dirichlet', face_mode: str = 'harmonic', kappa: torch.Tensor | float | None = None) -> torch.Tensor¶
Diagonal of A(σ) (the Jacobi preconditioner), shape of σ (broadcast with κ).
def face_conductances(sigma: torch.Tensor, spacing: SpacingSpec, bc: BCSpec = 'dirichlet', face_mode: str = 'harmonic', ndim: int | None = None) -> list[AxisConductance]¶
Per-axis face conductances σ̄/h² of the finite-volume operator (see module docstring).
def face_mean(a: torch.Tensor, b: torch.Tensor, mode: str = 'harmonic') -> torch.Tensor¶
Face coefficient from the two adjacent cell values.
def is_transformed(t: Any) -> bool¶
True for forward-mode dual tensors and tensors wrapped by torch.func transforms.
def masked_mean(x: torch.Tensor, weights: torch.Tensor, ndim: int, eps: float = 1e-30) -> torch.Tensor¶
Weighted mean over the trailing ndim axes (keepdim), weights broadcastable to x.
def normalize_bc(bc: BCSpec, ndim: int) -> tuple[tuple[str, str], ...]¶
Normalize a boundary-condition spec to ((lo, hi), ...), one pair per axis.
def point_source_rhs(domain: Domain, positions: torch.Tensor | Sequence[Sequence[float]], rates: torch.Tensor | Sequence[float], *, spread: str = 'linear', dtype: torch.dtype = ...¶
Source density of point sources/sinks (wells, optodes, electrodes) on domain's grid.
def solve_elliptic(sigma: torch.Tensor, rhs: torch.Tensor, spacing: SpacingSpec, bc: BCSpec = 'dirichlet', *, kappa: torch.Tensor | float | None = None, x0: torch.Tensor | Non...¶
Solve −∇·(σ∇u) + κu = f (matrix-free Jacobi-PCG) for one or many right-hand sides.
def weighted_area_downsample(data: torch.Tensor, weights: torch.Tensor, shape: Sequence[int]) -> tuple[torch.Tensor, torch.Tensor]¶
Weighted area-average of sparse observations onto a coarser (or equal) grid.
nefi.physics.magnetostatics¶
Magnetostatics of planar sources imaged on a parallel plane (NV / SQUID / Hall microscopy).
class BiotSavartOperator(obs_domain: Domain, z0: float, *, field: str = 'g', mu0: float = 1.2566370614359173e-06, components: str = 'z', chunk_elems: int = 4000000)¶
Direct real-space forward model g ↦ B (float64 data generation, inverse-crime guard).
forward(self, fields: Fields) -> torch.Tensoroutput_shape(self, shape: Sequence[int]) -> tuple[int, ...]
class CurrentDensityOperator(domain: Domain, z0: float, *, source: str = 'stream', field: str = 'g', current_fields: Sequence[str] = ('jx', 'jy'), mu0: float = 1.2566370614359173e-06, p...¶
Thin current sheet → stray field at standoff z0 (linear, homogeneity = 1).
at_resolution(self, shape: Sequence[int]) -> CurrentDensityOperatorbz_hat(self, fields, kgrid)required_fields(self) -> tuple[str, ...]transfer(self, device=None, dtype=None) -> torch.Tensor—B̂_z / ĝ = (μ0/2) k e^{−k z0}on the padded rfft grid (stream-function source).
class MagnetizationOperator(domain: Domain, z0: float, thickness: float, *, field: str = 'mz', direction: Sequence[float] = (0.0, 0.0, 1.0), mu0: float = 1.2566370614359173e-06, pad_fa...¶
Magnetized thin film (uniaxial, fixed direction d̂) → stray field at standoff z0.
at_resolution(self, shape: Sequence[int]) -> MagnetizationOperatorbz_hat(self, fields, kgrid)transfer(self, device=None, dtype=None) -> torch.Tensor— ComplexB̂_z / M̂on the padded rfft grid.
def biot_savart_segments(starts: torch.Tensor, ends: torch.Tensor, currents: torch.Tensor | float, points: torch.Tensor, *, mu0: float = 1.2566370614359173e-06) -> torch.Tensor¶
Exact field (M, 3) of straight line-current segments (thin wires, polygonal loops).
def biot_savart_sheet(jx: torch.Tensor, jy: torch.Tensor, src_xy: torch.Tensor, obs_xy: torch.Tensor, z0: float, *, cell_area: float, mu0: float = 1.2566370614359173e-06, compone...¶
Direct Biot–Savart summation of a discretized current sheet (midpoint rule).
def current_divergence(jx: torch.Tensor, jy: torch.Tensor, spacing: Sequence[float], method: str = 'central', pad_factor: float = 1.0) -> torch.Tensor¶
∂_x K_x + ∂_y K_y with the stencil matching :func:stream_to_current.
def fourier_inversion(bz: torch.Tensor, z0: float, spacing: Sequence[float], *, mu0: float = 1.2566370614359173e-06, reg: float = 0.001, window: str = 'hanning', cutoff_wavelengt...¶
Classical regularized k-space inversion of a B_z map (the standard NV baseline).
def magnetic_constant(length: str = 'm', current: str = 'A', field: str = 'T') -> float¶
μ0 expressed in [field]·[length]/[current].
def padded_shape(shape: Sequence[int], pad_factor: float = 2.0) -> tuple[int, ...]¶
FFT grid for zero-padded transforms: shape itself for pad_factor <= 1 (periodic), else the next FFT-friendly size >= pad_factor·n per axis.
def stream_to_current(g: torch.Tensor, spacing: Sequence[float], method: str = 'central', pad_factor: float = 1.0) -> tuple[torch.Tensor, torch.Tensor]¶
Sheet current K = ∇×(g ẑ) = (∂_y g, −∂_x g) of a stream function g (..., H, W).
def upward_continuation(b: torch.Tensor, dz: float, spacing: Sequence[float], *, pad_factor: float = 2.0, allow_downward: bool = False) -> torch.Tensor¶
Continue a potential-field map b (..., H, W) from its plane to dz further away.
def wavenumbers(shape: Sequence[int], spacing: Sequence[float], device=None, dtype=None) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]¶
Angular wavenumbers of a 2-D real FFT grid: (k_x (P0,1), k_y (1,P1//2+1), |k|).
nefi.physics.optics¶
Scalar coherent optics: angular-spectrum propagation, thin phase objects, inline holography.
class HolographyOperator(domain: Domain, distances: Sequence[float], wavelength: float, *, field: str = 'phase', absorption: str | float | None = None, illumination: float = 1.0, ba...¶
Multi-distance inline holography: phase (and absorption) → intensities |P_z t|².
at_resolution(self, shape)extra_repr(self) -> strfields_at_detectors(self, fields: Fields) -> torch.Tensor— Complex fields at every distance(n_distances, H, W).forward(self, fields: Fields) -> torch.Tensorn_distances(property)output_shape(self, shape)required_fields(self)transmission(self, fields: Fields) -> torch.Tensor— Complex sample transmission (:class:PhaseObject).
class PhaseObject(field: str = 'phase', absorption: str | float | None = None, illumination: float = 1.0)¶
Thin-object (projection-approximation) transmission t = A · exp(iφ − a).
forward(self, fields: Fields) -> torch.Tensoroutput_shape(self, shape)required_fields(self)
def gaussian_beam(domain: Domain, w0: float, center: Sequence[float] | None = None, dtype=torch.float64) -> torch.Tensor¶
Gaussian beam waist exp(−r²/w0²) (complex) sampled on domain (at z = 0).
def gerchberg_saxton(intensities: torch.Tensor, op: HolographyOperator, n_iter: int = 100, pure_phase: bool = True, init_phase: torch.Tensor | None = None) -> tuple[torch.Tensor...¶
Multi-distance Gerchberg–Saxton phase retrieval (sequential plane projections).
def phase_object(phase: torch.Tensor, absorption: torch.Tensor | float | None = None)¶
Thin-object transmission exp(iφ − a) (complex, same shape as phase).
def propagate(u: torch.Tensor, dz: float | Sequence[float] | torch.Tensor, wavelength: float, spacing: Sequence[float] | float, *, band_limit: bool = True, pad_factor: fl...¶
Angular-spectrum propagation of u (*batch, H, W) over dz.
def rayleigh_range(w0: float, wavelength: float) -> float¶
z_R = π w0² / λ of a Gaussian beam with waist w0 (1/e² intensity radius).
def transfer_function(shape: Sequence[int], spacing: Sequence[float], wavelength: float, distances: Sequence[float] | torch.Tensor | float, band_limit: bool = True, evanescent: s...¶
Angular-spectrum transfer functions H_z on an FFT grid: (n_z, *shape) complex.
nefi.physics.reaction_diffusion¶
Gray–Scott reaction–diffusion with spatially varying reaction parameters.
class GrayScottIC(u_base: float = 1.0, v_base: float = 0.0, a_u: float = 0.5, a_v: float = 0.25, g_mean: float = 0.75, g_amp: float = 0.25, modes: tuple[int, ...] = (2, 3), p...¶
Known smooth initial condition evaluated analytically at any resolution.
pattern(self, domain: Domain, dtype=torch.float64) -> torch.Tensor
class ReactionDiffusionOperator(domain: Domain, obs_times: Sequence[float], *, field: str = 'F', unknown: str = 'F', F: float = 0.045, k: float = 0.06, Du: float = 2e-05, Dv: float = 1e-05...¶
Gray–Scott forward model: parameter field (F or k) → snapshots of u/v.
at_resolution(self, shape)extra_repr(self) -> strforward(self, fields: Fields) -> torch.Tensoroutput_shape(self, shape)record(property)simulate(self, param: torch.Tensor, observe=None, record=None)— Run the explicit-Euler loop; returns(final_state, observations).
def diffusion_term(x: torch.Tensor, spacing: Sequence[float], boundary: str = 'neumann', coeff: torch.Tensor | None = None) -> torch.Tensor¶
Conservative second-order diffusion operator over the trailing len(spacing) dims.
def gray_scott_step(u: torch.Tensor, v: torch.Tensor, F_: torch.Tensor | float, k: torch.Tensor | float, Du: torch.Tensor | float, Dv: float, dt: float, spacing: Sequence[float...¶
One explicit-Euler Gray–Scott step. Du may be a spatial field (conservative form).
def stable_dt(d_max: float, spacing: Sequence[float], F_max: float, k_max: float, reaction_rate_bound: float = 1.0, safety: float = 0.9) -> float¶
Largest stable explicit-Euler step for Gray–Scott (see module docstring).
nefi.physics.scattering¶
2-D scalar diffraction tomography: Born / Rytov linearized scattering and Lippmann–Schwinger.
class BornOperator(domain: Domain, wavelength: float, **kw)¶
Linearized (Born / first-order Rytov) 2-D diffraction-tomography forward model.
at_resolution(self, shape)forward(self, fields: Fields) -> torch.Tensorscattered_field(self, x: torch.Tensor, kernel_method: str = 'truncated') -> torch.Tensor— Born scattered field inside the domain,(n_angles, N1, N2)complex.
class LippmannSchwingerOperator(domain: Domain, wavelength: float, *, n_iter: int = 200, tol: float = 1e-10, relax: float = 1.0, kernel_method: str = 'truncated', **kw)¶
Multiple-scattering forward model (independent data generator for diffraction tomography).
at_resolution(self, shape)forward(self, fields: Fields) -> torch.Tensortotal_field(self, x: torch.Tensor) -> torch.Tensor— Total fielduon the grid for every incidence(n_angles, N1, N2).
def adjoint_reconstruction(op: Operator, data: torch.Tensor, shape: Sequence[int]) -> torch.Tensor¶
Back-projection Aᵀ d of a linear operator via autograd (unnormalized).
def circle_receivers(n: int, radius: float, center: Sequence[float] = (0.0, 0.0)) -> torch.Tensor¶
n receivers equally spaced on a circle: (n, 2) (float64), angles 2πm/n.
def complex_dtype(dtype: torch.dtype) -> torch.dtype¶
Complex dtype matching a real one (float64 → complex128, else complex64).
def filtered_backpropagation(data: torch.Tensor, op: _ScatteringBase, shape: Sequence[int] | None = None, n_harmonics: int | None = None) -> torch.Tensor¶
Filtered backpropagation (Devaney 1982) for full-view circular-receiver Born/Rytov data.
def green_2d(r: torch.Tensor, k0: float) -> torch.Tensor¶
Outgoing 2-D Helmholtz Green's function (i/4) H₀⁽¹⁾(k0 r) ((Δ + k0²)G = −δ).
def green_convolve(q: torch.Tensor, kernel_fft: torch.Tensor) -> torch.Tensor¶
G ∗ q on the grid of q (*batch, N1, N2) with a kernel from :func:helmholtz_kernel_fft (zero-padded linear convolution, same-size output).
def hankel1(order: int, z: torch.Tensor) -> torch.Tensor¶
Hankel function of the first kind H_ν⁽¹⁾(z) = J_ν(z) + i Y_ν(z) for ν ∈ {0, 1}.
def helmholtz_kernel_fft(shape: Sequence[int], spacing: Sequence[float], k0: float, method: str = 'truncated', eps: float | None = None, device=None, dtype: torch.dtype = torch.floa...¶
FFT of the discrete Green's-function convolution kernel for a zero-padded 2N grid.
def lippmann_schwinger(f: torch.Tensor, u_inc: torch.Tensor, kernel_fft: torch.Tensor, n_iter: int = 100, tol: float = 1e-10, relax: float = 1.0) -> tuple[torch.Tensor, dict]¶
Solve u = u_inc + G ∗ (f u) by relaxed fixed-point iteration (the Born series).
def plane_waves(coords: torch.Tensor, k0: float, angles: torch.Tensor) -> torch.Tensor¶
exp(i k0 d̂_j·x) for incidence angles θ_j: (n_angles, *coords.shape[:-1]).
def receiver_green_matrix(receivers: torch.Tensor, coords: torch.Tensor, k0: float, cell_area: float) -> torch.Tensor¶
Quadrature matrix M_{r,j} = G(|x_r − y_j|) ΔA: (n_rec, n_pixels) complex128.
nefi.physics.timestep¶
Generic explicit time integration with optional per-block gradient checkpointing.
def auto_checkpoint_every(n_steps: int) -> int¶
Block length ceil(sqrt(n_steps)) minimizing n_steps/k + k (memory-optimal).
def laplacian_eig_max(spacing: Sequence[float], order: int = 2) -> float¶
Largest eigenvalue magnitude of the discrete Laplacian, Σ_a ρ_order / h_a².
def run_timestepping(step_fn: StepFn, state0: torch.Tensor | Sequence[torch.Tensor], params: torch.Tensor | Sequence[torch.Tensor] | None, n_steps: int, dt: float, observe: Obse...¶
Integrate state_{n+1} = step_fn(state_n, params, n, dt) for n_steps steps.
def stable_dt_diffusion(d_max: float, spacing: Sequence[float], rate_max: float = 0.0, safety: float = 0.9) -> float¶
Largest stable explicit-Euler step for u_t = D Δu − r u (5-point / 2nd-order Laplacian).
def stable_dt_wave(c_max: float, spacing: Sequence[float], order: int = 2, courant: float = 0.9) -> float¶
Largest stable leapfrog step for p_tt = c² Δp times a safety factor courant.
nefi.physics.wave¶
Scalar acoustic wave equation: differentiable time-domain simulation and wave-family operators.
class WaveGrid(domain: Domain, absorb_width: float = 0.0)¶
Simulation grid = field grid padded by an absorbing layer of n_abs cells per side.
cell_volume(property)crop(self, x: torch.Tensor) -> torch.Tensor— Inverse of :meth:pad(keeps the physical-domain cells).depth(self) -> torch.Tensor— Cell-center relative depth per axis, shape(d, *padded).layer_depth(self, ax: int, faces: bool = False) -> torch.Tensor— Relative depthξ ∈ [0, 1]into the layer along axisax(1-D, float64).nearest_index(self, points: torch.Tensor) -> torch.Tensor— Flattened index of the padded-grid cell containing each physical point.pad(self, x: torch.Tensor, mode: str = 'replicate') -> torch.Tensor— Pad the trailing spatial dims ("replicate"for material maps,"zeros"else).point_weights(self, points: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]— Multilinear interpolation stencils of physicalpoints (n, d)on the padded grid.source_density(self, points: torch.Tensor) -> torch.Tensor— Discrete δ-functions atpoints:(n, *padded)(adjoint weights / cell volume).
class WaveInitialConditionOperator(domain: Domain, receivers, *, n_t: int, dt_obs: float, c: float | torch.Tensor = 1.0, c_max: float | None = None, field: str = 'p0', order: int = 4, absorbi...¶
Photoacoustic forward model: initial pressure p0(x) (zero initial velocity) → traces.
at_resolution(self, shape)forward(self, fields: Fields) -> torch.Tensoroutput_shape(self, shape)simulate(self, p0: torch.Tensor, observe=None, record=None)— Leapfrog run fromp0(field grid); see :meth:WaveOperator.simulate.
class WaveOperator(domain: Domain, sources, receivers, *, n_t: int, dt_obs: float, f0: float, c_max: float, field: str = 'c', t0: float | None = None, amplitude: float = 1.0, ...¶
Acoustic FWI forward model: sound speed c(x) → receiver traces for every source.
at_resolution(self, shape)forward(self, fields: Fields) -> torch.Tensorn_receivers(property)n_sources(property)output_shape(self, shape)simulate(self, c: torch.Tensor, observe: Callable | None = None, record: Sequence[int] | None = None, sources: slice | None = None) -> tuple[tuple[torch.Tensor, ...]...— Run the leapfrog loop for sound speedc(field grid) and return(state, obs).wavelet(self) -> torch.Tensor— Source time function at the simulation stepst_n = n·dt(n < n_steps).
def laplacian(p: torch.Tensor, spacing: Sequence[float], order: int = 2, kernel: torch.Tensor | None = None) -> torch.Tensor¶
Centered finite-difference Laplacian over the trailing len(spacing) dims.
def laplacian_kernel(spacing: Sequence[float], order: int, device=None, dtype=None) -> torch.Tensor¶
Stencil of :func:laplacian as a (1, 1, k, …, k) convolution kernel.
def ricker(t: torch.Tensor, f0: float, t0: float | None = None) -> torch.Tensor¶
Ricker ("Mexican hat") wavelet (1 − 2a²) exp(−a²), a = π f0 (t − t0).
def stable_dt(c_max: float, spacing: Sequence[float], order: int = 4, courant: float = 0.9) -> float¶
CFL-limited leapfrog step (see :func:nefi.physics.timestep.stable_dt_wave).
def time_reversal(traces: torch.Tensor, domain: Domain, receivers, *, dt_obs: float, c: float | torch.Tensor = 1.0, mode: str = 'dirichlet', order: int = 4, absorbing: str = ...¶
Time-reversal reconstruction of an initial pressure from receiver traces (PAT baseline).
def wave_energy(p_prev: torch.Tensor, p: torch.Tensor, c: torch.Tensor, spacing: Sequence[float], dt: float, order: int = 2) -> torch.Tensor¶
Discrete leapfrog energy E^{n-1/2} (conserved exactly without absorption and sources).