Installation

Quick installation

The fastest way of installing this package is through PyPI or Conda:

TensorWaves can work with different computational backends. They are provided through optional dependencies. JAX is usually the fastest backend, so it’s recommended to install that as in the install examples here.

PyPI

python3 -m pip install tensorwaves[jax]

To install the package with support for amplitude analysis run:

python3 -m pip install tensorwaves[jax,pwa]

Conda

conda install -c conda-forge tensorwaves jax jaxlib

To install the package with support for amplitude analysis run:

conda install -c conda-forge ampform phasespace

This installs the latest release that you can find on the stable branch.

Optional dependencies can be installed as follows:

pip install tensorwaves[pwa]  # installs tensorwaves with ampform
pip install tensorwaves[jax,scipy,tf]
pip install tensorwaves[all]  # all runtime dependencies

GPU support

Computations with are fastest on a GPU. To get JAX and TensorFlow work with your graphics card, you need to install at least CUDA Toolkit 11.2 and cuDNN 8.1.

Below is an installation guide for installing TF and JAX with GPU support. These instructions may become outdated, so refer to this TF page an these JAX instructions if you run into trouble.

1. Download and install CUDA Toolkit 11.x by following these platform-dependent instructions. There may be dependency conflicts with existing NVIDIA packages. In that case, have a look here.

2. Tell the system where to find CUDA. In Ubuntu (or WSL-Ubuntu), this can be done by adding the following line to your .bashrc file:

   export PATH="$PATH:/usr/local/cuda/bin"
   

   But first, check if /usr/local/cuda/bin indeed exists!

3. Restart your machine. Then, launch a terminal to check whether CUDA is installed:

   nvcc --version
   

4. Download and install cuDNN following these instructions. Make sure that you download cuDNN for CUDA 11.x!

   In Ubuntu (Debian), there are two convenient options: (1) installing through apt or (2) using a local installer.

5. [Recommended] Install JAX with GPU binaries in your virtual environment following these instructions.

   pip install --upgrade jax[cuda] jaxlib -f https://storage.googleapis.com/jax-releases/jax_releases.html
   

   It doesn’t matter whether you do this before or after installing TensorWaves.

6. If TensorFlow can correctly find your GPU, the following should return a non-empty list:

   python3 -c 'import tensorflow as tf; print(tf.config.list_physical_devices("GPU"))'
   

   If JAX can correctly find your GPU, the following should return gpu:

   python3 -c 'from jax.lib import xla_bridge; print(xla_bridge.get_backend().platform)'
   

The latest version on the main branch (or any other branch, tag, or commit) can be installed as follows:

python3 -m pip install git+https://github.com/ComPWA/tensorwaves@main

Or, with optional dependencies:

python3 -m pip install "tensorwaves[jax,pwa] @ git+https://github.com/ComPWA/tensorwaves@main"

Developer installation

To contribute to the project, you need to install the package in a virtual environment. This can be done best with uv (see installation instructions here). For this, you first need to get the source code with Git:

git clone https://github.com/ComPWA/tensorwaves
cd tensorwaves

Now it’s simply a matter of creating and activating the virtual environment with uv sync. The dependencies for the project are ‘pinned’ in each commit through the uv.lock file.

uv sync --all-extras
source .venv/bin/activate

Formatting and linting checks are automatically performed when committing changes. This is done with pre-commit. To install the hooks in your local repository, run install pre-commit with uv:

uv tool install pre-commit --with pre-commit-uv --force-reinstall

and pre-commit install once:

pre-commit install --install-hooks

In addition, it may be handy to install tox:

uv tool install tox --with tox-uv

That’s all! Have a look at General examples to try out the package. You can also have a look at Help developing for tips on how to work with this ‘editable’ developer setup!