Neural Network API Reference

The opifex.neural package provides the building blocks for scientific machine learning models, built on top of Flax NNX.

Base Architectures

Standard MLP

opifex.neural.base.StandardMLP

StandardMLP(layer_sizes: list[int], activation: str = 'gelu', dropout_rate: float = 0.0, use_bias: bool = True, apply_final_dropout: bool = False, *, dtype: Any | None = None, param_dtype: Any = float32, rngs: Rngs, kernel_init: Callable = xavier_uniform(), bias_init: Callable = zeros)

Bases: Module

Modern Multi-Layer Perceptron implementation using FLAX NNX.

Fully compliant with Flax NNX best practices including: - Proper RNG handling with keyword-only rngs parameter - Modern activation functions (GELU default, configurable) - Efficient dropout strategies with deterministic control - Custom initialization strategies following NNX patterns - Automatic differentiation with JAX - Performance-optimized state management

Attributes:

Name	Type	Description
layer_sizes		List of layer sizes including input and output dimensions
activation		Name of the activation function to use
dropout_rate		Dropout probability (0.0 means no dropout)
use_bias		Whether to include bias terms in linear layers
apply_final_dropout		Whether to apply dropout after the final layer
layers		Sequence of linear transformation layers
activation_fn		The actual activation function
dropout	Dropout | None	Dropout layer (None if dropout_rate is 0)

Parameters:

Name	Type	Description	Default
layer_sizes	list[int]	List of layer sizes, e.g., [input_dim, hidden1, hidden2, output_dim]	required
activation	str	Activation function name ('gelu', 'tanh', 'relu', 'sigmoid', 'silu') Default is 'gelu' for modern neural networks	'gelu'
dropout_rate	float	Dropout probability for regularization (0.0 = no dropout)	0.0
use_bias	bool	Whether to use bias in linear projections	True
apply_final_dropout	bool	Whether to apply dropout after final layer (useful for some transformer-style architectures)	False
dtype	Any | None	Computation dtype for NNX linear layers. None preserves the Flax default promotion behavior.	None
param_dtype	Any	Parameter storage dtype for NNX linear layers.	float32
rngs	Rngs	FLAX NNX random number generator state (keyword-only)	required
kernel_init	Callable	Kernel initialization function (callable)	xavier_uniform()
bias_init	Callable	Bias initialization function (callable)	zeros

Source code in opifex/neural/base.py

def __init__(
    self,
    layer_sizes: list[int],
    activation: str = "gelu",
    dropout_rate: float = 0.0,
    use_bias: bool = True,
    apply_final_dropout: bool = False,
    *,
    dtype: Any | None = None,
    param_dtype: Any = jnp.float32,
    rngs: nnx.Rngs,
    kernel_init: Callable = nnx.initializers.xavier_uniform(),
    bias_init: Callable = nnx.initializers.zeros,
) -> None:
    """Initialize the StandardMLP following modern NNX patterns.

    Args:
        layer_sizes: List of layer sizes, e.g.,
            [input_dim, hidden1, hidden2, output_dim]
        activation: Activation function name
            ('gelu', 'tanh', 'relu', 'sigmoid', 'silu')
            Default is 'gelu' for modern neural networks
        dropout_rate: Dropout probability for regularization
            (0.0 = no dropout)
        use_bias: Whether to use bias in linear projections
        apply_final_dropout: Whether to apply dropout after final layer
            (useful for some transformer-style architectures)
        dtype: Computation dtype for NNX linear layers. ``None`` preserves
            the Flax default promotion behavior.
        param_dtype: Parameter storage dtype for NNX linear layers.
        rngs: FLAX NNX random number generator state (keyword-only)
        kernel_init: Kernel initialization function (callable)
        bias_init: Bias initialization function (callable)
    """
    super().__init__()

    # Store configuration
    self.layer_sizes = layer_sizes
    self.activation = activation
    self.dropout_rate = dropout_rate
    self.use_bias = use_bias
    self.apply_final_dropout = apply_final_dropout
    self.dtype = dtype
    self.param_dtype = param_dtype

    # Validate layer sizes
    if len(layer_sizes) < 2:
        raise ValueError("layer_sizes must have at least 2 elements (input and output)")

    # Create layers following NNX patterns (use nnx.List for Flax 0.12.0+)
    layers = []
    for i in range(len(layer_sizes) - 1):
        layer = nnx.Linear(
            in_features=layer_sizes[i],
            out_features=layer_sizes[i + 1],
            use_bias=use_bias,
            kernel_init=kernel_init,
            bias_init=bias_init,
            dtype=dtype,
            param_dtype=param_dtype,
            rngs=rngs,
        )
        layers.append(layer)
    self.layers = nnx.List(layers)

    # Set activation function using the activation library
    self.activation_fn = get_activation(activation)

    # Initialize dropout if needed - pass rngs directly
    if dropout_rate > 0.0:
        self.dropout: nnx.Dropout | None = nnx.Dropout(rate=dropout_rate, rngs=rngs)
    else:
        self.dropout = None

Atomistic Models

Machine-learning interatomic potentials live in opifex.neural.atomistic. They follow a backbone → typed property heads assembly: a backbone produces per-atom embeddings and named heads read them out into energy, forces and stress. See the Atomistic Potentials guide for the design, the three backbones (SchNet, PaiNN, NequIP) and a registry-driven build.

opifex.neural.atomistic.base.AtomisticModel

AtomisticModel(*, backbone: Backbone, heads: dict[str, PropertyHead], neighbor_list: NeighborList, max_edges: int)

Bases: Module

A backbone plus a dict of named property heads (the MLIP assembly).

Parameters:

Name	Type	Description	Default
backbone	Backbone	Embedding producer satisfying the Backbone protocol.	required
heads	dict[str, PropertyHead]	Mapping of head name to PropertyHead. Must include an "energy" head (:class:EnergyHead) -- the conservative force/stress heads differentiate its output.	required
neighbor_list	NeighborList	Edge builder satisfying the NeighborList protocol.	required
max_edges	int	Static upper bound on the number of edges (output size of the neighbour list under jit).	required

Raises:

Type	Description
ValueError	If no "energy" head is supplied.

implemented_properties property

implemented_properties: tuple[str, ...]

The union of every property emitted by the configured heads.

opifex.neural.atomistic.backbones

Concrete interatomic-potential backbones (embedding producers).

Each backbone is an :class:flax.nnx.Module satisfying :class:opifex.core.quantum.protocols.Backbone and self-registering in the opifex.core.quantum.registry backbone registry, so importing this package makes "schnet", "painn" and "nequip" discoverable by name:

• :class:SchNet -- invariant continuous-filter convolutions (Schuett 2018);
• :class:PaiNN -- equivariant scalar/vector message passing, l <= 1 (Schuett 2021);
• :class:NequIP -- E(3)-equivariant Clebsch-Gordan tensor-product message passing (Batzner 2022).

All three compose opifex's Q0 equivariant kit (:mod:opifex.neural.equivariant) via the shared :mod:opifex.neural.atomistic.backbones._message_passing helper, and emit per-atom invariant "node_features" consumed by the property heads.

SchNet

SchNet(*, config: SchNetConfig | None = None, rngs: Rngs)

Bases: Module

Invariant continuous-filter convolutional backbone (Schuett et al. 2018).

Satisfies :class:opifex.core.quantum.protocols.Backbone: maps a :class:~opifex.core.quantum.molecular_system.MolecularSystem and its padded edge index to {"node_features": (n_atoms, feature_dim)} invariant scalars.

Parameters:

Name	Type	Description	Default
config	SchNetConfig | None	Backbone hyper-parameters. Defaults to :class:SchNetConfig.	None
rngs	Rngs	Random number generators (keyword-only) seeding all weights.	required

SchNetConfig dataclass

SchNetConfig(*, feature_dim: int = 64, num_interactions: int = 3, num_radial_basis: int = 16, cutoff: float = 5.0, filter_hidden_dim: int = 64)

Hyper-parameters of a :class:SchNet backbone.

Attributes:

