Opifex

Research preview. Opifex is under rapid iteration and the API will change while we iterate toward v1.0. What that means concretely:

Area	Status	Impact
API	🔄 Unstable	Breaking changes are expected. Public interfaces may change without deprecation warnings. Pin to specific commits if stability is required.
Tests	🔄 In Flux	Test suite is being expanded. Some tests may fail or be skipped. Coverage metrics are improving but not yet full.
Documentation	🔄 Evolving	Docs may not reflect current implementation. Code examples might be outdated. Refer to source code and tests for accurate usage.

Pin a version if you need stability, and do not put it in production yet. For research and experimentation it is ready to use today, with the understanding that APIs will evolve.

This is public this early on purpose. Issues, questions and pull requests genuinely steer what gets built next, and a star tells us which layer to push on.

A JAX-native platform for scientific machine learning, built for unified excellence, probabilistic-first design, and high performance.

🎯 Core Vision

• 🔬 Unified Excellence: Single platform supporting all major Opifex paradigms with mathematical clarity
• 📊 Probabilistic-First: Built-in uncertainty quantification treating all computation as Bayesian inference
• ⚡ High Performance: Optimized for speed with JAX transformations and GPU acceleration
• 🏗️ Production-Oriented: Designed with benchmarking and deployment tools for future production use
• 🤝 Community-Driven: Open patterns for education, research collaboration, and industrial adoption

✨ Key Features

• Neural Operators: FNO, DeepONet, SFNO, U-FNO, UNO, TFNO, GNO, PINO, Local FNO, DISCO, and more (26 architectures)
• Physics-Informed Neural Networks: Standard PINNs plus domain decomposition (FBPINN, XPINN, CPINN)
• Atomistic Potentials: E(3)-equivariant SchNet, PaiNN, and NequIP backbones (with MACE-style higher body-order via symmetric contraction), energy/forces/stress heads, and an ASE calculator
• Quantum Chemistry: Differentiable Kohn-Sham DFT, neural exchange-correlation functionals, variational Monte Carlo, and equivariant Hamiltonian prediction (QH9)
• Equivariant Core: Native E(3) algebra - irreps, Clebsch-Gordan, Wigner-D, and spherical harmonics
• Uncertainty Quantification: Conformal prediction, calibration, Gaussian processes, Bayesian quadrature, probabilistic numerics, simulation-based inference, and a broad adapter suite (ensembles, last-layer, SNGP, evidential)
• Equation Discovery: SINDy, Ensemble SINDy, Weak SINDy, and Bayesian SINDy
• Field Operations: JAX-native differential operators, advection, and pressure projection on structured grids
• Data Loading: JAX-native pipelines for PDEBench tensors and VTK unstructured meshes on the datarax Source/Pipeline contract
• Advanced Training: NTK analysis, GradNorm loss balancing, adaptive sampling (RAR-D)
• Optimization: Learn-to-optimize, meta-optimization (MAML/Reptile), and second-order methods
• Unified SciML Solvers: Standardized protocol for PINNs, Neural Operators, and Hybrid solvers
• 59 Working Examples: Full coverage from getting started to advanced research workflows

For detailed feature documentation, see Features.

🚀 Quick Start

Prerequisites

• Python 3.12+
• CUDA-compatible GPU (optional but recommended)

Installation

# Clone the repository
git clone https://github.com/avitai/opifex.git
cd opifex

# Set up development environment
./setup.sh

# Activate environment
source ./activate.sh

# Run tests to verify installation
uv run pytest tests/ -v

For detailed installation instructions, see Installation Guide.

📚 Basic Usage

Fourier Neural Operator (FNO)

import jax
from flax import nnx
from opifex.neural.operators.fno import FourierNeuralOperator

# Create FNO for learning PDE solution operators
rngs = nnx.Rngs(jax.random.PRNGKey(0))

fno = FourierNeuralOperator(
    in_channels=1,
    out_channels=1,
    hidden_channels=32,
    modes=12,
    num_layers=4,
    rngs=rngs,
)

# Input: (batch, channels, *spatial_dims)
x = jax.random.normal(jax.random.PRNGKey(1), (4, 1, 64, 64))
y = fno(x)
print(f"FNO: {x.shape} -> {y.shape}")  # (4, 1, 64, 64) -> (4, 1, 64, 64)

Deep Operator Network (DeepONet)

import jax
from flax import nnx
from opifex.neural.operators.deeponet import DeepONet

# Create DeepONet for function-to-function mapping
rngs = nnx.Rngs(jax.random.PRNGKey(0))

deeponet = DeepONet(
    branch_sizes=[100, 64, 64, 32],  # 100 sensor locations
    trunk_sizes=[2, 64, 64, 32],     # 2D output coordinates
    activation="gelu",
    rngs=rngs,
)

# Branch input: function values at sensors (batch, num_sensors)
# Trunk input: evaluation coordinates (batch, n_locations, coord_dim)
branch_input = jax.random.normal(jax.random.PRNGKey(1), (8, 100))
trunk_input = jax.random.uniform(jax.random.PRNGKey(2), (8, 50, 2))  # 50 eval points

output = deeponet(branch_input, trunk_input)
print(f"DeepONet output: {output.shape}")  # (8, 50)

Equation Discovery (SINDy)

import jax.numpy as jnp
from opifex.discovery.sindy import SINDy, SINDyConfig

# Generate Lorenz trajectory (σ=10, ρ=28, β=8/3) with RK4
def lorenz(state, sigma=10.0, rho=28.0, beta=8.0 / 3.0):
    x, y, z = state
    return jnp.array([sigma * (y - x), x * (rho - z) - y, x * y - beta * z])

dt, state = 0.001, jnp.array([1.0, 1.0, 1.0])
trajectory, derivatives = [state], [lorenz(state)]
for _ in range(10000):
    k1 = lorenz(state)
    k2 = lorenz(state + 0.5 * dt * k1)
    k3 = lorenz(state + 0.5 * dt * k2)
    k4 = lorenz(state + dt * k3)
    state = state + (dt / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4)
    trajectory.append(state)
    derivatives.append(lorenz(state))

# Discover governing equations from data
model = SINDy(SINDyConfig(polynomial_degree=2, threshold=0.3))
model.fit(jnp.stack(trajectory), jnp.stack(derivatives))

for eq in model.equations(["x", "y", "z"]):
    print(eq)
# dx/dt = -9.999 x + 10.000 y            (true: -10 x + 10 y)
# dy/dt = 28.000 x + -1.000 y + -1.000 x z  (true: 28 x - y - x z)
# dz/dt = -2.667 z + 1.000 x y            (true: -8/3 z + x y)

For full examples and tutorials, see the Examples directory and Documentation.

🔧 Development

# Run tests
uv run pytest tests/ -v

# Code quality checks
uv run pre-commit run --all-files

For detailed development guidelines, see Development Guide.

📖 Documentation

• Getting Started: Installation and basic usage
• Features: Complete feature overview
• Architecture: Framework design and structure
• API Reference: Complete API documentation
• Examples: Working examples and tutorials
• Development: Contributing and development setup

🤝 Contributing

We welcome contributions! Please see our Contributing Guide for details.

📄 License

This project is licensed under the MIT License - see the LICENSE file for details.

Ready to get started? Check out our Quick Start Guide or explore the Examples directory!