Training API Documentation¶

Overview¶

The opifex.training module provides full training infrastructure for scientific machine learning, including physics-informed neural networks, optimization algorithms, and quantum-aware training workflows.

Module Structure:

opifex.core.training.trainer - Unified Trainer (Recommended)
opifex.core.training.config - Training configuration classes
opifex.core.training.physics_configs - Physics-specific configurations
opifex.training.basic_trainer - Core trainer implementations
opifex.core.training.monitoring.metrics - Metrics tracking and state management
opifex.core.training.components.recovery - Error recovery and stability handling
opifex.core.training.components - Modular training components

Core Classes¶

Trainer ⭐ RECOMMENDED¶

The unified, composable trainer architecture for all training workflows.

from opifex.core.training.trainer import Trainer
from opifex.core.training.config import TrainingConfig
from opifex.core.training.physics_configs import ConservationConfig, MultiScaleConfig

# Configure physics-aware training
conservation_config = ConservationConfig(
    laws=["energy", "momentum"],
    energy_tolerance=1e-6,
)

multiscale_config = MultiScaleConfig(
    scales=["molecular", "atomic"],
    weights={"molecular": 0.5, "atomic": 0.5},
)

config = TrainingConfig(
    num_epochs=100,
    learning_rate=1e-3,
    conservation_config=conservation_config,
    multiscale_config=multiscale_config,
)

# Create and use trainer
# Create a dummy model for demonstration
class SimpleModel(nnx.Module):
    def __init__(self, rngs: nnx.Rngs):
        self.linear = nnx.Linear(10, 1, rngs=rngs)
    def __call__(self, x):
        return self.linear(x)

model = SimpleModel(rngs=nnx.Rngs(0))
trainer = Trainer(model, config)
trained_model, history = trainer.train(train_data, val_data)

Key Features:

Composable Architecture: Mix and match physics configurations
Type-Safe: Full type hints and IDE support
Zero Runtime Overhead: Configuration at initialization only
Extensible: Add custom configs without modifying trainer
Production-Ready: Full testing and error handling

Supported Physics Configurations:

ConservationConfig: Energy, momentum, mass, and symmetry conservation
MultiScaleConfig: Multi-scale physics with coupling
QuantumTrainingConfig: Quantum chemistry and electronic structure
BoundaryConfig: Boundary condition enforcement
DFTConfig: Density functional theory workflows
SCFConfig: Self-consistent field convergence
MetricsTrackingConfig: Custom metrics tracking
LoggingConfig: Advanced logging and alerting

BasicTrainer¶

Standard training workflow with physics-informed capabilities.

from opifex.training.basic_trainer import BasicTrainer
from opifex.core.training.config import TrainingConfig

trainer = BasicTrainer(model, config)
trained_model, history = trainer.train(train_data, val_data)

Key Features:

Physics-informed neural network (PINN) training
Orbax-compatible checkpointing
JAX Array and automatic differentiation support
Type-safe with jaxtyping annotations

ModularTrainer ✅ NEW¶

Component-based training architecture with production-grade capabilities.

from opifex.training.basic_trainer import ModularTrainer
from opifex.core.training.config import TrainingConfig
from opifex.core.training.components.recovery import ErrorRecoveryManager
from opifex.core.training.components import FlexibleOptimizerFactory

trainer = ModularTrainer(
    model=model,
    config=config,
    rngs=rngs,
    components={
        "error_recovery": ErrorRecoveryManager(),
        "optimizer_factory": FlexibleOptimizerFactory()
    }
)

Key Features:

Component composition architecture
Pluggable training components
Production-grade error handling
Optimization strategies
Physics-aware metrics collection

Components¶

ErrorRecoveryManager ✅ NEW¶

Production-grade error handling with gradient stability and automatic recovery.

from opifex.core.training.components.recovery import ErrorRecoveryManager