Name	Type	Description
feature_dim	int	Width F of the per-atom scalar feature vector.
num_interactions	int	Number of continuous-filter interaction blocks T.
num_radial_basis	int	Number of Bessel radial-basis functions.
cutoff	float	Connection / cutoff radius r_c (in the system's length units).
filter_hidden_dim	int	Hidden width of the radial filter-generating MLP.

PaiNN

PaiNN(*, config: PaiNNConfig | None = None, rngs: Rngs)

Bases: Module

Equivariant scalar/vector message-passing backbone (Schuett et al. 2021).

Satisfies :class:opifex.core.quantum.protocols.Backbone: maps a :class:~opifex.core.quantum.molecular_system.MolecularSystem and its padded edge index to {"node_features": (n_atoms, feature_dim)} invariant scalars (the equivariant vector channels are internal state).

Parameters:

Name	Type	Description	Default
config	PaiNNConfig | None	Backbone hyper-parameters. Defaults to :class:PaiNNConfig.	None
rngs	Rngs	Random number generators (keyword-only) seeding all weights.	required

PaiNNConfig dataclass

PaiNNConfig(*, feature_dim: int = 64, num_interactions: int = 3, num_radial_basis: int = 16, cutoff: float = 5.0)

Hyper-parameters of a :class:PaiNN backbone.

Attributes:

Name	Type	Description
feature_dim	int	Width F of the scalar (and vector) feature channels.
num_interactions	int	Number of message + update interaction blocks T.
num_radial_basis	int	Number of Bessel radial-basis functions.
cutoff	float	Connection / cutoff radius r_c (in the system's length units).

NequIP

NequIP(*, config: NequIPConfig | None = None, rngs: Rngs)

Bases: Module

E(3)-equivariant tensor-product backbone (Batzner et al. 2022).

Satisfies :class:opifex.core.quantum.protocols.Backbone: maps a :class:~opifex.core.quantum.molecular_system.MolecularSystem and its padded edge index to {"node_features": (n_atoms, num_scalar_features)} invariant scalars (the 0e channels of the final steerable features).

Parameters:

Name	Type	Description	Default
config	NequIPConfig | None	Backbone hyper-parameters. Defaults to :class:NequIPConfig.	None
rngs	Rngs	Random number generators (keyword-only) seeding all weights.	required

Raises:

Type	Description
ValueError	If config.correlation < 1 or the hidden irreps carry no 0e scalar channel to read out.

NequIPConfig dataclass

NequIPConfig(*, hidden_irreps: str = '16x0e + 8x1o + 4x2e', sh_lmax: int = 2, num_interactions: int = 3, num_radial_basis: int = 8, radial_hidden_dim: int = 64, cutoff: float = 5.0, average_num_neighbors: float = 1.0, correlation: int = _DEFAULT_CORRELATION, sh_normalization: str = 'component', normalize_gate_act: bool = True, species: tuple[int, ...] = ())

Hyper-parameters of a :class:NequIP backbone.

Attributes:

Name	Type	Description
hidden_irreps	str	Steerable layout of the per-atom hidden features, e.g. "16x0e + 8x1o + 4x2e".
sh_lmax	int	Maximum spherical-harmonic degree of the edge embedding.
num_interactions	int	Number of tensor-product convolution layers T.
num_radial_basis	int	Number of Bessel radial-basis functions.
radial_hidden_dim	int	Hidden width of the radial network MLP.
cutoff	float	Connection / cutoff radius r_c (in the system's length units).
average_num_neighbors	float	Constant sqrt normaliser for the aggregated message (NequIP's n_neighbors internal normalisation).
correlation	int	Body-order correlation. 1 is the two-body edge tensor product with a gate nonlinearity; > 1 adds the MACE-style symmetric contraction (requires uniform-multiplicity hidden_irreps and species).
sh_normalization	str	Normalisation convention for the edge spherical-harmonic embedding, one of "component" (default; unit per-component variance, the NequIP convention that keeps the embedding at the unit scale the tensor-product weight init assumes), "integral" or "norm".
normalize_gate_act	bool	If True (default), the gate rescales each activation to unit second moment under a standard-normal input, so feature magnitudes do not drift across stacked gated interaction layers.
species	tuple[int, ...]	Sorted distinct atomic numbers in the dataset (e.g. (1, 6, 8) for an H/C/O system). When non-empty, each interaction's self-connection is species-indexed (a per-element residual, the NequIP convention: the skip is a tensor product of the node features with the one-hot atom type) instead of a single shared linear -- giving every element its own self-interaction. Empty (default) uses the shared linear self-connection.

opifex.neural.atomistic.heads

Typed property heads for atomistic models (backbone -> named outputs).

Each head owns exactly one property family (single responsibility) and satisfies the :class:opifex.core.quantum.protocols.PropertyHead protocol:

• :class:EnergyHead -- sum of per-atom scalar energies (invariant total energy);
• :class:ForcesHead -- conservative forces -grad(E) (autodiff);
• :class:StressHead -- virial / stress via strain-displacement autodiff.

Conservative force/stress are the default strategies; direct-readout variants plug into the same protocol later.

EnergyHead

EnergyHead(*, feature_dim: int, hidden_dim: int | None = None, scale_shift: AtomicScaleShift | None = None, rngs: Rngs)

Bases: Module

Sum-of-atomic-energies readout producing a scalar total energy.

Parameters:

Name	Type	Description	Default
feature_dim	int	Width of the backbone's "node_features" embedding.	required
hidden_dim	int | None	Hidden width of the per-atom MLP. Defaults to feature_dim.	None
scale_shift	AtomicScaleShift | None	Optional per-atom energy scale-shift (E0 + normaliser) applied to the summed energy. None (default) leaves the raw sum-of-atomic-energies unchanged.	None
rngs	Rngs	Random number generators (keyword-only) seeding the MLP weights.	required

implemented_properties property

implemented_properties: tuple[str, ...]

This head emits the total "energy".

ForcesHead

Forces as -grad(energy) w.r.t. positions (the conservative strategy).

Stateless: it differentiates the position-to-energy closure injected by the assembled model, so it owns no parameters.

implemented_properties property

implemented_properties: tuple[str, ...]

This head emits "forces".

StressHead

Stress as the strain-derivative of energy divided by cell volume.

Stateless: it differentiates the symmetric-strain-to-energy closure injected by the assembled model, so it owns no parameters.

implemented_properties property

implemented_properties: tuple[str, ...]

This head emits "stress".

Neural Quantum

opifex.neural.quantum

Neural quantum chemistry modules for scientific machine learning.

The public surface spans three integral-independent families:

• the learned exchange-correlation functional (:class:~opifex.neural.quantum.neural_xc.NeuralXCFunctional) and the differentiable Kohn-Sham density-functional theory solver in :mod:opifex.neural.quantum.dft (:class:~opifex.neural.quantum.dft.SCFSolver);
• the neural-wavefunction / variational Monte Carlo stack in :mod:opifex.neural.quantum.vmc (:class:~opifex.neural.quantum.vmc.FermiNet).

The Kohn-Sham DFT names (:class:SCFSolver, :class:SCFResult, :class:Functional, :class:SolverMode) are exposed lazily through :pep:562 __getattr__ so that integral-free subpackages -- notably the VMC family, which the task spec keeps free of any dependency on the Gaussian-integral engine -- can be imported without pulling in the DFT grid/SCF machinery (and its opifex.core.quantum backend). The names remain importable exactly as before; they are simply resolved on first access.

Functional

Bases: StrEnum

Supported exchange-correlation functionals.

SCFResult dataclass

SCFResult(*, total_energy: Array, orbital_energies: Array, density_matrix: Array, coefficients: Array, n_iterations: int, converged: bool)

Outcome of a restricted Kohn-Sham SCF calculation.

Attributes:

Name	Type	Description
total_energy	Array	Converged Kohn-Sham total energy (Hartree).
orbital_energies	Array	Molecular-orbital eigenvalues [Shape: (n_ao,)].
density_matrix	Array	Converged AO density matrix [Shape: (n_ao, n_ao)].
coefficients	Array	MO coefficients [Shape: (n_ao, n_ao)].
n_iterations	int	Number of SCF iterations performed.
converged	bool	Whether the density change fell below the tolerance.

SCFSolver

SCFSolver(system: MolecularSystem, basis: AtomicOrbitalBasis | None = None, *, functional: Functional | str = LDA, mode: SolverMode | str = DIIS, neural_functional: NeuralXCFunctional | None = None, grid_template: MolecularGridTemplate | None = None, max_iterations: int = 100, convergence_tolerance: float = 1e-08)

Restricted Kohn-Sham (RKS) self-consistent-field solver.

Parameters:

Name	Type	Description	Default
system	MolecularSystem	The molecular system to solve.	required
basis	AtomicOrbitalBasis | None	The AO basis (defaults to STO-3G built from the system).	None
functional	Functional | str	The exchange-correlation functional ("lda", "pbe" or "neural").	LDA
mode	SolverMode | str	"diis" for the Anderson-accelerated self-consistent SCF (Pulay DIIS on the density residual) or "direct" for direct minimisation.	DIIS
neural_functional	NeuralXCFunctional | None	A learned XC functional; required (and selects the "neural" functional) when functional == "neural".	None
grid_template	MolecularGridTemplate | None	A pre-built molecular-grid template; defaults to the standard Becke grid for system. Supplying a coarser template trades XC-integration accuracy for speed (e.g. in tests).	None
max_iterations	int	Maximum SCF / fixed-point / minimisation iterations.	100
convergence_tolerance	float	RMS density-change convergence threshold.	1e-08

functional property

functional: Functional

The exchange-correlation functional in use.

energy_from_positions

energy_from_positions(positions: Array, *, differentiable: str = 'implicit') -> Array

Converged Kohn-Sham total energy as a function of nuclear positions.

The self-consistent density is found as an implicit fixed point of the Roothaan step regardless of the forward :class:SolverMode (direct minimisation and the DIIS/fixed-point iteration converge to the same Kohn-Sham density). Differentiating the implicit fixed point gives exact, memory-cheap gradients via the implicit function theorem and avoids the gauge-singular Hessian of the direct-minimisation parametrisation, so :meth:compute_forces is robust for both modes.

Parameters:

Name	Type	Description	Default
positions	Array	Nuclear positions in Bohr [Shape: (n_atoms, 3)].	required
differentiable	str	"implicit" (default) finds the density as an implicit fixed point (IFT gradient); "unroll" runs a fixed number of differentiable SCF steps (gradient cross-check).	'implicit'

Returns:

Type	Description
Array	The scalar converged total energy (Hartree).

energy_from_state

energy_from_state(state: State, positions: Array | None = None) -> Array

Converged total energy as a differentiable function of the XC state.

The entry point for learned-XC training: jax.grad of this with respect to state gives the exact dE/dtheta through the implicit-diff SCF (the implicit function theorem differentiates the converged fixed point, not the iterations).

Parameters:

Name	Type	Description	Default
state	State	The neural XC parameter state (an nnx.State pytree, as produced by :func:flax.nnx.split).	required
positions	Array | None	Geometry to evaluate at (defaults to the system geometry).	None

Returns:

Type	Description
Array	The scalar converged total energy (Hartree).

energy

energy() -> Array

Converged total energy at the system's nuclear geometry.

compute_forces

compute_forces(positions: Array | None = None) -> Array

Analytic nuclear forces :math:F = -\partial E/\partial R.

Computed by :func:jax.grad of the implicit-diff total energy with respect to the nuclear coordinates.

Parameters:

Name	Type	Description	Default
positions	Array | None	Geometry to evaluate at (defaults to the system geometry).	None

Returns:

Type	Description
Array	Forces in Hartree/Bohr [Shape: (n_atoms, 3)].

energy_and_forces

energy_and_forces(positions: Array | None = None) -> tuple[Array, Array]

Converged total energy and the analytic forces :math:-\partial E/\partial R.

Parameters:

Name	Type	Description	Default
positions	Array | None	Geometry to evaluate at (defaults to the system geometry).	None

Returns:

Type	Description
tuple[Array, Array]	A pair (energy, forces) with forces in Hartree/Bohr.

solve

solve(*, initial_density: Array | None = None) -> SCFResult

Run the forward SCF (DIIS or direct minimisation) to convergence.

Parameters:

Name	Type	Description	Default
initial_density	Array | None	Optional closed-shell density to seed the Anderson/DIIS iteration. A high-quality guess (e.g. reconstructed from a neural-network predicted Fock via :func:density_from_fock) reaches the fixed point in fewer iterations; the converged result is unchanged. Ignored by the direct-minimisation mode, which is seeded internally.	None

Returns:

Name	Type	Description
The	SCFResult	class:SCFResult with the converged total energy and orbitals.

SolverMode

Bases: StrEnum

How the self-consistent density is found.

NeuralXCFunctional

NeuralXCFunctional(hidden_sizes: Sequence[int] = (128, 128, 64), activation: Callable = gelu, use_attention: bool = True, num_attention_heads: int = 8, use_advanced_features: bool = True, dropout_rate: float = 0.0, *, rngs: Rngs)

Bases: Module

Neural exchange-correlation functional for DFT calculations.

Implements a modern neural XC functional with attention mechanisms for capturing non-local correlations, enhanced physics constraints, and chemical accuracy optimization.

Parameters:

Name	Type	Description	Default
hidden_sizes	Sequence[int]	Sequence of hidden layer sizes	(128, 128, 64)
activation	Callable	Activation function to use	gelu
use_attention	bool	Whether to use attention mechanism for non-local correlations	True
num_attention_heads	int	Number of attention heads	8
use_advanced_features	bool	Whether to include advanced physics features	True
dropout_rate	float	Dropout rate for regularization	0.0
rngs	Rngs	Random number generators	required

energy_density_from_sigma

energy_density_from_sigma(density: Array, sigma: Array, *, deterministic: bool = True) -> Array

XC energy per particle as a function of rho and sigma=|grad rho|^2.

The GGA-native interface used on a real molecular grid and for the AD XC potential: the gradient direction is irrelevant to a (semi-)local functional, so the dimensionless features depend only on :math:(\rho,\sigma). The Cartesian gradient is reconstructed along a single axis with magnitude :math:\sqrt\sigma purely to reuse the feature extractor; the resulting energy density is identical for any direction.

Parameters:

Name	Type	Description	Default
density	Array	Electron density [Shape: (n_points,)].	required
sigma	Array	Squared density gradient |grad rho|^2 [Shape: (n_points,)].	required
deterministic	bool	Whether to use deterministic computation.	True

Returns:

Type	Description
Array	XC energy per particle [Shape: (n_points,)].

xc_potential_components

xc_potential_components(density: Array, sigma: Array, *, deterministic: bool = True) -> tuple[Array, Array]

GGA XC potential pair :math:(v_\rho, v_\sigma) by autodiff.

Returns both functional derivatives of the XC energy density :math:\rho\,\varepsilon_{xc}(\rho,\sigma):

.. math:: v_\rho = \frac{\partial(\rho\varepsilon_{xc})}{\partial\rho},\qquad v_\sigma = \frac{\partial(\rho\varepsilon_{xc})}{\partial\sigma}.

Both channels are live -- the density-gradient (:math:\sigma) channel is differentiated, not zeroed -- so the GGA potential is correct.

Parameters:

Name	Type	Description	Default
density	Array	Electron density [Shape: (n_points,)].	required
sigma	Array	Squared density gradient |grad rho|^2 [Shape: (n_points,)].	required
deterministic	bool	Whether to use deterministic computation.	True

Returns:

Type	Description
tuple[Array, Array]	The pair (v_rho, v_sigma) each [Shape: (n_points,)].

compute_functional_derivative

compute_functional_derivative(density: Array, gradients: Array, *, deterministic: bool = False) -> Array

Density-channel functional derivative d(rho eps_xc)/d rho.

Computes the live GGA density-channel potential at fixed :math:\sigma=|\nabla\rho|^2. The full GGA potential additionally needs the :math:\sigma channel; use :meth:xc_potential_components for both.

Parameters:

Name	Type	Description	Default
density	Array	Electron density [batch, grid_points] or [grid_points].	required
gradients	Array	Density gradients [..., 3].	required
deterministic	bool	Whether to use deterministic computation.	False

Returns:

Type	Description
Array	d(rho eps_xc)/d rho with the same shape as density.

assess_chemical_accuracy

assess_chemical_accuracy(density: Array, gradients: Array, reference_energy: Array | None = None, *, deterministic: bool = False) -> dict[str, float]

Assess chemical accuracy of XC functional predictions.

Parameters:

Name	Type	Description	Default
density	Array	Electron density	required
gradients	Array	Density gradients	required
reference_energy	Array | None	Reference XC energy for comparison (optional)	None
deterministic	bool	Whether to use deterministic computation	False

Returns:

Type	Description
dict[str, float]	Dictionary containing accuracy metrics

Kohn-Sham DFT solver

Restricted Kohn-Sham (RKS) self-consistent-field solver.

A closed-shell RKS driver built on the native McMurchie-Davidson integral backend (:class:~opifex.core.quantum.backend.JaxGaussianBackend) and the exchange-correlation functionals in :mod:opifex.neural.quantum.dft.xc (LDA Slater+VWN5 and the PBE GGA).

Forward SCF

The Kohn-Sham equations are solved by symmetric-orthogonalisation fixed-point iteration with Anderson acceleration (Pulay DIIS on the density residual -- :class:~opifex.neural.quantum.dft._fixed_point.AndersonAcceleration):

. Orthogonalise with Lowdin's :math:S^{-1/2}.

. Build the Fock matrix :math:F(D) = h_\text{core} + J[D] + V_{xc}[D] with the

Coulomb matrix :math:J_{\mu\nu} = \sum_{\lambda\sigma} (\mu\nu|\lambda\sigma) D_{\lambda\sigma} and the LDA/GGA :math:V_{xc} on a real molecular grid.

. Solve :math:F' C' = C' \varepsilon, back-transform, occupy the lowest

:math:n_\text{occ} orbitals, form the Roothaan step :math:D' = 2 C_\text{occ} C_\text{occ}^\top.

. Anderson-mix a short history of densities to converge the residual

:math:D' - D; plain Roothaan iteration charge-sloshes and stalls.

The forward solve and the differentiable energy path share this one fixed-point engine, so both are jit-compatible and converge identically.

A direct-minimisation (SCF-free) mode is available behind the same interface: the Kohn-Sham energy is minimised directly over a QR-orthonormalised coefficient matrix (jrystal / DWD, arXiv:2411.05033) -- intended for the learned-XC path.

Differentiable energy and analytic forces

:meth:SCFSolver.energy_from_positions returns the converged total energy as a pure, differentiable function of the nuclear coordinates: the integrals, grid and XC matrix are rebuilt from positions and the self-consistent density is found as an implicit fixed point (:mod:opifex.neural.quantum.dft._energy). Optimistix's :class:~optimistix.ImplicitAdjoint differentiates the converged fixed point by the implicit function theorem, so :meth:SCFSolver.compute_forces / :meth:SCFSolver.energy_and_forces -- the analytic forces :math:F = -\partial E/\partial R from :func:jax.grad -- are exact and avoid backprop through the SCF iterations (the PySCFAD rationale, Zhang & Chan 2022).

The reported total energy is the proper Kohn-Sham energy :math:E = \operatorname{Tr}[D\,h_\text{core}] + \tfrac12 \operatorname{Tr}[D\,J] + E_{xc} + E_{nn}.

References

• P. Pulay, Chem. Phys. Lett. 73, 393 (1980) -- DIIS; D. G. Anderson, J. ACM 12, 547 (1965) -- Anderson acceleration (the density-space DIIS used here).
• X. Zhang, G. K.-L. Chan, J. Chem. Phys. 157, 204801 (2022), arXiv:2207.13836 -- implicit differentiation of the SCF fixed point (PySCFAD).
• L. Y. Yao et al., arXiv:2411.05033 (jrystal / DWD) -- direct minimisation.
• R. G. Parr, W. Yang, Density-Functional Theory of Atoms and Molecules, Oxford (1989), Ch. 7 -- the Kohn-Sham total-energy expression.
• A. Szabo, N. S. Ostlund, Modern Quantum Chemistry, Dover (1996), Ch. 3 -- Roothaan equations and Lowdin symmetric orthogonalisation.

Functional

Bases: StrEnum

Supported exchange-correlation functionals.

SolverMode

Bases: StrEnum

How the self-consistent density is found.

SCFResult dataclass

SCFResult(*, total_energy: Array, orbital_energies: Array, density_matrix: Array, coefficients: Array, n_iterations: int, converged: bool)

Outcome of a restricted Kohn-Sham SCF calculation.

Attributes:

Name	Type	Description
total_energy	Array	Converged Kohn-Sham total energy (Hartree).
orbital_energies	Array	Molecular-orbital eigenvalues [Shape: (n_ao,)].
density_matrix	Array	Converged AO density matrix [Shape: (n_ao, n_ao)].
coefficients	Array	MO coefficients [Shape: (n_ao, n_ao)].
n_iterations	int	Number of SCF iterations performed.
converged	bool	Whether the density change fell below the tolerance.

SCFSolver

SCFSolver(system: MolecularSystem, basis: AtomicOrbitalBasis | None = None, *, functional: Functional | str = LDA, mode: SolverMode | str = DIIS, neural_functional: NeuralXCFunctional | None = None, grid_template: MolecularGridTemplate | None = None, max_iterations: int = 100, convergence_tolerance: float = 1e-08)

Restricted Kohn-Sham (RKS) self-consistent-field solver.

Parameters:

Name	Type	Description	Default
system	MolecularSystem	The molecular system to solve.	required
basis	AtomicOrbitalBasis | None	The AO basis (defaults to STO-3G built from the system).	None
functional	Functional | str	The exchange-correlation functional ("lda", "pbe" or "neural").	LDA
mode	SolverMode | str	"diis" for the Anderson-accelerated self-consistent SCF (Pulay DIIS on the density residual) or "direct" for direct minimisation.	DIIS
neural_functional	NeuralXCFunctional | None	A learned XC functional; required (and selects the "neural" functional) when functional == "neural".	None
grid_template	MolecularGridTemplate | None	A pre-built molecular-grid template; defaults to the standard Becke grid for system. Supplying a coarser template trades XC-integration accuracy for speed (e.g. in tests).	None
max_iterations	int	Maximum SCF / fixed-point / minimisation iterations.	100
convergence_tolerance	float	RMS density-change convergence threshold.	1e-08

functional property

functional: Functional

The exchange-correlation functional in use.

energy_from_positions

energy_from_positions(positions: Array, *, differentiable: str = 'implicit') -> Array

Converged Kohn-Sham total energy as a function of nuclear positions.

The self-consistent density is found as an implicit fixed point of the Roothaan step regardless of the forward :class:SolverMode (direct minimisation and the DIIS/fixed-point iteration converge to the same Kohn-Sham density). Differentiating the implicit fixed point gives exact, memory-cheap gradients via the implicit function theorem and avoids the gauge-singular Hessian of the direct-minimisation parametrisation, so :meth:compute_forces is robust for both modes.

Parameters:

Name	Type	Description	Default
positions	Array	Nuclear positions in Bohr [Shape: (n_atoms, 3)].	required
differentiable	str	"implicit" (default) finds the density as an implicit fixed point (IFT gradient); "unroll" runs a fixed number of differentiable SCF steps (gradient cross-check).	'implicit'

Returns:

Type	Description
Array	The scalar converged total energy (Hartree).

energy_from_state

energy_from_state(state: State, positions: Array | None = None) -> Array

Converged total energy as a differentiable function of the XC state.

The entry point for learned-XC training: jax.grad of this with respect to state gives the exact dE/dtheta through the implicit-diff SCF (the implicit function theorem differentiates the converged fixed point, not the iterations).

Parameters:

Name	Type	Description	Default
state	State	The neural XC parameter state (an nnx.State pytree, as produced by :func:flax.nnx.split).	required
positions	Array | None	Geometry to evaluate at (defaults to the system geometry).	None

Returns:

Type	Description
Array	The scalar converged total energy (Hartree).

energy

energy() -> Array

Converged total energy at the system's nuclear geometry.

compute_forces

compute_forces(positions: Array | None = None) -> Array

Analytic nuclear forces :math:F = -\partial E/\partial R.

Computed by :func:jax.grad of the implicit-diff total energy with respect to the nuclear coordinates.

Parameters:

Name	Type	Description	Default
positions	Array | None	Geometry to evaluate at (defaults to the system geometry).	None

Returns:

Type	Description
Array	Forces in Hartree/Bohr [Shape: (n_atoms, 3)].

energy_and_forces

energy_and_forces(positions: Array | None = None) -> tuple[Array, Array]

Converged total energy and the analytic forces :math:-\partial E/\partial R.

Parameters:

Name	Type	Description	Default
positions	Array | None	Geometry to evaluate at (defaults to the system geometry).	None

Returns:

Type	Description
tuple[Array, Array]	A pair (energy, forces) with forces in Hartree/Bohr.

solve

solve(*, initial_density: Array | None = None) -> SCFResult

Run the forward SCF (DIIS or direct minimisation) to convergence.

Parameters:

Name	Type	Description	Default
initial_density	Array | None	Optional closed-shell density to seed the Anderson/DIIS iteration. A high-quality guess (e.g. reconstructed from a neural-network predicted Fock via :func:density_from_fock) reaches the fixed point in fewer iterations; the converged result is unchanged. Ignored by the direct-minimisation mode, which is seeded internally.	None

Returns:

Name	Type	Description
The	SCFResult	class:SCFResult with the converged total energy and orbitals.

density_from_fock

density_from_fock(fock: Array, overlap: Array, n_occupied: int) -> Array

Closed-shell density from a Fock matrix by solving FC = SCe.

Reconstructs an initial-guess density from a Fock matrix (such as one predicted by a neural Hamiltonian model) in the same AO basis as overlap: it Lowdin-orthogonalises with S^{-1/2}, diagonalises the orthonormal Fock, back-transforms the lowest n_occupied orbitals and forms D = 2 C_occ C_occ^T. Pair the result with :meth:SCFSolver.solve(initial_density=...)<SCFSolver.solve> to seed the SCF.

Parameters:

Name	Type	Description	Default
fock	Array	The Fock matrix (n_ao, n_ao).	required
overlap	Array	The AO overlap matrix S (n_ao, n_ao) in the same basis.	required
n_occupied	int	Number of doubly-occupied orbitals (electrons // 2).	required

Returns:

Type	Description
Array	The closed-shell density matrix (n_ao, n_ao).

SCF acceleration from a predicted Fock

Measure SCF iteration reduction from a high-quality initial guess.

A neural Hamiltonian model that predicts a Fock matrix close to the self-consistent one can seed the SCF with a near-converged density, so the Anderson/DIIS iteration reaches the fixed point in fewer steps than the default core-Hamiltonian guess. This module quantifies that reduction: it runs the same :class:~opifex.neural.quantum.dft.scf.SCFSolver from the default guess (baseline) and from a supplied initial_density (guided) and reports the iteration counts, having checked that both reach the same converged energy.

The guess must be a closed-shell density in the solver's own AO basis. Use :func:~opifex.neural.quantum.dft.scf.density_from_fock to turn a predicted Fock (in that basis) into a density first. Wiring a QH9-trained spherical def2-SVP B3LYP predictor additionally requires a matching spherical-def2-SVP solver path; that basis bridge is tracked separately and is not assumed here.

SCFAccelerationResult dataclass

SCFAccelerationResult(*, baseline_iterations: int, guided_iterations: int, energy_hartree: float, converged: bool)

Outcome of comparing a guided SCF solve against the default guess.

Attributes:

Name	Type	Description
baseline_iterations	int	SCF cycles from the default core-Hamiltonian guess.
guided_iterations	int	SCF cycles from the supplied initial density.
energy_hartree	float	The converged total energy (Hartree); identical for both.
converged	bool	Whether the guided solve reached the convergence tolerance.

iteration_reduction property

iteration_reduction: int

Number of SCF cycles saved by the guided initial guess.

measure_scf_acceleration

measure_scf_acceleration(solver: SCFSolver, initial_density: Array, *, energy_tolerance: float = 1e-06) -> SCFAccelerationResult

Compare a guided SCF solve against the default-guess baseline.

Runs solver once from the default core-Hamiltonian guess and once from initial_density, then reports the iteration counts. The two solves must reach the same converged energy (the seed only changes the path, not the fixed point); a mismatch beyond energy_tolerance indicates an inconsistent guess (e.g. a density in the wrong AO basis) and raises.

Parameters:

Name	Type	Description	Default
solver	SCFSolver	The configured SCF solver (Anderson/DIIS mode).	required
initial_density	Array	Closed-shell density seed in the solver's AO basis.	required
energy_tolerance	float	Maximum allowed energy difference (Hartree) between the baseline and guided solves.	1e-06

Returns:

Name	Type	Description
The	SCFAccelerationResult	class:SCFAccelerationResult with both iteration counts.

Raises:

Type	Description
ValueError	If the guided solve converges to a different energy than the baseline (a sign the seed is inconsistent with the solver's basis).

spherical_fock_to_cartesian_density

spherical_fock_to_cartesian_density(spherical_fock: Array, cartesian_overlap: Array, angular_momenta: tuple[int, ...], n_occupied: int) -> Array

Build a Cartesian SCF seed density from a spherical-basis Fock matrix.

Bridges the predictor's spherical def2-SVP Fock (the standard 2l+1-per-shell basis) to an initial density in the SCF's Cartesian basis ((l+1)(l+2)/2 per shell -- e.g. 6 d components, with the extra contaminant). With the validated Cartesian->spherical block transform T (:func:~opifex.core.quantum._spherical.build_block_transform, columns in the spherical AO order), the Cartesian overlap is mapped to spherical (S_sph = T^T S_cart T), the closed-shell density is solved there (:func:~opifex.neural.quantum.dft.scf.density_from_fock), and embedded back as D_cart = T D_sph T^T. This congruence preserves the electron count Tr(D_cart S_cart) = 2 n_occ and overlap-metric idempotency D_cart S_cart D_cart = 2 D_cart exactly, so D_cart is a valid closed-shell seed for :meth:SCFSolver.solve(initial_density=...)<...solve>.

The seed lives in the spherical subspace of the Cartesian basis (the d contaminant starts at zero and the SCF relaxes it), so it is an approximate guess, not the exact Cartesian fixed point. spherical_fock must be in the same spherical AO order as T's columns; a QH9-predictor Fock (the pyscf_def2svp p-order) needs :func:~opifex.neural.quantum.hamiltonian.qh9_eval.to_pyscf_internal_ordering applied first.

Parameters:

Name	Type	Description	Default
spherical_fock	Array	The Fock matrix in the spherical AO basis (n_sph, n_sph).	required
cartesian_overlap	Array	The SCF's Cartesian AO overlap (n_cart, n_cart).	required
angular_momenta	tuple[int, ...]	The angular momentum l of each shell, in AO order (the SCF basis's shell.angular_momentum sequence).	required
n_occupied	int	Number of doubly-occupied orbitals (electrons // 2).	required

Returns:

Type	Description
Array	The Cartesian closed-shell seed density (n_cart, n_cart).

Exchange-correlation functionals

LDA and PBE (GGA) exchange-correlation functionals.

Provides the spin-unpolarised exchange-correlation energy densities used by the restricted Kohn-Sham solver in :mod:opifex.neural.quantum.dft.scf.

Local density approximation (LDA)

• Slater (Dirac) exchange -- the uniform-electron-gas exchange energy per particle :math:\varepsilon_x(\rho) = -C_x\,\rho^{1/3} with :math:C_x = \tfrac34 (3/\pi)^{1/3}.
• VWN5 correlation -- the Vosko-Wilk-Nusair (1980) parametrisation of the Ceperley-Alder uniform-gas correlation energy (their fit V; libxc/PySCF code 7, selected by 'lda,vwn').

Generalised gradient approximation (PBE)

• PBE exchange -- the uniform exchange times the enhancement factor :math:F_x(s) = 1 + \kappa - \kappa/(1+\mu s^2/\kappa) with :math:\kappa=0.804, :math:\mu=0.2195149727645171 and the reduced gradient :math:s = |\nabla\rho|/(2 k_F\rho), :math:k_F=(3\pi^2\rho)^{1/3} (PRL 77, 3865 (1996), eq. 13-14).
• PBE correlation -- the PW92 uniform correlation plus the gradient correction :math:H(r_s,t) (eq. 7-8). The uniform part is the Perdew-Wang (1992) lda_c_pw_mod fit (:math:A=0.0310907), which is what libxc's gga_c_pbe uses internally.

Everything is written in JAX so the functionals are differentiable and the exchange-correlation potential is obtained by automatic differentiation rather than a hand-coded derivative. For the LDA the potential is :math:v_{xc} = d(\rho\varepsilon_{xc})/d\rho; for the GGA it is the pair :math:(\partial(\rho\varepsilon)/\partial\rho,\; \partial(\rho\varepsilon)/\partial\sigma) with :math:\sigma=|\nabla\rho|^2.

References

• P. A. M. Dirac, Proc. Cambridge Philos. Soc. 26, 376 (1930) (exchange).
• S. H. Vosko, L. Wilk, M. Nusair, Can. J. Phys. 58, 1200 (1980), eq. 4.4 and Table 5 (paramagnetic fit) -- the VWN5 correlation parametrisation.
• J. P. Perdew, Y. Wang, Phys. Rev. B 45, 13244 (1992), Table I -- the uniform-gas correlation fit (A=0.0310907 pw_mod variant).
• J. P. Perdew, K. Burke, M. Ernzerhof, Phys. Rev. Lett. 77, 3865 (1996), eq. 7-8 (correlation H) and eq. 13-14 (exchange factor F_x).

slater_exchange_energy_density

slater_exchange_energy_density(density: Array) -> Array

Slater exchange energy per particle :math:\varepsilon_x(\rho).

Parameters:

Name	Type	Description	Default
density	Array	Total electron density rho [any shape].	required

Returns:

Type	Description
Array	Exchange energy per particle -C_x rho^(1/3) (same shape).

vwn_correlation_energy_density

vwn_correlation_energy_density(density: Array) -> Array

VWN5 correlation energy per particle :math:\varepsilon_c(\rho).

Implements the Vosko-Wilk-Nusair (1980) eq. 4.4 closed form for the paramagnetic (spin-unpolarised) electron gas:

.. math:: \varepsilon_c = A\Big[ \ln\frac{x^2}{X(x)} + \frac{2b}{Q}\arctan\frac{Q}{2x+b} - \frac{b x_0}{X(x_0)}\Big( \ln\frac{(x-x_0)^2}{X(x)} + \frac{2(b+2x_0)}{Q}\arctan\frac{Q}{2x+b}\Big)\Big],

with :math:x=\sqrt{r_s}, :math:X(x)=x^2+bx+c, :math:Q=\sqrt{4c-b^2} and the Wigner-Seitz radius :math:r_s=(3/4\pi\rho)^{1/3}.

Parameters:

Name	Type	Description	Default
density	Array	Total electron density rho [any shape].	required

Returns:

Type	Description
Array	Correlation energy per particle (same shape).

lda_energy_density

lda_energy_density(density: Array) -> Array

LDA exchange-correlation energy per particle :math:\varepsilon_{xc}.

Parameters:

Name	Type	Description	Default
density	Array	Total electron density rho [any shape].	required

Returns:

Type	Description
Array	epsilon_x + epsilon_c (same shape).

lda_exchange_correlation_potential

lda_exchange_correlation_potential(density: Array) -> Array

LDA XC potential :math:v_{xc} = d(\rho\,\varepsilon_{xc})/d\rho.

Computed by automatic differentiation of the XC energy density :math:\rho\,\varepsilon_{xc}(\rho).

Parameters:

Name	Type	Description	Default
density	Array	Total electron density rho [Shape: (n_points,)].	required

Returns:

Type	Description
Array	XC potential at each point [Shape: (n_points,)].

pw92_correlation_energy_density

pw92_correlation_energy_density(density: Array) -> Array

Perdew-Wang (1992) uniform-gas correlation energy per particle.

Implements the closed-form fit (Perdew & Wang 1992, eq. 10)

.. math:: \varepsilon_c^{\text{unif}}(r_s) = -2A(1+\alpha_1 r_s)\, \ln!\Big(1 + \frac{1}{2A(\beta_1 r_s^{½}+\beta_2 r_s +\beta_3 r_s^{3/2}+\beta_4 r_s^2)}\Big),

with the spin-unpolarised pw_mod constants (A=0.0310907) used by libxc's gga_c_pbe. Provides the uniform reference for PBE correlation.

Parameters:

Name	Type	Description	Default
density	Array	Total electron density rho [any shape].	required

Returns:

Type	Description
Array	Uniform-gas correlation energy per particle (same shape).

pbe_exchange_energy_density

pbe_exchange_energy_density(density: Array, sigma: Array) -> Array

PBE exchange energy per particle :math:\varepsilon_x^{\text{PBE}}.

The uniform-gas exchange is enhanced by the gradient-dependent factor :math:F_x(s) (PRL 77, 3865 (1996), eq. 13-14):

.. math:: \varepsilon_x^{\text{PBE}} = \varepsilon_x^{\text{unif}}\,F_x(s),\quad F_x(s) = 1 + \kappa - \frac{\kappa}{1 + \mu s^2/\kappa},\quad s = \frac{|\nabla\rho|}{2 k_F \rho},\; k_F = (3\pi^2\rho).

Parameters:

Name	Type	Description	Default
density	Array	Total electron density rho [any shape].	required
sigma	Array	Squared density gradient |grad rho|^2 [same shape].	required

Returns:

Type	Description
Array	PBE exchange energy per particle (same shape).

pbe_correlation_energy_density

pbe_correlation_energy_density(density: Array, sigma: Array) -> Array

PBE correlation energy per particle :math:\varepsilon_c^{\text{PBE}}.

Adds the gradient correction :math:H to the PW92 uniform correlation (PRL 77, 3865 (1996), eq. 7-8; here for the unpolarised case :math:\phi=1):

.. math:: H = \gamma\,\ln!\Big[1 + \frac{\beta}{\gamma} t^2 \frac{1+At^2}{1+At2+A^2t4}\Big],\quad A = \frac{\beta}{\gamma}\Big[e^{{-\varepsilon_c}-1 \Big]^{-1},\quad t = \frac{|\nabla\rho|}{2 k_s \rho},\; k_s = \sqrt{4 k_F/\pi}.}}/\gamma

Parameters:

Name	Type	Description	Default
density	Array	Total electron density rho [any shape].	required
sigma	Array	Squared density gradient |grad rho|^2 [same shape].	required

Returns:

Type	Description
Array	PBE correlation energy per particle (same shape).

pbe_energy_density

pbe_energy_density(density: Array, sigma: Array) -> Array

PBE exchange-correlation energy per particle :math:\varepsilon_{xc}.

Parameters:

Name	Type	Description	Default
density	Array	Total electron density rho [any shape].	required
sigma	Array	Squared density gradient |grad rho|^2 [same shape].	required

Returns:

Type	Description
Array	epsilon_x^PBE + epsilon_c^PBE (same shape).

pbe_exchange_correlation_potential

pbe_exchange_correlation_potential(density: Array, sigma: Array) -> tuple[Array, Array]

GGA XC potential components for PBE, by automatic differentiation.

Returns the two functional derivatives of the XC energy density :math:\rho\,\varepsilon_{xc}(\rho,\sigma) needed to assemble the GGA Fock contribution:

.. math:: v_\rho = \frac{\partial(\rho\varepsilon_{xc})}{\partial\rho},\qquad v_\sigma = \frac{\partial(\rho\varepsilon_{xc})}{\partial\sigma}.

Both are obtained with :func:jax.grad rather than a hand-coded derivative.

Parameters:

Name	Type	Description	Default
density	Array	Total electron density rho [Shape: (n_points,)].	required
sigma	Array	Squared density gradient |grad rho|^2 [Shape: (n_points,)].	required

Returns:

Type	Description
tuple[Array, Array]	A pair (v_rho, v_sigma) each [Shape: (n_points,)].

Hamiltonian prediction

opifex.neural.quantum.hamiltonian

Equivariant DFT Hamiltonian prediction in QHNet block form (native JAX/NNX).

A jax/flax.nnx implementation of equivariant electronic-structure matrix prediction (Yu et al. 2023, "QHNet", arXiv:2306.04922) built on opifex's Q0 equivariant kit (:mod:opifex.neural.equivariant) and the NequIP steerable trunk (:mod:opifex.neural.atomistic.backbones.nequip). Rather than assembling one dense matrix per fixed composition, the predictor emits a fixed (14, 14) diagonal block per atom and (14, 14) off-diagonal block per directed edge, so heterogeneous molecules concatenate into a single flat batch (the def2-SVP FULL_ORBITALS = 14 AO slots per second-row atom, masked per element).

The public surface:

• :class:HamiltonianBlockExpansion -- the QHNet expansion head (last-index Clebsch-Gordan contraction) turning a steerable feature plus an invariant embedding into a dense (14, 14) block.
• :class:BlockHamiltonianPredictor / :class:BlockHamiltonianConfig -- the heterogeneous-batchable per-atom / per-edge block predictor; its :meth:~...block_predictor.BlockHamiltonianPredictor.assemble_matrix scatters the blocks into a single molecule's symmetric dense Fock matrix.
• The orbital-layout primitives (:data:BLOCK_IRREPS, :data:FULL_ORBITALS, :data:ORBITAL_MASK, :func:atom_orbital_counts, :func:block_validity_mask) fixing the def2-SVP AO slots each block occupies.
• The GPU-fused block training surface (:class:BlockTrainConfig, :func:per_molecule_block_loss, :func:make_fused_block_train_step / :func:make_fused_block_eval_step) used by scripts/train_qh9_blocks.py.

BLOCK_IRREPS module-attribute

BLOCK_IRREPS: Irreps = Irreps('3x0e + 2x1e + 1x2e')

The 14-dim row/col representation of a Fock block (3x0e + 2x1e + 1x2e).

FULL_ORBITALS module-attribute

FULL_ORBITALS: int = 14

def2-SVP full second-row AO count (3 s + 2 p + 1 d = 3 + 6 + 5).

ORBITAL_MASK module-attribute

ORBITAL_MASK: dict[int, tuple[int, ...]] = {1: _HYDROGEN_AO_INDICES, 6: tuple(range(FULL_ORBITALS)), 7: tuple(range(FULL_ORBITALS)), 8: tuple(range(FULL_ORBITALS)), 9: tuple(range(FULL_ORBITALS))}

Atomic number -> valid AO indices into the 14-slot irrep-ordered block.

HamiltonianBlockExpansion

HamiltonianBlockExpansion(*, feature_irreps: Irreps | str = _DEFAULT_FEATURE_IRREPS, embed_dim: int = 64, mlp_hidden_dim: int = 128, rngs: Rngs)

Bases: Module

Expand a bottleneck feature + invariant embedding into a (14, 14) block.

Implements QHNet's Expansion (reference OpenDFT/QHBench/QH9/models/Expanson.py) over the output shell grid of :data:~opifex.neural.quantum.hamiltonian._orbital_layout.BLOCK_IRREPS (3x0e + 2x1e + 1x2e). Per-sample path weights (and scalar-block biases) are produced by an MLP on a provided invariant embedding, so the same module builds diagonal blocks (from a node embedding) and off-diagonal blocks (from a concatenated pair embedding). The Clebsch-Gordan contraction reuses :func:opifex.geometry.algebra.wigner.clebsch_gordan (no reimplementation).

Parameters:

Name	Type	Description	Default
feature_irreps	Irreps | str	Layout of the incoming steerable bottleneck feature. Defaults to 8x0e + 8x1e + 8x2e + 8x3e + 8x4e -- every degree reachable by a d-d (l = 2, 2) shell pair (L up to 4).	_DEFAULT_FEATURE_IRREPS
embed_dim	int	Width of the invariant embedding driving the weight/bias MLP.	64
mlp_hidden_dim	int	Hidden width of the weight/bias MLP.	128
rngs	Rngs	Random number generators (keyword-only) seeding the MLP.	required

BlockHamiltonianConfig dataclass

BlockHamiltonianConfig(*, hidden_irreps: str = '16x0e + 16x1o + 16x2e + 16x3o + 16x4e', sh_lmax: int = 4, num_interactions: int = 3, start_refinement_layer: int = 0, bottleneck_multiplicity: int = 16, num_radial_basis: int = 8, radial_hidden_dim: int = 64, cutoff: float = 20.0, average_num_neighbors: float = 1.0, embed_dim: int = 64, block_mlp_hidden_dim: int = 128, pair_weight_hidden_dim: int = 64)

Hyper-parameters of a :class:BlockHamiltonianPredictor.

Defaults sit well below the QHNet reference (hidden multiplicity ~128, sh_lmax 4, 5 interactions) so the documented defaults stay test-fast; production / training configs should raise hidden_irreps to a uniform-multiplicity Hx0e + Hx1o + Hx2e + Hx3o + Hx4e (sh_lmax 4) toward the reference.

Attributes:

Name	Type	Description
hidden_irreps	str	Steerable layout of the per-atom hidden / message-passing features (QHNet's hidden_irrep). Must be uniform multiplicity across all degrees (the channel-wise refinement tensor products require it) and contain a 0e scalar channel.
sh_lmax	int	Maximum spherical-harmonic degree of the edge embedding.
num_interactions	int	Number of NequIP convolution layers (QHNet's num_gnn_layers ConvNetLayer stack).
start_refinement_layer	int	Convolution index after which the self / pair refinement layers run (QHNet's start_layer); refinement happens for every layer with index strictly greater than it, so there are num_interactions - 1 - start_refinement_layer refinement layers.
bottleneck_multiplicity	int	Multiplicity of the even bottleneck feeding the block heads (QHNet's bottle_hidden_size).
num_radial_basis	int	Number of Bessel radial-basis functions.
radial_hidden_dim	int	Hidden width of the radial network MLP.
cutoff	float	Connection / cutoff radius r_c (Bohr). Defaults large so the complete within-molecule graph is retained.
average_num_neighbors	float	Constant sqrt normaliser for the aggregated message (NequIP's internal normalisation).
embed_dim	int	Width of the invariant embedding driving the block head's per-sample weight/bias MLP.
block_mlp_hidden_dim	int	Hidden width of the block head's weight/bias MLP.
pair_weight_hidden_dim	int	Hidden width of the pair layer's per-edge weight MLPs.

to_nequip

to_nequip() -> NequIPConfig

Return the matching :class:NequIPConfig for the reused conv layers.

BlockHamiltonianPredictor

BlockHamiltonianPredictor(*, config: BlockHamiltonianConfig | None = None, rngs: Rngs)

Bases: Module

Heterogeneous-batchable per-atom / per-edge QHNet Fock block predictor.

Consumes a flat concatenated batch (atomic_numbers, positions, within-molecule edge_index) and emits a fixed (14, 14) diagonal block per atom and (14, 14) off-diagonal block per directed edge. Reuses the NequIP convolution trunk (segment-based, hence batch-transparent), the QHNet self / pair interaction refinement layers and the shared :class:HamiltonianBlockExpansion head (reference divelab/AIRS OpenDFT/QHBench/QH9/models/QHNet.py).

Parameters:

Name	Type	Description	Default
config	BlockHamiltonianConfig | None	Hyper-parameters. Defaults to :class:BlockHamiltonianConfig.	None
rngs	Rngs	Random number generators (keyword-only) seeding all weights.	required

Raises:

Type	Description
ValueError	If config.hidden_irreps carries no 0e scalar channel, is not uniform multiplicity, or leaves no room for a refinement layer.

assemble_matrix

assemble_matrix(diagonal_blocks: Float[Array, 'n_atoms 14 14'], off_diagonal_blocks: Float[Array, 'n_edges 14 14'], atomic_numbers: Int[Array, ' n_atoms'], edge_index: Int[Array, '2 n_edges']) -> Float[Array, 'n_ao n_ao']

Assemble a single molecule's dense, symmetric (n_ao, n_ao) matrix.

Masks each block to its element's valid AO slots (:func:~...._orbital_layout.block_validity_mask) and scatters it into the dense matrix at the per-atom AO offsets (:func:~...._orbital_layout.atom_orbital_counts). The off-diagonal blocks are written at both (i, j) and (j, i); the directed graph carries both edges, so the QHNet off-diagonal law H[i, j] = B_ij + B_ji^T is realised by the final symmetrisation H = H~ + H~^T.

This is a host-side inference helper for a single molecule (it builds a dense matrix and uses Python sizing), not part of the batched forward.

Parameters:

Name	Type	Description	Default
diagonal_blocks	Float[Array, 'n_atoms 14 14']	(A, 14, 14) per-atom blocks from :meth:__call__.	required
off_diagonal_blocks	Float[Array, 'n_edges 14 14']	(E, 14, 14) per-edge blocks.	required
atomic_numbers	Int[Array, ' n_atoms']	(A,) atomic numbers of the single molecule.	required
edge_index	Int[Array, '2 n_edges']	(2, E) directed edge index of the single molecule.	required

Returns:

Type	Description
Float[Array, 'n_ao n_ao']	The symmetric dense AO matrix of shape (n_ao, n_ao).

BlockTrainConfig dataclass

BlockTrainConfig(*, learning_rate: float = 0.0005, beta1: float = 0.99, beta2: float = 0.999, weight_decay: float = 0.0, warmup_steps: int = 1000, total_steps: int = 300000, lr_end: float = 1e-07, power: float = 1.0, grad_clip_norm: float = 5.0)

QHNet training hyper-parameters for the block Hamiltonian predictor.

Defaults reproduce the QH9/QHNet reference setup (OpenDFT/QHBench/QH9): AdamW with lr = 5e-4 and betas = (0.99, 0.999), a polynomial (power = 1) decay schedule with a 1000-step warmup over 300000 total steps to lr_end = 1e-7, and global-norm gradient clipping at 5.0.

Attributes:

Name	Type	Description
learning_rate	float	Peak AdamW learning rate (post-warmup).
beta1	float	AdamW first-moment decay.
beta2	float	AdamW second-moment decay.
weight_decay	float	AdamW decoupled weight decay.
warmup_steps	int	Linear warmup steps to the peak learning rate.
total_steps	int	Total schedule steps (decay horizon).
lr_end	float	Final (floor) learning rate after polynomial decay.
power	float	Polynomial-decay power (1 = linear decay).
grad_clip_norm	float	Global gradient-norm clip threshold.

schedule

schedule() -> Schedule

Return the warmup + polynomial-decay learning-rate schedule.

Mirrors HuggingFace get_polynomial_decay_schedule_with_warmup used by the QHNet reference: a linear warmup from 0 to learning_rate over warmup_steps, then a polynomial decay to lr_end over the remaining total_steps - warmup_steps steps.

optimizer

optimizer() -> GradientTransformation

Return the AdamW + global-norm-clip optax transform for this config.

QH9TestSetMetrics dataclass

QH9TestSetMetrics(*, n_molecules: int, orbital_energy_mae: float, orbital_energy_mae_occ: float, coefficient_similarity: float, homo_lumo_gap_mae: float, hamiltonian_mae: float)

Aggregated QH9 test-set metrics over the evaluated molecules.

Attributes:

Name	Type	Description
n_molecules	int	Number of molecules evaluated.
orbital_energy_mae	float	Mean ε-MAE over all orbitals (Hartree).
orbital_energy_mae_occ	float	Mean ε-MAE over occupied orbitals (Hartree).
coefficient_similarity	float	Mean occupied-orbital ψ-cosine similarity.
homo_lumo_gap_mae	float	Mean HOMO-LUMO-gap MAE (Hartree).
hamiltonian_mae	float	Mean Fock-matrix MAE (Hartree).

as_dict

as_dict() -> dict[str, float | int]

Return the metrics as a plain dict (for JSON/logging).

atom_orbital_counts

atom_orbital_counts(atomic_numbers: Int[Array, '...']) -> Int[Array, '...']

Return the number of populated AOs per atom (5 for H/He, 14 for C/N/O/F).

Parameters:

Name	Type	Description	Default
atomic_numbers	Int[Array, '...']	Integer atomic numbers Z (any leading shape).	required

Returns:

Type	Description
Int[Array, '...']	Integer AO counts of shape atomic_numbers.shape.

block_validity_mask

block_validity_mask(row_atomic_numbers: Int[Array, '...'], col_atomic_numbers: Int[Array, '...'] | None = None) -> Bool[Array, '... 14 14']

Return the (..., 14, 14) AO validity mask of an atom or directed pair.

For a single atom (col_atomic_numbers is None) the mask is the outer product of the atom's per-AO validity with itself (the diagonal Fock block). For a directed edge it is the outer product of the row element's mask and the column element's mask -- mask[i, j] = row_valid[i] & col_valid[j] -- matching QHNet's per-pair matrix_block_mask (reference OpenDFT/QHBench/QH9/datasets.py).

Parameters:

Name	Type	Description	Default
row_atomic_numbers	Int[Array, '...']	Atomic numbers Z of the row (receiver) atoms.	required
col_atomic_numbers	Int[Array, '...'] | None	Atomic numbers Z of the column (sender) atoms; if None, the row atoms are reused (diagonal block).	None

Returns:

Type	Description
Bool[Array, '... 14 14']	Boolean mask of shape row_atomic_numbers.shape + (14, 14).

make_fused_block_eval_step

make_fused_block_eval_step(decode_op: OperatorModule, cut_op: OperatorModule, *, swap_edges: bool = True) -> Callable[..., Float[Array, '']]

Build a fused decode + cut + predict Hamiltonian-MAE eval step.

The evaluation analogue of :func:make_fused_block_train_step: it reuses the same operators to produce the target blocks on device, runs the per-molecule predictor and returns the Hamiltonian MAE (Hartree) without an optimizer.update.

Parameters:

Name	Type	Description	Default
decode_op	OperatorModule	The Fock spherical-decode operator.	required
cut_op	OperatorModule	The Fock block-cut operator.	required
swap_edges	bool	Edge-orientation flag forwarded to :func:predict_blocks_vmapped.	True

Returns:

Type	Description
Callable[..., Float[Array, '']]	A jitted (predictor, raw_batch) -> hamiltonian_mae step (Hartree).

make_fused_block_train_step

make_fused_block_train_step(decode_op: OperatorModule, cut_op: OperatorModule, *, num_molecules: int, swap_edges: bool = True) -> Callable[..., tuple[Float[Array, ''], Float[Array, '']]]

Build the fused decode + cut + predict + loss + update train step.

The returned nnx.jit closure runs, inside one compiled graph over a per-molecule padded batch (leading molecule axis):

1. the Fock spherical decode and block cut operators, vmapped over the molecule axis Batch-free via :meth:~datarax.core.operator.OperatorModule._apply_on_raw (no apply_batch, no Batch object);
2. the single-molecule predictor vmapped per molecule (:func:predict_blocks_vmapped);
3. one nnx.value_and_grad (has_aux=True) of :func:per_molecule_block_loss against the operator-produced target blocks -- a single forward yielding (loss, mae);
4. one optimizer.update.

No per-step host sync (no float() / block_until_ready) happens here; the caller syncs at log cadence. The operators carry no parameters, so the optimizer differentiates only the predictor.

Parameters:

Name	Type	Description	Default
decode_op	OperatorModule	The :class:~...qh9_fock_operators.FockSphericalDecodeOperator.	required
cut_op	OperatorModule	The :class:~...qh9_fock_operators.FockBlockCutOperator.	required
num_molecules	int	The fixed per-batch molecule count (the leading axis size).	required
swap_edges	bool	Edge-orientation flag forwarded to :func:predict_blocks_vmapped.	True

Returns:

Type	Description
Callable[..., tuple[Float[Array, ''], Float[Array, '']]]	A jitted (predictor, optimizer, raw_batch) -> (loss, mae) step.

per_molecule_block_loss

per_molecule_block_loss(predictions: dict[str, Float[Array, 'b ... 14 14']], batch: dict[str, Array]) -> tuple[Float[Array, ''], dict[str, Float[Array, '']]]

Combine the QHNet block loss (MSE + MAE) over a per-molecule batch.

The batch is (b, max_atoms, ...) / (b, max_edges, ...) per-molecule padded arrays (from :class:~opifex.data.sources.qh9_padded_source.QH9PaddedSource after the Fock operators), not one flat segment-concatenation, so each molecule's masked squared / absolute error is a plain per-molecule reduction (:func:_per_molecule_block_error). Each molecule's combined diagonal + off-diagonal error is normalised by its combined valid-element count and averaged over molecules.

Parameters:

Name	Type	Description	Default
predictions	dict[str, Float[Array, 'b ... 14 14']]	{"diagonal_blocks" (b, max_atoms, 14, 14), "off_diagonal_blocks" (b, max_edges, 14, 14)} (per-molecule vmapped predictor outputs).	required
batch	dict[str, Array]	A per-molecule padded batch dict carrying the operator-produced targets/masks and the node/edge pad masks.	required

Returns:

Type	Description
Float[Array, '']	(loss, metrics) with loss = mse + mae and metrics carrying
dict[str, Float[Array, '']]	"mae", "mse", "rmse", "hamiltonian_mae" (Hartree) and
tuple[Float[Array, ''], dict[str, Float[Array, '']]]	"hamiltonian_mae_micro" (micro-Hartree).

predict_blocks_vmapped

predict_blocks_vmapped(predictor: BlockHamiltonianPredictor, batch: dict[str, Array], *, swap_edges: bool = True) -> dict[str, Float[Array, 'b ... 14 14']]

Run the predictor per molecule over a leading-axis padded batch.

The batch carries a leading molecule axis (atomic_numbers (b, max_atoms), positions (b, max_atoms, 3), edge_index (b, 2, max_edges)), so the single-molecule predictor is :func:nnx.vmap-ed over that axis. The data path stores edge_index = (receiver, sender) while the predictor reads (sender, receiver); with swap_edges (default) each molecule's edge index is row-swapped before the call (see the module docstring's edge orientation note).

Parameters:

Name	Type	Description	Default
predictor	BlockHamiltonianPredictor	The single-molecule block Hamiltonian predictor.	required
batch	dict[str, Array]	A per-molecule padded batch dict (leading molecule axis).	required
swap_edges	bool	Whether to present the predictor its native (sender, receiver) edge order. Defaults to True.	True

Returns:

Type	Description
dict[str, Float[Array, 'b ... 14 14']]	``{"diagonal_blocks" (b, max_atoms, 14, 14),
dict[str, Float[Array, 'b ... 14 14']]	"off_diagonal_blocks" (b, max_edges, 14, 14)}``.

cal_orbital_and_energies

cal_orbital_and_energies(overlap: Float[Array, 'n_ao n_ao'], hamiltonian: Float[Array, 'n_ao n_ao']) -> tuple[Float[Array, ' n_ao'], Float[Array, 'n_ao n_ao']]

Solve the generalized eigenproblem F C = S C diag(eps) via Löwdin.

Faithful JAX port of the reference cal_orbital_and_energies (OpenDFT/QHBench/QH9/test.py lines 112-120): symmetric (Löwdin) orthogonalization S^{-1/2} = U diag(1/sqrt(s)) U^T (built as U / sqrt(s)), transform Fs = (S^{-1/2})^T F S^{-1/2}, eigendecompose Fs -> (orbital_energies, C_orth) and rotate the coefficients back to the AO basis C = S^{-1/2} C_orth. Eigenvalues of S are floored at 1e-8 before the inverse square root (numerical guard, as in the reference).

Parameters:

Name	Type	Description	Default
overlap	Float[Array, 'n_ao n_ao']	The AO overlap matrix S (symmetric positive-definite).	required
hamiltonian	Float[Array, 'n_ao n_ao']	The Fock matrix F in the same AO ordering as S.	required

Returns:

Type	Description
Float[Array, ' n_ao']	(orbital_energies, orbital_coefficients) -- ascending orbital energies
Float[Array, 'n_ao n_ao']	(n_ao,) and AO-basis coefficients (n_ao, n_ao) whose column k
tuple[Float[Array, ' n_ao'], Float[Array, 'n_ao n_ao']]	is orbital k.

evaluate_examples

evaluate_examples(predictor: BlockHamiltonianPredictor, examples: Iterable[QH9Example]) -> QH9TestSetMetrics

Aggregate :func:evaluate_fock over decoded QH9 examples.

For each example the predicted Fock is assembled from the predictor (:func:predict_fock) and compared with the example's target Fock; the per-molecule metrics are averaged (unweighted) over the molecules.

Parameters:

Name	Type	Description	Default
predictor	BlockHamiltonianPredictor	The block Hamiltonian predictor (trained or fresh).	required
examples	Iterable[QH9Example]	Decoded :class:~opifex.data.sources.qh9_source.QH9Example records to evaluate.	required

Returns:

Type	Description
QH9TestSetMetrics	The aggregated :class:QH9TestSetMetrics.

Raises:

Type	Description
ValueError	If examples is empty.

evaluate_fock

evaluate_fock(predicted_fock: Float[Array, 'n_ao n_ao'], target_fock: Float[Array, 'n_ao n_ao'], atomic_numbers: Int[NDArray[int32], ' n_atoms'], positions: Float[NDArray[float64], 'n_atoms 3'], n_electrons: int) -> dict[str, Float[Array, '']]

QH9 evaluation metrics for one molecule's predicted vs. target Fock.

Both Fock matrices are reordered into PySCF's internal spherical AO ordering (:func:to_pyscf_internal_ordering), paired with the PySCF overlap (:func:overlap_matrix_def2svp) and diagonalized via :func:cal_orbital_and_energies. The ε-MAE is reported over all orbitals and over the occupied set, the ψ-similarity over the occupied orbitals, and the HOMO-LUMO-gap MAE and Fock MAE complete the dict.

Parameters:

Name	Type	Description	Default
predicted_fock	Float[Array, 'n_ao n_ao']	Predicted Fock in opifex spherical AO ordering.	required
target_fock	Float[Array, 'n_ao n_ao']	Target Fock in opifex spherical AO ordering.	required
atomic_numbers	Int[NDArray[int32], ' n_atoms']	Nuclear charges, shape (n_atoms,).	required
positions	Float[NDArray[float64], 'n_atoms 3']	Atom positions in Bohr, shape (n_atoms, 3).	required
n_electrons	int	Total electron count sum(Z) of the closed-shell molecule; the occupied-orbital count is n_electrons // 2.	required

Returns:

Type	Description
dict[str, Float[Array, '']]	``{"orbital_energy_mae", "orbital_energy_mae_occ", "coefficient_similarity",
dict[str, Float[Array, '']]	"homo_lumo_gap_mae", "hamiltonian_mae"}`` of scalar JAX arrays.

evaluate_qh9_test_set

evaluate_qh9_test_set(predictor: BlockHamiltonianPredictor, db_path: Path, *, checkpoint_path: Path | None = None, limit: int | None = None) -> QH9TestSetMetrics

Evaluate the QH9 benchmark metrics over the QH9-Stable test split.

Reads only the deterministic 0.8/0.1/0.1 test-split molecules (:func:~opifex.data.sources.qh9_source.read_qh9_test_split, which computes the split from the cheap id count and decodes the test subset lazily -- no full 130k-molecule decode), optionally restores a best-val orbax checkpoint into predictor (:func:load_predictor_checkpoint), then aggregates :func:evaluate_fock over the test molecules (:func:evaluate_examples).

Parameters:

Name	Type	Description	Default
predictor	BlockHamiltonianPredictor	The block Hamiltonian predictor, built with the config matching checkpoint_path when one is given.	required
db_path	Path	Path to QH9Stable.db.	required
checkpoint_path	Path | None	Optional best-val checkpoint to restore before evaluating (e.g. <run>/checkpoints/best_epoch_N).	None
limit	int | None	Optional cap on the number of test-split molecules evaluated (the split itself is computed over the full database for fidelity).	None

Returns:

Type	Description
QH9TestSetMetrics	The aggregated :class:QH9TestSetMetrics over the evaluated test molecules.

hamiltonian_mae

hamiltonian_mae(predicted_fock: Float[Array, 'n_ao n_ao'], target_fock: Float[Array, 'n_ao n_ao']) -> Float[Array, '']

Mean absolute error between predicted and target Fock matrices (Hartree).

Parameters:

Name	Type	Description	Default
predicted_fock	Float[Array, 'n_ao n_ao']	Predicted Fock matrix.	required
target_fock	Float[Array, 'n_ao n_ao']	Target Fock matrix (same ordering and shape).	required

Returns:

Type	Description
Float[Array, '']	Scalar mean absolute Fock-element error (Hartree).

homo_lumo_gap

homo_lumo_gap(orbital_energies: Float[Array, ' n_ao'], n_occupied: int) -> Float[Array, '']

HOMO-LUMO gap eps[n_occ] - eps[n_occ - 1] from ascending energies.

Parameters:

Name	Type	Description	Default
orbital_energies	Float[Array, ' n_ao']	Ascending orbital energies (n_ao,).	required
n_occupied	int	Number of doubly-occupied orbitals.	required

Returns:

Type	Description
Float[Array, '']	The scalar HOMO-LUMO gap (Hartree).

latest_checkpoint

latest_checkpoint(checkpoint_dir: Path) -> Path | None

Return the newest best_epoch_* checkpoint under checkpoint_dir.

Parameters:

Name	Type	Description	Default
checkpoint_dir	Path	The run's checkpoints/ directory.	required

Returns:

Type	Description
Path | None	The path of the highest-epoch checkpoint, or None if none exist.

load_predictor_checkpoint

load_predictor_checkpoint(predictor: BlockHamiltonianPredictor, checkpoint_path: Path) -> BlockHamiltonianPredictor

Restore a best-val orbax checkpoint into predictor (in place) and return it.

Mirrors scripts/train_qh9_blocks.py's save format: the checkpoint is the nnx.to_pure_dict(nnx.state(predictor, nnx.Param)) pure-dict written by an :class:orbax.checkpoint.StandardCheckpointer. Restoration reads back into the same pure-dict structure and replaces the predictor's parameter state.

Parameters:

Name	Type	Description	Default
predictor	BlockHamiltonianPredictor	A predictor built with the same config as the checkpoint.	required
checkpoint_path	Path	Path to the saved orbax checkpoint directory.	required

Returns:

Type	Description
BlockHamiltonianPredictor	The same predictor with restored parameters.

Raises:

Type	Description
FileNotFoundError	If checkpoint_path does not exist.

occupied_orbital_count

occupied_orbital_count(atomic_numbers: Int[NDArray[int32], ' n_atoms']) -> int

Number of doubly-occupied orbitals of a closed-shell neutral molecule.

n_occ = sum(Z) / 2 for these closed-shell neutral QH9 molecules (reference test.py num_orb = int(batch.atoms.sum() / 2)).

Parameters:

Name	Type	Description	Default
atomic_numbers	Int[NDArray[int32], ' n_atoms']	Nuclear charges, shape (n_atoms,).	required

Returns:

Type	Description
int	The integer occupied-orbital count.

orbital_coefficient_similarity

orbital_coefficient_similarity(predicted_coefficients: Float[Array, 'n_ao n'], target_coefficients: Float[Array, 'n_ao n']) -> Float[Array, '']

Mean sign-invariant per-orbital cosine similarity of orbital coefficients.

Ports the reference ψ-similarity (test.py lines 57-59): cosine_similarity(pred, target, dim=0).abs().mean() -- the cosine similarity is taken per orbital (over the AO axis, dim=0), made sign-invariant via abs (orbital coefficients are defined up to a global sign), and averaged over orbitals.

Parameters:

Name	Type	Description	Default
predicted_coefficients	Float[Array, 'n_ao n']	Predicted AO-basis coefficients (n_ao, n) (column k is orbital k).	required
target_coefficients	Float[Array, 'n_ao n']	Target AO-basis coefficients (n_ao, n).	required

Returns:

Type	Description
Float[Array, '']	Scalar mean absolute per-orbital cosine similarity in [0, 1].

orbital_energy_mae

orbital_energy_mae(predicted_energies: Float[Array, ' n'], target_energies: Float[Array, ' n']) -> Float[Array, '']

Mean absolute error between predicted and target orbital energies.

Parameters:

Name	Type	Description	Default
predicted_energies	Float[Array, ' n']	Predicted orbital energies (any matching shape).	required
target_energies	Float[Array, ' n']	Target orbital energies.	required

Returns:

Type	Description
Float[Array, '']	Scalar mean absolute orbital-energy error (Hartree).

overlap_matrix_def2svp

overlap_matrix_def2svp(atomic_numbers: Int[NDArray[int32], ' n_atoms'] | Sequence[int], positions_bohr: Float[NDArray[float64], 'n_atoms 3']) -> Float[Array, 'n_ao n_ao']

Return the PySCF def2-SVP spherical AO overlap S at a QH9 geometry.

Builds the molecule with pyscf.gto.M(basis='def2svp', unit='Bohr') at the QH9 positions (Bohr, opifex convention) and reads int1e_ovlp_sph. Cached per geometry (:func:_overlap_cached) since the eval revisits geometries; the PySCF call is host-side only (this is an eval, not a training, path).

Parameters:

Name	Type	Description	Default
atomic_numbers	Int[NDArray[int32], ' n_atoms'] | Sequence[int]	Nuclear charges, shape (n_atoms,).	required
positions_bohr	Float[NDArray[float64], 'n_atoms 3']	Atom positions in Bohr, shape (n_atoms, 3).	required

Returns:

Type	Description
Float[Array, 'n_ao n_ao']	The overlap matrix S of shape (n_ao, n_ao) (PySCF internal
Float[Array, 'n_ao n_ao']	spherical AO ordering).

predict_fock

predict_fock(predictor: BlockHamiltonianPredictor, atomic_numbers: Int[NDArray[int32], ' n_atoms'], positions_bohr: Float[NDArray[float64], 'n_atoms 3']) -> Float[Array, 'n_ao n_ao']

Assemble a single molecule's predicted dense Fock from the block predictor.

Runs the predictor on the complete directed graph and assembles the symmetric dense Fock via :meth:BlockHamiltonianPredictor.assemble_matrix (reused). The edge index is in the predictor's (sender, receiver) convention so the assembled off-diagonal law matches the predictor's training orientation.

Parameters:

Name	Type	Description	Default
predictor	BlockHamiltonianPredictor	The trained (or fresh) block Hamiltonian predictor.	required
atomic_numbers	Int[NDArray[int32], ' n_atoms']	Nuclear charges, shape (n_atoms,).	required
positions_bohr	Float[NDArray[float64], 'n_atoms 3']	Atom positions in Bohr, shape (n_atoms, 3).	required

Returns:

Type	Description
Float[Array, 'n_ao n_ao']	The assembled symmetric dense Fock (n_ao, n_ao) in opifex spherical
Float[Array, 'n_ao n_ao']	AO ordering.

to_pyscf_internal_ordering

to_pyscf_internal_ordering(fock: Float[Array, 'n_ao n_ao'], atomic_numbers: Sequence[int] | Int[NDArray[int32], ' n_atoms']) -> Float[Array, 'n_ao n_ao']

Reorder an opifex-spherical Fock into PySCF's internal spherical AO order.

Ports the reference matrix_transform(..., convention='back2pyscf') applied to the spherical data/predicted Fock before pairing it with the PySCF overlap (test.py lines 149-167). With (I, s) from :func:_back2pyscf_indices this is the symmetric congruence F'[i, j] = F[I[i], I[j]] * s[i] * s[j], aligning F with mol.intor('int1e_ovlp_sph') so :func:cal_orbital_and_energies is in one consistent basis.

Parameters:

Name	Type	Description	Default
fock	Float[Array, 'n_ao n_ao']	Fock matrix in opifex spherical (pyscf_def2svp) AO ordering.	required
atomic_numbers	Sequence[int] | Int[NDArray[int32], ' n_atoms']	Nuclear charges of the molecule (H, C, N, O, F only).	required

Returns:

Type	Description
Float[Array, 'n_ao n_ao']	The Fock matrix in PySCF's internal spherical AO ordering.

Neural Operators

opifex.neural.operators

Opifex Neural Operators: Full Operator Learning Library

This module provides the most complete collection of neural operators for scientific machine learning, including all major variants from the neuraloperator repository and advanced architectures.

The library includes:

• Fourier Neural Operators (FNO, TFNO, U-FNO, SFNO, Local FNO, AM-FNO)
• Deep Operator Networks (DeepONet and variants)
• Specialized operators (GINO, MGNO, UQNO, LNO, WNO, GNO)
• Physics-informed operators (PINO)
• Graph-based operators
• Uncertainty quantification operators

All operators are built with JAX/FLAX NNX for high performance and support automatic differentiation, just-in-time compilation, and multi-device parallelization.

AdaptiveDeepONet

AdaptiveDeepONet(branch_input_dim: int, trunk_input_dim: int, base_latent_dim: int, *, num_resolution_levels: int = 3, adaptive_latent_scaling: bool = True, use_residual_connections: bool = True, activation: str = 'tanh', rngs: Rngs)

Bases: Module

Adaptive DeepONet with dynamic architecture adjustment.

This variant can adapt its architecture based on problem complexity and provides multiple resolution levels for different accuracy requirements.

Parameters:

Name	Type	Description	Default
branch_input_dim	int	Branch network input dimension	required
trunk_input_dim	int	Trunk network input dimension	required
base_latent_dim	int	Base latent dimension (scaled for different levels)	required
num_resolution_levels	int	Number of resolution levels	3
adaptive_latent_scaling	bool	Whether to scale latent dimensions adaptively	True
use_residual_connections	bool	Whether to use residual connections	True
activation	str	Activation function name	'tanh'
rngs	Rngs	Random number generators	required

DeepONet

DeepONet(branch_sizes: list[int], trunk_sizes: list[int], *, activation: str = 'gelu', output_activation: str | None = None, use_bias: bool = True, rngs: Rngs)

Bases: Module

Deep Operator Network for learning function-to-function mappings.

DeepONet learns to approximate nonlinear operators G that map functions to functions: G: u → G(u), where u and G(u) are functions.

The architecture consists of: - Branch network: Processes input function u evaluated at sensors - Trunk network: Processes evaluation locations y - Dot product combination of branch and trunk outputs

Fully compliant with modern Flax NNX patterns.

Parameters:

Name	Type	Description	Default
branch_sizes	list[int]	Layer sizes for branch network [input_sensors, hidden1, hidden2, ..., output_dim]	required
trunk_sizes	list[int]	Layer sizes for trunk network [location_dim, hidden1, hidden2, ..., output_dim] Note: output_dim should match branch output_dim	required
activation	str	Activation function name for hidden layers	'gelu'
output_activation	str | None	Optional activation for final output (None means no activation on output)	None
use_bias	bool	Whether to use bias in linear layers	True
rngs	Rngs	Random number generators (keyword-only)	required

get_branch_output

get_branch_output(branch_input: Array, *, deterministic: bool = True) -> Array

Get branch network output for analysis purposes.

Parameters:

Name	Type	Description	Default
branch_input	Array	Function values at sensor locations	required
deterministic	bool	Whether to use deterministic mode	True

Returns:

Type	Description
Array	Branch network output

get_trunk_output

get_trunk_output(trunk_input: Array, *, deterministic: bool = True) -> Array

Get trunk network output for analysis purposes.

Parameters:

Name	Type	Description	Default
trunk_input	Array	Evaluation locations	required
deterministic	bool	Whether to use deterministic mode	True

Returns:

Type	Description
Array	Trunk network output

FourierEnhancedDeepONet

FourierEnhancedDeepONet(branch_sizes: list[int], trunk_sizes: list[int], *, fourier_modes: int = 16, use_spectral_branch: bool = True, use_spectral_trunk: bool = False, activation: str = 'tanh', rngs: Rngs)

Bases: Module

Fourier-Enhanced DeepONet combining spectral and operator learning.

This variant integrates Fourier Neural Operator concepts into DeepONet architecture for improved performance on problems with spectral structure.

Parameters:

Name	Type	Description	Default
branch_sizes	list[int]	Branch network layer sizes [input, hidden..., output]	required
trunk_sizes	list[int]	Trunk network layer sizes [input, hidden..., output]	required
fourier_modes	int	Number of Fourier modes for spectral layers	16
use_spectral_branch	bool	Whether to use spectral convolution in branch	True
use_spectral_trunk	bool	Whether to use spectral convolution in trunk	False
activation	str	Activation function name	'tanh'
rngs	Rngs	Random number generators	required

MultiPhysicsDeepONet

MultiPhysicsDeepONet(branch_input_dim: int, trunk_input_dim: int, branch_hidden_dims: list[int], trunk_hidden_dims: list[int], latent_dim: int, *, num_physics_systems: int = 1, use_attention: bool = True, attention_heads: int = 8, physics_constraints: list[str] | None = None, sensor_optimization: bool = False, num_sensors: int | None = None, activation: Callable[[Array], Array] = tanh, rngs: Rngs)

Bases: Module

Enhanced DeepONet with multi-physics support and attention mechanisms.

Extends the basic DeepONet architecture with physics-aware attention, multi-physics coupling, and sensor optimization for improved operator learning.

Parameters:

Name	Type	Description	Default
branch_input_dim	int	Branch network input dimension	required
trunk_input_dim	int	Trunk network input dimension	required
branch_hidden_dims	list[int]	Branch network hidden dimensions	required
trunk_hidden_dims	list[int]	Trunk network hidden dimensions	required
latent_dim	int	Latent dimension for inner product	required
num_physics_systems	int	Number of physics systems to handle	1
use_attention	bool	Whether to use physics-aware attention	True
attention_heads	int	Number of attention heads	8
physics_constraints	list[str] | None	List of physics constraints to enforce	None
sensor_optimization	bool	Whether to use sensor optimization	False
num_sensors	int | None	Number of sensors (required if sensor_optimization=True)	None
activation	Callable[[Array], Array]	Activation function	tanh
rngs	Rngs	Random number generators	required

branch_nets property

branch_nets: list[Module]

Get branch networks from all physics operators.

get_sensor_positions

get_sensor_positions() -> Array | None

Get current sensor positions if sensor optimization is enabled.

set_physics_constraints

set_physics_constraints(constraints: list[str]) -> None

Update physics constraints for attention mechanism.

AmortizedFourierNeuralOperator

AmortizedFourierNeuralOperator(in_channels: int, out_channels: int, hidden_channels: int = 32, modes: Sequence[int] = (16, 16), num_layers: int = 4, kernel_hidden_dim: int = 128, kernel_layers: int = 3, max_frequency: float = 10.0, activation: Callable = gelu, use_layer_norm: bool = False, use_kernel_regularization: bool = True, *, rngs: Rngs)

Bases: Module

Amortized Fourier Neural Operator with neural kernel parameterization.

get_regularization_loss

get_regularization_loss(x: Array) -> Array

Compute regularization loss on demand.

get_kernel_analysis

get_kernel_analysis(freq_range: tuple[float, float], num_points: int = 100) -> dict[str, Array]

Analyze learned kernel functions.

AmortizedSpectralConvolution

AmortizedSpectralConvolution(in_channels: int, out_channels: int, modes: Sequence[int], kernel_hidden_dim: int = 128, kernel_layers: int = 3, max_frequency: float = 10.0, use_kernel_regularization: bool = True, *, rngs: Rngs)

Bases: Module

Amortized spectral convolution with neural kernel parameterization.

KernelNetwork

KernelNetwork(freq_dim: int, output_dim: int, hidden_dim: int = 128, num_layers: int = 3, activation: Callable = gelu, use_frequency_encoding: bool = True, max_frequency: float = 10.0, *, rngs: Rngs)

Bases: Module

Neural network to parameterize Fourier kernels.

FourierLayer

FourierLayer(in_channels: int, out_channels: int, modes: int, *, activation: Callable[[Array], Array] = gelu, spatial_dims: int = 2, factorization: str | None = None, factorization_rank: float | None = None, rngs: Rngs)

Bases: Module

Fourier layer combining spectral convolution with activation.

This layer performs: 1. FFT to transform input to spectral domain 2. Spectral convolution 3. IFFT to transform back to spatial domain 4. Linear transformation and activation with proper residual connection

Fully compliant with modern Flax NNX patterns.

Parameters:

Name	Type	Description	Default
in_channels	int	Number of input channels	required
out_channels	int	Number of output channels	required
modes	int	Number of Fourier modes	required
activation	Callable[[Array], Array]	Activation function	gelu
spatial_dims	int	Number of spatial dimensions (1, 2, or 3). Controls which spectral weights are allocated — avoids dead parameters.	2
factorization	str | None	Optional low-rank factorization of the spectral weight ('tucker', 'cp', or 'tt'); None uses a dense weight.	None
factorization_rank	float | None	Compression ratio for the factorization (per-mode Tucker ratio, or ratio of min(shape) for CP/TT); defaults to 0.5.	None
rngs	Rngs	Random number generators (keyword-only)	required

get_compression_stats

get_compression_stats() -> dict[str, float]

Report factorized-vs-dense parameter compression for this layer.

Returns:

Type	Description
dict[str, float]	Mapping with the factorized parameter count, the equivalent dense
dict[str, float]	spectral-weight count, their ratio, and the fractional reduction.

Raises:

Type	Description
ValueError	If the layer uses dense (non-factorized) spectral weights.

FourierNeuralOperator

FourierNeuralOperator(in_channels: int, out_channels: int, hidden_channels: int, modes: int, num_layers: int, *, activation: Callable[[Array], Array] = gelu, factorization_type: str | None = None, factorization_rank: float | None = None, positional_embedding: bool = False, use_mixed_precision: bool = False, domain_padding: float = 0.0, spatial_dims: int = 2, rngs: Rngs)

Bases: Module

Fourier Neural Operator for learning solution operators of PDEs.

Implements the complete FNO architecture with optional tensor factorization and mixed precision training capabilities. Fully compliant with modern Flax NNX patterns.

Parameters:

Name	Type	Description	Default
in_channels	int	Number of input channels	required
out_channels	int	Number of output channels	required
hidden_channels	int	Number of hidden channels	required
modes	int	Number of Fourier modes	required
num_layers	int	Number of Fourier layers	required
activation	Callable[[Array], Array]	Activation function	gelu
factorization_type	str | None	Optional tensor factorization ('tucker', 'cp', 'tt')	None
factorization_rank	float | None	Rank for tensor factorization	None
positional_embedding	bool	If True, append normalised grid-coordinate channels to the input before lifting (needed for boundary-value problems such as Darcy flow).	False
use_mixed_precision	bool	Whether to use mixed precision	False
domain_padding	float	Fraction of each spatial dimension to zero-pad before the spectral layers (reduces the Gibbs phenomenon for non-periodic problems such as Darcy flow). Specified as a fraction (e.g. 0.25), NOT pixels, so the padding scales with resolution and preserves the FNO's discretisation-invariance / zero-shot super-resolution property. 0 disables.	0.0
spatial_dims	int	Number of spatial dimensions (1, 2, or 3). Determines which spectral weights are allocated per layer.	2
rngs	Rngs	Random number generators (keyword-only)	required

get_compression_stats

get_compression_stats() -> dict[str, float]

Aggregate factorized-vs-dense spectral compression across all layers.

Returns:

Type	Description
dict[str, float]	Mapping with summed factorized and equivalent-dense spectral parameter
dict[str, float]	counts, their ratio, and the fractional reduction.

Raises:

Type	Description
ValueError	If the operator uses dense (non-factorized) spectral weights.

count_parameters

count_parameters() -> int

Count total number of trainable parameters in the model.

FactorizedFourierLayer

FactorizedFourierLayer(in_channels: int, out_channels: int, modes: int, factorization_type: str, factorization_rank: int, *, activation: Callable[[Array], Array] = gelu, rngs: Rngs)

Bases: Module

Fourier layer with tensor factorization for parameter reduction.

Implements Tucker or CP factorization of the spectral convolution weights to achieve significant parameter reduction (up to 95%) while maintaining performance.

Parameters:

Name	Type	Description	Default
in_channels	int	Number of input channels	required
out_channels	int	Number of output channels	required
modes	int	Number of Fourier modes	required
factorization_type	str	Type of factorization ("tucker" or "cp")	required
factorization_rank	int	Rank for factorization	required
activation	Callable[[Array], Array]	Activation function	gelu
rngs	Rngs	Random number generators	required

get_parameter_count

get_parameter_count() -> dict[str, int | float]

Get parameter count breakdown for analysis.

LocalFourierLayer

LocalFourierLayer(in_channels: int, out_channels: int, modes: Sequence[int], kernel_size: int = 3, activation: Callable = gelu, mixing_weight: float = 0.5, *, rngs: Rngs)

Bases: Module

Fourier layer with local convolution for capturing short-range interactions.

Combines global spectral convolution with local spatial convolution for full feature extraction.

Parameters:

Name	Type	Description	Default
in_channels	int	Number of input channels	required
out_channels	int	Number of output channels	required
modes	Sequence[int]	Fourier modes for spectral convolution	required
kernel_size	int	Kernel size for local convolution	3
activation	Callable	Activation function	gelu
mixing_weight	float	Weight for combining spectral and local branches	0.5
rngs	Rngs	Random number generator state	required

get_mixing_analysis

get_mixing_analysis(x: Array) -> tuple[Array, Array, Array]

Analyze global vs local contributions for this layer.

Parameters:

Name	Type	Description	Default
x	Array	Input tensor (batch, in_channels, *spatial).	required

Returns:

Type	Description
Array	Tuple of (global_features, local_features, mixing_weights)
Array	where mixing_weights is a scalar array of the spectral weight.

LocalFourierNeuralOperator

LocalFourierNeuralOperator(in_channels: int, out_channels: int, hidden_channels: int, modes: Sequence[int], num_layers: int = 4, kernel_size: int = 3, use_adaptive_mixing: bool = True, use_residual_connections: bool = True, activation: Callable = gelu, *, rngs: Rngs)

Bases: Module

Local Fourier Neural Operator combining global and local operations.

This operator is designed for problems that require both: - Long-range dependencies (captured by Fourier operations) - Local features and fine details (captured by convolutions)

Examples include: - Turbulent flows with both large-scale structures and small eddies - Wave propagation with local scattering and global modes - Multi-physics problems with different characteristic scales

Parameters:

Name	Type	Description	Default
in_channels	int	Number of input channels	required
out_channels	int	Number of output channels	required
hidden_channels	int	Hidden layer width	required
modes	Sequence[int]	Fourier modes for global operations	required
num_layers	int	Number of Local Fourier layers	4
kernel_size	int	Kernel size for local convolutions	3
use_adaptive_mixing	bool	Whether to use adaptive feature mixing	True
use_residual_connections	bool	Whether to use residual connections	True
activation	Callable	Activation function	gelu
rngs	Rngs	Random number generator state	required

analyze_global_local_contributions

analyze_global_local_contributions(x: Array) -> dict[str, list[Array]]

Analyze global vs local contributions at each layer.

Returns:

Type	Description
dict[str, list[Array]]	Dictionary with global and local feature maps

MultiScaleFourierNeuralOperator

MultiScaleFourierNeuralOperator(in_channels: int, out_channels: int, hidden_channels: int, modes_per_scale: list[int], num_layers_per_scale: list[int], *, spatial_dims: int = 2, activation: Callable[[Array], Array] = gelu, use_cross_scale_attention: bool = True, attention_heads: int = 8, dropout_rate: float = 0.0, use_gradient_checkpointing: bool = True, rngs: Rngs)

Bases: Module

Multi-Scale Fourier Neural Operator for hierarchical resolution handling.

This operator learns operators across multiple scales simultaneously, enabling efficient handling of multi-scale physics problems like turbulence, multi-phase flows, and hierarchical material structures.

Features: - Hierarchical spectral convolutions at different resolution levels - Adaptive scale selection based on input characteristics - Cross-scale information exchange through attention mechanisms - Memory-efficient implementation with gradient checkpointing

Parameters:

Name	Type	Description	Default
in_channels	int	Number of input channels	required
out_channels	int	Number of output channels	required
hidden_channels	int	Hidden channel dimension	required
modes_per_scale	list[int]	List of Fourier modes for each scale	required
num_layers_per_scale	list[int]	List of layer counts for each scale	required
spatial_dims	int	Number of spatial dimensions (1 or 2)	2
activation	Callable[[Array], Array]	Activation function	gelu
use_cross_scale_attention	bool	Whether to use cross-scale attention	True
attention_heads	int	Number of attention heads	8
dropout_rate	float	Dropout rate for regularization	0.0
use_gradient_checkpointing	bool	Whether to use gradient checkpointing	True
rngs	Rngs	Random number generators	required

SphericalFourierNeuralOperator

SphericalFourierNeuralOperator(in_channels: int, out_channels: int, hidden_channels: int, lmax: int, mmax: int | None = None, num_layers: int = 4, activation: Callable = gelu, grid: str = 'legendre-gauss', *, rngs: Rngs)

Bases: Module

Spherical Fourier Neural Operator for data on spherical domains.

Uses spherical harmonic transforms instead of regular FFTs, making it ideal for: - Global atmospheric modeling - Ocean circulation - Planetary science - Any data naturally defined on spheres

Parameters:

Name	Type	Description	Default
in_channels	int	Number of input channels.	required
out_channels	int	Number of output channels.	required
hidden_channels	int	Hidden layer width.	required
lmax	int	Maximum spherical harmonic degree (controls spectral resolution).	required
mmax	int | None	Maximum azimuthal order (if None, uses lmax).	None
num_layers	int	Number of SFNO layers.	4
activation	Callable	Activation function.	gelu
grid	str	Latitude quadrature grid for the real SHT ("legendre-gauss").	'legendre-gauss'
rngs	Rngs	Random number generator state.	required

get_spherical_modes

get_spherical_modes(x: Array) -> Array

Get spherical harmonic coefficients for analysis.

Parameters:

Name	Type	Description	Default
x	Array	Input tensor on sphere	required

Returns:

Type	Description
Array	Spherical harmonic coefficients

compute_power_spectrum

compute_power_spectrum(x: Array) -> Array

Compute the spherical harmonic power spectrum per degree l.

The real SHT stores only non-negative orders m; the negative orders of a real field are their conjugates, so the angular power at degree l is |c_l^0|^2 + 2 * sum_{m>0} |c_l^m|^2.

Parameters:

Name	Type	Description	Default
x	Array	Input tensor on sphere (batch, channels, nlat, nlon).	required

Returns:

Type	Description
Array	Power spectrum (batch, channels, lmax) as a function of degree.

SphericalHarmonicConvolution

SphericalHarmonicConvolution(in_channels: int, out_channels: int, lmax: int, mmax: int | None = None, *, rngs: Rngs)

Bases: Module

Spherical harmonic convolution for spherical domains.

Operates in spherical harmonic space analogous to how standard FNO operates in Fourier space, but adapted for spherical geometry. The coefficient layout is (batch, channels, lmax, mmax) with non-negative orders m only, matching the real SHT of torch-harmonics / neuralop SphericalConv. A learnable complex weight contracts the channel axis per spherical mode.

Parameters:

Name	Type	Description	Default
in_channels	int	Number of input channels.	required
out_channels	int	Number of output channels.	required
lmax	int	Maximum spherical harmonic degree (controls resolution).	required
mmax	int | None	Maximum azimuthal order (if None, uses lmax).	None
rngs	Rngs	Random number generator state.	required

TensorizedFourierNeuralOperator

TensorizedFourierNeuralOperator(in_channels: int, out_channels: int, hidden_channels: int = 64, modes: Sequence[int] | int = (16, 16), num_layers: int = 4, factorization: Literal['tucker', 'cp', 'tt'] = 'tucker', rank: float = 0.1, *, rngs: Rngs)

Bases: FourierNeuralOperator

Tensorized FNO — a Fourier Neural Operator with low-rank spectral weights.

Thin specialisation of :class:~opifex.neural.operators.fno.base.FourierNeuralOperator that stores each spectral-convolution weight as a CP / Tucker / Tensor-Train factorization (Kossaifi et al., "Multi-Grid Tensorized Fourier Neural Operator"). It inherits the full, correct FNO forward pass — lifting, the activation(spectral + skip) Fourier blocks, grid positional embedding, and the two-layer projection head — so the only difference from a dense FNO is the factorized weight. This deletes the previously duplicated (and incorrect) spectral/forward implementation in favour of the single shared one (Rule 1).

TensorizedSpectralConvolution

TensorizedSpectralConvolution(in_channels: int, out_channels: int, modes: Sequence[int], decomposition_type: Literal['tucker', 'cp', 'tt'] = 'tucker', rank: float = 0.1, *, rngs: Rngs)

Bases: Module

Spectral convolution whose weight is a low-rank CP / Tucker / TT factorization.

Transforms a real spatial field to the Fourier domain, contracts the centered low-frequency band against the factorized weight (keeping both positive and negative low frequencies) and transforms back — see :func:opifex.neural.operators.fno._factorized.factorized_spectral_conv.

get_compression_stats

get_compression_stats() -> dict[str, float]

Get compression statistics.

UFNODecoderBlock

UFNODecoderBlock(in_channels: int, skip_channels: int, out_channels: int, modes: Sequence[int], upsample_factor: int = 2, activation: Callable = gelu, *, rngs: Rngs)

Bases: Module

Clean U-FNO decoder block with standardized tensor operations.

Performs: upsampling + skip fusion + spectral convolution

UFNOEncoderBlock

UFNOEncoderBlock(in_channels: int, out_channels: int, modes: Sequence[int], downsample_factor: int = 2, activation: Callable = gelu, *, rngs: Rngs)

Bases: Module

Clean U-FNO encoder block with standardized tensor operations.

Performs: spectral convolution + skip connection + downsampling

UFourierNeuralOperator

UFourierNeuralOperator(in_channels: int, out_channels: int, hidden_channels: int, modes: Sequence[int], num_levels: int = 3, downsample_factor: int = 2, activation: Callable = gelu, *, rngs: Rngs)

Bases: Module

U-Net style Fourier Neural Operator with clean, standardized architecture.

Features: - Consistent tensor dimension handling - Standardized spectral operations - Clean encoder-decoder structure - Proper channel management throughout

GraphNeuralOperator

GraphNeuralOperator(node_dim: int, hidden_dim: int, num_layers: int, *, edge_dim: int = 0, activation: Callable[[Array], Array] = gelu, rngs: Rngs)

Bases: Module

Graph Neural Operator for learning operators on irregular domains.

Implements message passing neural networks with geometric awareness for learning operators on graph-structured data. Suitable for irregular meshes, molecular systems, and other graph-based scientific computing applications.

Parameters:

Name	Type	Description	Default
node_dim	int	Dimension of node features	required
hidden_dim	int	Hidden dimension for message passing	required
num_layers	int	Number of message passing layers	required
edge_dim	int	Dimension of edge features (0 for no edge features)	0
activation	Callable[[Array], Array]	Activation function	gelu
rngs	Rngs	Random number generators	required

MessagePassingLayer

MessagePassingLayer(node_dim: int, edge_dim: int, hidden_dim: int, *, activation: Callable[[Array], Array] = gelu, rngs: Rngs)

Bases: Module

Message passing layer for graph neural networks.

Implements the message passing paradigm: 1. Compute messages between connected nodes 2. Aggregate messages at each node 3. Update node features based on aggregated messages

Parameters:

Name	Type	Description	Default
node_dim	int	Dimension of node features	required
edge_dim	int	Dimension of edge features	required
hidden_dim	int	Hidden dimension for message computation	required
activation	Callable[[Array], Array]	Activation function	gelu
rngs	Rngs	Random number generators	required

PhysicsAwareAttention

PhysicsAwareAttention(embed_dim: int, num_heads: int, *, physics_constraints: list[str] | None = None, dropout_rate: float = 0.0, rngs: Rngs)

Bases: Module

Physics-aware attention mechanism with constraint enforcement.

Integrates physics constraints into the attention mechanism to ensure physically meaningful attention patterns.

Parameters:

Name	Type	Description	Default
embed_dim	int	Embedding dimension	required
num_heads	int	Number of attention heads	required
physics_constraints	list[str] | None	List of physics constraints to enforce	None
dropout_rate	float	Dropout rate for attention weights	0.0
rngs	Rngs	Random number generators	required

PhysicsCrossAttention

PhysicsCrossAttention(embed_dim: int, num_heads: int, physics_constraints: list[str], num_physics_systems: int, *, conservation_weight: float = 0.1, adaptive_weighting: bool = True, cross_system_coupling: bool = True, dropout_rate: float = 0.0, rngs: Rngs)

Bases: Module

Physics-Cross-Attention mechanism for enhanced multi-physics coupling.

Implements cross-attention between different physics systems with conservation law enforcement and adaptive weighting based on physics constraints.

Parameters:

Name	Type	Description	Default
embed_dim	int	Embedding dimension	required
num_heads	int	Number of attention heads	required
physics_constraints	list[str]	List of physics constraints to enforce	required
num_physics_systems	int	Number of different physics systems	required
conservation_weight	float	Weight for conservation law enforcement	0.1
adaptive_weighting	bool	Whether to use adaptive constraint weighting	True
cross_system_coupling	bool	Whether to enable cross-system coupling	True
dropout_rate	float	Dropout rate for attention weights	0.0
rngs	Rngs	Random number generators	required

forward_with_conservation

forward_with_conservation(x: Array, *, physics_info: Array | None = None, training: bool = False) -> tuple[Array, Array]

Forward pass returning the output and its conservation-law residual.

The conservation loss is the squared flux divergence of the predicted field (see :meth:_compute_conservation_loss); it is always computed from the output, independent of physics_info.

Parameters:

Name	Type	Description	Default
x	Array	Input tensor.	required
physics_info	Array | None	Optional physics constraint information passed to the forward pass.	None
training	bool	Whether in training mode.	False

Returns:

Type	Description
tuple[Array, Array]	Tuple of (output, conservation_loss).

PhysicsInformedOperator

PhysicsInformedOperator(layer_sizes: list[int], physics_type: str = 'pde', *, activation: str = 'gelu', physics_weight: float = 1.0, data_weight: float = 1.0, use_bias: bool = True, rngs: Rngs)

Bases: Module

Physics-Informed Neural Operator with embedded physical constraints.

This operator combines standard neural operator architectures with physics-based constraints and differential operators to ensure physically consistent solutions.

Fully compliant with modern Flax NNX patterns.

Parameters:

Name	Type	Description	Default
layer_sizes	list[int]	Layer sizes for the neural network [input_dim, hidden1, hidden2, ..., output_dim]	required
physics_type	str	Type of physics constraint ('pde', 'conservation', 'symmetry')	'pde'
activation	str	Activation function name	'gelu'
physics_weight	float	Weight for physics loss component	1.0
data_weight	float	Weight for data loss component	1.0
use_bias	bool	Whether to use bias in linear layers	True
rngs	Rngs	Random number generators (keyword-only)	required

compute_physics_loss

compute_physics_loss(coordinates: Array, *, deterministic: bool = True) -> Array

Compute physics-based loss components.

Parameters:

Name	Type	Description	Default
coordinates	Array	Space-time coordinates	required
deterministic	bool	Whether to use deterministic mode	True

Returns:

Type	Description
Array	Physics loss value

compute_total_loss

compute_total_loss(coordinates: Array, target_solution: Array | None = None, *, deterministic: bool = True) -> dict[str, Array]

Compute total loss combining data and physics components.

Parameters:

Name	Type	Description	Default
coordinates	Array	Space-time coordinates	required
target_solution	Array | None	Target solution (optional, for supervised learning)	None
deterministic	bool	Whether to use deterministic mode	True

Returns:

Type	Description
dict[str, Array]	Dictionary containing individual loss components and total loss

GeometryAttention

GeometryAttention(feature_dim: int, geometry_dim: int, num_heads: int = 8, use_distance_attention: bool = True, *, rngs: Rngs)

Bases: Module

Geometry-aware attention mechanism.

Computes attention weights based on both feature similarity and geometric relationships with proper dimension handling.

Parameters:

Name	Type	Description	Default
feature_dim	int	Dimension of feature vectors	required
geometry_dim	int	Dimension of geometry embeddings	required
num_heads	int	Number of attention heads	8
use_distance_attention	bool	Whether to include distance-based attention	True
rngs	Rngs	Random number generator state	required

GeometryEncoder

GeometryEncoder(coord_dim: int, hidden_dim: int, output_dim: int, use_positional_encoding: bool = True, num_frequencies: int = 8, max_position: float = 10000.0, embedding_type: str = 'transformer', *, rngs: Rngs)

Bases: Module

Encoder mapping point coordinates to geometry embeddings.

Coordinates are first lifted into a sinusoidal positional embedding and then passed through an MLP to produce per-point geometry features. This is the geometry-embedding component of the Geometry-Informed Neural Operator.

References

Li, Z. et al. (2023). "Geometry-Informed Neural Operator for Large-Scale 3D PDEs." NeurIPS 2023, arXiv:2309.00583. The transformer-style sinusoidal coordinate embedding mirrors the input/output GNO positional embedding in neuraloperator/neuralop/layers/gno_block.py (self.pos_embedding, L154-L160 / L239-L241) built from neuralop/layers/embeddings.py::SinusoidalEmbedding.

Parameters:

Name	Type	Description	Default
coord_dim	int	Dimension of input coordinates.	required
hidden_dim	int	Hidden layer dimension of the encoding MLP.	required
output_dim	int	Output geometry-embedding dimension.	required
use_positional_encoding	bool	Whether to lift coordinates with a sinusoidal positional embedding before the MLP.	True
num_frequencies	int	Number of sinusoidal frequencies per coordinate.	8
max_position	float	max_positions for transformer-style embedding.	10000.0
embedding_type	str	Sinusoidal embedding style, "transformer" or "nerf".	'transformer'
rngs	Rngs	Random number generator state.	required

GeometryInformedNeuralOperator

GeometryInformedNeuralOperator(in_channels: int, out_channels: int, hidden_channels: int = 64, modes: Sequence[int] = (16, 16), num_layers: int = 4, geometry_dim: int = 32, coord_dim: int = 2, use_geometry_attention: bool = True, use_spectral_conv: bool = True, *, rngs: Rngs)

Bases: Module

Complete Geometry-Informed Neural Operator.

Advanced neural operator that incorporates geometric information throughout the network for improved performance on spatially complex problems.

Parameters:

Name	Type	Description	Default
in_channels	int	Number of input channels	required
out_channels	int	Number of output channels	required
hidden_channels	int	Hidden channel dimension	64
modes	Sequence[int]	Fourier modes for spectral convolution	(16, 16)
num_layers	int	Number of GINO blocks	4
geometry_dim	int	Dimension of geometry embeddings	32
coord_dim	int	Coordinate dimension	2
use_geometry_attention	bool	Whether to use geometry attention	True
use_spectral_conv	bool	Whether to use spectral convolution	True
rngs	Rngs	Random number generator state	required

GINOBlock

GINOBlock(in_channels: int, out_channels: int, modes: Sequence[int], geometry_dim: int, coord_dim: int = 2, use_geometry_attention: bool = True, use_spectral_conv: bool = True, activation: Callable[[Array], Array] = gelu, *, rngs: Rngs)

Bases: Module

Single GINO block with spectral convolution and geometry attention.

Combines spectral convolutions with geometry-aware processing for enhanced spatial understanding.

Parameters:

Name	Type	Description	Default
in_channels	int	Number of input channels	required
out_channels	int	Number of output channels	required
modes	Sequence[int]	Fourier modes for spectral convolution	required
geometry_dim	int	Dimension of geometry embeddings	required
coord_dim	int	Dimension of coordinates	2
use_geometry_attention	bool	Whether to use geometry attention	True
use_spectral_conv	bool	Whether to use spectral convolution	True
activation	Callable[[Array], Array]	Activation function	gelu
rngs	Rngs	Random number generator state	required

LatentNeuralOperator

LatentNeuralOperator(in_channels: int, out_channels: int, latent_dim: int, num_latent_tokens: int, *, num_attention_heads: int = 8, num_encoder_layers: int = 4, num_decoder_layers: int = 4, physics_constraints: list[str] | None = None, dropout_rate: float = 0.0, activation: Callable[[Array], Array] = gelu, rngs: Rngs)

Bases: Module

Latent Neural Operator with attention-based latent representations.

This operator learns compact latent representations of function spaces using attention mechanisms, enabling efficient learning of complex operator mappings with reduced computational overhead.

Features: - Learnable latent space for function representation - Multi-head attention for function-to-latent and latent-to-function mappings - Physics-aware attention constraints - Efficient inference through latent space operations

Parameters:

Name	Type	Description	Default
in_channels	int	Number of input channels	required
out_channels	int	Number of output channels	required
latent_dim	int	Dimension of latent space	required
num_latent_tokens	int	Number of latent tokens	required
num_attention_heads	int	Number of attention heads	8
num_encoder_layers	int	Number of encoder layers	4
num_decoder_layers	int	Number of decoder layers	4
physics_constraints	list[str] | None	List of physics constraints	None
dropout_rate	float	Dropout rate	0.0
activation	Callable[[Array], Array]	Activation function	gelu
rngs	Rngs	Random number generators	required

MGNOLayer

MGNOLayer(channels: int, max_multipole_order: int = 4, use_local_messages: bool = True, dropout_rate: float = 0.1, *, rngs: Rngs)

Bases: Module

MGNO layer with numerical stability and robust message passing.

Combines multipole expansion with local graph neural network operations for handling both long-range and short-range interactions.

Parameters:

Name	Type	Description	Default
channels	int	Number of feature channels	required
max_multipole_order	int	Maximum multipole expansion order	4
use_local_messages	bool	Whether to use local message passing	True
dropout_rate	float	Dropout rate for regularization	0.1
rngs	Rngs	Random number generator state	required

MultipoleExpansion

MultipoleExpansion(channels: int, max_order: int = 4, epsilon: float = 1e-08, stabilization_factor: float = 0.1, *, rngs: Rngs)

Bases: Module

Numerically stable multipole expansion layer.

Computes multipole moments with proper numerical stability to prevent overflow and NaN generation in hierarchical computations.

Parameters:

Name	Type	Description	Default
channels	int	Number of feature channels	required
max_order	int	Maximum multipole order	4
epsilon	float	Small constant for numerical stability	1e-08
stabilization_factor	float	Factor for moment normalization	0.1
rngs	Rngs	Random number generator state	required

MultipoleGraphNeuralOperator

MultipoleGraphNeuralOperator(in_features: int, out_features: int, hidden_features: int = 64, num_layers: int = 3, max_degree: int = 4, use_local_messages: bool = True, dropout_rate: float = 0.1, *, rngs: Rngs)

Bases: Module

Complete Multipole Graph Neural Operator with numerical stability.

Neural operator for systems with long-range interactions such as molecular dynamics, N-body simulations, and plasma physics.

Parameters:

Name	Type	Description	Default
in_features	int	Number of input feature channels	required
out_features	int	Number of output feature channels	required
hidden_features	int	Hidden layer width	64
num_layers	int	Number of MGNO layers	3
max_degree	int	Maximum multipole expansion order	4
use_local_messages	bool	Whether to use local message passing	True
dropout_rate	float	Dropout rate for regularization	0.1
rngs	Rngs	Random number generator state	required

OperatorNetwork

OperatorNetwork(operator_type: str, config: dict[str, Any], *, rngs: Rngs)

Bases: Module

Unified interface for different operator network types.

This class provides a common interface for different neural operator architectures (FNO, DeepONet, etc.) to enable easy experimentation and comparison.

Parameters:

Name	Type	Description	Default
operator_type	str	Type of operator ('fno', 'deeponet', 'fourier_deeponet', 'adaptive_deeponet', etc.)	required
config	dict[str, Any]	Configuration dictionary for the operator	required
rngs	Rngs	Random number generators	required

BayesianLinear

BayesianLinear(in_features: int, out_features: int, prior_std: float = 1.0, deterministic: bool = False, *, rngs: Rngs)

Bases: Module

Variational diagonal-Gaussian dense layer.

Reference: Blundell et al. 2015, "Weight Uncertainty in Neural Networks" (arXiv:1505.05424). Per-parameter diagonal-Gaussian posterior with reparameterization-trick sampling and analytic KL against an isotropic Gaussian prior. ../bayesian-torch and ../blitz-bayesian-deep-learning serve as PyTorch reference implementations of the same variational layer family.

Weight and bias each carry a (mean, log-variance) posterior; sampling uses the reparameterization trick.

Mode handling follows the :class:nnx.Dropout convention: the module holds a self.deterministic flag that the NNX train() and inference-mode methods flip via set_attributes recursion. Sampling is enabled when the resolved mode is non-deterministic AND rngs is supplied. A per-call deterministic keyword overrides the module flag for one call site (mirrors nnx.Dropout.__call__).

deterministic defaults to False so the module ships in training (sampling) mode; switch the module to inference mode to disable sampling globally.

kl_divergence

kl_divergence() -> Array

Total KL divergence (weights + bias) under the layer's diagonal Gaussian prior.

BayesianSpectralConvolution

BayesianSpectralConvolution(in_channels: int, out_channels: int, modes: tuple[int, ...], prior_std: float = 1.0, deterministic: bool = False, *, rngs: Rngs)

Bases: Module

Variational Fourier-spectral convolution with complex Gaussian weights.

Implements the canonical Zongyi Li Fourier Neural Operator spectral block (Li et al. 2021, arXiv:2010.08895; reference implementation: ../deeponet-fno/src/darcy_rectangular_pwc/fourier_2d.py:SpectralConv2d) with a variational diagonal-Gaussian posterior over each complex weight.

The trainable Fourier weights split into real and imaginary parts; each part carries a diagonal-Gaussian posterior (mean, log-variance). Sampling uses the reparameterization trick and combines the parts into a complex weight tensor for the spectral convolution.

Mode handling. For jnp.fft.rfft-style transforms:

• 1D: only the real-FFT axis exists, low-frequency modes are [:modes[0]]. One weight tensor of shape (out, in, modes[0]).
• 2D: jnp.fft.rfftn(x, axes=(-2, -1)) is a full FFT on the H axis and a real FFT on the W axis. The H axis therefore carries BOTH positive [:modes[0]] and negative [-modes[0]:] low-frequency modes; the W axis carries only [:modes[1] // 2 + 1]. Following Li, TWO weight tensors of shape (out, in, modes[0], modes[1] // 2 + 1) are used — one for each H-frequency band — so the spectral kernel captures the full low-frequency response rather than a single quadrant.

Output spatial shape matches the input; only in_channels becomes out_channels. Aleatoric / epistemic uncertainty extraction is the caller's responsibility — this layer returns only the convolved tensor.

deterministic follows the :class:nnx.Dropout convention; ships in non-deterministic (sampling) mode and is flipped by the NNX inference-mode toggle via set_attributes recursion.

kl_divergence

kl_divergence() -> Array

Sum diagonal-Gaussian KL across every (real/imag, pos/neg-H) weight posterior.

UncertaintyQuantificationNeuralOperator

UncertaintyQuantificationNeuralOperator(*, base: UQNOBaseSolutionOperator, residual: UQNOResidualOperator, calibrator: UQNOConformalCalibrator | None = None)

Bases: Module

Three-stage conformal UQNO orchestrator.

Holds a base solution operator, a residual quantile operator, and an optional fitted :class:UQNOConformalCalibrator. Use:

1. Train self.base to convergence on the regression task with any standard FNO training loop.
2. Train self.residual against :class:opifex.uncertainty.losses.PointwiseQuantileLoss on base(x) - y_true residuals (gradients through :meth:__call__ are stopped at the base via jax.lax.stop_gradient so residual-stage updates do not contaminate the base).
3. Call :meth:calibrate on a held-out calibration set to obtain a :class:UQNOConformalCalibrator; attach it via :meth:with_calibrator.
4. Call :meth:predict_with_bands at test time.

The class never claims native Bayesian or distributional support; the matching capability declaration is :class:opifex.uncertainty.adapters.operators.FNOConformalAdapterSpec.

predict_base

predict_base(x: Array) -> Array

Apply the base solution operator only.

predict_residual

predict_residual(x: Array) -> Array

Apply the residual quantile operator (non-negative).

calibrate

calibrate(x_calib: Array, y_calib: Array, *, alpha: float, delta: float, eps: float = 1e-12) -> UQNOConformalCalibrator

Derive a scalar uncertainty scaling factor on a calibration set.

Mirrors ../neuraloperator/scripts/train_uqno_darcy.py: for every calibration sample, compute per-grid ratios |y - base(x)| / (residual(x) + eps); take the domain_idx-th largest ratio per function (per-batch); then the function_idx-th largest of those across the batch is the scalar scaling factor.

Parameters:

Name	Type	Description	Default
x_calib	Array	Calibration inputs, shape (n_samples, ...).	required
y_calib	Array	Calibration targets, same shape as the base model output.	required
alpha	float	Target pointwise miscoverage in (0, 1).	required
delta	float	Target function-level miscoverage in (0, 1).	required
eps	float	Floor added to residual(x) before division to avoid divide-by-zero on near-zero predicted widths.	1e-12

with_calibrator

with_calibrator(calibrator: UQNOConformalCalibrator) -> UncertaintyQuantificationNeuralOperator

Attach calibrator to this operator and return self.

NNX modules support in-place mutation; with_* is the fluent-attach name (matches the canonical neuraloperator uqno_data_proc.set_scale_factor pattern in spirit).

predict_with_bands

predict_with_bands(x: Array) -> PredictiveDistribution

Return PredictiveDistribution with bands base ± E * scaling_factor.

Requires a fitted :class:UQNOConformalCalibrator (attach via :meth:with_calibrator or by passing calibrator= at construction). The metadata records ("method", "conformal"), ("alpha", alpha), ("delta", delta); epistemic and samples stay None (conformal is not Bayesian).

WaveletNeuralOperator

WaveletNeuralOperator(in_channels: int, out_channels: int, hidden_channels: int, num_levels: int, *, wavelet_type: str = 'db4', mode: str = 'symmetric', activation: Callable[[Array], Array] = gelu, use_learnable_wavelets: bool = False, rngs: Rngs)

Bases: Module

Wavelet Neural Operator for multi-scale wavelet-based learning.

This operator uses wavelet transforms to capture multi-scale features in the input functions, enabling efficient learning of operators with multi-scale characteristics like turbulence and material heterogeneity.

Features: - Discrete Wavelet Transform (DWT) for multi-scale decomposition - Learnable wavelet coefficients processing - Multi-resolution reconstruction - Adaptive wavelet basis selection

Parameters:

Name	Type	Description	Default
in_channels	int	Number of input channels	required
out_channels	int	Number of output channels	required
hidden_channels	int	Hidden channel dimension	required
num_levels	int	Number of wavelet decomposition levels	required
wavelet_type	str	Type of wavelet (e.g., 'db4', 'haar')	'db4'
mode	str	Boundary condition mode	'symmetric'
activation	Callable[[Array], Array]	Activation function	gelu
use_learnable_wavelets	bool	Whether to use learnable wavelet bases	False
rngs	Rngs	Random number generators	required

DeepONetConformalAdapterSpec dataclass

DeepONetConformalAdapterSpec(*, operator_family: str = _DEEPONET_FAMILY, default_strategy: DefaultStrategy = CONFORMAL, source_package: str = 'opifex', spatial_axes: tuple[int, ...] = _DEEPONET_SPATIAL_AXES, spectral_axes: tuple[int, ...] | None = None, supported_metrics: tuple[str, ...] = _CONFORMAL_METRICS, required_capabilities: tuple[str, ...] = _OPERATOR_REQUIRED_CAPABILITIES)

Bases: OperatorAdapterSpec

DeepONet + conformal-calibration adapter spec.

DeepONetDeepEnsembleAdapterSpec dataclass

DeepONetDeepEnsembleAdapterSpec(*, operator_family: str = _DEEPONET_FAMILY, default_strategy: DefaultStrategy = ENSEMBLE, source_package: str = 'opifex', spatial_axes: tuple[int, ...] = _DEEPONET_SPATIAL_AXES, spectral_axes: tuple[int, ...] | None = None, supported_metrics: tuple[str, ...] = _ENSEMBLE_METRICS, required_capabilities: tuple[str, ...] = _OPERATOR_REQUIRED_CAPABILITIES)

Bases: OperatorAdapterSpec

DeepONet + deep-ensemble adapter spec.

DeepONetMCDropoutAdapterSpec dataclass

DeepONetMCDropoutAdapterSpec(*, operator_family: str = _DEEPONET_FAMILY, default_strategy: DefaultStrategy = MC_DROPOUT, source_package: str = 'opifex', spatial_axes: tuple[int, ...] = _DEEPONET_SPATIAL_AXES, spectral_axes: tuple[int, ...] | None = None, supported_metrics: tuple[str, ...] = _MCDROPOUT_METRICS, required_capabilities: tuple[str, ...] = _OPERATOR_REQUIRED_CAPABILITIES)

Bases: OperatorAdapterSpec

DeepONet + MC-dropout adapter spec (caller-owned rngs at predict-time).

FNOConformalAdapterSpec dataclass

FNOConformalAdapterSpec(*, operator_family: str = _FNO_FAMILY, default_strategy: DefaultStrategy = CONFORMAL, source_package: str = 'opifex', spatial_axes: tuple[int, ...] = _FNO_SPATIAL_AXES, spectral_axes: tuple[int, ...] | None = _FNO_SPECTRAL_AXES, supported_metrics: tuple[str, ...] = _CONFORMAL_METRICS, required_capabilities: tuple[str, ...] = _OPERATOR_REQUIRED_CAPABILITIES)

Bases: OperatorAdapterSpec

FNO + conformal-calibration adapter spec (pre-UQNO-rewrite path).

FNODeepEnsembleAdapterSpec dataclass

FNODeepEnsembleAdapterSpec(*, operator_family: str = _FNO_FAMILY, default_strategy: DefaultStrategy = ENSEMBLE, source_package: str = 'opifex', spatial_axes: tuple[int, ...] = _FNO_SPATIAL_AXES, spectral_axes: tuple[int, ...] | None = _FNO_SPECTRAL_AXES, supported_metrics: tuple[str, ...] = _ENSEMBLE_METRICS, required_capabilities: tuple[str, ...] = _OPERATOR_REQUIRED_CAPABILITIES)

Bases: OperatorAdapterSpec

FNO + deep-ensemble adapter spec (member tuple via DeepEnsembleState).

FNOMCDropoutAdapterSpec dataclass

FNOMCDropoutAdapterSpec(*, operator_family: str = _FNO_FAMILY, default_strategy: DefaultStrategy = MC_DROPOUT, source_package: str = 'opifex', spatial_axes: tuple[int, ...] = _FNO_SPATIAL_AXES, spectral_axes: tuple[int, ...] | None = _FNO_SPECTRAL_AXES, supported_metrics: tuple[str, ...] = _MCDROPOUT_METRICS, required_capabilities: tuple[str, ...] = _OPERATOR_REQUIRED_CAPABILITIES)

Bases: OperatorAdapterSpec

FNO + MC-dropout adapter spec (caller-owned rngs at predict-time).

OperatorAdapterSpec dataclass

OperatorAdapterSpec(*, operator_family: str, default_strategy: DefaultStrategy, source_package: str = 'opifex', spatial_axes: tuple[int, ...] = (), spectral_axes: tuple[int, ...] | None = None, supported_metrics: tuple[str, ...] = (), required_capabilities: tuple[str, ...] = ())

Base class for operator-family UQ adapter specs.

Fields:

• operator_family — lowercase family name ("fno", "deeponet").
• default_strategy — :class:DefaultStrategy enum value advertising which adapter the spec configures (CONFORMAL, ENSEMBLE, MC_DROPOUT).
• source_package — owning package name (always "opifex" here).
• spatial_axes — tuple of input/output axes treated as spatial by the operator (used by function-space metrics and the calibrator).
• spectral_axes — subset of spatial_axes that participate in a Fourier-spectral kernel; None for non-spectral operators (e.g. DeepONet).
• supported_metrics — tuple of metric names the eventual calibrator can compute against this spec ("l2", "h1", "spatial_coverage", "spectral_coverage").
• required_capabilities — capability tags the operator must satisfy before the spec is wired (e.g. ("native_nnx_module",)).

recommended_capability

recommended_capability() -> UQCapability

Return an honest :class:UQCapability for this spec.

native_bayesian is always False — adapter-mediated UQ on a deterministic operator is not native Bayesian. The matching strategy capability flag (supports_conformal / supports_ensemble) is set to True; supports_function_space is always True for operator adapters.

function_space_metadata

function_space_metadata() -> MetadataItems

Return the operator function-space provenance as metadata pairs.

Records operator_family, spatial_axes, the supported metric tuple, and — for spectral operators only — spectral_axes. This provenance is merged into every wrapped predictive distribution so downstream consumers can identify the output topology (and which function-space metrics are admissible) without re-deriving it from the spec.

wrap

wrap(model: Any, capability: UQCapability) -> Any

Wire the spec to its concrete adapter, dispatching on default_strategy.

ENSEMBLE packages model (the operator-member tuple) as a :class:DeepEnsembleState and delegates to :class:DeepEnsembleAdapter; MC_DROPOUT delegates an :class:MCDropoutState to :class:MCDropoutAdapter. The wrapped object's predict_distribution output is enriched with this spec's :meth:function_space_metadata. CONFORMAL raises an actionable redirect to the dedicated conformal calibrators (the conformal contract takes calibration data, not a model). A capability falsely claiming native_bayesian=True is rejected.

create_high_frequency_amfno

create_high_frequency_amfno(in_channels: int, out_channels: int, modes: Sequence[int] = (128, 128), **kwargs) -> AmortizedFourierNeuralOperator

Create AM-FNO optimized for high-frequency problems.

create_shock_amfno

create_shock_amfno(in_channels: int = 3, out_channels: int = 3, modes: Sequence[int] = (96, 96), **kwargs) -> AmortizedFourierNeuralOperator

Create AM-FNO for problems with shocks/discontinuities.

create_wave_amfno

create_wave_amfno(in_channels: int = 2, out_channels: int = 2, modes: Sequence[int] = (64, 64), **kwargs) -> AmortizedFourierNeuralOperator

Create AM-FNO for wave propagation problems.

create_multiphysics_local_fno

create_multiphysics_local_fno(in_channels: int = 5, out_channels: int = 5, modes: Sequence[int] = (24, 24), **kwargs) -> LocalFourierNeuralOperator

Create Local FNO for multi-physics problems.

create_turbulence_local_fno

create_turbulence_local_fno(in_channels: int = 3, out_channels: int = 3, modes: Sequence[int] = (32, 32), **kwargs) -> LocalFourierNeuralOperator

Create Local FNO optimized for turbulent flow modeling.

create_wave_local_fno

create_wave_local_fno(in_channels: int = 2, out_channels: int = 2, modes: Sequence[int] = (64, 64), **kwargs) -> LocalFourierNeuralOperator

Create Local FNO for wave propagation with scattering.

create_climate_sfno

create_climate_sfno(in_channels: int = 5, out_channels: int = 5, lmax: int = 32, **kwargs) -> SphericalFourierNeuralOperator

Create SFNO optimized for global climate modeling.

create_ocean_sfno

create_ocean_sfno(in_channels: int = 4, out_channels: int = 4, lmax: int = 48, **kwargs) -> SphericalFourierNeuralOperator

Create SFNO for global ocean circulation modeling.

create_planetary_sfno

create_planetary_sfno(in_channels: int = 3, out_channels: int = 3, lmax: int = 16, **kwargs) -> SphericalFourierNeuralOperator

Create SFNO for planetary-scale phenomena.

create_weather_sfno

create_weather_sfno(in_channels: int = 7, out_channels: int = 7, lmax: int = 64, **kwargs) -> SphericalFourierNeuralOperator

Create SFNO for high-resolution weather prediction.

create_cp_fno

create_cp_fno(in_channels: int, out_channels: int, hidden_channels: int = 64, modes: Sequence[int] = (16, 16), rank: float = 0.1, num_layers: int = 4, *, rngs: Rngs) -> TensorizedFourierNeuralOperator

Create CP factorized FNO.

create_tt_fno

create_tt_fno(in_channels: int, out_channels: int, hidden_channels: int = 64, modes: Sequence[int] = (16, 16), rank: float = 0.1, num_layers: int = 4, *, rngs: Rngs) -> TensorizedFourierNeuralOperator

Create Tensor Train factorized FNO.

create_tucker_fno

create_tucker_fno(in_channels: int, out_channels: int, hidden_channels: int = 64, modes: Sequence[int] = (16, 16), rank: float = 0.1, num_layers: int = 4, *, rngs: Rngs) -> TensorizedFourierNeuralOperator

Create Tucker factorized FNO.

create_deep_ufno

create_deep_ufno(in_channels: int, out_channels: int, hidden_channels: int = 32, modes: Sequence[int] = (32, 32), **kwargs) -> UFourierNeuralOperator

Create deep U-FNO (5 levels) for complex multi-scale problems.

create_shallow_ufno

create_shallow_ufno(in_channels: int, out_channels: int, hidden_channels: int = 64, modes: Sequence[int] = (16, 16), **kwargs) -> UFourierNeuralOperator

Create shallow U-FNO (2 levels) for simple multi-scale problems.

create_turbulence_ufno

create_turbulence_ufno(in_channels: int = 4, out_channels: int = 3, **kwargs) -> UFourierNeuralOperator

Create U-FNO optimized for turbulent flow modeling.

create_3d_gino

create_3d_gino(in_channels: int, out_channels: int, *, rngs: Rngs) -> GeometryInformedNeuralOperator

Create GINO optimized for 3D problems.

create_adaptive_mesh_gino

create_adaptive_mesh_gino(in_channels: int, out_channels: int, *, rngs: Rngs) -> GeometryInformedNeuralOperator

Create GINO for adaptive mesh refinement.

create_cad_gino

create_cad_gino(in_channels: int, out_channels: int, *, rngs: Rngs) -> GeometryInformedNeuralOperator

Create GINO optimized for CAD geometries.

create_multiscale_gino

create_multiscale_gino(in_channels: int, out_channels: int, *, rngs: Rngs) -> GeometryInformedNeuralOperator

Create GINO for multiscale problems.

create_molecular_mgno

create_molecular_mgno(in_features: int, out_features: int, *, rngs: Rngs) -> MultipoleGraphNeuralOperator

Create MGNO optimized for molecular dynamics simulations.

create_nbody_mgno

create_nbody_mgno(in_features: int, out_features: int, *, rngs: Rngs) -> MultipoleGraphNeuralOperator

Create MGNO for N-body gravitational simulations.

create_plasma_mgno

create_plasma_mgno(in_features: int, out_features: int, *, rngs: Rngs) -> MultipoleGraphNeuralOperator

Create MGNO for plasma physics simulations.

get_operator_capability

get_operator_capability(operator_type: str) -> UQCapability

Return the :class:UQCapability for operator_type.

Raises:

Type	Description
KeyError	If operator_type isn't in :data:OPERATOR_CAPABILITY_REGISTRY.

create_operator

create_operator(operator_type: str, **kwargs: Any) -> Any

Factory function to create any operator by name.

Parameters:

Name	Type	Description	Default
operator_type	str	Type of operator to create	required
**kwargs	Any	Arguments for operator initialization	{}

Returns:

Type	Description
Any	Initialized operator instance

Raises:

Type	Description
ValueError	If operator_type is not recognized

Example

Create a Tensorized FNO

tfno = create_operator("TFNO", ... in_channels=3, out_channels=1, ... hidden_channels=64, modes=(16, 16), ... factorization="tucker", rank=0.1, ... rngs=rngs)

recommend_operator

recommend_operator(application: str) -> dict[str, Any]

Recommend the best operator for a specific application.

Parameters:

Name	Type	Description	Default
application	str	Application domain	required

Returns:

Type	Description
dict[str, Any]	Dictionary with recommendations

Example

rec = recommend_operator("turbulent_flow") print(f"Recommended: {rec['primary']}") print(f"Reason: {rec['reason']}")

list_operators

list_operators(category: str | None = None) -> dict[str, Sequence[str]]

List available operators by category.

Parameters:

Name	Type	Description	Default
category	str | None	Optional category filter	None

Returns:

Type	Description
dict[str, Sequence[str]]	Dictionary of operators by category

get_operator_info

get_operator_info(operator_type: str) -> dict[str, Any]

Get detailed information about a specific operator.

Parameters:

Name	Type	Description	Default
operator_type	str	Type of operator	required

Returns:

Type	Description
dict[str, Any]	Dictionary with operator information

Bayesian Networks

opifex.neural.bayesian

Bayesian neural network components with uncertainty quantification.

CalibrationTools

CalibrationTools(*, rngs: Rngs)

Bases: Module

Enhanced tools for uncertainty calibration assessment and improvement.

Parameters:

Name	Type	Description	Default
rngs	Rngs	Random number generators	required

assess_calibration

assess_calibration(predictions: Array, uncertainties: Array, true_values: Array, num_bins: int = 10) -> dict[str, float | dict[str, Array]]

Assess calibration quality of uncertainty estimates.

Parameters:

Name	Type	Description	Default
predictions	Array	Model predictions	required
uncertainties	Array	Predicted uncertainties	required
true_values	Array	Ground truth values	required
num_bins	int	Number of bins for reliability diagram	10

Returns:

Type	Description
dict[str, float | dict[str, Array]]	Dictionary with calibration metrics

compute_reliability_diagram

compute_reliability_diagram(confidences: Array, accuracies: Array, num_bins: int = 10) -> dict[str, Array]

Compute reliability diagram data.

Parameters:

Name	Type	Description	Default
confidences	Array	Predicted confidence values	required
accuracies	Array	Binary accuracy indicators	required
num_bins	int	Number of bins for the diagram	10

Returns:

Type	Description
dict[str, Array]	Dictionary with binned confidence and accuracy data

platt_scaling

platt_scaling(logits: Array, labels: Array, validation_logits: Array) -> tuple[float, float]

Fit Platt scaling and return its (slope, intercept) parameters.

Delegates to :class:PlattScaling (the single source of truth for Platt calibration) so there is exactly one fitting implementation. The returned slope / intercept are the fitted sigmoid parameters a / b from P(y=1|f) = sigmoid(a * f + b).

Parameters:

Name	Type	Description	Default
logits	Array	Training logits for fitting scaling parameters.	required
labels	Array	Training labels.	required
validation_logits	Array	Validation logits (accepted for API compatibility; the fitted parameters are independent of them).	required

Returns:

Type	Description
tuple[float, float]	Tuple of (slope, intercept) scaling parameters.

isotonic_regression_calibration

isotonic_regression_calibration(confidences: Array, accuracies: Array) -> Array

Fit isotonic regression and return calibrated confidences.

Delegates to :class:IsotonicRegression (the single source of truth, which uses a convergent pool-adjacent-violators fit) so there is exactly one isotonic implementation.

Parameters:

Name	Type	Description	Default
confidences	Array	Predicted confidence values.	required
accuracies	Array	Binary accuracy indicators.	required

Returns:

Type	Description
Array	Calibrated confidence values, aligned with confidences.

IsotonicRegression

IsotonicRegression(n_bins: int = 100, *, rngs: Rngs)

Bases: Module

Isotonic regression for calibration.

Non-parametric calibration method that learns a monotonic mapping from confidence scores to calibrated probabilities.

Parameters:

Name	Type	Description	Default
n_bins	int	Number of bins for isotonic regression	100
rngs	Rngs	Random number generators	required

fit

fit(confidences: Array, labels: Array) -> None

Fit isotonic regression using pool adjacent violators algorithm.

Parameters:

Name	Type	Description	Default
confidences	Array	Training confidence scores	required
labels	Array	Binary labels (0 or 1)	required

PlattScaling

PlattScaling(*, rngs: Rngs)

Bases: Module

Platt scaling for probabilistic calibration.

Applies a sigmoid function to logits to improve calibration of binary classification problems.

Parameters:

Name	Type	Description	Default
rngs	Rngs	Random number generators	required

fit

fit(logits: Array, labels: Array, max_iterations: int = 100) -> None

Fit Platt scaling parameters using maximum likelihood.

Parameters:

Name	Type	Description	Default
logits	Array	Training logits	required
labels	Array	Binary labels (0 or 1)	required
max_iterations	int	Maximum number of optimization iterations	100

TemperatureScaling

TemperatureScaling(physics_constraints: Sequence[str] = (), adaptive: bool = False, learning_rate: float = 0.01, constraint_strength: float = 1.0, *, rngs: Rngs)

Bases: Module

Temperature scaling for uncertainty calibration.

Applies learnable temperature scaling to improve calibration of probabilistic predictions while respecting physics constraints.

Parameters:

Name	Type	Description	Default
physics_constraints	Sequence[str]	List of physics constraints to enforce	()
adaptive	bool	Whether to use adaptive temperature learning	False
learning_rate	float	Learning rate for temperature optimization	0.01
constraint_strength	float	Strength of physics constraint enforcement	1.0
rngs	Rngs	Random number generators	required

apply_physics_aware_calibration

apply_physics_aware_calibration(predictions: Array, inputs: Array) -> tuple[Array, float]

Apply physics-aware temperature scaling with constraint enforcement.

Parameters:

Name	Type	Description	Default
predictions	Array	Model predictions to calibrate	required
inputs	Array	Input data for constraint evaluation	required

Returns:

Type	Description
tuple[Array, float]	Tuple of (calibrated_predictions, physics_constraint_penalty)

optimize_temperature

optimize_temperature(logits: Array, labels: Array) -> float

Optimize temperature parameter for calibration.

Parameters:

Name	Type	Description	Default
logits	Array	Model logits for validation data	required
labels	Array	True labels for validation data	required

Returns:

Type	Description
float	Optimized temperature value

optimize_temperature_with_physics_constraints

optimize_temperature_with_physics_constraints(predictions: Array, targets: Array, inputs: Array) -> float

Optimize temperature parameter with physics constraint awareness.

Parameters:

Name	Type	Description	Default
predictions	Array	Model predictions	required
targets	Array	Target values	required
inputs	Array	Input data for constraint evaluation	required

Returns:

Type	Description
float	Optimized temperature value

adaptive_temperature_scaling

adaptive_temperature_scaling(predictions: Array, uncertainties: Array, true_values: Array) -> Array

Apply adaptive temperature scaling based on uncertainty quality.

Parameters:

Name	Type	Description	Default
predictions	Array	Model predictions	required
uncertainties	Array	Predicted uncertainties	required
true_values	Array	Ground truth values	required

Returns:

Type	Description
Array	Adaptively calibrated temperatures

AmortizedVariationalFramework

AmortizedVariationalFramework(base_model: Module, prior_config: PriorConfig, variational_config: VariationalConfig, *, rngs: Rngs)

Bases: Module

Variational framework with amortized uncertainty estimation.

This framework combines a base neural network model with variational Bayesian inference capabilities, enabling uncertainty quantification through amortized variational inference.

Parameters:

Name	Type	Description	Default
base_model	Module	Base neural network model to augment with uncertainty.	required
prior_config	PriorConfig	Configuration for physics-informed priors.	required
variational_config	VariationalConfig	Configuration for variational inference.	required
rngs	Rngs	Random number generator state.	required

predict_with_uncertainty

predict_with_uncertainty(x: Float[Array, 'batch input_dim'], num_samples: int | None = None, *, rngs: Rngs) -> tuple[Float[Array, 'batch output_dim'], Float[Array, 'batch output_dim']]

Forward pass with uncertainty quantification.

Parameters:

Name	Type	Description	Default
x	Float[Array, 'batch input_dim']	Input tensor of shape (batch_size, input_dim).	required
num_samples	int | None	Number of Monte Carlo samples for uncertainty estimation.	None
rngs	Rngs	Random number generator state.	required

Returns:

Type	Description
tuple[Float[Array, 'batch output_dim'], Float[Array, 'batch output_dim']]	Tuple of (mean_prediction, uncertainty) both of shape (batch_size, output_dim).

compute_elbo

compute_elbo(x: Float[Array, 'batch input_dim'], y: Float[Array, 'batch output_dim'], num_samples: int | None = None, *, rngs: Rngs) -> Float[Array, '']

Compute Evidence Lower BOund (ELBO).

Parameters:

Name	Type	Description	Default
x	Float[Array, 'batch input_dim']	Input tensor of shape (batch_size, input_dim).	required
y	Float[Array, 'batch output_dim']	Target tensor of shape (batch_size, output_dim).	required
num_samples	int | None	Number of Monte Carlo samples for ELBO estimation.	None
rngs	Rngs	Random number generator state.	required

Returns:

Type	Description
Float[Array, '']	ELBO scalar value (higher is better).

sample_predictive_distribution

sample_predictive_distribution(x: Float[Array, 'batch input_dim'], num_samples: int | None = None, *, rngs: Rngs) -> Float[Array, 'samples batch output_dim']

Sample from predictive distribution.

Parameters:

Name	Type	Description	Default
x	Float[Array, 'batch input_dim']	Input tensor of shape (batch_size, input_dim).	required
num_samples	int | None	Number of predictive samples to generate.	None
rngs	Rngs	Random number generator state.	required

Returns:

Type	Description
Float[Array, 'samples batch output_dim']	Predictive samples of shape (num_samples, batch_size, output_dim).

MeanFieldGaussian

MeanFieldGaussian(num_params: int, *, rngs: Rngs, prior_mean: float = 0.0, prior_std: float = 1.0, observation_noise: float = 0.1)

Bases: Module

Mean-field Gaussian variational posterior over a weight vector.

The factorized posterior q(w) = N(mu, diag(sigma^2)) over a weight vector w in R^num_params is the variational object injected into a base network by :class:AmortizedVariationalFramework.

On its own it is also a complete Bayesian linear model (Bishop, PRML 3.3): for an input x in R^num_params the prediction f(x) = w . x has the closed-form predictive f(x) ~ N(mu . x, sum_i x_i^2 sigma_i^2) -- Gaussian because the map is linear and q(w) is Gaussian. A homoscedastic Gaussian observation noise y ~ N(f(x), sigma_y^2) (the learnable log_observation_std) completes the likelihood, so the layer exposes the platform UQ protocol surfaces (:meth:predict_distribution, :meth:loss_components, :meth:negative_elbo, :meth:kl_divergence) directly, with the expected NLL available in closed form (no sampling).

Parameters:

Name	Type	Description	Default
num_params	int	Number of weights w the posterior factorizes over.	required
rngs	Rngs	Random number generator state.	required
prior_mean	float	Mean of the factorized Gaussian prior p(w) used by :meth:kl_divergence.	0.0
prior_std	float	Standard deviation of the prior p(w); must be positive.	1.0
observation_noise	float	Initial homoscedastic observation-noise standard deviation sigma_y of the Gaussian likelihood; learnable via log_observation_std. Must be positive.	0.1

sample

sample(num_samples: int, *, rngs: Rngs) -> Float[Array, 'samples params']

Sample from variational posterior.

Parameters:

Name	Type	Description	Default
num_samples	int	Number of samples to draw.	required
rngs	Rngs	Random number generator state.	required

Returns:

Type	Description
Float[Array, 'samples params']	Array of shape (num_samples, num_params) containing parameter samples.

log_prob

log_prob(samples: Float[Array, 'samples params']) -> Float[Array, 'samples']

Compute log probability of samples.

Parameters:

Name	Type	Description	Default
samples	Float[Array, 'samples params']	Parameter samples of shape (num_samples, num_params).	required

Returns:

Type	Description
Float[Array, 'samples']	Log probabilities for each sample of shape (num_samples,).

kl_divergence

kl_divergence(prior_mean: float | None = None, prior_std: float | None = None) -> Float[Array, '']

Compute KL(q(w) || p(w)) from the factorized Gaussian prior.

Parameters:

Name	Type	Description	Default
prior_mean	float | None	Prior mean; defaults to the value supplied at construction (prior_mean=0.0 unless overridden).	None
prior_std	float | None	Prior standard deviation; defaults to the value supplied at construction (prior_std=1.0 unless overridden).	None

Returns:

Type	Description
Float[Array, '']	KL divergence scalar value.

predict_distribution

predict_distribution(x: Float[Array, 'batch params'], *, rngs: Rngs | None = None) -> PredictiveDistribution

Return the closed-form Bayesian-linear predictive for inputs x.

The predictive f(x) ~ N(mu . x, x^2 . sigma^2) plus the homoscedastic observation noise sigma_y^2 gives epistemic = x^2 . sigma^2, aleatoric = sigma_y^2, and total = epistemic + aleatoric -- all in closed form, so no Monte-Carlo rngs are needed.

Parameters:

Name	Type	Description	Default
x	Float[Array, 'batch params']	Inputs of shape (batch, num_params) (or (num_params,)).	required
rngs	Rngs | None	Unused -- the predictive is exact; accepted for protocol conformance with stochastic models.	None

Returns:

Name	Type	Description
A	PredictiveDistribution	class:PredictiveDistribution with mean, variance, and the
	PredictiveDistribution	epistemic / aleatoric / total decomposition.

loss_components

loss_components(batch: Mapping[str, Any], *, config: ObjectiveConfig, rngs: Rngs | None = None) -> UQLossComponents

Return the per-batch negative-ELBO decomposition.

The expected negative log-likelihood under q(w) is available in closed form for the Gaussian likelihood::

E_q[NLL] = 0.5 log(2 pi sigma_y^2)
           + (mean((y - mu.x)^2) + mean(Var_q[f])) / (2 sigma_y^2)

and is combined with KL(q || p) by :meth:UQLossComponents.from_components using the weights / dataset scaling in config.

Parameters:

Name	Type	Description	Default
batch	Mapping[str, Any]	Mapping with required fields x ((batch, num_params)) and y ((batch,)).	required
config	ObjectiveConfig	Loss weights and dataset metadata.	required
rngs	Rngs | None	Unused -- the expected NLL is exact; accepted for protocol conformance.	None

Returns:

Type	Description
UQLossComponents	The optimizer-facing :class:UQLossComponents decomposition.

negative_elbo

negative_elbo(batch: Mapping[str, Any], *, config: ObjectiveConfig, rngs: Rngs | None = None) -> Float[Array, '']

Return the scalar negative-ELBO objective for one batch.

Parameters:

Name	Type	Description	Default
batch	Mapping[str, Any]	Mapping with required fields x and y.	required
config	ObjectiveConfig	Loss weights and dataset metadata.	required
rngs	Rngs | None	Forwarded to :meth:loss_components (unused there).	None

Returns:

Type	Description
Float[Array, '']	The scalar total of :meth:loss_components -- the value passed
Float[Array, '']	to jax.value_and_grad / optimizer.update.

PriorConfig dataclass

PriorConfig(conservation_laws: Sequence[str] = (), boundary_conditions: Sequence[str] = (), physics_constraints: Sequence[str] = (), prior_scale: float = 1.0)

Configuration for physics-informed priors.

Attributes:

Name	Type	Description
conservation_laws	Sequence[str]	List of conservation laws to enforce (e.g., ['energy', 'momentum']).
boundary_conditions	Sequence[str]	List of boundary conditions to incorporate.
physics_constraints	Sequence[str]	List of physics constraints to respect.
prior_scale	float	Scale parameter for the prior distribution.

UncertaintyEncoder

UncertaintyEncoder(input_dim: int, hidden_dims: Sequence[int], output_dim: int, *, rngs: Rngs)

Bases: Module

Neural network for amortized uncertainty estimation.

This encoder network predicts the parameters of the variational posterior directly from input data, enabling amortized variational inference.

Parameters:

Name	Type	Description	Default
input_dim	int	Dimensionality of input features.	required
hidden_dims	Sequence[int]	Sequence of hidden layer dimensions.	required
output_dim	int	Dimensionality of output (typically 2 * num_params for mean and log_std).	required
rngs	Rngs	Random number generator state.	required

VariationalConfig dataclass

VariationalConfig(input_dim: int, hidden_dims: Sequence[int] = (64, 32), num_samples: int = 10, kl_weight: float = 1.0, temperature: float = 1.0)

Configuration for variational inference.

Attributes:

Name	Type	Description
input_dim	int	Dimensionality of input features.
hidden_dims	Sequence[int]	Tuple of hidden layer dimensions for the encoder.
num_samples	int	Number of samples to draw during inference.
kl_weight	float	Weight for the KL divergence term in ELBO.
temperature	float	Temperature parameter for variational distribution.

register_bayesian_capabilities

register_bayesian_capabilities(registry: UQRegistry) -> None

Register the Task 7.2 Bayesian model capabilities into registry.

Explicit registration — called from a composition root rather than at import time (Rule 13: no mutable side effects on import). The shared singleton :class:UQRegistry is populated with the ProbabilisticPINN and MultiFidelityPINN model declarations.

Idempotent: names already present are skipped, so repeated calls (and the re-entrancy of :func:bayesian_uq_registry) never trip CalibraX's duplicate-registration rejection.

Parameters:

Name	Type	Description	Default
registry	UQRegistry	Target :class:UQRegistry. Almost always the singleton instance, but any registry is accepted for test isolation.	required

bayesian_uq_registry

bayesian_uq_registry() -> UQRegistry

Return the shared singleton UQRegistry with bayesian models registered.

Lazy composition-root accessor: callers that need the registry already holding the Task 7.2 model capabilities use this instead of relying on an import-time side effect. Registration is idempotent, so this is safe to call repeatedly.

ProbabilisticPINN shared-objective surface

opifex.neural.bayesian.ProbabilisticPINN is an nnx.Module that implements the canonical VariationalModule protocol from opifex.uncertainty.protocols:

• kl_divergence() -> jax.Array — total KL across every Bayesian layer in the network.
• predict_distribution(x, *, rngs, mode) -> PredictiveDistribution — returns the canonical Phase-1 contract. mode is a PredictiveMode value (deterministic / single-sample / monte-carlo ensemble); unknown modes raise ValueError.
• loss_components(batch, *, rngs, objective) -> UQLossComponents — returns the data / KL / physics / boundary / initial-condition terms as the canonical pattern-B container.
• negative_elbo(batch, *, rngs, objective) -> UQLossComponents — UQLossComponents.from_components evaluated with sign flipped for optimisers that maximise ELBO.

All four methods take a traced rngs: nnx.Rngs argument; no module holds a hidden fallback RNG. The shared objective API replaces hand-rolled data_loss + kl_weight * kl assembly in the example notebooks.

RobustPINNOptimizer (uncertainty-guided training)

RobustPINNOptimizer.compute_loss_components(batch, *, rngs, objective) returns the same UQLossComponents pattern-B container as ProbabilisticPINN so a robust-PINN training loop can plug into the shared objective surface without diverging. uncertainty_guided_sampling(x_candidates, num_samples, *, rngs) selects the highest-uncertainty samples for the next training batch.

ComputationAwareSpectralConvolution (CASpec)

opifex.neural.operators.fno.bayesian.ComputationAwareSpectralConvolution is a sibling of BayesianSpectralConvolution whose uncertainty over the flattened spectral weights is maintained as a low-rank CAKF posterior — the implicit posterior_cov = prior_cov - factor @ factor^T representation of Pförtner+ 2024 (arXiv:2405.08971) and the CAGP precursor Wenger+ 2023 (arXiv:2306.07879). The constructor mirrors BayesianSpectralConvolution; __call__ runs the deterministic spectral conv using the BSC posterior-mean weights, and cakf_refine(observation=, observation_matrix=, observation_cov=, max_iter=) returns a _CAKFSpectralRefinement carrying the updated (cakf_mean, cakf_factor) pair (rank gained per call == max_iter). The same module also re-exports BayesianSpectralConvolution from its canonical home at opifex.uncertainty.layers.bayesian so callers can import either sibling from a single namespace.

gp_pinn_predictive_posterior (GP-PINN)

opifex.neural.pinns.gp_pinn.gp_pinn_predictive_posterior(*, pinn_forward, laplace_posterior, coordinates, gp_adapter_spec) returns a function-valued GP predictive over a trained PINN via the linearised-Laplace equivalence (Immer, Korzepa, Bauer 2021, AISTATS, arXiv:2008.08400 §3). The math is identical to LUNO (opifex.uncertainty.curvature.linearized_neural_operator_posterior) and is reused directly; what differs is the context: the input is a PINN forward consuming spatial / spatio-temporal coordinates, and the gp_adapter_spec parameter (a GP adapter spec such as TinygpAdapterSpec or GPJaxAdapterSpec) is recorded in the predictive metadata so consumers can resolve the linearised-Laplace ↔ GP correspondence. Concrete GP fit / predict is available through the opifex.uncertainty.gp subpackage.

ProbabilisticFourierNeuralOperator (PNO)

opifex.neural.operators.fno.probabilistic.ProbabilisticFourierNeuralOperator equips a standard FNO backbone with twin pointwise heads — a mean head and a log-variance head — producing a per-location heteroscedastic-Gaussian PredictiveDistribution (Kendall & Gal 2017, arXiv:1703.04977 §3.1; companion to the Magnani+ 2024 LUNO function-uncertainty thread, arXiv:2406.04317). The training objective is the elementwise heteroscedastic-Gaussian negative log-likelihood, exposed as probabilistic_fno_negative_log_likelihood(model, x, y); the predictive uncertainty is aleatoric by construction. Epistemic uncertainty is supplied orthogonally by wrapping a fitted PNO with the existing LaplaceAdapterSpec (opifex.uncertainty.curvature) or a deep-ensemble adapter (FNODeepEnsembleAdapterSpec). The log-variance head is clipped to [log_variance_floor, log_variance_ceiling] (defaults [-10, 10]) for numerical stability.

UncertaintyQuantificationNeuralOperator (UQNO)

The conformal neural operator under opifex.neural.operators.specialized.uqno is composed of three NNX modules:

• UQNOBaseSolutionOperator — the underlying FNO that produces point predictions of the PDE solution field.
• UQNOResidualOperator — a Bayesian residual-magnitude operator built on shared BayesianSpectralConvolution layers; predicts per-pixel calibrated uncertainty.
• UQNOConformalCalibrator — applies pointwise conformal calibration to the residual output so the resulting bands carry the requested empirical coverage.

The three-stage pipeline (predict_base → calibrate → predict_with_bands) is documented end-to-end in the examples/uncertainty/uqno_darcy example. UQNO exposes the conformal contract only — no Bayesian-objective surface (predict_distribution / loss_components / negative_elbo are intentionally absent on UQNO itself; those live on ProbabilisticPINN and the shared layers).

Domain Decomposition PINNs

Domain decomposition methods for physics-informed neural networks, enabling efficient training on complex geometries.

Base Classes

opifex.neural.pinns.domain_decomposition.base

Base classes for Domain Decomposition PINNs.

This module provides the foundational classes for domain decomposition approaches to physics-informed neural networks.

Key Classes

• Subdomain: Represents a subdomain region in the computational domain
• Interface: Represents the interface between adjacent subdomains
• DomainDecompositionPINN: Abstract base class for DD-PINN variants

Design Principles

• Each subdomain has its own neural network
• Interfaces enforce continuity and flux matching
• Window functions provide smooth blending (for FBPINN variants)

References

• Survey Section 8.3: Domain Decomposition Methods

Subdomain dataclass

Subdomain(*, id: int, bounds: Float[Array, 'dim 2'], overlap: float = 0.0)

Representation of a subdomain in the computational domain.

A subdomain is a rectangular region defined by its bounds in each spatial dimension.

Attributes:

Name	Type	Description
id	int	Unique identifier for this subdomain
bounds	Float[Array, 'dim 2']	Array of shape (dim, 2) with [min, max] for each dimension
overlap	float	Optional overlap with neighboring subdomains (for Schwarz methods)

center property

center: Float[Array, ' dim']

Compute the center of the subdomain.

volume property

volume: Float[Array, '']

Compute the volume (area in 2D, length in 1D) of the subdomain.

contains

contains(x: Float[Array, ' dim']) -> Array

Check if a point is inside this subdomain.

Parameters:

Name	Type	Description	Default
x	Float[Array, ' dim']	Point coordinates of shape (dim,)	required

Returns:

Type	Description
Array	Boolean array (scalar) indicating if point is inside subdomain

Interface dataclass

Interface(*, subdomain_ids: tuple[int, int], points: Float[Array, 'num_points dim'], normal: Float[Array, ' dim'])

Representation of an interface between two subdomains.

The interface stores sample points for enforcing continuity conditions between adjacent subdomains.

Attributes:

Name	Type	Description
subdomain_ids	tuple[int, int]	Tuple of (left_id, right_id) for adjacent subdomains
points	Float[Array, 'num_points dim']	Sample points on the interface, shape (num_points, dim)
normal	Float[Array, ' dim']	Outward normal vector from first subdomain, shape (dim,)

DomainDecompositionPINN

DomainDecompositionPINN(input_dim: int, output_dim: int, subdomains: Sequence[Subdomain], interfaces: Sequence[Interface], hidden_dims: Sequence[int], *, activation: Callable[[Array], Array] = tanh, rngs: Rngs)

Bases: Module

Base class for Domain Decomposition PINNs.

This class provides the infrastructure for training separate networks on subdomains with interface coupling conditions.

Attributes:

Name	Type	Description
input_dim		Input spatial dimension
output_dim		Output dimension (solution fields)
subdomains		List of subdomain definitions
interfaces		List of interface definitions
networks		List of subdomain networks

Parameters:

Name	Type	Description	Default
input_dim	int	Input spatial dimension	required
output_dim	int	Output dimension	required
subdomains	Sequence[Subdomain]	List of subdomain definitions	required
interfaces	Sequence[Interface]	List of interface definitions	required
hidden_dims	Sequence[int]	Hidden layer dimensions (shared across subdomains)	required
activation	Callable[[Array], Array]	Activation function	tanh
rngs	Rngs	Random number generators	required

get_subdomain_outputs

get_subdomain_outputs(x: Float[Array, ...]) -> list[Float[Array, 'batch out']]

Get outputs from all subdomain networks.

Parameters:

Name	Type	Description	Default
x	Float[Array, ...]	Input coordinates	required

Returns:

Type	Description
list[Float[Array, 'batch out']]	List of outputs from each subdomain network

compute_interface_residual

compute_interface_residual() -> Float[Array, '']

Compute interface continuity residual.

Enforces u_left = u_right at interface points.

Returns:

Type	Description
Float[Array, '']	Scalar interface residual (MSE of discontinuity)

compute_flux_residual

compute_flux_residual(derivative_fn: Callable[[Module, Float[Array, ...]], Float[Array, ...]]) -> Float[Array, '']

Compute interface flux continuity residual.

Enforces (du/dn)_left = (du/dn)_right at interface points.

Parameters:

Name	Type	Description	Default
derivative_fn	Callable[[Module, Float[Array, ...]], Float[Array, ...]]	Function to compute gradient of network output	required

Returns:

Type	Description
Float[Array, '']	Scalar flux residual

SubdomainNetwork

SubdomainNetwork(input_dim: int, output_dim: int, hidden_dims: Sequence[int], *, activation: Callable[[Array], Array] = tanh, rngs: Rngs)

Bases: Module

Neural network for a single subdomain.

A simple MLP that processes inputs for a specific subdomain.

Parameters:

Name	Type	Description	Default
input_dim	int	Input dimension	required
output_dim	int	Output dimension	required
hidden_dims	Sequence[int]	List of hidden layer dimensions	required
activation	Callable[[Array], Array]	Activation function	tanh
rngs	Rngs	Random number generators	required

uniform_partition

uniform_partition(bounds: Float[Array, 'dim 2'], num_partitions: tuple[int, ...], interface_points: int = 10) -> tuple[list[Subdomain], list[Interface]]

Create a uniform N-D partition of a rectangular (hyperrectangular) domain.

The domain is tiled into a tensor-product grid of axis-aligned subdomains, one per grid cell, with subdomain ids enumerated in row-major (C) order. Internal faces between axis-adjacent cells become :class:Interface objects with an axis-aligned unit normal and a grid of sample points on the shared face. The construction is dimension-agnostic and works for 1D, 2D, 3D and higher (no per-dimension special-casing).

Reference

Moseley, Markham, Nissen-Meyer (2023), "Finite Basis Physics-Informed Neural Networks", arXiv:2107.07871. The FBPINN subdomain tiling is a tensor product across dimensions; see RectangularDecompositionND in the reference implementation (https://github.com/benmoseley/FBPINNs), which lays out subdomains via np.meshgrid(*subdomain_xs).

Parameters:

Name	Type	Description	Default
bounds	Float[Array, 'dim 2']	Domain bounds, shape (dim, 2) with [min, max] per axis	required
num_partitions	tuple[int, ...]	Number of partitions in each dimension (length dim)	required
interface_points	int	Target number of sample points per interface face	10

Returns:

Type	Description
tuple[list[Subdomain], list[Interface]]	Tuple of (subdomains, interfaces)

XPINN (Extended PINN)

opifex.neural.pinns.domain_decomposition.xpinn

Extended Physics-Informed Neural Network (XPINN).

XPINN extends the PINN framework to handle domain decomposition with explicit interface conditions for continuity and flux matching.

Key Features

• Separate networks for each subdomain
• Interface continuity conditions (u_left = u_right)
• Flux continuity conditions (du/dn_left = du/dn_right)
• Weighted loss combination for interface enforcement

References

• Jagtap & Karniadakis (2020): Extended Physics-Informed Neural Networks
• Survey Section 8.3.1: XPINNs
• GitHub: https://github.com/AmeyaJagtap/XPINNs

XPINN

XPINN(input_dim: int, output_dim: int, subdomains: Sequence[Subdomain], interfaces: Sequence[Interface], hidden_dims: Sequence[int], *, config: XPINNConfig | None = None, activation: Callable[[Array], Array] = tanh, rngs: Rngs)

Bases: DomainDecompositionPINN

Extended Physics-Informed Neural Network.

XPINN decomposes the computational domain into non-overlapping subdomains, training a separate neural network for each subdomain. Interface conditions enforce solution continuity and flux matching between adjacent subdomains.

The total loss includes

• Data loss (if available)
• PDE residual loss (per subdomain)
• Interface continuity loss: ||u_left - u_right||²
• Interface flux loss: ||∂u/∂n_left - ∂u/∂n_right||²

Attributes:

Name	Type	Description
config		XPINN configuration with loss weights
input_dim		Spatial dimension
output_dim		Solution dimension
subdomains		List of subdomain definitions
interfaces		List of interface definitions
networks		List of subdomain networks

Example

subdomains = [ ... Subdomain(id=0, bounds=jnp.array([[0.0, 0.5]])), ... Subdomain(id=1, bounds=jnp.array([[0.5, 1.0]])), ... ] interfaces = [ ... Interface(subdomain_ids=(0, 1), points=jnp.array([[0.5]]), ... normal=jnp.array([1.0])) ... ] model = XPINN( ... input_dim=1, output_dim=1, ... subdomains=subdomains, interfaces=interfaces, ... hidden_dims=[32, 32], rngs=nnx.Rngs(0) ... )

Parameters:

Name	Type	Description	Default
input_dim	int	Spatial dimension	required
output_dim	int	Solution dimension	required
subdomains	Sequence[Subdomain]	List of subdomain definitions	required
interfaces	Sequence[Interface]	List of interface definitions	required
hidden_dims	Sequence[int]	Hidden layer dimensions for subdomain networks	required
config	XPINNConfig | None	XPINN configuration. Uses defaults if None.	None
activation	Callable[[Array], Array]	Activation function	tanh
rngs	Rngs	Random number generators	required

compute_continuity_loss

compute_continuity_loss() -> Float[Array, '']

Compute interface continuity loss.

Delegates to base class compute_interface_residual (DRY).

Returns:

Type	Description
Float[Array, '']	Scalar continuity loss (MSE of discontinuity)

compute_flux_loss

compute_flux_loss() -> Float[Array, '']

Compute interface flux continuity loss.

Enforces ∂u/∂n_left = ∂u/∂n_right at all interface points, where n is the interface normal direction.

Returns:

Type	Description
Float[Array, '']	Scalar flux loss (MSE of flux discontinuity)

compute_interface_loss

compute_interface_loss() -> Float[Array, '']

Compute total weighted interface loss.

Combines continuity and flux losses with configured weights.

Returns:

Type	Description
Float[Array, '']	Scalar total interface loss

compute_subdomain_residual

compute_subdomain_residual(subdomain_id: int, residual_fn: Callable[[Callable[[Float[Array, ...]], Float[Array, 'batch out']], Float[Array, ...]], Float[Array, ' batch']], collocation_points: Float[Array, ...]) -> Float[Array, '']

Compute PDE residual for a specific subdomain.

Parameters:

Name	Type	Description	Default
subdomain_id	int	ID of the subdomain	required
residual_fn	Callable[[Callable[[Float[Array, ...]], Float[Array, 'batch out']], Float[Array, ...]], Float[Array, ' batch']]	Function that computes PDE residual given network and points	required
collocation_points	Float[Array, ...]	Points where to evaluate residual	required

Returns:

Type	Description
Float[Array, '']	Scalar residual loss for this subdomain

compute_total_residual

compute_total_residual(residual_fn: Callable[[Callable[[Float[Array, ...]], Float[Array, 'batch out']], Float[Array, ...]], Float[Array, ' batch']], collocation_points_per_subdomain: Sequence[Float[Array, ...]]) -> Float[Array, '']

Compute total PDE residual across all subdomains.

Parameters:

Name	Type	Description	Default
residual_fn	Callable[[Callable[[Float[Array, ...]], Float[Array, 'batch out']], Float[Array, ...]], Float[Array, ' batch']]	Function that computes PDE residual	required
collocation_points_per_subdomain	Sequence[Float[Array, ...]]	Collocation points for each subdomain	required

Returns:

Type	Description
Float[Array, '']	Scalar total residual loss

XPINNConfig dataclass

XPINNConfig(continuity_weight: float = 1.0, flux_weight: float = 1.0, residual_weight: float = 1.0, average_residual_weight: float = 0.0)

Configuration for XPINN training.

Attributes:

Name	Type	Description
continuity_weight	float	Weight for interface continuity loss (u_left = u_right)
flux_weight	float	Weight for interface flux continuity loss (du/dn matching)
residual_weight	float	Weight for PDE residual loss in each subdomain
average_residual_weight	float	Weight for residual averaging at interfaces

FBPINN (Finite Basis PINN)

opifex.neural.pinns.domain_decomposition.fbpinn

Finite Basis Physics-Informed Neural Network (FBPINN).

FBPINN uses smooth window functions to create a partition of unity, enabling smooth blending of subdomain solutions without explicit interface conditions.

Key Features

• Smooth window functions (cosine, Gaussian)
• Partition of unity through normalization
• No explicit interface conditions needed
• Naturally handles overlapping subdomains

References

• Moseley et al. (2023): Finite Basis Physics-Informed Neural Networks
• Survey Section 8.3.2: FBPINNs
• GitHub: https://github.com/benmoseley/FBPINNs

FBPINN

FBPINN(input_dim: int, output_dim: int, subdomains: Sequence[Subdomain], interfaces: Sequence, hidden_dims: Sequence[int], *, config: FBPINNConfig | None = None, activation: Callable[[Array], Array] = tanh, rngs: Rngs)

Bases: DomainDecompositionPINN

Finite Basis Physics-Informed Neural Network.

FBPINN decomposes the computational domain into overlapping subdomains, using smooth window functions to blend subdomain network outputs. This creates a partition of unity that ensures smooth global solutions.

The output is computed as

u(x) = Σᵢ wᵢ(x) * uᵢ(x) / Σⱼ wⱼ(x)

where wᵢ(x) is the window function for subdomain i and uᵢ(x) is the network output for subdomain i.

Attributes:

Name	Type	Description
config		FBPINN configuration
windows		List of window functions for each subdomain

Example

subdomains = [ ... Subdomain(id=0, bounds=jnp.array([[0.0, 0.6]])), ... Subdomain(id=1, bounds=jnp.array([[0.4, 1.0]])), ... ] model = FBPINN( ... input_dim=1, output_dim=1, ... subdomains=subdomains, interfaces=[], ... hidden_dims=[32, 32], rngs=nnx.Rngs(0) ... )

Parameters:

Name	Type	Description	Default
input_dim	int	Spatial dimension	required
output_dim	int	Solution dimension	required
subdomains	Sequence[Subdomain]	List of subdomain definitions (should overlap)	required
interfaces	Sequence	List of interface definitions (optional for FBPINN)	required
hidden_dims	Sequence[int]	Hidden layer dimensions for subdomain networks	required
config	FBPINNConfig | None	FBPINN configuration. Uses defaults if None.	None
activation	Callable[[Array], Array]	Activation function	tanh
rngs	Rngs	Random number generators	required

compute_window_weights

compute_window_weights(x: Float[Array, ...]) -> Float[Array, 'batch num_subdomains']

Compute window weights for all subdomains.

Parameters:

Name	Type	Description	Default
x	Float[Array, ...]	Input coordinates	required

Returns:

Type	Description
Float[Array, 'batch num_subdomains']	Window weights, shape (batch, num_subdomains)

FBPINNConfig dataclass

FBPINNConfig(window_type: Literal['cosine', 'gaussian'] = 'cosine', normalize_windows: bool = True, overlap_factor: float = 0.2, gaussian_sigma: float = 0.25)

Configuration for FBPINN training.

Attributes:

Name	Type	Description
window_type	Literal['cosine', 'gaussian']	Type of window function ("cosine" or "gaussian")
normalize_windows	bool	Whether to normalize window weights to sum to 1
overlap_factor	float	Factor controlling subdomain overlap (for auto-partitioning)
gaussian_sigma	float	Sigma parameter for Gaussian windows

WindowFunction

WindowFunction(subdomain: Subdomain)

Bases: ABC

Abstract base class for window functions.

Window functions define the influence region of each subdomain network. They should be smooth, have compact support within the subdomain, and enable partition of unity when combined.

Parameters:

Name	Type	Description	Default
subdomain	Subdomain	The subdomain this window is associated with	required

CosineWindow

CosineWindow(subdomain: Subdomain)

Bases: WindowFunction

Cosine-based window function.

w(x) = 0.5 * (1 + cos(π * r)) for r < 1, else 0

where r is the normalized distance from the subdomain center, scaled by the subdomain half-width.

This creates a smooth bump function that is 1 at the center and 0 at the boundary.

GaussianWindow

GaussianWindow(subdomain: Subdomain, sigma: float = 0.25)

Bases: WindowFunction

Gaussian-based window function.

w(x) = exp(-||x - center||² / (2 * σ²))

where σ controls the width of the Gaussian.

Parameters:

Name	Type	Description	Default
subdomain	Subdomain	The subdomain this window is associated with	required
sigma	float	Standard deviation of the Gaussian (relative to subdomain size)	0.25

CPINN (Conservative PINN)

opifex.neural.pinns.domain_decomposition.cpinn

Conservative Physics-Informed Neural Network (cPINN).

cPINN extends XPINN with explicit flux conservation at interfaces, enforcing strong conservation properties required for conservation laws.

Key Features

• Explicit flux computation at interfaces
• Strong conservation enforcement
• Weighted combination of continuity and flux losses

References

• Jagtap et al. (2020): Conservative physics-informed neural networks
• Survey Section 8.3.2: Conservative PINNs

CPINN

CPINN(input_dim: int, output_dim: int, subdomains: Sequence[Subdomain], interfaces: Sequence[Interface], hidden_dims: Sequence[int], *, config: CPINNConfig | None = None, activation: Callable[[Array], Array] = tanh, rngs: Rngs)

Bases: DomainDecompositionPINN

Conservative Physics-Informed Neural Network.

cPINN enforces strong conservation at subdomain interfaces by explicitly computing and matching fluxes across boundaries.

The total interface loss includes

• Continuity loss: ||u_left - u_right||²
• Flux conservation loss: ||F_left · n - F_right · n||²

where F = ∇u is the flux (gradient) of the solution.

Attributes:

Name	Type	Description
config		cPINN configuration with loss weights
input_dim		Spatial dimension
output_dim		Solution dimension
subdomains		List of subdomain definitions
interfaces		List of interface definitions
networks		List of subdomain networks

Example

subdomains = [ ... Subdomain(id=0, bounds=jnp.array([[0.0, 0.5]])), ... Subdomain(id=1, bounds=jnp.array([[0.5, 1.0]])), ... ] interfaces = [ ... Interface(subdomain_ids=(0, 1), points=jnp.array([[0.5]]), ... normal=jnp.array([1.0])) ... ] model = CPINN( ... input_dim=1, output_dim=1, ... subdomains=subdomains, interfaces=interfaces, ... hidden_dims=[32, 32], rngs=nnx.Rngs(0) ... )

Parameters:

Name	Type	Description	Default
input_dim	int	Spatial dimension	required
output_dim	int	Solution dimension	required
subdomains	Sequence[Subdomain]	List of subdomain definitions	required
interfaces	Sequence[Interface]	List of interface definitions	required
hidden_dims	Sequence[int]	Hidden layer dimensions for subdomain networks	required
config	CPINNConfig | None	cPINN configuration. Uses defaults if None.	None
activation	Callable[[Array], Array]	Activation function	tanh
rngs	Rngs	Random number generators	required

compute_continuity_loss

compute_continuity_loss() -> Float[Array, '']

Compute interface continuity loss.

Delegates to base class compute_interface_residual (DRY).

Returns:

Type	Description
Float[Array, '']	Scalar continuity loss (MSE of discontinuity)

compute_flux_conservation_loss

compute_flux_conservation_loss() -> Float[Array, '']

Compute flux conservation loss at interfaces.

Enforces F_left · n = F_right · n at all interface points, where F = ∇u is the flux.

Returns:

Type	Description
Float[Array, '']	Scalar flux conservation loss

compute_interface_loss

compute_interface_loss() -> Float[Array, '']

Compute total weighted interface loss.

Combines continuity and flux conservation losses with configured weights.

Returns:

Type	Description
Float[Array, '']	Scalar total interface loss

CPINNConfig dataclass

CPINNConfig(flux_weight: float = 1.0, continuity_weight: float = 1.0, conservation_weight: float = 0.1)

Configuration for cPINN training.

Attributes:

Name	Type	Description
flux_weight	float	Weight for flux conservation loss at interfaces
continuity_weight	float	Weight for solution continuity loss
conservation_weight	float	Weight for global conservation enforcement

APINN (Augmented PINN)

opifex.neural.pinns.domain_decomposition.apinn

Augmented Physics-Informed Neural Network (APINN).

APINN uses a learnable gating network to smoothly blend subdomain solutions, allowing the model to learn optimal subdomain selection.

Key Features

• Learnable gating network for subdomain weighting
• Temperature-controlled softmax for soft/hard selection
• Differentiable blending for end-to-end training

References

• Survey Section 8.3.3: Augmented PINNs

APINN

APINN(input_dim: int, output_dim: int, subdomains: Sequence[Subdomain], interfaces: Sequence[Interface], hidden_dims: Sequence[int], *, config: APINNConfig | None = None, activation: Callable[[Array], Array] = tanh, rngs: Rngs)

Bases: DomainDecompositionPINN

Augmented Physics-Informed Neural Network.

APINN uses a learnable gating network to determine how to blend solutions from different subdomains. Unlike FBPINN which uses fixed window functions, APINN learns the optimal blending.

The output is computed as

u(x) = Σᵢ gᵢ(x) * uᵢ(x)

where gᵢ(x) are the learned gating weights (sum to 1) and uᵢ(x) are the subdomain network outputs.

Attributes:

Name	Type	Description
config		APINN configuration
gating_network		Network that produces blending weights
input_dim		Spatial dimension
output_dim		Solution dimension
subdomains		List of subdomain definitions
interfaces		List of interface definitions
networks		List of subdomain networks

Example

subdomains = [ ... Subdomain(id=0, bounds=jnp.array([[0.0, 0.5]])), ... Subdomain(id=1, bounds=jnp.array([[0.5, 1.0]])), ... ] interfaces = [ ... Interface(subdomain_ids=(0, 1), points=jnp.array([[0.5]]), ... normal=jnp.array([1.0])) ... ] model = APINN( ... input_dim=1, output_dim=1, ... subdomains=subdomains, interfaces=interfaces, ... hidden_dims=[32, 32], rngs=nnx.Rngs(0) ... )

Parameters:

Name	Type	Description	Default
input_dim	int	Spatial dimension	required
output_dim	int	Solution dimension	required
subdomains	Sequence[Subdomain]	List of subdomain definitions	required
interfaces	Sequence[Interface]	List of interface definitions	required
hidden_dims	Sequence[int]	Hidden layer dimensions for subdomain networks	required
config	APINNConfig | None	APINN configuration. Uses defaults if None.	None
activation	Callable[[Array], Array]	Activation function	tanh
rngs	Rngs	Random number generators	required

get_gating_weights

get_gating_weights(x: Float[Array, 'batch dim']) -> Float[Array, 'batch num_subdomains']

Get gating weights for given points.

Parameters:

Name	Type	Description	Default
x	Float[Array, 'batch dim']	Input coordinates	required

Returns:

Type	Description
Float[Array, 'batch num_subdomains']	Gating weights for each subdomain

compute_interface_loss

compute_interface_loss() -> Float[Array, '']

Compute weighted interface continuity loss.

Delegates continuity computation to base class compute_interface_residual and applies the configured continuity weight (DRY).

Returns:

Type	Description
Float[Array, '']	Scalar interface loss

APINNConfig dataclass

APINNConfig(temperature: float = 1.0, gating_hidden_dims: list[int] = (lambda: [16, 16])(), continuity_weight: float = 1.0)

Configuration for APINN training.

Attributes:

Name	Type	Description
temperature	float	Softmax temperature for gating. Lower values give sharper (more discrete) weights, higher values give smoother (more uniform) weights.
gating_hidden_dims	list[int]	Hidden dimensions for the gating network
continuity_weight	float	Weight for interface continuity loss

GatingNetwork

GatingNetwork(input_dim: int, num_subdomains: int, hidden_dims: Sequence[int], *, activation: Callable[[Array], Array] = tanh, rngs: Rngs)

Bases: Module

Gating network for subdomain selection.

This network takes spatial coordinates and outputs weights for blending subdomain solutions.

Attributes:

Name	Type	Description
layers		List of linear layers
activation		Activation function

Parameters:

Name	Type	Description	Default
input_dim	int	Input spatial dimension	required
num_subdomains	int	Number of subdomains to gate	required
hidden_dims	Sequence[int]	Hidden layer dimensions	required
activation	Callable[[Array], Array]	Activation function	tanh
rngs	Rngs	Random number generators	required

For usage examples and best practices, see the Domain Decomposition PINNs Guide.

Activations

opifex.neural.activations

Activation functions optimized for scientific neural networks.

This module provides a full collection of activation functions specifically optimized for scientific machine learning applications. All functions are fully compatible with Flax NNX patterns and JAX transformations.

MODERNIZATION APPLIED: - Full Flax NNX compliance with proper type annotations - Enhanced activation function selection with error handling - Optimized implementations for scientific computing - Support for both standard and specialized activation patterns

get_activation

get_activation(name: str | Callable) -> Any

Get activation function by name or return function if already callable.

Parameters:

Name	Type	Description	Default
name	str | Callable	Name of the activation function (case-insensitive) or callable function	required

Returns:

Type	Description
Any	JAX activation function or callable

Raises:

Type	Description
ValueError	If activation function is not found

list_activations

list_activations() -> list[str]

List all available activation functions.

Returns:

Type	Description
list[str]	List of activation function names

Examples:

>>> activations = list_activations()
>>> print(f"Available activations: {', '.join(activations)}")

register_activation

register_activation(name: str, func: Callable) -> None

Register a custom activation function.

Parameters:

Name	Type	Description	Default
name	str	Name of the activation function	required
func	Callable	The activation function (should accept and return JAX arrays)	required

Examples:

>>> def my_activation(x):
...     return x ** 3
>>> register_activation("cubic", my_activation)
>>> cubic_fn = get_activation("cubic")

mish

mish(x: Array) -> Array

Mish activation function: x * tanh(softplus(x)).

Mish is a self-gated activation function that has shown excellent performance in deep networks. It's smooth and non-monotonic.

Mathematical definition: f(x) = x * tanh(ln(1 + exp(x)))

Parameters:

Name	Type	Description	Default
x	Array	Input array	required

Returns:

Type	Description
Array	Output array with Mish activation applied

Note

This implementation uses softplus(x) = ln(1 + exp(x)) for numerical stability.

snake_activation

snake_activation(x: Array, a: float = 1.0) -> Array

Snake activation function: x + sin²(αx)/α.

Snake activation has been shown to work well for certain scientific applications, particularly those involving periodic patterns.

Mathematical definition: f(x) = x + (1/α) * sin²(αx)

Parameters:

Name	Type	Description	Default
x	Array	Input array	required
a	float	Frequency parameter (default: 1.0)	1.0

Returns:

Type	Description
Array	Output array with Snake activation applied

Note

The frequency parameter α controls the oscillation frequency. Higher values create more frequent oscillations.

gaussian_activation

gaussian_activation(x: Array, sigma: float = 1.0) -> Array

Gaussian activation function: exp(-x²/(2σ²)).

Gaussian activation can be useful for radial basis function networks and certain scientific applications where localized responses are desired.

Mathematical definition: f(x) = exp(-x²/(2σ²))

Parameters:

Name	Type	Description	Default
x	Array	Input array	required
sigma	float	Standard deviation parameter (default: 1.0)	1.0

Returns:

Type	Description
Array	Output array with Gaussian activation applied

Note

The σ parameter controls the width of the Gaussian. Smaller values create sharper peaks.

normalized_tanh

normalized_tanh(x: Array) -> Array

Normalized tanh activation: 1.7159 * tanh(2x/3).

This is a normalized version of tanh that has unit variance for normalized inputs, which can help with training stability.

Parameters:

Name	Type	Description	Default
x	Array	Input array	required

Returns:

Type	Description
Array	Output array with normalized tanh applied

soft_exponential

soft_exponential(x: Array, alpha: float = 0.0) -> Array

Soft exponential activation function.

This is a parameterized activation that interpolates between different behaviors based on the alpha parameter.

Mathematical definition: - If α < 0: -ln(1 - α(x + α)) / α - If α = 0: x - If α > 0: (exp(αx) - 1) / α + α

Parameters:

Name	Type	Description	Default
x	Array	Input array	required
alpha	float	Shape parameter	0.0

Returns:

Type	Description
Array	Output array with soft exponential applied

get_derivative_activation

get_derivative_activation(name: str) -> Any

Get the derivative of an activation function.

This is useful for implementations that need explicit derivatives rather than relying on automatic differentiation.

Parameters:

Name	Type	Description	Default
name	str	Name of the activation function	required

Returns:

Type	Description
Any	Derivative function of the specified activation

Raises:

Type	Description
ValueError	If activation name is not recognized or derivative not available