Skip to content

Technology Stack

Opifex is built on a carefully curated technology stack that provides high performance, reliability, and modern development practices for scientific machine learning applications.

🛠️ Core Technologies

Core JAX Ecosystem

The foundation of Opifex is built on the JAX ecosystem, providing high-performance numerical computing with automatic differentiation and GPU acceleration.

  • JAX 0.8.0: Core framework with CUDA support

    • Automatic differentiation for gradient-based optimization
    • Just-in-time (JIT) compilation for performance
    • Vectorization and parallelization support
    • GPU and TPU acceleration

  • FLAX 0.12.0: Modern neural network framework (exclusive)

    • Stateful neural network transformations
    • Modular and composable neural network components
    • Integration with JAX transformations
    • Type-safe parameter handling

  • Optax 0.2.6+: Optimization algorithms

    • Gradient-based optimization algorithms
    • Learning rate scheduling
    • Gradient clipping and normalization
    • Composable optimization transformations

  • BlackJAX 1.2.5+: MCMC sampling

    • Hamiltonian Monte Carlo (HMC)
    • No-U-Turn Sampler (NUTS)
    • Metropolis-Hastings algorithms
    • Bayesian inference support

  • Diffrax 0.4.0+: Differential equations

    • Ordinary differential equation (ODE) solvers
    • Stochastic differential equation (SDE) solvers
    • Neural differential equations
    • Adaptive step size control

Quantum Chemistry Stack

Specialized components for quantum mechanical calculations and molecular systems.

  • Neural DFT: Chemical accuracy (<1 kcal/mol) quantum calculations

    • Neural exchange-correlation functionals
    • Self-consistent field (SCF) acceleration
    • Density functional theory implementations
    • Chemical accuracy validation

  • Molecular Systems: 3D geometry with periodic boundary conditions

    • Molecular structure representation
    • Periodic boundary condition handling
    • Symmetry operations and point groups
    • Force field integration

  • Electronic Structure: Quantum mechanical problem definitions

    • Wavefunction representations
    • Basis set management
    • Quantum mechanical operators
    • Many-body theory support

  • Physics Constraints: Conservation laws and quantum mechanical principles

    • Particle number conservation
    • Energy conservation
    • Symmetry enforcement
    • Quantum mechanical constraints

Advanced Training Stack

Infrastructure for physics-aware training and meta-optimization.

  • Physics-Informed Losses: Multi-physics composition with adaptive weighting

    • Hierarchical loss composition
    • Adaptive weight scheduling
    • Conservation law enforcement
    • Multi-physics problem support

  • Conservation Law Enforcement: Mass, momentum, energy, quantum conservation

    • Automatic conservation law detection
    • Soft and hard constraint enforcement
    • Physics-aware regularization
    • Constraint satisfaction monitoring

  • Meta-Optimization: L2O algorithms with neural meta-learning

    • Learn-to-optimize (L2O) framework
    • Meta-learning algorithms (MAML, Reptile)
    • Multi-objective optimization
    • Reinforcement learning for optimization

  • Adaptive Scheduling: Performance-based learning rate and weight adaptation

    • Performance monitoring
    • Adaptive learning rate scheduling
    • Dynamic weight adjustment
    • Convergence detection

  • Quantum-Aware Training: SCF acceleration and quantum constraint handling

    • SCF convergence acceleration
    • Quantum constraint enforcement
    • Chemical accuracy monitoring
    • Quantum mechanical validation

Development Infrastructure

Modern development tools ensuring code quality, testing, and documentation.

  • uv: Package management (exclusive)

    • Fast Python package installation
    • Dependency resolution
    • Virtual environment management
    • Lock file generation

  • ruff + pyright: Code quality (exclusive)

    • Fast Python linting with ruff
    • Type checking with pyright
    • Code formatting and style enforcement
    • Import sorting and organization

  • pytest: Testing framework

    • Parametrized testing
    • Fixture management
    • Coverage reporting

  • MkDocs: Documentation system

    • Markdown-based documentation
    • Material theme for modern UI
    • Mathematical notation support
    • API documentation generation

🔧 Supporting Libraries

Numerical Computing

  • Optimistix: Root finding & minimization
  • Lineax: Linear solvers
  • Distrax: Probabilistic programming
  • Orbax: Checkpointing system

Data Management

  • NumPy: Numerical array operations
  • SQLAlchemy: Database integration with type safety (optional, in platform extras)
  • HDF5 (h5py): Large-scale data storage (optional, in neural-dft/scientific-data extras)
  • Pandas: Data manipulation and analysis (optional, in pdebench extras)

Visualization

  • Matplotlib: Scientific plotting
  • Plotly: Interactive visualizations

Security & Quality

  • Bandit: Security analysis
  • pydocstyle: Documentation standards (via pre-commit hook, not a direct dependency)
  • pre-commit: Git hook management
  • pyright: Static type checking

🚀 Performance Characteristics

Computational Performance

  • GPU Acceleration: Native CUDA support through JAX
  • JIT Compilation: Automatic optimization of computational graphs
  • Vectorization: SIMD operations for array computations
  • Memory Efficiency: Optimized memory usage patterns

Scalability

  • Distributed Computing: Multi-GPU and multi-node support
  • Batch Processing: Efficient batch operations
  • Streaming: Large dataset processing
  • Cloud Integration: Kubernetes-native deployment

Quality Metrics

  • Test Coverage: Extensive test coverage
  • Type Safety: Full type annotation coverage
  • Documentation: Full API documentation
  • Security: Zero known vulnerabilities

🔄 Version Management

Dependency Pinning

All dependencies are pinned to specific versions to ensure reproducibility and stability across different environments.

Compatibility Matrix

Component Version Python CUDA Notes
JAX 0.8.0 3.11+ 12.0+ Core framework
FLAX 0.12.0 3.11+ - Neural networks
Optax 0.2.6+ 3.11+ - Optimization
BlackJAX 1.2.5+ 3.11+ - MCMC sampling
Diffrax 0.4.0+ 3.11+ - Differential equations

Update Policy

  • Major versions: Carefully evaluated for breaking changes
  • Minor versions: Regular updates for new features
  • Patch versions: Automatic updates for bug fixes
  • Security updates: Immediate updates for security patches

Learn More