Metadata-Version: 2.4
Name: jax-morph
Version: 0.4.0
Summary: Differentiable particle-based physics for proliferating cells and active matter, in JAX.
Project-URL: Homepage, https://github.com/fmottes/jax-morph
Project-URL: Repository, https://github.com/fmottes/jax-morph
Project-URL: Documentation, https://fmottes.github.io/jax-morph/
Project-URL: Issues, https://github.com/fmottes/jax-morph/issues
Author-email: Francesco Mottes <fmottes@seas.harvard.edu>, Ramya Deshpande <rdeshpande@seas.harvard.edu>
License: Apache-2.0
License-File: LICENSE
Keywords: active matter,cell mechanics,differentiable programming,jax,morphogenesis,particle-based physics
Classifier: Development Status :: 4 - Beta
Classifier: Programming Language :: Python
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: Programming Language :: Python :: Implementation :: CPython
Requires-Python: >=3.11
Requires-Dist: diffrax>=0.6.0
Requires-Dist: equinox>=0.11.7
Requires-Dist: jax>=0.4.35
Provides-Extra: neighbors
Requires-Dist: jax-md>=0.2.8; extra == 'neighbors'
Provides-Extra: viz
Requires-Dist: matplotlib>=3.8; extra == 'viz'
Description-Content-Type: text/markdown

> [!IMPORTANT]
> The code to reproduce the results in
> [*Engineering morphogenesis of cell clusters with differentiable programming*](https://doi.org/10.1038/s43588-025-00851-4)
> is available on the [`paper-natcompsci-2025` branch](https://github.com/fmottes/jax-morph/tree/paper-natcompsci-2025).

<p align="center">
  <img src="docs/_static/images/jax-morph-wordmark-white.png" alt="jax-morph" width="720">
</p>

Differentiable particle-based physics for proliferating cells and active matter, built with JAX and Equinox.

- **Differentiable end-to-end** — pathwise gradients through the continuous physics, score-function gradients through discrete events.
- **Composable physics steps** — mechanics, diffusion, Brownian and active-Brownian dynamics, growth, division, and death.
- **Principled time-stepping** — steps compose into one macro-step by Lie-Trotter operator splitting: each advances the shared state over `dt`, giving a consistent, first-order-accurate integrator of the coupled dynamics.
- **Continuous-time dynamics and control** — neural-ODE and gene-network controllers that model cell decision making from local cues.
- **Core abstractions built for easy extension** — add your own physics steps to the pipeline. Guides provided with the library for coding agents reference.
- **JAX and Equinox native** — `jit`, `vmap`, Equinox filtered transformations and neural-network modules work throughout the pipeline; GPU support out of the box.

See [Key Concepts](docs/concepts.md) and usage guides for more.

---
## Installation

jax-morph requires Python 3.11 or later.

```bash
pip install jax-morph
```

Or with uv:

```bash
uv add jax-morph
```

The visualization API is an optional matplotlib-backed extra. Install it when using
`jxm.viz.draw`, `jxm.viz.animate`, or `jxm.viz.plot_timeseries`:

```bash
pip install 'jax-morph[viz]'
```

Or with uv:

```bash
uv add 'jax-morph[viz]'
```

The base package does not install matplotlib; `import jax_morph.viz` remains available, and a
rendering call explains how to add the extra if it is missing.

---
## Quickstart

This model seeds a single cell that diffuses under a pairwise interaction and stochastically divides:

```python
import jax

import jax_morph as jxm
from jax_morph.physics import BrownianDynamics, Division, SoftSphere

# Cells diffuses under a soft-sphere interaction and divide stochastically
model = jxm.Model([
    BrownianDynamics(SoftSphere(), n_space_dim=2, kT=0.05),
    Division(n_space_dim=2),
])

# Build state class and initialize with a single cell
StateClass = jxm.build_state_from_model(model)
state = StateClass.init_empty(capacity=32, n_space_dim=2, n_types=1)
state = state.update(
    alive=state.alive.at[0].set(True),
    radius=state.radius.at[0].set(0.5),
    celltype=state.celltype.at[0, 0].set(1.0),
    division_rate=state.division_rate.at[0].set(1.0),
)

# Simulate
sim_key = jax.random.PRNGKey(0)
final_state = jxm.simulate(model, state, n_steps=20, dt=0.1, key=sim_key)
```

See the [documentation](docs/index.md) and [example notebooks](examples/README.md) for more.

---
## Installed usage guides

Usage guides ship with the library so they remain available with a PyPI installation and match the
installed API version:

```python
import jax_morph as jxm

jxm.guides.list_guides()
jxm.guides.guide('extending')
jxm.guides.guide('optimization/pathwise')
```

---
## Reference

If you use Jax-Morph, please cite:

```
@article{deshpandemottes2025,
  title={Engineering morphogenesis of cell clusters with differentiable programming},
  author={Deshpande, Ramya and Mottes, Francesco and Vlad, Ariana-Dalia and Brenner, Michael P and Dal Co, Alma},
  journal   = {Nature Computational Science},
  year      = {2025},
  doi       = {10.1038/s43588-025-00851-4},
  url       = {https://doi.org/10.1038/s43588-025-00851-4}
}
```

The published paper can be read [here](https://rdcu.be/eAyqo).
