Installation

Requirements

Python 3.11+
A CUDA-capable GPU is required for training; CPU-only is sufficient for statistics, sampling evaluation, and running the tutorials.

From source (recommended)

git clone https://github.com/AlexBM173/cmb_foregrounds_diffusion
cd cmb_foregrounds_diffusion
python -m venv .venv
source .venv/bin/activate

# CPU-only PyTorch (avoids pulling in the large CUDA build on machines
# without a GPU):
pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu

# Install the package with development extras:
pip install -e ".[dev]"

Optional dependencies

Some tutorial notebooks require additional packages not installed by default:

# Minkowski functionals (tutorial 12):
pip install quantimpy

# Scattering transforms (tutorial 11) — choose one:
pip install kymatio           # CPU/GPU via PyTorch
# or clone the Cheng et al. (2021) implementation

# Documentation build (local):
pip install -e ".[docs]"

GPU training

For training on a CUDA GPU, install the CUDA-enabled PyTorch wheel matching your driver version before running pip install -e .:

# Example for CUDA 12.1:
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121
pip install -e .
accelerate config     # select GPU count and mixed precision
accelerate launch foregrounds_diffusion/train.py --run-name my_run_v1

SLURM cluster

Pre-built SLURM scripts are provided at the repo root:

# Edit RUN_NAME and USE_WANDB at the top of each script, then:
sbatch train_slurm.sh    # single GPU, 1–12 h wall time
sbatch sample_slurm.sh   # 4 GPUs, 2 h wall time