adasurrog 0.5.6

Automatic high-precision surrogate modeling for parameterized ODE systems.

AdaSurroG

AdaSurroG trains fast, differentiable surrogate models for parameterized ordinary differential equation systems. It is designed for repeated-query scientific workflows such as parameter studies, uncertainty quantification, optimization, and systems-biology simulation.

The default high-precision pipeline uses a fixed high-fidelity dataset and combines:

• automatic state transforms for multiscale variables;
• adaptive fixed output-time selection from pilot trajectories;
• proper orthogonal decomposition (POD) trajectory compression;
• a Chebyshev polynomial coefficient backbone;
• a residual PyTorch network;
• time/accuracy-bounded AdamW training;
• deterministic CPU float64 full-batch L-BFGS refinement;
• resumable atomic checkpoints and explicit accuracy diagnostics.

AdaSurroG trains one surrogate per mechanistic ODE model. The pipeline is generic: users provide the ODE definition, parameter bounds, initial conditions, and time interval rather than manually designing a new neural architecture for each model.

Installation

Install the universal ODE pipeline:

python -m pip install "adasurrog[universal]"

Install plotting utilities and all runtime extras:

python -m pip install "adasurrog[all]"

For development from a source checkout:

python -m pip install -e ".[dev]"

Python 3.10–3.13 is supported. PyTorch is a required dependency. SciPy is required for generating ODE reference trajectories, and PyYAML is required for YAML-based training configurations.

Quick start: built-in benchmark

List the included benchmark models:

adasurrog-benchmark --list

Train a surrogate for the SEIR benchmark:

adasurrog-benchmark \
  --problem seir \
  --train 192 \
  --validation 48 \
  --test 48 \
  --target-rmse 0.002 \
  --training-seconds 900 \
  --lbfgs-seconds 300

The pipeline writes its dataset, checkpoints, pilot diagnostics, and final certification report under the selected run directory.

YAML training

problem:
  factory: benchmark:seir

sampling:
  pilot_trajectories: 24
  train_trajectories: 192
  validation_trajectories: 48
  test_trajectories: 48
  time_points: 128
  seed: 17

accuracy:
  target_rmse: 0.002
  target_p99: 0.01
  pod_floor_fraction: 0.30

budget:
  training_seconds: 900
  lbfgs_seconds: 300

surrogate:
  weighted_pod: auto
  rank_candidates: [4, 8, 12, 16, 24, 32, 48, 64]
  polynomial_degrees: [1, 2, 3, 4]

runtime:
  run_directory: runs/universal_seir
  device: mps
  lbfgs_device: cpu
  resume: none

storage:
  mode: compact
  dataset_dtype: float32
  checkpoint_dtype: float32
  keep_dataset: true
  minimum_free_disk_gb: 5.0

Run it with:

adasurrog-train examples/universal/seir.yaml

On Apple Silicon, AdamW may run on MPS while float64 L-BFGS runs on CPU.

Define a custom ODE problem

import numpy as np

from adasurrog import ODEProblem


def create_problem() -> ODEProblem:
    def rhs(time, state, parameters):
        del time
        production, degradation = parameters
        return np.array([production - degradation * state[0]])

    def jacobian_state(time, state, parameters):
        del time, state
        return np.array([[-parameters[1]]])

    return ODEProblem(
        name="production_degradation",
        state_names=("concentration",),
        parameter_names=("production", "degradation"),
        parameter_lower=np.array([0.1, 0.05]),
        parameter_upper=np.array([10.0, 2.0]),
        time_span=(0.0, 20.0),
        rhs_function=rhs,
        initial_condition_function=lambda parameters: np.array([0.0]),
        jacobian_state_function=jacobian_state,
    )

Reference the factory from YAML:

problem:
  factory: my_models.production_degradation:create_problem

Parameters currently require strictly positive lower bounds because the default pipeline uses a normalized log-parameter box.

Python API

from pathlib import Path

from adasurrog import (
    AccuracyConfig,
    AutoPipelineConfig,
    BudgetConfig,
    SamplingConfig,
    UniversalODEPipeline,
)
from my_models.production_degradation import create_problem

config = AutoPipelineConfig(
    sampling=SamplingConfig(
        pilot_trajectories=16,
        train_trajectories=128,
        validation_trajectories=32,
        test_trajectories=32,
        time_points=96,
    ),
    accuracy=AccuracyConfig(target_rmse=1e-3),
    budget=BudgetConfig(training_seconds=600, lbfgs_seconds=180),
    run_directory=Path("runs/production_degradation"),
)

surrogate, report = UniversalODEPipeline(create_problem(), config).fit()
print(report.status, report.validation_rmse, report.test_rmse)

Load and query a trained surrogate

import numpy as np

from adasurrog import AdaSurrogate

surrogate = AdaSurrogate.load(
    "runs/production_degradation/checkpoints/universal_best.pt",
    device="cpu",
)

trajectory = surrogate.predict(
    parameters=np.array([2.0, 0.4]),
)

