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.

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

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 .

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.

module load cuda/11.8.0 cudnn/8.9.7.29-cuda

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]

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)

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