API — nefi.operators¶
Generated by scripts/gen_api_docs.py — do not edit by hand. Index
nefi.operators¶
Differentiable forward operators.
nefi.operators.base¶
Forward-operator base class and generic wrappers.
class LambdaOperator(fn: Callable[[Fields], torch.Tensor], primary: str = 'x', homogeneity: float | None = None, output_shape_fn=None)¶
Wrap a plain function f(fields) -> Tensor (prototyping / tests).
forward(self, fields: Fields) -> torch.Tensoroutput_shape(self, shape)
class Nuisance(inner: Operator, gain: bool = True, offset: bool = False, init_gain: float = 1.0, init_offset: float = 0.0)¶
Learnable global gain / offset around an inner operator: exp(log_gain)·F(x) + offset.
at_resolution(self, shape)batchable(property)forward(self, fields: Fields) -> torch.Tensorgain(property)output_shape(self, shape)required_fields(self)traceable(property)
class Operator(*args: Any, **kwargs: Any)¶
Base class for differentiable forward models.
apply(self, x: torch.Tensor, **others: torch.Tensor) -> torch.Tensorat_resolution(self, shape: Sequence[int]) -> Operator— Operator evaluating fields sampled atshape.forward(self, fields: Fields) -> torch.Tensorget_field(self, fields: Fields, name: str | None = None) -> torch.Tensoroutput_shape(self, shape: Sequence[int]) -> tuple[int, ...] | None— Prediction shape for fields sampled atshape(Noneif not statically known).required_fields(self) -> tuple[str, ...]
class Sequential(*ops: Operator)¶
Compose operators: the output of each becomes the primary field of the next.
at_resolution(self, shape)batchable(property)forward(self, fields: Fields) -> torch.Tensorrequired_fields(self)traceable(property)
nefi.operators.conv¶
Generic N-D linear convolution via FFT with per-resolution kernel caches.
class FFTConvolution(kernel: torch.Tensor | KernelFn, domain: Domain, field: str = 'x', periodic: bool = False, kernel_extent: str = 'full', post: Callable[[torch.Tensor], torch...¶
Linear convolution y = k ∗ x on the domain grid, rebuilt at any resolution.
at_resolution(self, shape)clear_cache(self)forward(self, fields: Fields) -> torch.Tensorkernel(self, shape: Sequence[int] | None = None, device=None, dtype=None) -> torch.Tensor— Kernel sampled for a field grid ofshape(default the operator's domain shape).output_shape(self, shape)
def fft_convolve(x: torch.Tensor, kernel: torch.Tensor, periodic: bool = False, kernel_fft: torch.Tensor | None = None, fft_shape: Sequence[int] | None = None)¶
Convolve the trailing kernel.ndim dims of x with kernel (center at shape//2).
def gaussian_kernel_fn(sigma: float | Sequence[float], normalize: bool = True) -> KernelFn¶
Kernel function for an isotropic/anisotropic Gaussian blur with physical std sigma.
def kernel_offsets(spacing: Sequence[float], shape: Sequence[int], device=None, dtype=None)¶
Physical offset grid (*shape, d) centered at the kernel center shape // 2.
def prepare_kernel(kernel: torch.Tensor, n: Sequence[int], periodic: bool = False)¶
Precompute the kernel transform for inputs with trailing shape n.
nefi.operators.function¶
Wrap any differentiable Python callable as a nefi :class:~nefi.operators.Operator.
class FunctionOperator(fn: Callable, *, field: str = 'x', takes_dict: bool = False, fields: Sequence[str] | None = None, homogeneity: float | None = None, output_shape: ShapeFn | ...¶
A differentiable callable as a forward operator.
at_resolution(self, shape)extra_repr(self) -> strforward(self, fields: Fields) -> torch.Tensormultiscale_capable(property) — True when coarse grids are supported (at_resolution/output_shape/upsample).output_shape(self, shape)required_fields(self)
class UpsampleToNative(op: Operator, native_shape: Sequence[int], mode: str = 'linear')¶
Evaluate op on fields upsampled from a coarse grid to native_shape.
at_resolution(self, shape)forward(self, fields: Fields) -> torch.Tensoroutput_shape(self, shape)required_fields(self)
def as_operator(forward: Operator | Callable, field: str = 'x', **kw: object) -> Operator¶
Return forward if it is an :class:Operator, else wrap it in :class:FunctionOperator.
def field_shape(fields: Fields, name: str, ndim: int | None) -> tuple[int, ...] | None¶
Trailing ndim dims of fields[name] (None if unknown).
def upsample_fields(fields: Fields, shape: Sequence[int], mode: str = 'linear', names: Sequence[str] | None = None) -> dict[str, torch.Tensor]¶
Resample the trailing len(shape) dims of (selected) fields to shape.
nefi.operators.linear¶
Linear measurement operators: resolution loss, sampling, Fourier (MRI-style) sampling, algebra.
class Downsample(factor: int | Sequence[int] = 2, mode: str = 'area', field: str = 'x')¶
Resolution loss by an integer factor per axis (detector binning / low-res imaging).
forward(self, fields: Fields) -> torch.Tensoroutput_shape(self, shape)
class FourierSampling(mask: torch.Tensor, norm: str = 'ortho', real_output: bool = True, centered: bool = True, compact: bool = False, field: str = 'x', imag_field: str | None = ...¶
Undersampled Fourier measurements (MRI / interferometry): y = M ⊙ F x.
adjoint(self, y: torch.Tensor, real: bool = True) -> torch.Tensor— Zero-filled inverseF⁻¹ Mᵀ y(the classic baseline / a good initial guess).estimate_noise(self, measurement, outer: float = 0.5) -> float | None— Per-component noise std from the MAD of observed outer-k-space coefficients.fft(self, x: torch.Tensor) -> torch.Tensorforward(self, fields: Fields) -> torch.Tensorifft(self, k: torch.Tensor) -> torch.Tensormeasurement_mask(self) -> torch.Tensor | None— Mask for :class:~nefi.measurement.Measurement(zero-filled layout), elseNone.noise_region(self, outer: float = 0.5) -> torch.Tensor— Boolean mask of sampled coefficients beyondouter ×the maximal k-space radius.output_shape(self, shape)required_fields(self)
class Identity(field: str = 'x')¶
y = x (denoising / inpainting with Measurement.mask). Homogeneity 1.
forward(self, fields: Fields) -> torch.Tensoroutput_shape(self, shape)
class Pointwise(fn: Callable[[torch.Tensor], torch.Tensor], homogeneity: float | None = None, field: str = 'x')¶
Elementwise map y = fn(x) (detector response, nonlinearity, unit conversion).
forward(self, fields: Fields) -> torch.Tensoroutput_shape(self, shape)
class Sampling(mask: torch.Tensor | None = None, indices: torch.Tensor | None = None, field: str = 'x', native_shape: Sequence[int] | None = None)¶
Point sampling of the field: inpainting, sparse sensors, scattered probes.
forward(self, fields: Fields) -> torch.Tensormeasurement(self, data: torch.Tensor, noise_std: float | None = None)—Measurement(data, mask=self.measurement_mask(), noise_std=noise_std).measurement_mask(self) -> torch.Tensor | None— Mask to put in :class:~nefi.measurement.Measurement(mask mode), elseNone.output_shape(self, shape)
class Stack(*ops: Operator, mode: str = 'stack')¶
Several measurements of the same unknown(s) (multi-modal / multi-view data).
forward(self, fields: Fields) -> torch.Tensoroutput_shape(self, shape)
class Sum(*ops: Operator)¶
y = Σ_i op_i(fields) (superposition of contributions, e.g. two field components).
batchable(property)forward(self, fields: Fields) -> torch.Tensoroutput_shape(self, shape)
def random_kspace_mask(shape: Sequence[int], fraction: float = 0.3, center_fraction: float = 0.08, seed: int = 0, variable_density: bool = True) -> torch.Tensor¶
Random (variable-density) k-space sampling mask in centered layout (DC at shape//2).
nefi.operators.nonlinear¶
Nonlinear measurement operators: phase retrieval, Beer–Lambert attenuation, detector saturation.
class BeerLambert(axis: int = 0, path_op: Operator | Callable[[torch.Tensor], torch.Tensor] | None = None, I0: float | torch.Tensor = 1.0, domain: Domain | None = None, mode:...¶
Beer–Lambert attenuation I = I0 · exp(−∫ μ dl) (X-ray / optical transmission).
at_resolution(self, shape)forward(self, fields: Fields) -> torch.Tensorline_integral(self, mu: torch.Tensor) -> torch.Tensoroutput_shape(self, shape)
class PhaseRetrieval(oversample: int | Sequence[int] = 2, norm: str = 'ortho', native_shape: Sequence[int] | None = None, field: str = 'x')¶
Fourier phase retrieval (coherent diffraction imaging): y = |F pad(x)|².
forward(self, fields: Fields) -> torch.Tensoroutput_shape(self, shape)
class Saturation(kind: str = 'tanh', level: float = 1.0, gain: float = 1.0, threshold: float = 0.0, field: str = 'x')¶
Saturating detector response (pointwise), to compose after a linear operator.
forward(self, fields: Fields) -> torch.Tensoroutput_shape(self, shape)
nefi.operators.pde¶
Differentiable PDE operators: finite-volume heat conduction with a discrete adjoint (NeFTY).
nefi.operators.pde.adjoint¶
Implicit-Euler time stepping and its discrete adjoint (NeFTY §4.3, App. D.3).
class HeatSolveConfig(spacing: tuple[float, ...], dt: float, n_steps: int, obs_steps: tuple[int, ...], bc: BoundarySpec, face_mode: str = 'harmonic', solver: str = 'jacobi', inne...¶
Static description of an implicit-Euler heat solve (hashable, no tensors).
ndim(property)stencil(self, alpha: torch.Tensor) -> DiffusionStencil
class ImplicitEulerAdjoint(*args, **kwargs)¶
Discrete adjoint of the implicit-Euler heat solve (NeFTY Eq. 10–11, App. D.3 Eq. 25–27).
backward(ctx, grad_frames: torch.Tensor)forward(ctx, alpha: torch.Tensor, T0: torch.Tensor, cfg: HeatSolveConfig) -> torch.Tensor
def implicit_euler_frames(alpha: torch.Tensor, T0: torch.Tensor, cfg: HeatSolveConfig, grad_mode: str = 'adjoint') -> torch.Tensor¶
Observed surface frames of the implicit-Euler heat solve with the chosen gradient mode.
def rollout(alpha: torch.Tensor, T0: torch.Tensor, cfg: HeatSolveConfig, *, keep_states: bool = False, residuals: bool = False, step_fn: Callable | None = None, n_steps...¶
Run the implicit-Euler recurrence (NeFTY Eq. 8) and collect the observed surface frames.
def solve_step(system: ImplicitSystem, b: torch.Tensor, x0: torch.Tensor | None, cfg: HeatSolveConfig) -> torch.Tensor¶
Solve A x = b with the configured inner solver, warm-started from x0.
def surface(T: torch.Tensor) -> torch.Tensor¶
Observed front surface: the first slice along the last (through-thickness) axis.
nefi.operators.pde.heat¶
Differentiable 3-D heat conduction: implicit-Euler forward model and an explicit data simulator.
class ExplicitHeatSimulator(domain: Domain, dt: float = 0.05, n_steps: int = 100, obs_frames: Sequence[int] | int | None = None, initial=None, bc: BoundarySpec | str | Sequence[str] | ...¶
Independent explicit (forward-Euler, substepped, float64) heat simulator — data only.
at_resolution(self, shape: Sequence[int]) -> ExplicitHeatSimulatordivergence(self, T: torch.Tensor, faces: list[torch.Tensor]) -> torch.Tensor—∇·(α∇T)by flux differences with ghost cells (plus the Robin sink).extra_repr(self) -> strforward(self, fields: Fields) -> torch.Tensorn_substeps(self, alpha: torch.Tensor) -> int
class GaussianFlash(amplitude: float = 100.0, width_xy: float = 2.5, width_z: float = 0.2, center: Sequence[float] | None = None, periodic: bool = True, cell_average: bool = True)¶
Gaussian post-flash profile of the synthetic benchmark (NeFTY App. A.5)::
to_dict(self) -> dict
class HeatOperator(domain: Domain, dt: float = 0.05, n_steps: int = 100, obs_frames: Sequence[int] | int | None = None, initial=None, bc: BoundarySpec | str | Sequence[str] | ...¶
NeFTY hard-constrained forward model: diffusivity → observed front-surface frames.
at_resolution(self, shape: Sequence[int]) -> HeatOperator— Same physics (Δt,N_t, frames, BCs, solver) on gridshapeof the same slab; spacing and the initial condition are rebuilt (a learnable IC module is shared).extra_repr(self) -> strforward(self, fields: Fields) -> torch.Tensorsimulate(self, alpha: torch.Tensor, T0: torch.Tensor | None = None, return_states: bool = False, residuals: bool = False) -> dict[str, torch.Tensor | list[float] | N...— Diagnostic forward solve (no gradients).
class InitialCondition()¶
Callable (domain, shape=None) -> T0 evaluated at cell centres (ambient-shifted).
to_dict(self) -> dict
class UniformFlash(amplitude: float = 100.0, width_z: float = 0.2, cell_average: bool = True)¶
Near-uniform front-face flash of the real-PVC setting (NeFTY App. A.5 / H.2)::
to_dict(self) -> dict
def explicit_substeps(dt: float, spacing: Sequence[float], alpha_max: float, min_substeps: int = 10, safety: float = 2.0) -> int¶
Substeps per frame, NeFTY App. E.1: N_sub = max(10, ⌈Δt/Δt_stable × 2⌉) with Δt_stable = Δx_min² / (2 D α_max) (D the spatial dimension).
nefi.operators.pde.linear_solvers¶
Matrix-free linear solvers for the implicit-Euler system A(α) T^{n+1} = T^n.
def conjugate_gradient(apply_A: MatVec, b: torch.Tensor, x0: torch.Tensor | None = None, tol: float = 1e-06, max_iter: int = 200, *, atol: float = 0.0, precond: torch.Tensor | Mat...¶
(Preconditioned) conjugate gradients for a symmetric positive-definite A.
def jacobi(apply_A: MatVec, diag: torch.Tensor, b: torch.Tensor, x0: torch.Tensor | None = None, iters: int = 50, *, omega: float = 1.0, apply_offdiag: MatVec | None =...¶
Unrolled (damped) Jacobi iteration, NeFTY Eq. (24).
nefi.operators.pde.stencil¶
Finite-volume diffusion stencil L(α) with harmonic-mean face coefficients.
class BoundarySpec(kinds: tuple[str, ...], robin_h: float = 0.0)¶
Per-axis boundary conditions of the diffusion operator (NeFTY App. A.5 / D.2).
has_robin(property)ndim(property)parse(bc: BoundarySpec | str | Sequence[str] | None, ndim: int, robin_h: float | None = None) -> BoundarySpec— Normalize a boundary-condition spec.periodic(self, axis: int) -> bool
class DiffusionStencil(alpha: torch.Tensor, spacing: Sequence[float], bc: BoundarySpec | str | Sequence[str] | None = None, face_mode: str = 'harmonic', robin_h: float | None = No...¶
Precomputed 2·ndim+1-point stencil of L(α) (NeFTY Eq. 7 / App. D.2).
apply(self, T: torch.Tensor) -> torch.Tensor—L(α) TforTof shape(*batch, *grid).bilinear(self, mu: torch.Tensor, T: torch.Tensor) -> torch.Tensor—μᵀ L(α) Tvia the symmetric face form-Σ_faces a (Δμ)(ΔT) - Σ β μ T.layout(property) — The :class:FlatLayoutof this stencil's grid (on its device).system(self, dt: float) -> ImplicitSystem— The implicit-Euler matrixA = I - Δt L(α)(NeFTY Eq. 8 / 23).
class FlatLayout(grid: Sequence[int], device: torch.device | str | None = None)¶
Flat padded layout of a grid in which all 2·ndim neighbours are contiguous 1-D views.
center(self, P: torch.Tensor) -> torch.Tensor— The cells themselves (range layout) as a view of the padded array.from_range(self, xr: torch.Tensor) -> torch.Tensor— Range layout(*batch, L)→ grid layout(*batch, *grid)(one gather).neighbor(self, P: torch.Tensor, j: int) -> torch.Tensor— Neighbourj(order of :func:neighbors) of every range cell, as a view ofP.pad_grid(self, x: torch.Tensor) -> torch.Tensor— Padded array(*batch, n_padded)from a grid-layout tensor (one gather).pad_range(self, xr: torch.Tensor) -> torch.Tensor— Padded array(*batch, n_padded)from a range-layout tensor (one gather).to_range(self, x: torch.Tensor) -> torch.Tensor— Grid layout(*batch, *grid)→ range layout(*batch, L)(one gather).weighted_sum(self, base: torch.Tensor, weights: Sequence[torch.Tensor], P: torch.Tensor) -> torch.Tensor—base + Σ_j w_j ⊙ neighbour_j(range layout;2·ndimfused multiply-adds).
class ImplicitSystem(stencil: DiffusionStencil, dt: float)¶
A(α) = I - Δt L(α): symmetric, strictly diagonally dominant, SPD (NeFTY App. D.2).
apply(self, T: torch.Tensor) -> torch.Tensor—A T.apply_offdiag(self, T: torch.Tensor) -> torch.Tensor—R TwithA = D + R(the Jacobi splitting of NeFTY Eq. 24).chebyshev_omegas(self, iters: int) -> list[torch.Tensor]— :func:chebyshev_omegasfor this system (cached periters; constants for autograd, i.e. thegrad_mode="autograd"gradient treats the iteration's weights as fixed).chebyshev_sweeps(self, b: torch.Tensor, x: torch.Tensor, iters: int) -> torch.Tensor—itersChebyshev-accelerated Jacobi iterations (semi-iterative method).jacobi_sweep(self, b_scaled: torch.Tensor, x: torch.Tensor) -> torch.Tensor— One fused Jacobi sweepD⁻¹(b - R x)givenb_scaled = D⁻¹ b(Eq. 24).jacobi_sweep_weights(self) -> list[torch.Tensor]—W_j = -R_j / Dso that one Jacobi sweep isx ← b/D + Σ_j W_j ⊙ shift_j(x).jacobi_sweeps(self, b: torch.Tensor, x: torch.Tensor, iters: int) -> torch.Tensor—itersfused Jacobi sweepsx ← D⁻¹b + Σ_j W_j ⊙ shift_j(x)(NeFTY Eq. 24).layout(property)spectral_bound(self) -> torch.Tensor— Gershgorin boundρ = max_i Σ_j W_j[i] < 1for the spectral radius of the Jacobi map (0-dim tensor, no gradient; the eigenvalues ofD⁻¹Rare real sinceAis symmetric).
def apply_A(T: torch.Tensor, alpha: torch.Tensor, dt: float, spacing: Sequence[float], bc: BoundarySpec | str | Sequence[str] | None = None, face_mode: str = 'harmonic'...¶
A(α) T = T - Δt L(α) T (NeFTY Eq. 8).
def apply_diffusion(T: torch.Tensor, alpha: torch.Tensor, spacing: Sequence[float], bc: BoundarySpec | str | Sequence[str] | None = None, face_mode: str = 'harmonic', robin_h: ...¶
L(α) T — second-order finite-volume diffusion (NeFTY Eq. 7, App. D.2).
def chebyshev_omegas(rho: torch.Tensor, iters: int) -> list[torch.Tensor]¶
Chebyshev semi-iteration weights ω_2 … ω_K for a Jacobi map of spectral radius ≤ rho.
def diagonal_A(alpha: torch.Tensor, dt: float, spacing: Sequence[float], bc: BoundarySpec | str | Sequence[str] | None = None, face_mode: str = 'harmonic', robin_h: float ...¶
Diagonal D of A(α) (the Jacobi preconditioner, NeFTY Eq. 24).
def face_coefficients(alpha: torch.Tensor, axis: int, mode: str = 'harmonic', periodic: bool = True, ndim: int | None = None) -> torch.Tensor¶
Face-centred diffusivity ᾱ_{i+1/2} between cell i and i+1 along axis.
def face_conductances(alpha: torch.Tensor, spacing: Sequence[float], bc: BoundarySpec | str | Sequence[str] | None = None, mode: str = 'harmonic') -> list[torch.Tensor]¶
Per-axis face conductances a_{i+1/2} = ᾱ_{i+1/2} / Δ_d² (App. D.2 notation a_pq).
def flat_layout(grid: Sequence[int], device: torch.device | str | None = None) -> FlatLayout¶
Cached :class:FlatLayout for grid with index tensors on device.
def forward_difference(x: torch.Tensor, axis: int, ndim: int) -> torch.Tensor¶
x_{i+1} - x_i along grid axis (circular wrap; the boundary entry is masked by the zero conductance for non-periodic axes).
def neighbors(T: torch.Tensor, ndim: int) -> list[torch.Tensor]¶
Circularly shifted copies [T_{i+1}, T_{i-1}] for each of the trailing ndim axes.
nefi.operators.timestepping¶
Generic differentiable time integration with memory-controlled adjoints.
class TimeStepper(step: StepFn, init: Callable[[Fields], State], n_steps: int, dt: float, observe: Callable[[State, int], torch.Tensor | None] | None = None, *, params: Calla...¶
Differentiable time integration of a user-written scheme (see module docstring).
at_resolution(self, shape)extra_repr(self) -> strforward(self, fields: Fields) -> torch.Tensoroutput_shape(self, shape)required_fields(self)
def advection_diffusion_step(state: State, params: Mapping[str, Any], dt: float, n: int, *, spacing: Sequence[float], velocity: Any = None, diffusivity: Any = 0.0, boundary: str = 'peri...¶
One explicit step of u_t + v·∇u = ∇·(D ∇u) + s (first-order upwind advection).
def laplacian(u: torch.Tensor, spacing: Sequence[float], boundary: str = 'neumann') -> torch.Tensor¶
Second-order finite-difference Laplacian over the trailing len(spacing) dims.
def stable_dt(spacing: Sequence[float], *, c_max: float | None = None, velocity_max: float | Sequence[float] | None = None, diffusivity_max: float | None = None, cfl: flo...¶
Largest stable explicit time step (times a safety factor cfl ≤ 1).
def wave_initial_state(u0: torch.Tensor, v0: torch.Tensor | None = None, *, c: Any, dt: float, spacing: Sequence[float], boundary: str = 'neumann') -> tuple[torch.Tensor, torch.Te...¶
Leapfrog start (u_{-1}, u_0) for initial displacement u0 and velocity v0.
def wave_step(state: State, params: Mapping[str, Any], dt: float, n: int, *, spacing: Sequence[float], c: Any = 'c', boundary: str = 'absorbing', source: Any = None, damp...¶
One leapfrog step of u_tt = c² Δu + s (optionally damped − γ u_t).