jaxhybridmodelsTrain functions inside ODE models

A JAX library for combining user-written differential equations with trainable predictors, bounded physical quantities, and irregular time-series data.

Get started

Browse examples

API reference

User-owned dynamics

Write one simulate_fn for a single experiment. The library applies batching, JIT compilation, and differentiation around it.

Irregular observations

Give each channel its own timestamps. The data layer builds masks and groups experiments by timestamp-axis length.

Physical bounds

Declare a range for a parameter or a network's inputs and output. Training works on an unbounded latent and can never leave the range.

Two training loops

Use Optax for gradient-based fitting or Evosax for population-based search over small parameter sets.

PyTree composition

Predictors are any PyTree of Equinox modules, usually a dict keyed by role. HybridModel binds them to the simulator; the functional layer takes them directly.

Explicit extension points

Replace the predictor, loss, regulariser, solver, UI, or training loop with a callable or PyTree of your own.

Start here

From a repository checkout:

bash

uv sync --extra examples
uv run python examples/pendulum/train_harmonic.py

It fits the frequency of a harmonic oscillator and predicts with the fitted model. Getting started walks through the same code, and Examples lists the larger hybrid models.

In a nutshell

You write the ODE solve for one experiment. The library batches it over irregular data, differentiates through it, and trains the predictors.

Part	Responsibility
`predictors`	A PyTree of trainable values and networks.
`simulate_fn`	Integrates one experiment and returns the full state.
`state_to_output`	Selects the observed channels from the full state.
`SolverConfig`	Stores the Diffrax solver and its settings.
`Dataset`	Stores bucketed observations and masks.

HybridModel binds the first four; the data is passed per call, as a list of experiments or a Dataset:

python

import jaxhybridmodels as hm

model = hm.HybridModel(
    predictors=predictors,
    simulate_fn=simulate_fn,
    state_to_output=state_to_output,
    output_channel_names=("conc",),
)
result = model.fit(experiments, hm.OptaxTrainingConfig(steps=500, lr=1e-2), key=key)
predictions = result.model.predict(experiments)
hm.print_metrics(predictions)

Getting started defines every variable here. Concepts explains the data and model interfaces.

Find your path

If you want to...	Start here	Main API
Build irregular experiments and datasets	Data and buckets	`make_experiment`, `experiments_from_dataframe`
Write the ODE and observation map	Model interface	`HybridModel`
Choose parameters, networks and bounds	Predictors and bounds	`Parameter`, `BoundedPredictor`
Train with gradients or population search	Training	`Objective`, `OptaxTrainingConfig`, `EvosaxTrainingConfig`
Add time-varying inputs or schedules	Profiles and schedules	`ramp_profile`, `annealing_schedule`
Save, evaluate, or ensemble models	Saving and loading, Ensembles	`save`, `compute_metrics`

Install

Python 3.11 or newer: