Installation#

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#

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

2. Create an environment#

The repository supports two environment managers. Pick one.

uv reads pyproject.toml and uv.lock:

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:

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

pixi reads pixi.toml (conda-forge packages plus the JAX wheels from PyPI):

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:

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, or the fast part of the test suite:

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.