error_manager = ErrorRecoveryManager(
    config={
        "max_retries": 5,
        "gradient_clip_threshold": 1.0,
        "loss_explosion_threshold": 100.0,
        "checkpoint_on_error": True
    }
)

Features:

Gradient clipping with automatic threshold adaptation
NaN detection and recovery mechanisms
Loss explosion detection and mitigation
Multiple recovery strategies (gradient clipping, learning rate reduction, parameter reinitialization)
Full error logging and analytics

FlexibleOptimizerFactory ✅ NEW¶

Optimizer creation with scheduling support.

from opifex.core.training.components import FlexibleOptimizerFactory

optimizer_factory = FlexibleOptimizerFactory(
    config={
        "optimizer_type": "adamw",  # "adam", "adamw", "sgd"
        "learning_rate": 1e-3,
        "schedule_type": "cosine",  # "cosine", "exponential", "linear"
        "total_steps": 1000,
        "cosine_alpha": 0.0
    }
)

Supported Optimizers:

Adam: Adaptive moment estimation
AdamW: Adam with weight decay
SGD: Stochastic gradient descent with momentum

Supported Schedules:

Cosine: Cosine annealing learning rate decay
Exponential: Exponential decay
Linear: Linear decay

MetricsCollector ✅ NEW¶

Physics-aware metrics collection with convergence tracking.

from opifex.core.training.monitoring.metrics import AdvancedMetricsCollector

collector = AdvancedMetricsCollector()
collector.start_training()
metrics = collector.collect_physics_metrics(model, x, y_true)

Collected Metrics:

Training loss and validation metrics
Gradient norms and stability indicators
Physics-specific metrics (energy conservation, mass conservation)
Convergence rates and training diagnostics
Real-time performance analytics

TrainingComponent ✅ NEW¶

Base class for creating custom training components.

from opifex.core.training.components import TrainingComponent

class CustomComponent(TrainingComponent):
    def setup(self, model, training_state):
        # Initialize component
        pass

    def step(self, model, training_state):
        # Update component state
        pass

Purpose:

Enables modular component development
Provides common interface for training components
Supports pluggable architecture patterns

Configuration Classes¶

TrainingConfig¶

Training configuration with full parameter control.

from opifex.core.training.config import TrainingConfig

config = TrainingConfig(
    num_epochs=1000,
    batch_size=64,
    learning_rate=1e-3,
    validation_frequency=100,
    checkpoint_frequency=500,
    early_stopping=True,
    patience=50
)

Parameters:

num_epochs: Number of training epochs
batch_size: Training batch size
learning_rate: Learning rate (can be overridden by optimizer factory)
validation_frequency: Validation evaluation frequency
checkpoint_frequency: Model checkpointing frequency
early_stopping: Enable early stopping
patience: Early stopping patience

TrainingState¶

Enhanced training state with full tracking.

from opifex.core.training.monitoring.metrics import TrainingState

# Automatically managed by trainers
state = trainer.training_state
print(f"Current epoch: {state.epoch}")
print(f"Best validation loss: {state.best_val_loss}")

Tracked Information:

Current epoch and step counters
Best validation metrics
Model and optimizer states
Recovery attempt history
Training diagnostics

Physics-Informed Training¶

PhysicsInformedLoss¶

Hierarchical multi-physics loss composition with adaptive weighting.

from opifex.core.physics.losses import PhysicsInformedLoss, PhysicsLossConfig

physics_loss = PhysicsInformedLoss(
    config=PhysicsLossConfig(
        physics_weight=1.0,
        boundary_weight=1.0,
        data_weight=1.0,
        adaptive_weighting=True
    )
)

# Use with BasicTrainer
trainer.set_physics_loss(physics_loss)

Supported Physics:

Partial differential equations (PDEs)
Conservation laws (mass, momentum, energy)
Quantum mechanical constraints
Boundary condition enforcement

Uncertainty-Aware Trainers¶

Three trainer surfaces in opifex.training.basic_trainer extend BasicTrainer with uncertainty-driven behaviours; each is independent and opt-in.

