Clone the repository
git clone https://github.com/FarrOutLab/GWInferno.git
Recommended to use conda to set up your environment with python>=3.9:
For CPU usage only create an environment and install requirements and GWInferno with:
cd gwinferno
conda create -n gwinferno python=3.12
conda activate gwinferno
conda install -c conda-forge numpyro h5py
pip install --upgrade pip
pip install -r pip_requirements.txt
python -m pip install .
To enable JAX access to CUDA enabled GPUs we need to specify specific versions to install. The following procedure will only work for Linux x86_64 and Linux aarch64; for other platforms see Jax documentation
Jax recommends installing Nvidia CUDA and cuDNN with pip wheels. If you use local installations of CUDA and cuDNN, which could be the case for a remote cluster, then you'll need to install jax from the single CUDA wheel variant it offers. As of writing, this wheel is only compatible with CUDA >= 12.1 and cuDNN >= 9.1 < 10.0. See JAX installation instructions for more details.
cd gwinferno
conda create -n gwinferno_gpu python=3.12
conda activate gwinferno_gpu
pip install --upgrade pip
Next, install CUDA and cuDNN pip wheels. See Nvidia's CUDA quickstart guide for the CUDA installation procedure and the cuDNN documentation for the cuDNN installation procedure. Once that has finished, continue with these steps:
pip install --upgrade "jax[cuda12]"
pip install numpyro[cuda]
pip install -r pip_requirements.txt
python -m pip install .
Ensure Nvidia CUDA and cuDNN are installed locally.
cd gwinferno
conda create -n gwinferno_gpu python=3.12
conda activate gwinferno_gpu
pip install --upgrade pip
pip install --upgrade "jax[cuda12_local]"
pip install numpyro[cuda]
pip install -r pip_requirements.txt
python -m pip install .
Released under MIT by @FarrOutLab.
If GWInferno is used please cite this paper
We make use of and build upon other open source software projects, please cite them if using GWInferno.
- Bruce Edelman -- [email protected]
- Jaxen Godfrey -- [email protected]
- Gino Carrillo -- [email protected]
- Ben Farr -- [email protected]