Installation

Follow these instructions to install DESC and its dependencies.

We only test installation with conda environments, so your mileage may vary with other managers. For information on using conda, see here.

Attention

DESC requires python>=3.10. If you have python2 also locally installed, replace all pip commands with pip3 and all python commands with python3 to ensure the correct python version is used.

Attention

We do not support DESC on Windows OS. There have been instances where numerical discrepancies on Windows cause incorrect results. For these reasons, we recommend using Windows Subsystem for Linux (WSL) if you have a Windows machine. For instructions to install WSL see here. To use WSL in VS code see here.

On Your Local Machine

CPU

CPU+GPU (without FINUFFT)

For GPU support, you must install the JAX library as discussed in JAX installation docs. For example, below are the instructions to install on compatible devices with an NVIDIA GPU. These instructions do not install GPU support for jax-finufft which is used in Bounce2D computations such as GammaC and EffectiveRipple objectives. Set nufft_eps=0 in these objectives if necessary.

git clone https://github.com/PlasmaControl/DESC.git
cd DESC
conda create --name desc-env 'python>=3.10, <=3.14'
conda activate desc-env
sed -i '1 s/^jax/jax[cuda12]/' requirements.txt
pip install --editable .

Note that on BSD systems, the sed command that replaces jax with jax[cuda12] in the requirements.txt file is sed -i '' '1 s/^jax/jax[cuda12]/' requirements.txt.

You may optionally install developer requirements if you want to run tests.

pip install -r devtools/dev-requirements.txt

CPU+GPU (with FINUFFT)

git clone https://github.com/PlasmaControl/DESC.git
cd DESC
conda create --name desc-env -c conda-forge 'python>=3.10, <=3.13' 'fftw' 'gxx<12'
conda activate desc-env

sed -i '1 s/^jax/jax[cuda12]/' requirements.txt
sed -i '/^jax-finufft/d' requirements.txt

export CMAKE_PREFIX_PATH=$CONDA_PREFIX:$CMAKE_PREFIX_PATH
pip install --editable .
pip install -Ccmake.define.JAX_FINUFFT_USE_CUDA=ON --no-binary=jax-finufft 'jax-finufft >= 1.1.0'

Note that on BSD systems, the sed command that replaces jax with jax[cuda12] in the requirements.txt file is sed -i '' '1 s/^jax/jax[cuda12]/' requirements.txt.

You may optionally install developer requirements if you want to run tests.

pip install -r devtools/dev-requirements.txt

On Most Linux Computing Clusters

CPU

CPU+GPU

For GPU support, you must install the JAX library as discussed in JAX installation docs. We will show instructions that work for the clusters we have tested. Most of these clusters are x86_64 architectures. You may try the instructions for the cluster most resembling your own, or see if your cluster has JAX GPU installation instructions, as that is the main cause for installation differences.

Attention

DESC does not always guarantee support for the latest version of JAX (which does not have a stable 1.0 release yet). Older versions of jax may need to be installed, which may in turn require lower versions of jaxlib, as well as CUDA and CuDNN.

Verifying your Installation

To verify your installation works, try the following.

CPU

The following command should show an output stating the DESC version, the JAX version, and your device.

from desc.backend import print_backend_info
print_backend_info()

You can try running an example equilibrium solve. (The filepath shown here is from the DESC folder if you have cloned the git repo. Otherwise the file can be downloaded here.)

python -m desc -vv desc/examples/SOLOVEV

CPU+GPU

The following command should show an output stating the DESC version, the JAX version, and your device.

from desc import set_device
set_device('gpu')

from desc.backend import print_backend_info
print_backend_info()

You can try running an example equilibrium solve. (The filepath shown here is from the DESC folder if you have cloned the git repo. Otherwise the file can be downloaded here.)

python -m desc -vv desc/examples/SOLOVEV -g

If you installed DESC with FINUFFT on GPU, the following code should complete without error.

from desc import set_device
set_device("gpu")

from desc.examples import get
from desc.objectives import ObjectiveFunction, GammaC

obj = ObjectiveFunction(GammaC(get("W7-X"), num_transit=1, num_pitch=1))
obj.build()
x = obj.x()
obj.compute_scaled_error(x).block_until_ready()

Troubleshooting

We list common problems and their possible solutions. If you encounter other problems, please make an issue on Github and we will help.

Tip

Problem: My installation yields the error ModuleNotFoundError: No module named 'desc'.

Solution: First ensure you have activated the conda environment where DESC is installed (conda activate desc-env). If the issue persists, it is possible that DESC has not been added to your PYTHONPATH. Try adding the DESC directory to your PYTHONPATH manually by adding export PYTHONPATH="$PYTHONPATH:path/to/DESC" to the end of your ~/.bashrc (or other shell configuration) file. Note that /path/to/DESC should be replaced with the actual path to the DESC directory on your machine. You will also need to run source ~/.bashrc after making that change to ensure that your path updates for your current terminal session.

Tip

Problem: My installation yields the error ModuleNotFoundError: No module named 'termcolor' (or another module which is not desc).

Solution: Ensure you have activated the conda environment where DESC is installed (conda activate desc-env).

Tip

Problem: Attempts to install yield ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behavior is the source of the following dependency conflicts. desc-opt ... with a list of incompatibilities.

Solution: This may be due to another version of DESC or jax that is installed in the conda base environment. Try deleting the DESC folder, ensuring that pip list in the conda base environment no longer lists desc-opt or jax, then repeating the installation.

Tip

Problem: I am using JAX version 0.6.1 and getting errors like XlaRuntimeError: INTERNAL: cuSolver internal error

Solution: It is recommended to upgrade JAX to a newer version where these issues are resolved. If you must use version 0.6.1, then you must install the following package.

pip install 'nvidia-cublas-cu12==12.9.0.13'

Tip

Problem: Using pytest to run tests leads to import errors as discussed here.

Solution: This issue occurs because pip is an imperfect package manager, and the packages it installs have a tendency to leak out of the environment when pip thinks it can cache files globally to share among local environments. One way to resolve the issue is to prepend python -m to any command with pytest. Alternatively one can fix the broken pytest as follows. Since pytest has leaked out of the environment, first remove it globally. If you use conda it should suffice to remove it from the base environment, then install in the local environment as follows.

conda deactivate
conda activate base
pip uninstall pytest
conda activate desc-env
pip install pytest