`UncertaintyGuidedTrainer`¶

Adaptive sampling driven by the model's epistemic uncertainty: candidate points with the highest predictive variance get re-weighted into the next training batch. Plugs into any UncertaintyAwareModule.predict_distribution(...) surface.

from flax import nnx
from opifex.training.basic_trainer import UncertaintyGuidedTrainer

trainer = UncertaintyGuidedTrainer(
    model=bayesian_model,
    uncertainty_quantifier=uq,
    rngs=nnx.Rngs(0),
    uncertainty_threshold=0.1,
)
state = trainer.train(dataset, num_epochs=100)

`MultiFidelityUncertaintyTrainer`¶

Multi-fidelity training that interleaves a low-fidelity solver and a high-fidelity solver while tracking per-fidelity uncertainty. Routes expensive high-fidelity calls only when the low-fidelity uncertainty exceeds a configurable threshold.

from flax import nnx
from opifex.training.basic_trainer import MultiFidelityUncertaintyTrainer

trainer = MultiFidelityUncertaintyTrainer(
    low_fidelity_model=fast_model,
    high_fidelity_model=accurate_model,
    uncertainty_quantifier=uq,
    rngs=nnx.Rngs(0),
)

`ActiveUncertaintyLearner`¶

Active-learning loop that queries the most informative samples from a candidate pool using the same epistemic-uncertainty proxy. Designed for small-data regimes where labelling is expensive.

from flax import nnx
from opifex.training.basic_trainer import ActiveUncertaintyLearner

learner = ActiveUncertaintyLearner(
    model=bayesian_model,
    uncertainty_quantifier=uq,
    rngs=nnx.Rngs(0),
    acquisition_size=32,
)

All three trainers consume the canonical PredictiveDistribution contract from opifex.uncertainty.types so they slot in alongside the calibration and conformal post-processing pipelines.

Usage Examples¶

Basic Training Workflow¶

import jax
import jax.numpy as jnp
import flax.nnx as nnx
from opifex.neural.base import StandardMLP
from opifex.training.basic_trainer import BasicTrainer
from opifex.core.training.config import TrainingConfig

# Create model
model = StandardMLP([1, 32, 32, 1], activation="tanh", rngs=nnx.Rngs(42))

# Configure training
config = TrainingConfig(num_epochs=1000, batch_size=64, learning_rate=1e-3)

# Create trainer and train
trainer = BasicTrainer(model, config)
trained_model, history = trainer.train(train_data, val_data)

Modular Training¶

from opifex.training.basic_trainer import ModularTrainer
from opifex.core.training.config import TrainingConfig
from opifex.core.training.components.recovery import ErrorRecoveryManager
from opifex.core.training.components import FlexibleOptimizerFactory

# Configure components
error_recovery = ErrorRecoveryManager(config={"max_retries": 5, "gradient_clip_threshold": 1.0})
optimizer_factory = FlexibleOptimizerFactory(config={"optimizer_type": "adamw", "schedule_type": "cosine"})

# Create modular trainer
# Note: AdvancedMetricsCollector is automatically created by ModularTrainer
trainer = ModularTrainer(
    model=model,
    config=config,
    rngs=rngs,
    components={
        "error_recovery": error_recovery,
        "optimizer_factory": optimizer_factory
    }
)

# Train with capabilities
trained_model, history = trainer.train(train_data, val_data)

Physics-Informed Training¶

from opifex.training.basic_trainer import BasicTrainer
from opifex.core.physics.losses import PhysicsInformedLoss

# Define PDE residual
def pde_residual(model_fn, x, t):
    u = model_fn(jnp.array([x, t]).reshape(1, -1))
    # Compute PDE residual (example: heat equation)
    u_t = jax.grad(lambda t: model_fn(jnp.array([x, t]).reshape(1, -1)))(t)
    u_xx = jax.grad(jax.grad(lambda x: model_fn(jnp.array([x, t]).reshape(1, -1))))(x)
    return u_t - 0.1 * u_xx

