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

Backend precision

TensorWaves uses 64-bit precision for JAX and TensorFlow by default. Configure 32-bit precision before creating backend arrays or TensorWaves functions with:

import tensorwaves

tensorwaves.configure(
    jax_precision="float32",
    tensorflow_precision="float32",
)

The two options are independent and can be specified separately.

Precision is a property of the backend, not of TensorWaves, so it applies from the moment it is set: arrays that were created earlier keep the dtype they were created with. A backend that has already been imported is reconfigured immediately, and one that has not is configured when TensorWaves first uses it. Calling tensorwaves.configure() right after your imports therefore covers both cases.

TensorWaves also respects JAX’s JAX_ENABLE_X64 environment variable. An explicit call to tensorwaves.configure() takes precedence over it. Note that JAX itself only reads that variable when it is imported, whereas TensorWaves reads it when it first uses JAX. TensorFlow has no equivalent environment variable, so tensorflow_precision is the only way to select its precision.

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 --python=3.13

and pre-commit install once:

pre-commit install --install-hooks

Poe the Poet is used as a task runner. Install it globally (within your home folder) with uv:

uv tool install poethepoet --force-reinstall --python=3.13

You can see which local CI checks it defines by running

poe

For instance, all style checks can be run with

poe style

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!