Metadata-Version: 2.5
Name: jax-hybridmodels
Version: 0.2.0b1
Summary: Hybrid ODE and neural models in JAX
Project-URL: Homepage, https://github.com/DanielePessina/jax-hybridmodels
Project-URL: Documentation, https://danielepessina.github.io/jax-hybridmodels/
Project-URL: Repository, https://github.com/DanielePessina/jax-hybridmodels
Project-URL: Issues, https://github.com/DanielePessina/jax-hybridmodels/issues
Author-email: Daniele Pessina <daniele.pessina18@imperial.ac.uk>
License-Expression: BSD-3-Clause
License-File: LICENSE
Keywords: diffrax,equinox,hybrid models,jax,neural ODE
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Topic :: Scientific/Engineering :: Mathematics
Classifier: Typing :: Typed
Requires-Python: >=3.11
Requires-Dist: diffrax>=0.7.2
Requires-Dist: equinox>=0.13.7
Requires-Dist: evosax>=0.2.0
Requires-Dist: flax<0.12.9,>=0.12.1
Requires-Dist: jax<0.11,>=0.10.0
Requires-Dist: jaxkan>=0.3.7
Requires-Dist: jaxlib<0.11,>=0.10.0
Requires-Dist: jaxtyping>=0.3.9
Requires-Dist: numpy>=2.4.4
Requires-Dist: optax>=0.2.8
Requires-Dist: rich>=15.0.0
Requires-Dist: scipy>=1.17.1
Provides-Extra: examples
Requires-Dist: matplotlib>=3.10.9; extra == 'examples'
Requires-Dist: openpyxl>=3.1.5; extra == 'examples'
Requires-Dist: pandas>=3.0.2; extra == 'examples'
Requires-Dist: python-libsbml>=5.20.4; extra == 'examples'
Requires-Dist: rlax>=0.1.9; extra == 'examples'
Requires-Dist: sympy>=1.14; extra == 'examples'
Description-Content-Type: text/markdown

# 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](https://danielepessina.github.io/jax-hybridmodels/guide/getting-started)
page, use the [Examples](https://danielepessina.github.io/jax-hybridmodels/examples/)
chooser to find a modelling pattern, and use the generated [API reference](https://danielepessina.github.io/jax-hybridmodels/api/)
for exact signatures.

| Need | Start here |
| --- | --- |
| Understand the model contract | [Model interface](https://danielepessina.github.io/jax-hybridmodels/guide/model-interface) |
| Build irregular data | [Data and buckets](https://danielepessina.github.io/jax-hybridmodels/guide/data) |
| Choose predictors and bounds | [Predictors and bounds](https://danielepessina.github.io/jax-hybridmodels/guide/predictors) |
| Train or freeze parameters | [Training](https://danielepessina.github.io/jax-hybridmodels/guide/training) |
| Extend the library | [Custom predictors](https://danielepessina.github.io/jax-hybridmodels/guide/custom-predictors) |

The site source lives under [`docs/`](./docs).

## First run

The examples are plain scripts. From a checkout:

```bash
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](https://danielepessina.github.io/jax-hybridmodels/examples/)
for the other workflows.

## Installation

### From PyPI

```bash
uv add jax-hybridmodels==0.2.0b1
```

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

### Development

```bash
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:

```bash
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`](./docs/README.md) for local preview, API generation,
and Pages deployment details.

## License

BSD-3-Clause. See [LICENSE](./LICENSE).
