jax-morpho
Status: alpha (0.1.x). The API is being designed in the open and may change without notice until 0.2.
Differentiable, GPU-scale developmental morphogenesis in JAX. Cells are points with a Morse adhesion/repulsion potential; relaxation is autodiff gradient descent, and growth is cell division. An evolvable genome grows a tissue form — and gradients reach all the way back to the genome.
Why
Mature tissue simulators (tyssue, Chaste, CompuCell3D) are CPU vertex- or
Potts-based and forward-only. The differentiable-morphogenesis work (Growing
NCA, VertAX, Deshpande et al. 2024) is autodiff-native but tops out at ~10²
cells. jax-morpho occupies the gap: a differentiable, GPU-scale,
genome-driven center-based engine.
- Reproduces real-epithelium statistics. Coupling proliferation to relaxation lands the polygon-side distribution on the Gibson-2006 epithelium (~50% hexagons), tunable between random (Poisson–Voronoi) and crystalline packings by the relaxation-per-division ratio.
- Scales. The same Morse energy over a
jax_mdcell-list reaches 1–2 million cells on a single GPU (~32 ms/step on an RTX 4090), verified to the same energy minimum as the dense O(N²) path. - Differentiable to the genome. Autodiff through the whole relaxation recovers the mechanical program — or the genome — that produces a target form (gradient-based inverse design / evo-devo), beating gradient-free search at equal budget.
Install
pip install jax-morpho # CPU jax
pip install 'jax[cuda12]' # add GPU support (CUDA 12)
pip install 'jax-morpho[scale]' # + jax_md neighbor lists for 10^5–10^6 cells
Quickstart
import jax_morpho as jm
# Grow a tissue by coupled division + relaxation; measure its packing.
pos, alive = jm.grow_relax(n_max=1400, n_start=20, target=1000, relax_steps=8)
sides = jm.interior_side_counts(pos, alive)
print("hexagon fraction:", jm.side_distribution(sides)[6]) # ~0.5, near Gibson
# Invert a target form into the genome that grows it.
# (see examples/demo_genome_mechanics.py)
Runnable demos:
python examples/demo_center_based.py # emergent Gibson statistics + differentiability
python examples/demo_inverse_design.py # gradient-based shape inverse design
python examples/demo_genome_mechanics.py # genome -> form, gradients to the genome
python examples/demo_gpu_scaling.py # 1–2M cells on GPU (needs [scale] + jax[cuda12])
Modules
| Module | What |
|---|---|
jax_morpho.center_based |
Morse engine: relaxation, division, growth, topology + shape descriptors |
jax_morpho.scale |
jax_md cell-list relaxation for O(N) scaling (needs [scale]) |
jax_morpho.inverse_design |
gradient-based inverse design of tissue shape |
jax_morpho.genome |
genome → mechanics → form (differentiable to the genome) |
jax_morpho.stats |
epithelial topology references (Poisson–Voronoi, Gibson) + L1 distance |
Development
pip install -e '.[test]'
pytest -q -n auto --cov=jax_morpho
Citing
See CITATION.cff. A Zenodo DOI is minted on the first
tagged release.
License
MIT — see LICENSE.
1 maintainer
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file jax_morpho-0.3.0.tar.gz.
File metadata
- Download URL: jax_morpho-0.3.0.tar.gz
- Upload date:
- Size: 172.6 kB
- Tags: Source
- Uploaded using Trusted Publishing? Yes
- Uploaded via: twine/7.0.0 CPython/3.13.14
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
3ea2eafdfd3ea18ca80f8e8b15493350009a29f2a17a751e6e756c1e572761c4
|
|
| MD5 |
04f9ed3ca184b38e48ce43f20bade97f
|
|
| BLAKE2b-256 |
046cf93a66d3a67ae1a242ad7ff3b168dcd1d5543c50c3ebf837a24ed0d368b0
|
Provenance
The following attestation bundles were made for jax_morpho-0.3.0.tar.gz:
Publisher:
publish-pypi.yml on JimGalasyn/jax-morpho
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
jax_morpho-0.3.0.tar.gz -
Subject digest:
3ea2eafdfd3ea18ca80f8e8b15493350009a29f2a17a751e6e756c1e572761c4 - Sigstore transparency entry: 2342042778
- Sigstore integration time:
-
Permalink:
JimGalasyn/jax-morpho@e7e85fe63dac7fa48fe87de926bdfddba18c1303 -
Branch / Tag:
refs/tags/v0.3.0 - Owner: https://github.com/JimGalasyn
-
Access:
public
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
publish-pypi.yml@e7e85fe63dac7fa48fe87de926bdfddba18c1303 -
Trigger Event:
release
-
Statement type:
File details
Details for the file jax_morpho-0.3.0-py3-none-any.whl.
File metadata
- Download URL: jax_morpho-0.3.0-py3-none-any.whl
- Upload date:
- Size: 87.1 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? Yes
- Uploaded via: twine/7.0.0 CPython/3.13.14
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
758bf4b79bf0a20cc6ec3ac2f24ef6e8b49b2efe5b8ffbe6870e17ecbc8b8a63
|
|
| MD5 |
1aa0485860b736df32de00cf10a9c770
|
|
| BLAKE2b-256 |
c0706aee85c624d9049c01d56113ebd789d15f006db0557b86c81ac1e2552449
|
Provenance
The following attestation bundles were made for jax_morpho-0.3.0-py3-none-any.whl:
Publisher:
publish-pypi.yml on JimGalasyn/jax-morpho
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
jax_morpho-0.3.0-py3-none-any.whl -
Subject digest:
758bf4b79bf0a20cc6ec3ac2f24ef6e8b49b2efe5b8ffbe6870e17ecbc8b8a63 - Sigstore transparency entry: 2342042782
- Sigstore integration time:
-
Permalink:
JimGalasyn/jax-morpho@e7e85fe63dac7fa48fe87de926bdfddba18c1303 -
Branch / Tag:
refs/tags/v0.3.0 - Owner: https://github.com/JimGalasyn
-
Access:
public
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
publish-pypi.yml@e7e85fe63dac7fa48fe87de926bdfddba18c1303 -
Trigger Event:
release
-
Statement type:
