Contribute#

TRACMASS-JAX is developed on GitHub at TRACMASS/tracmass_jax. Bug reports, questions and pull requests are welcome. For the Fortran code, see the TRACMASS contribution guide.

Development setup#

git clone https://github.com/TRACMASS/tracmass_jax.git
cd tracmass_jax
uv sync                       # or: pixi install
uv run pytest tests/          # fast tests
uv run pytest tests/ --slow   # also the benchmarks and data-dependent tests

Tests marked slow need GLORYS data or run long benchmarks. tests/conftest.py neutralises the run-output settings (monitor backend, automatic recording, move_vertically), so a tracmass.toml in your working directory cannot redirect test output.

Coding standards#

The code follows NumPy conventions:

  • NumPy-style docstrings for all public functions, methods and classes, with Parameters, Returns, Raises, Notes and Examples sections. They are rendered in the API reference.

  • Type annotations on function signatures.

  • snake_case for functions, variables and modules, PascalCase for classes.

  • At most 88 characters per line (Black-compatible).

  • jnp inside JAX-compiled functions, np for host-side setup and tests.

  • Array shape comments next to field declarations, e.g. uflux: jnp.ndarray  # (imt, jmt, km, 2).

JAX notes#

  • jax.tree_util.register_dataclass does not add _replace, so GridFields and FluxFields define it with dataclasses.replace.

  • no_wflux is a Python bool fixed when the kernel is built (piecewise(grid, no_wflux=...)), so it selects a compiled variant.

  • The kernels close over grid. Anything in GridFields becomes a compile-time constant, and a different grid means a new compilation. Flux data is a dynamic argument.

  • Keep shapes stable inside loops. A shape that changes from window to window (for example a growing particle array) recompiles the kernel.

Building the documentation#

The documentation is built with Sphinx and MyST from the Markdown files in docs/. The API pages are generated from the docstrings.

uv run --with-requirements docs/requirements.txt \
    sphinx-build -b html docs docs/_build/html
# open docs/_build/html/index.html

or, with the requirements installed, make -C docs html. make -C docs livehtml rebuilds on save, if sphinx-autobuild is installed.

docs/conf.py imports the package from src/, so the build does not install it and needs only CPU JAX.

Hosting#

The same sources build on both hosting services:

  • GitHub Pages: .github/workflows/docs.yml builds the site on every push to main and deploys it with GitHub Actions. Turn it on once under Settings → Pages → Build and deployment → Source: GitHub Actions. The site is then at https://tracmass.github.io/tracmass_jax/.

  • Read the Docs: .readthedocs.yaml configures the build. Import the repository at https://readthedocs.org/dashboard/import/, and Read the Docs builds every push, with versioned docs for tags and branches.

Writing pages#

  • Pages are Markdown (MyST). Math uses $...$ and $$...$$, notes use :::{note} fences, and cross-references are relative links such as [Seeding](tutorials/seeding.md).

  • Keep examples runnable. The quick start and the adapter tutorial were run when they were written; rerun them when the API changes.

  • Working notes that should not be published can stay in docs/. Add them to exclude_patterns in docs/conf.py.