# Contribute

TRACMASS-JAX is developed on GitHub at
<https://github.com/TRACMASS/tracmass_jax>. Bug reports, questions and pull
requests are welcome. For the Fortran code, see the
[TRACMASS contribution guide](https://www.tracmass.org/docs/contribute/contribute.html).

## Development setup

```bash
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](api/index.md).
- 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](https://www.sphinx-doc.org) and
[MyST](https://myst-parser.readthedocs.io) from the Markdown files in
`docs/`. The API pages are generated from the docstrings.

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