Because Veros should be usable by both beginners and experts, it has a number of optional dependencies that are not strictly required to run Veros, but highly recommended in high-performance contexts.
Veros notifies you when a slower code path has to be taken due to a missing dependency. So unless you are getting a warning, there is usually no need to install optional dependencies (except JAX and MPI).
When to use
Linux, OSX, Windows
To run 4x faster on CPU, and for GPU support
Veros Cython extensions
20% speedup when using JAX
MPI + mpi4py
To run in parallel
To run in parallel with JAX
PETSc + petsc4py
Faster linear solver when using
more than 10 processes (or GPUs)
On this page, we give all installation instructions via
pip install. If you used conda to install Veros, consider replacing them with
Using the JAX backend, Veros is typically about 4x faster than with NumPy, so this should be the first thing to try if you want to get more performance.
JAX is available on all major platforms and can be installed via:
$ pip install jax jaxlib
To use JAX on GPU, you have to install a CUDA-enabled version of jaxlib, e.g.:
$ pip install jax jaxlib==0.1.67+cuda111 -f https://storage.googleapis.com/jax-releases/jax_releases.html
(see also the JAX installation guide).
Veros also supplies Cython extensions that optimize certain bottlenecks in JAX. You can make sure they are installed by running:
$ pip install cython
$ python setup.py build_ext --inplace
in the Veros repository root.
To run Veros on more than one process you need to use MPI. This requires that you install an MPI implementation (such as OpenMPI) on your system. Additionally, you have to install
mpi4py to interface with it:
$ pip install mpi4py
Then, you can run Veros in parallel via MPI.
For optimal performance on many processes, Veros supports using PETSc as a linear solver. To use it, you will have to install the PETSc library and
petsc4py Python package:
$ PETSC_DIR=/path/to/petsc3.12 pip install petsc4py==3.12
Note that the versions of PETSc and
petsc4py have to match.
Using JAX + MPI#
To use JAX together with MPI, you need to install
mpi4jax after installing
$ pip install mpi4jax