jaxhybridmodels

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.

Bounded predictors

Declare physical input and output ranges. The predictor works in latent space and maps its output back into those ranges.

Two training loops

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

PyTree composition

Pass one predictor or a nested tuple, dictionary, or NamedTuple. There is no required model wrapper class.

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, run the known-answer sanity check first:

bash

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

It recovers the frequency of a synthetic harmonic oscillator from noisy position measurements. The Getting started page then builds the same pipeline step by step, and the Examples page shows where to go for a real hybrid ODE.

In a nutshell

jaxhybridmodels keeps the ODE simulation function in user code. The Python package is imported as jaxhybridmodels. The library adds the surrounding data, predictor, and training machinery.

Part	Responsibility
`predictors`	An Equinox PyTree of trainable functions.
`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.

Part

Responsibility

predictors

An Equinox PyTree of trainable functions.

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.

The core training call is:

python

import jaxhybridmodels as hm

dataset = hm.make_dataset(experiments, output_channel_names=("value",))
history, trained = hm.train_with_optax(
    predictors,
    dataset,
    config,
    simulate_fn=simulate_fn,
    state_to_output=state_to_output,
    solver=solver,
    key=key,
)

The variables in this short call are defined in the complete Getting started example. Read Concepts for the data and model interfaces.

If you want to...	Start here	Main API
Build irregular experiments and datasets	Data and buckets	`make_experiment`, `make_dataset`
Write the ODE and observation map	Model interface	`simulate_fn`, `state_to_output`
Choose a bounded predictor	Predictors and bounds	`BoundedPredictor`, `BoundScaler`
Train with gradients or population search	Training	`train_with_optax`, `train_with_evosax`
Add time-varying inputs or schedules	Profiles and schedules	`ramp_profile`, `annealing_schedule`
Save, evaluate, or ensemble models	Saving and loading, Ensembles	`save_run`, `predict_dataset`

If you want to...

Start here

Main API

Build irregular experiments and datasets

Data and buckets

make_experiment, make_dataset

Write the ODE and observation map

Model interface

simulate_fn, state_to_output

Choose a bounded predictor

Predictors and bounds

BoundedPredictor, BoundScaler