Installation and Setup

JaxDEM is not yet published on PyPI, so install it from the GitHub repository (or clone the repository and install locally). pyproject.toml defines the optional extras.

Create a Virtual Environment (Recommended)

python -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip

On Windows (PowerShell), activation is:

.\.venv\Scripts\Activate.ps1

Install JaxDEM from GitHub

Basic install:

python -m pip install "git+https://github.com/cdelv/JaxDEM.git"

The basic install pulls in everything needed to run core simulations (jax, numpy, vtk, orbax, and optax) — import jaxdem works out of the box. Extras only add GPU/TPU backends or reinforcement learning dependencies (flax, distrax, tqdm, scipy, tensorflow).

With CUDA 13 support:

python -m pip install "git+https://github.com/cdelv/JaxDEM.git#egg=JaxDEM[cuda13]"

With RL capabilities:

python -m pip install "git+https://github.com/cdelv/JaxDEM.git#egg=JaxDEM[rl]"

You can combine extras as needed:

python -m pip install "git+https://github.com/cdelv/JaxDEM.git#egg=JaxDEM[cuda13,rl]"

Install by Cloning (Alternative)

You can also clone first, then install locally:

git clone https://github.com/cdelv/JaxDEM.git
cd JaxDEM
python -m pip install .

With extras:

python -m pip install ".[rl]"
python -m pip install ".[cuda13,rl]"

Next Steps

After installation succeeds, continue with:

• Introduction

• System Guide

• Materials Guide