# Configure physics loss
physics_loss = PhysicsInformedLoss(pde_residual=pde_residual)

# Set up PINN training
trainer = BasicTrainer(model, config)
trainer.set_physics_loss(physics_loss)

# Train with physics constraints
trained_model, history = trainer.train(
    train_data=(domain_points, None),  # No target data for domain points
    boundary_data=(boundary_points, boundary_values)
)

Integration¶

With Neural Networks¶

from opifex.neural.base import StandardMLP

# Standard networks
standard_model = StandardMLP([3, 64, 64, 1], activation="swish", rngs=rngs)

For molecular and materials property prediction (energies, forces, stress), build a machine-learning interatomic potential from opifex.neural.atomistic instead; see the Atomistic Potentials guide.

Best Practices¶

Production Training¶

Use ModularTrainer for production workflows with error recovery
Configure appropriate error recovery strategies for your problem
Monitor training metrics with MetricsCollector
Use learning rate scheduling for better convergence
Enable checkpointing for long training runs

Physics-Informed Training¶

Balance loss weights between physics, boundary, and data terms
Use adaptive weighting for complex multi-physics problems
Monitor conservation laws during training
Validate physics constraints on test data

Performance Optimization¶

Choose appropriate batch sizes for your hardware
Use JAX transformations (vmap, jit) for efficiency
Profile training with JAX profiling tools
Monitor gradient health and stability

Troubleshooting¶

Common Issues¶

NaN losses: Enable NaN detection in ErrorRecoveryManager
Gradient explosions: Use gradient clipping with appropriate thresholds
Slow convergence: Try different optimizers and learning rate schedules
Physics constraint violations: Increase physics loss weights or improve residual computation

Debug Features¶

Full logging of training metrics and errors
Recovery attempt tracking for debugging stability issues
Gradient norm monitoring for optimization health
Physics constraint validation for PINN problems

Multilevel Training¶

Coarse-to-fine training hierarchies for accelerated convergence.

Width-Based Hierarchy (MLPs)¶

opifex.core.training.strategies.multilevel.coarse_to_fine ¶

Coarse-to-Fine Multilevel Training Framework.

This module implements multilevel training strategies where models are trained from coarse (fewer parameters) to fine (more parameters) levels.

Key Features

Network hierarchy with configurable coarsening
Transfer operators (prolongation/restriction)
Cascade training (sequential level training)

The key insight is that training a coarse network first provides a good initialization for the fine network, accelerating convergence.

References

Survey Section 8.2: Multilevel Training
Multigrid methods in numerical analysis

MultilevelConfig `dataclass` ¶

MultilevelConfig(num_levels: int = 3, coarsening_factor: float = 0.5, level_epochs: list[int] = (lambda: [100, 200, 300])(), warmup_epochs: int = 0)

Configuration for multilevel training.

Attributes:

Name	Type	Description
`num_levels`	`int`	Number of levels in the hierarchy
`coarsening_factor`	`float`	Factor to reduce width at each coarser level
`level_epochs`	`list[int]`	Number of epochs to train at each level
`warmup_epochs`	`int`	Extra epochs at the finest level

create_network_hierarchy ¶

create_network_hierarchy(input_dim: int, output_dim: int, base_hidden_dims: Sequence[int], num_levels: int, coarsening_factor: float = 0.5, *, activation: Callable[[Array], Array] = tanh, rngs: Rngs) -> list[MultilevelMLP]

Create hierarchy of networks from coarse to fine.

The finest level (highest index) uses the base_hidden_dims. Coarser levels use progressively smaller networks.

Parameters:

Name	Type	Description	Default
`input_dim`	`int`	Input dimension	required
`output_dim`	`int`	Output dimension	required
`base_hidden_dims`	`Sequence[int]`	Hidden dimensions for the finest level	required
`num_levels`	`int`	Number of levels in hierarchy	required
`coarsening_factor`	`float`	Factor to reduce width at each coarser level	`0.5`
`activation`	`Callable[[Array], Array]`	Activation function	`tanh`
`rngs`	`Rngs`	Random number generators	required

