jaxhybridmodels

A JAX/Equinox library for combining user-written ODE dynamics with trainable predictors. It supports bounded physical quantities, regular and irregular time-series experiments, and Optax or Evosax training. Crystallisation kinetics is the canonical example, not the scope.

Documentation

The documentation site is built with VitePress and deployed to GitHub Pages at https://danielepessina.github.io/jax-hybridmodels/.

Start with the Getting started page, use the Examples chooser to find a modelling pattern, and use the generated API reference for exact signatures.

The site source lives under docs/.

First run

The examples are plain scripts. From a checkout:

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

The harmonic oscillator has a known optimum, so it is a useful wiring check. See the examples overview for the other workflows.

Installation

From PyPI

uv add jax-hybridmodels==0.2.0b1

The PyPI distribution is named jax-hybridmodels; import it in Python as jaxhybridmodels.

Development

uv sync --extra examples

This installs jaxhybridmodels in editable mode together with jax, equinox, diffrax, optax, evosax, jaxkan, and the small CLI/UI dependencies, plus the optional dependencies used by the examples.

For the library only, use uv sync.

Contributing

Run the verification checks before pushing:

uv run ruff check .        # lint
uv run ty check src        # typecheck
uv run pytest -q           # test suite
npm --prefix docs run docs:check   # API sync + site build

See docs/README.md for local preview, API generation, and Pages deployment details.

License

BSD-3-Clause. See LICENSE.