# Installation

:::{admonition} Technical note
TRACMASS-JAX is a Python package (Python ≥ 3.12) built on JAX, `xarray`,
Dask and Zarr. It runs on Linux and macOS. On Linux, the default install
uses JAX's CUDA 13 build and runs on an NVIDIA GPU when one is available,
falling back to the CPU otherwise. On macOS it runs on the CPU.
:::

## 1. Get the code

```bash
git clone https://github.com/TRACMASS/tracmass_jax.git
cd tracmass_jax
```

## 2. Create an environment

The repository supports two environment managers. Pick one.

::::{tab-set}

:::{tab-item} uv (recommended)
[uv](https://docs.astral.sh/uv/) reads `pyproject.toml` and `uv.lock`:

```bash
uv sync                 # creates .venv with tracmass installed (editable)
uv run python -c "import tracmass, jax; print(jax.devices())"
```

Run scripts and the test suite through `uv run`:

```bash
uv run python my_run.py
uv run pytest tests/
uv run pytest tests/test_ned_grid.py -v
```
:::

:::{tab-item} pixi
[pixi](https://pixi.sh) reads `pixi.toml` (conda-forge packages plus the
JAX wheels from PyPI):

```bash
pixi install
pixi run python -c "import tracmass, jax; print(jax.devices())"
pixi run -e test pytest tests/
```

The pixi environment sets `JAX_ENABLE_X64=True`, so JAX uses 64-bit floats
by default.
:::

::::

When `tracmass` is imported it prints the free GPU memory, if it can query
the GPU:

```text
GPU: 11870 MiB free / 12282 MiB total
```

If the NVIDIA management library is missing or cannot talk to the driver,
you see `GPU: NVML unavailable (...)` and the import carries on.

## 3. GPU notes

- **CUDA comes from the JAX wheels.** `jax[cuda13]` pulls in its own CUDA
  runtime from PyPI, so you need neither a system CUDA toolkit nor `nvcc`.
  You only need an NVIDIA driver that is new enough for CUDA 13.
- **Driver and library mismatch.** If `import tracmass` reports
  `NVML/RM version mismatch`, or `jax.devices()` only lists the CPU after a
  driver upgrade, the kernel module and the user-space driver libraries
  (`libcuda`, `libnvidia-ml`) have different versions. Reinstall the driver
  package that matches the loaded kernel module (compare
  `cat /proc/driver/nvidia/version` with
  `ls -l /usr/lib/x86_64-linux-gnu/libnvidia-ml.so.1`), or reboot after a
  driver update.
- **Memory.** `tracmass` sets `XLA_PYTHON_CLIENT_PREALLOCATE=false`, so JAX
  allocates GPU memory as it needs it instead of reserving most of the card
  at start-up. Call `tracmass.free_device_memory()` between runs in a long
  interactive session.

## 4. Check the installation

Run the synthetic test case described in the [Quick start](quickstart.md),
or the fast part of the test suite:

```bash
uv run pytest tests/test_ned_grid.py tests/test_fortran_comparison.py
```

`tests/test_fortran_comparison.py` reruns the *Theoretical* project of the
Fortran code and compares the trajectory with the Fortran output. The long
benchmarks and the tests that need GLORYS data are marked `slow` and are
skipped unless you pass `--slow` (from `pytest-skip-slow`).

## Data for the reanalysis adapters

The synthetic adapters (`theoretical`, `fabboni`) need no data. The others
read data from local caches or remote servers:

| Adapter | Data source | Local cache |
|---|---|---|
| `glorys` | Copernicus Marine GLORYS12 reanalysis, daily Zarr stores | `/data/{disk}/mercator/{product}` |
| `ecco2` | ECCO2 cube92 daily fields over OPeNDAP (APDRC) | `/data/{disk}/ecco2/cube92_daily` |
| `scrippsgog` | Scripps Gulf of Guinea MITgcm run over HTTPS | `/data/raid/scrippsgog/PRELIM/daily/` |

The cache root is set by `settings["disk"]` in each `tracmass.seagrid`
module (default `"raid"`). See [GCM adapters](../components/adapters.md).