selected_times = surrogate.predict(
    parameters=np.array([2.0, 0.4]),
    time=np.linspace(0.0, 20.0, 50),
)

check = surrogate.check_domain(np.array([2.0, 0.4]))
print(check.inside)

Physical query times are accepted directly. A checkpoint contains the POD representation, coefficient model, state transforms, parameter domain, time grid, and training report; the original training dataset is not required for inference.

Stopping, checkpoint, and disk behavior

Training is controlled by a validation target and cumulative wall-clock budget instead of a fixed epoch count. Since version 0.5.2, storage.mode: compact is the default. Compact mode keeps only the float32 deployment checkpoint, fixed dataset, and diagnostics:

runs/<problem>/
├── data/fixed_dataset.npz
├── pilot_diagnostics.json
├── universal_diagnostics.json
└── checkpoints/
    └── universal_best.pt

It does not retain AdamW/L-BFGS optimizer histories, periodic checkpoints, terminal duplicates, or exception checkpoints. This is the recommended setting for benchmark matrices. Interrupted compact runs restart from the same fixed Sobol design.

Use storage.mode: resumable when mid-run resume is required. It keeps one latest optimizer-state checkpoint but no periodic history. storage.mode: full preserves the legacy debugging behavior.

Remove redundant checkpoints produced by older releases with a dry run followed by deletion:

adasurrog-clean-runs runs --show-files
adasurrog-clean-runs runs --apply

Add --remove-datasets only when regenerated high-fidelity data are not needed.

Inspect a checkpoint:

adasurrog checkpoint-info runs/my_model/checkpoints/universal_best.pt

Inspect the installed environment:

adasurrog info

Accuracy diagnostics

The final report classifies a run as one of:

• certified (legacy label meaning the configured validation target was met; it is not a formal certificate);
• representation_limited;
• regression_limited;
• data_limited;
• optimization_limited;
• tail_error_limited.

The report includes validation and test errors, explicit validation_target_met and test_target_met booleans, the POD representation floor, selected POD rank, selected polynomial model, and regression amplification over the POD floor. The legacy status is derived from validation diagnostics only; it is not a statistical guarantee or a certificate of held-out test accuracy. A failure to reach a target is therefore reported diagnostically rather than hidden behind a final loss value.

Included benchmark models

• Robertson kinetics: severe stiffness, conservation, and scale separation;
• Lotka–Volterra: nonlinear oscillation and phase sensitivity;
• Brusselator: limit-cycle dynamics and nearby bifurcation behavior;
• SEIR: nonnegative compartment dynamics, conservation, and a transient peak.

Legacy and experimental commands

The distribution retains the gene-feedback examples and the earlier pool-based active-learning prototype for reproducibility. The recommended production path is adasurrog-train or the UniversalODEPipeline; active learning is not part of the default automatic pipeline.

Experimental trajectory CorrectNet

AdaSurroG 0.5.6 provides a second-stage GatedCorrectNet for structured errors that remain after BaseNet training. The BaseNet remains frozen. CorrectNet receives the projected BaseNet trajectory, its finite-difference ODE defect, normalized parameters, and time coordinates. It predicts a gated residual trajectory.

State geometry is explicit rather than hard-coded:

• signed: negative and positive states are allowed;
• positive: nonnegative states;
• box: fixed lower and upper bounds;
• mixed: per-state combinations of signed, positive, and box domains;
• none or linear invariants of the form C y = b.

Correction geometry is selected explicitly: simplex_logit for equal-total nonnegative systems, positive_log_projected for nonnegative non-conserved systems, and additive_projected for signed/mixed states and general linear invariants. Candidate trajectories are projected to the declared convex feasible set and blended through the learned gate. A positive zero-lower-bound domain with C = [1, ..., 1] reproduces the conserved-simplex use case. Nonlinear invariants are not included in 0.5.6. See CORRECTNET.md and examples/correctnet/.

Current scope

Version 0.5.6 supports continuous parameterized ODE initial-value problems with strictly positive parameter bounds and fixed simulation horizons. Global unweighted or metric-weighted POD is used for trajectory representation. Event resets, discontinuous controls, automatic state/time-blocked POD, nonlinear latent representations, derivative-supervised training, and nonlinear invariant projection are not yet automatic. The experimental CorrectNet helper currently accepts autonomous batched RHS functions on a fixed shared output grid and supports signed or bounded states with optional static linear invariants. Deployable target="initial" invariants require an exact linear_invariant_target_torch(parameters) problem function.

Reproducibility

sampling.seed now seeds Python, NumPy, PyTorch, the Sobol design, and shuffled training batches in the universal pipeline. Set sampling.deterministic_algorithms: true to require deterministic PyTorch algorithms where they are available; this can reject unsupported accelerator operations.

Citation

See CITATION.cff for citation metadata. Release history is recorded in CHANGELOG.md.

License

AdaSurroG is distributed under the MIT License.