Metadata-Version: 2.4
Name: foregrounds_diffusion
Version: 0.1.0
Summary: DDPM pipeline for generating correlated CIB and tSZ extragalactic CMB foregrounds
Author-email: Alex Blake Martin <alexbm173@gmail.com>
License: MIT
Project-URL: Homepage, https://github.com/AlexBM173/cmb_foregrounds_diffusion
Project-URL: Documentation, https://cmb-foregrounds-diffusion.readthedocs.io
Project-URL: Repository, https://github.com/AlexBM173/cmb_foregrounds_diffusion
Project-URL: Bug Tracker, https://github.com/AlexBM173/cmb_foregrounds_diffusion/issues
Keywords: cosmology,CMB,diffusion models,deep learning,astrophysics
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Topic :: Scientific/Engineering :: Astronomy
Classifier: Programming Language :: Python :: 3.11
Requires-Python: >=3.11
Description-Content-Type: text/markdown
Requires-Dist: numpy>=1.26
Requires-Dist: scipy>=1.11
Requires-Dist: matplotlib>=3.7
Requires-Dist: cycler>=0.10
Requires-Dist: torch>=2.10
Requires-Dist: torchvision>=0.25
Requires-Dist: accelerate>=1.12
Requires-Dist: denoising-diffusion-pytorch>=2.2.5
Requires-Dist: healpy>=1.19
Requires-Dist: astropy>=7.2
Requires-Dist: einops>=0.8
Requires-Dist: ema-pytorch>=0.7
Requires-Dist: pillow>=12.0
Requires-Dist: pytorch-fid>=0.3
Requires-Dist: tqdm>=4.67
Provides-Extra: dev
Requires-Dist: pytest>=7; extra == "dev"
Requires-Dist: pytest-cov; extra == "dev"
Requires-Dist: joblib>=1.3; extra == "dev"
Provides-Extra: docs
Requires-Dist: sphinx<8,>=7; extra == "docs"
Requires-Dist: furo; extra == "docs"
Requires-Dist: nbsphinx; extra == "docs"
Requires-Dist: sphinx-copybutton; extra == "docs"
Requires-Dist: sphinx-autodoc-typehints; extra == "docs"
Requires-Dist: ipykernel; extra == "docs"
Provides-Extra: fast
Requires-Dist: numba>=0.60; extra == "fast"
Requires-Dist: quantimpy>=0.5.6; extra == "fast"

# Denoising Diffusion Probabilistic Models for Extragalactic Foregrounds from AGORA

This repository implements a denoising diffusion probabilistic model (DDPM) pipeline to generate realistic AGORA map patches, incorporating point-source masked Cosmic Infrared Background (CIB) and cluster masked thermal Sunyaev-Zeldovich (tSZ) maps. The model is trained to reproduce statistical features of simulated sky patches.

## Overview

- **Data**: AGORA maps with point sources masked at 2mJy threshold. Zeroed-out pixels represent masked regions.
- **Preprocessing**:
  - High-frequency suppression via sharp mode cutoff (`l > 7000`) to avoid aliasing.
  - Negative pixel values from filtering artifacts are zeroed out.
- **Patching**:
  - Patches of size 6°×6° projected to 256×256 pixel Cartesian grids.
  - Centered on a grid defined by step size of 6° adjusted for equal angular separation in galactic coordinates.

## Data Location

Maps are produced by Srini and are located at: /sptlocal/analysis/ymap/sims/mdpl2/data/v0.7/bahamas80_scal1.000/mask_radio_cib_2.0mjy/cib(tsz)

## Training
Training is handled using `huggingface-accelerate` by running the script `train.py`:
accelerate launch train.py

The training script:
Loads preprocessed maps from data/low_pass/{ptsrc}mJy/
Stacks CIB and tSZ maps into a 2-channel tensor: (N, 2, 256, 256)
Augments with 90°, 180°, 270° rotations and horizontal flips
Trains a U-Net-based DDPM model with flash attention

## Sampling
The trained model generates synthetic CIB and tSZ map pairs that resemble the original astrophysical simulations and preserve the correct cross-correlations. These samples are useful for data augmentation, uncertainty estimation, and testing cosmological inference pipelines.
New samples can be generated using `sample.py`, which loads a trained checkpoint and produces batches of correlated CIB–tSZ pairs:
accelerate launch sample.py

## Requirements
* Python 3+
** Denoising-diffusion (https://github.com/lucidrains/denoising-diffusion-pytorch)
