Furax

Furax: a Framework for Unified and Robust data Analysis with JAX.

This framework provides building blocks for solving inverse problems, in particular in the astrophysical and cosmological domains.

Installation

You should always use a virtual environment to install packages (e.g. venv, conda environment, etc.).

Start by installing JAX for the target architecture.

Furax is available as furax on PyPI, and can be installed with:

pip install furax

Development version

Clone the repository, and navigate to the root directory of the project. For example:

git clone [email protected]:CMBSciPol/furax.git
cd furax

Then, install the package with:

pip install .

Developing Furax

After cloning, install in editable mode and with development dependencies:

pip install -e .[dev]

We use pytest for testing. You can run the tests with:

pytest

To ensure that your code passes the quality checks, you can use our pre-commit configuration:

1. Install the pre-commit hooks with

pre-commit install

2. That's it! Every commit will trigger the code quality checks.

Running on JeanZay

Load cuda and and cudnn for JAX

module load cuda/11.8.0 cudnn/8.9.7.29-cuda

Create Python env (only the first time)

module load python/3.10.4 && conda deactivate
python -m venv venv
source venv/bin/activate
# install jax
pip install --upgrade "jax[cuda11_local]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
# install furax
pip install -e .[dev]

launch script

To launch only the pytests

sbatch slurms/astro-sim-v100-testing.slurm

To launch your own script

sbatch slurms/astro-sim-v100-run.slurm yourscript.py

You can also allocate ressources and go into bash mode

srun --pty --account=nih@v100 --nodes=1 --ntasks-per-node=1 --cpus-per-task=10 --gres=gpu:1 --hint=nomultithread bash
module purge
module load python/3.10.4
source venv/bin/activate
module load cuda/11.8.0  cudnn/8.9.7.29-cuda
# Then do your thing
python my_script.py
pytest

Don't leave the bash running !! (I would suggest running script with sbatch)

Specific for nih / SciPol project

The repo is already in the commun WORK folder, the data is downloaded and the environment is ready.

You only need to do this

cd $ALL_CCFRWORK/furax-main

Then launch scripts as you see fit