Returns:

Type	Description
`list[MultilevelMLP]`	List of networks from coarsest to finest

prolongate ¶

prolongate(coarse_model: MultilevelMLP, fine_model: MultilevelMLP) -> MultilevelMLP

Transfer (prolongate) parameters from coarse to fine model.

This copies the coarse model parameters to the corresponding subset of the fine model parameters. Additional fine model parameters are left at their initialized values.

Parameters:

Name	Type	Description	Default
`coarse_model`	`MultilevelMLP`	Coarse level model	required
`fine_model`	`MultilevelMLP`	Fine level model (will be modified in place)	required

Returns:

Type	Description
`MultilevelMLP`	Fine model with prolongated parameters

restrict ¶

restrict(fine_model: MultilevelMLP, coarse_model: MultilevelMLP) -> MultilevelMLP

Transfer (restrict) parameters from fine to coarse model.

This copies a subset of the fine model parameters to the coarse model.

Parameters:

Name	Type	Description	Default
`fine_model`	`MultilevelMLP`	Fine level model	required
`coarse_model`	`MultilevelMLP`	Coarse level model (will be modified in place)	required

Returns:

Type	Description
`MultilevelMLP`	Coarse model with restricted parameters

Mode-Based Hierarchy (FNOs)¶

opifex.core.training.strategies.multilevel.multilevel_fno ¶

Multilevel Training for Fourier Neural Operators.

This module implements multilevel training strategies specifically designed for Fourier Neural Operators, where the hierarchy is based on Fourier modes.

Key Features

Mode-based coarsening (fewer modes = coarser level)
Transfer operators for spectral weights
Cascade training from few modes to many modes

The key insight is that training with fewer Fourier modes first captures low-frequency features quickly, then finer modes capture high-frequency details.

References

Survey Section 8.2: Multilevel Training
Li et al. (2020): Fourier Neural Operator

MultilevelFNOTrainer ¶

MultilevelFNOTrainer(width: int, input_dim: int, output_dim: int, config: MultilevelFNOConfig | None = None, *, num_layers: int = 4, activation: Callable[[Array], Array] = gelu, rngs: Rngs)

Trainer for multilevel FNO.

Trains FNOs from coarse to fine modes, transferring learned spectral weights between levels.

Attributes:

Name	Type	Description
`config`		Multilevel configuration
`hierarchy`		List of FNOs from coarse to fine
`current_level`		Current training level

Parameters:

Name	Type	Description	Default
`width`	`int`	Hidden channel width	required
`input_dim`	`int`	Input channels	required
`output_dim`	`int`	Output channels	required
`config`	`MultilevelFNOConfig \| None`	Multilevel configuration	`None`
`num_layers`	`int`	FNO layers per network	`4`
`activation`	`Callable[[Array], Array]`	Activation function	`gelu`
`rngs`	`Rngs`	Random number generators	required

get_current_model ¶

get_current_model() -> SimpleFNO

Get FNO at current level.

Returns:

Type	Description
`SimpleFNO`	Current level FNO

advance_level ¶

advance_level() -> bool

Advance to next finer level.

Transfers learned weights from current level to next level.

Returns:

Type	Description
`bool`	True if advanced successfully, False if at finest level

is_at_finest ¶

is_at_finest() -> bool

Check if at finest level.

Returns:

Type	Description
`bool`	True if at finest level

get_epochs_for_current_level ¶

get_epochs_for_current_level() -> int

Get epochs for current level from config.

Returns:

Type	Description
`int`	Number of epochs to train at current level

MultilevelFNOConfig `dataclass` ¶

MultilevelFNOConfig(num_levels: int = 3, base_modes: int = 12, mode_reduction_factor: int = 2, level_epochs: list[int] = (lambda: [50, 100, 150])())

