Opifex: Unified Scientific Machine Learning Framework

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
• 🏗️ Research Infrastructure: Flexible design with integrated benchmarking and experimental tools
• 🤝 Community-Driven: Open patterns for education, research collaboration, and industrial adoption

🚀 Quick Start

Prerequisites

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

Installation

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

# Set up unified development environment (auto-detects GPU/CPU)
./setup.sh

# Activate environment
source ./activate.sh

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

For detailed setup instructions, troubleshooting, and configuration options, see the Environment Setup Guide.

Basic Example

import jax
import jax.numpy as jnp
import flax.nnx as nnx
from opifex.neural.operators.fno import FourierNeuralOperator

# Create FNO for PDE solving
key = jax.random.PRNGKey(42)
rngs = nnx.Rngs(key)

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

# Forward pass with 2D spatial data
x = jax.random.normal(key, (4, 1, 64, 64))  # (batch, channels, height, width)
y = fno(x)
print(f"FNO: {x.shape} -> {y.shape}")  # (4, 1, 64, 64) -> (4, 1, 64, 64)

📚 Documentation

Getting Started

• Installation Guide - Setup instructions and configuration
• Quick Start Tutorial - Get up and running quickly
• Environment Setup - Development environment configuration
• GPU Setup - GPU acceleration setup

Core Documentation

• Features - Overview of Opifex paradigms and capabilities
• Architecture - Framework design and 6-layer architecture
• Technology Stack - Dependencies, tools, and infrastructure

User Guides

• Concepts - Core concepts and terminology
• Neural Networks - Neural network implementations
• Geometry - Geometric operations and representations
• Training - Training infrastructure and workflows
• Problems - Problem definition and setup

Methods & Tutorials

• Neural Operators - FNO, DeepONet, and advanced architectures
• Physics-Informed Networks - PINNs and physics-constrained learning
• Neural DFT - Quantum chemistry with neural networks
• Learn-to-Optimize - Meta-optimization and adaptive algorithms
• Probabilistic Methods - Uncertainty quantification
• Advanced Benchmarking - Evaluation and validation

API Reference

• Core Package - Mathematical abstractions and numerical framework
• Neural Package - Neural operators and physics-informed networks
• Geometry Package - Geometric operations and manifolds
• Training Package - Training infrastructure and optimization
• Optimization Package - Meta-optimization and advanced algorithms
• Benchmarking Package - Evaluation and validation tools
• Bayesian Package - Uncertainty quantification and Bayesian methods

Examples

• Examples Overview - Runnable examples and demonstrations
• Neural Operators - FNO, TFNO, DeepONet, and more
• PINNs - Physics-informed neural networks
• Quantum Chemistry - Neural DFT and molecular examples

Development

• Contributing - How to contribute to Opifex
• Development Setup - Development environment and guidelines
• Code Quality - Standards and best practices
• Testing - Testing framework and guidelines
• GPU Development - GPU development guidelines

Deployment

• Local Development - Local deployment setup
• AWS Deployment - Amazon Web Services deployment
• GCP Deployment - Google Cloud Platform deployment
• Troubleshooting - Common deployment issues

🤝 Community

Getting Help

• FAQ - Frequently asked questions
• GitHub Issues - Report bugs and request features
• Discussions - Community Q&A and collaboration
• Research Partnerships - Academic collaboration opportunities

Contributing

We welcome contributions! Please see our Contributing Guide for details on:

• Code style and standards
• Testing requirements
• Documentation guidelines
• Pull request process

📄 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 Features to learn about Opifex's capabilities!