Configuration for multilevel FNO training.

Attributes:

Name	Type	Description
`num_levels`	`int`	Number of levels in the hierarchy
`base_modes`	`int`	Number of Fourier modes at finest level
`mode_reduction_factor`	`int`	Factor to reduce modes at each coarser level
`level_epochs`	`list[int]`	Epochs to train at each level

create_fno_hierarchy ¶

create_fno_hierarchy(base_modes: int, width: int, input_dim: int, output_dim: int, num_levels: int, reduction_factor: int = 2, *, num_layers: int = 4, activation: Callable[[Array], Array] = gelu, rngs: Rngs) -> list[SimpleFNO]

Create hierarchy of FNOs from coarse to fine.

Parameters:

Name	Type	Description	Default
`base_modes`	`int`	Modes at finest level	required
`width`	`int`	Hidden channel width (same for all levels)	required
`input_dim`	`int`	Input channels	required
`output_dim`	`int`	Output channels	required
`num_levels`	`int`	Number of levels	required
`reduction_factor`	`int`	Mode reduction per level	`2`
`num_layers`	`int`	FNO layers per network	`4`
`activation`	`Callable[[Array], Array]`	Activation function	`gelu`
`rngs`	`Rngs`	Random number generators	required

Returns:

Type	Description
`list[SimpleFNO]`	List of FNOs from coarsest to finest

create_mode_hierarchy ¶

create_mode_hierarchy(base_modes: int, num_levels: int, reduction_factor: int = 2) -> list[int]

Create hierarchy of mode counts from coarse to fine.

The finest level (highest index) uses base_modes. Coarser levels use progressively fewer modes.

Parameters:

Name	Type	Description	Default
`base_modes`	`int`	Number of modes at finest level	required
`num_levels`	`int`	Number of levels in hierarchy	required
`reduction_factor`	`int`	Factor to reduce modes at each coarser level	`2`

Returns:

Type	Description
`list[int]`	List of mode counts from coarsest to finest

prolongate_fno_modes ¶

prolongate_fno_modes(coarse_fno: SimpleFNO, fine_fno: SimpleFNO) -> SimpleFNO

Transfer spectral weights from coarse to fine FNO.

Copies the lower-frequency modes from the coarse network to the corresponding modes in the fine network.

Parameters:

Name	Type	Description	Default
`coarse_fno`	`SimpleFNO`	Coarse FNO (fewer modes)	required
`fine_fno`	`SimpleFNO`	Fine FNO (more modes, modified in place)	required

Returns:

Type	Description
`SimpleFNO`	Fine FNO with prolongated weights

restrict_fno_modes ¶

restrict_fno_modes(fine_fno: SimpleFNO, coarse_fno: SimpleFNO) -> SimpleFNO

Transfer spectral weights from fine to coarse FNO.

Copies the lower-frequency modes from the fine network to the coarse network (truncation).

Parameters:

Name	Type	Description	Default
`fine_fno`	`SimpleFNO`	Fine FNO (more modes)	required
`coarse_fno`	`SimpleFNO`	Coarse FNO (fewer modes, modified in place)	required

Returns:

Type	Description
`SimpleFNO`	Coarse FNO with restricted weights

For usage examples and best practices, see the Multilevel Training Guide.

Adaptive Sampling¶

Residual-based sampling strategies for efficient PINN training.

opifex.core.training.components.adaptive_sampling ¶

Residual-based Adaptive Distribution (RAD) sampling for PINNs.

This module implements adaptive collocation point sampling strategies for physics-informed neural networks, based on PDE residual magnitudes.

Key Features

RAD: Residual-weighted sampling distribution
RAR-D: Residual-based Adaptive Refinement with Distribution

The key insight is that collocation points should be concentrated in regions where the PDE residual is high, as these are the areas where the neural network needs more training signal.

References

Survey Section 5.2: Residual-based Adaptive Sampling
Wu et al. (2023): A full study of non-adaptive and residual-based adaptive sampling for physics-informed neural networks

RADSampler ¶

RADSampler(config: RADConfig | None = None)

Residual-based Adaptive Distribution sampler.

This sampler draws collocation points from the domain with probability proportional to the PDE residual magnitude.

Attributes:

Name	Type	Description
`config`		RAD configuration

Example

sampler = RADSampler() domain_points = jnp.linspace(0, 1, 100).reshape(-1, 1) residuals = compute_pde_residual(model, domain_points) key = jax.random.key(0) sampled = sampler.sample(domain_points, residuals, batch_size=32, key=key)

Parameters:

Name	Type	Description	Default
`config`	`RADConfig \| None`	RAD configuration. Uses defaults if None.	`None`

sample ¶

sample(domain_points: Float[Array, ...], residuals: Float[Array, ...], batch_size: int, key: PRNGKeyArray) -> Float[Array, ...]

Sample collocation points based on residual distribution.

Parameters:

Name	Type	Description	Default
`domain_points`	`Float[Array, ...]`	Candidate points in the domain	required
`residuals`	`Float[Array, ...]`	PDE residual magnitudes at each point	required
`batch_size`	`int`	Number of points to sample	required
`key`	`PRNGKeyArray`	JAX random key	required

Returns:

Type	Description
`Float[Array, ...]`	Sampled collocation points

compute_weights ¶

compute_weights(residuals: Float[Array, ...]) -> Float[Array, ...]

Compute importance weights for residual-weighted loss.

These weights can be used to weight the loss function instead of resampling the collocation points.

Parameters:

Name	Type	Description	Default
`residuals`	`Float[Array, ...]`	PDE residual magnitudes	required

Returns:

Type	Description
`Float[Array, ...]`	Importance weights

RADConfig `dataclass` ¶

RADConfig(beta: float = 1.0, resample_frequency: int = 100, min_probability: float = 1e-06, temperature: float = 1.0)

Configuration for Residual-based Adaptive Distribution sampling.

Attributes:

Name	Type	Description
`beta`	`float`	Exponent for residual weighting. Higher values concentrate sampling more strongly on high-residual regions. ξ_j = \|r_j\|^β / Σ_k \|r_k\|^β
`resample_frequency`	`int`	Number of training steps between resampling
`min_probability`	`float`	Minimum sampling probability to ensure coverage
`temperature`	`float`	Temperature for probability smoothing

RARDRefiner ¶

RARDRefiner(config: RARDConfig | None = None, num_new_points: int | None = None, noise_scale: float | None = None)

Residual-based Adaptive Refinement with Distribution.

This refiner adds new collocation points near regions with high PDE residual, adaptively increasing resolution where needed.

Attributes:

Name	Type	Description
`config`		RAR-D configuration

Example

refiner = RARDRefiner(num_new_points=20) refined_points = refiner.refine(points, residuals, bounds, key)

Parameters:

Name	Type	Description	Default
`config`	`RARDConfig \| None`	RAR-D configuration. Uses defaults if None.	`None`
`num_new_points`	`int \| None`	Override for number of new points	`None`
`noise_scale`	`float \| None`	Override for noise scale	`None`

refine ¶

refine(current_points: Float[Array, ...], residuals: Float[Array, ...], bounds: Float[Array, 'dim 2'], key: PRNGKeyArray) -> Float[Array, ...]

Add new points near high-residual regions.

Parameters:

Name	Type	Description	Default
`current_points`	`Float[Array, ...]`	Current collocation points	required
`residuals`	`Float[Array, ...]`	PDE residual magnitudes at each point	required
`bounds`	`Float[Array, 'dim 2']`	Domain bounds, shape (dim, 2) with [min, max]	required
`key`	`PRNGKeyArray`	JAX random key	required

Returns:

Type	Description
`Float[Array, ...]`	Refined point set including new points

identify_refinement_regions ¶

identify_refinement_regions(residuals: Float[Array, ...]) -> Float[Array, ...]

Identify which points are in refinement regions.

Parameters:

Name	Type	Description	Default
`residuals`	`Float[Array, ...]`	PDE residual magnitudes	required

Returns:

Type	Description
`Float[Array, ...]`	Boolean mask indicating refinement regions

RARDConfig `dataclass` ¶

RARDConfig(num_new_points: int = 10, percentile_threshold: float = 90.0, noise_scale: float = 0.1)

Configuration for RAR-D refinement.

Attributes:

Name	Type	Description
`num_new_points`	`int`	Number of new points to add per refinement
`percentile_threshold`	`float`	Only refine near points above this percentile
`noise_scale`	`float`	Scale of random perturbation for new points

compute_sampling_distribution ¶

compute_sampling_distribution(residuals: Float[Array, ...], beta: float = 1.0, min_probability: float = 1e-06) -> Float[Array, ...]

Compute sampling distribution from residual magnitudes.

The sampling probability for each point is proportional to the residual magnitude raised to the power beta: p_j = |r_j|^β / Σ_k |r_k|^β

Parameters:

Name	Type	Description	Default
`residuals`	`Float[Array, ...]`	PDE residual magnitudes at each collocation point	required
`beta`	`float`	Exponent for residual weighting	`1.0`
`min_probability`	`float`	Minimum probability to ensure all points have some chance of being sampled	`1e-06`

Returns:

Type	Description
`Float[Array, ...]`	Sampling probabilities that sum to 1

For detailed algorithms and best practices, see the Adaptive Sampling Guide.

Training API Documentation¶

Overview¶

Core Classes¶

Trainer ⭐ RECOMMENDED¶

BasicTrainer¶

ModularTrainer ✅ NEW¶

Components¶

ErrorRecoveryManager ✅ NEW¶

FlexibleOptimizerFactory ✅ NEW¶

MetricsCollector ✅ NEW¶

TrainingComponent ✅ NEW¶

Configuration Classes¶

TrainingConfig¶

TrainingState¶

Physics-Informed Training¶

PhysicsInformedLoss¶

Uncertainty-Aware Trainers¶

UncertaintyGuidedTrainer¶

MultiFidelityUncertaintyTrainer¶

ActiveUncertaintyLearner¶

Usage Examples¶

Basic Training Workflow¶

Modular Training¶

Physics-Informed Training¶

Integration¶

With Neural Networks¶

Best Practices¶

Production Training¶

Physics-Informed Training¶

Performance Optimization¶

Troubleshooting¶

Common Issues¶

Debug Features¶

Multilevel Training¶

Width-Based Hierarchy (MLPs)¶

opifex.core.training.strategies.multilevel.coarse_to_fine ¶

MultilevelConfig dataclass ¶

create_network_hierarchy ¶

prolongate ¶

restrict ¶

Mode-Based Hierarchy (FNOs)¶

opifex.core.training.strategies.multilevel.multilevel_fno ¶

MultilevelFNOTrainer ¶

get_current_model ¶

advance_level ¶

is_at_finest ¶

get_epochs_for_current_level ¶

MultilevelFNOConfig dataclass ¶

create_fno_hierarchy ¶

create_mode_hierarchy ¶

prolongate_fno_modes ¶

restrict_fno_modes ¶

Adaptive Sampling¶

opifex.core.training.components.adaptive_sampling ¶

RADSampler ¶

sample ¶

compute_weights ¶

RADConfig dataclass ¶

RARDRefiner ¶

refine ¶

identify_refinement_regions ¶

RARDConfig dataclass ¶

compute_sampling_distribution ¶

`UncertaintyGuidedTrainer`¶

`MultiFidelityUncertaintyTrainer`¶

`ActiveUncertaintyLearner`¶

MultilevelConfig `dataclass` ¶

MultilevelFNOConfig `dataclass` ¶

RADConfig `dataclass` ¶

RARDConfig `dataclass` ¶