{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "dd5dcc75",
   "metadata": {},
   "source": [
    "# 02 — Masking\n",
    "\n",
    "**Purpose:** Apply point-source and cluster masks to the full-sky HEALPix CIB and tSZ maps.\n",
    "\n",
    "This notebook implements the two-step masking pipeline described in §2 of the paper:\n",
    "\n",
    "1. **Point-source masking** — uses `get_point_source_mask_in_healpix` to identify pixels\n",
    "   above the 2 mJy flux threshold in the CIB map at NSIDE=8192, inpaints them with\n",
    "   Gaussian noise, then degrades to NSIDE=2048.\n",
    "\n",
    "2. **Cluster masking** — uses `get_apodised_mdpl2_cluster_mask` with the halo catalogue\n",
    "   from `01_halo_catalogue.ipynb` to build apodised circular masks around clusters with\n",
    "   M500c ≥ 3×10¹⁴ M☉ at NSIDE=2048. Masked pixels are inpainted with Gaussian noise\n",
    "   matching the map mean and standard deviation via `inpaint_masked_regions`.\n",
    "\n",
    "**Inputs:**\n",
    "- Full-sky CIB map (Jy/sr): `data/agora_len_mag_cibmap_act_150ghz.fits`\n",
    "- Full-sky tSZ map (Compton-y): `data/agora_ltszNG_bahamas80_bnd_unb_1.0e+12_1.0e+18_lensed.fits`\n",
    "- Halo catalogue: `data/halo_catalogue/halo_catalogue_m500gt3e14.npy.npz` (from `01_halo_catalogue.ipynb`)\n",
    "\n",
    "**Outputs:**\n",
    "- Masked CIB map (μK): `data/cib_150_masked.fits`\n",
    "- Masked tSZ map (μK): `data/tsz_150_masked.fits`\n",
    "\n",
    "**Key module functions:**\n",
    "- `foregrounds_diffusion.masking.get_point_source_mask_in_healpix`\n",
    "- `foregrounds_diffusion.masking.get_apodised_mdpl2_cluster_mask`\n",
    "- `foregrounds_diffusion.preprocessing.inpaint_masked_regions`\n",
    "\n",
    "**Paper reference:** §2 (point-source and cluster masking descriptions)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "a6d03532-edbc-45db-b2b7-1ceca1b42e15",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Obtaining file:///home/apb86/cmb_foregrounds_diffusion\n",
      "  Installing build dependencies ... \u001b[?25ldone\n",
      "\u001b[?25h  Checking if build backend supports build_editable ... \u001b[?25ldone\n",
      "\u001b[?25h  Getting requirements to build editable ... \u001b[?25ldone\n",
      "\u001b[?25h  Preparing editable metadata (pyproject.toml) ... \u001b[?25ldone\n",
      "\u001b[?25hRequirement already satisfied: numpy>=1.26 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from foregrounds_diffusion==0.1.0) (1.26.4)\n",
      "Requirement already satisfied: scipy>=1.11 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from foregrounds_diffusion==0.1.0) (1.17.0)\n",
      "Requirement already satisfied: torch>=2.10 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from foregrounds_diffusion==0.1.0) (2.10.0)\n",
      "Requirement already satisfied: torchvision>=0.25 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from foregrounds_diffusion==0.1.0) (0.25.0)\n",
      "Requirement already satisfied: accelerate>=1.12 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from foregrounds_diffusion==0.1.0) (1.12.0)\n",
      "Requirement already satisfied: denoising-diffusion-pytorch>=2.2.5 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from foregrounds_diffusion==0.1.0) (2.2.6)\n",
      "Requirement already satisfied: healpy>=1.19 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from foregrounds_diffusion==0.1.0) (1.19.0)\n",
      "Requirement already satisfied: astropy>=7.2 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from foregrounds_diffusion==0.1.0) (7.2.0)\n",
      "Requirement already satisfied: einops>=0.8 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from foregrounds_diffusion==0.1.0) (0.8.2)\n",
      "Requirement already satisfied: ema-pytorch>=0.7 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from foregrounds_diffusion==0.1.0) (0.7.9)\n",
      "Requirement already satisfied: pillow>=12.0 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from foregrounds_diffusion==0.1.0) (12.1.1)\n",
      "Requirement already satisfied: pytorch-fid>=0.3 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from foregrounds_diffusion==0.1.0) (0.3.0)\n",
      "Requirement already satisfied: tqdm>=4.67 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from foregrounds_diffusion==0.1.0) (4.67.3)\n",
      "Requirement already satisfied: packaging>=20.0 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from accelerate>=1.12->foregrounds_diffusion==0.1.0) (26.0)\n",
      "Requirement already satisfied: psutil in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from accelerate>=1.12->foregrounds_diffusion==0.1.0) (7.2.2)\n",
      "Requirement already satisfied: pyyaml in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from accelerate>=1.12->foregrounds_diffusion==0.1.0) (6.0.3)\n",
      "Requirement already satisfied: huggingface_hub>=0.21.0 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from accelerate>=1.12->foregrounds_diffusion==0.1.0) (1.4.1)\n",
      "Requirement already satisfied: safetensors>=0.4.3 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from accelerate>=1.12->foregrounds_diffusion==0.1.0) (0.7.0)\n",
      "Requirement already satisfied: pyerfa>=2.0.1.1 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from astropy>=7.2->foregrounds_diffusion==0.1.0) (2.0.1.5)\n",
      "Requirement already satisfied: astropy-iers-data>=0.2025.10.27.0.39.10 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from astropy>=7.2->foregrounds_diffusion==0.1.0) (0.2026.2.9.0.50.33)\n",
      "Requirement already satisfied: filelock in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from huggingface_hub>=0.21.0->accelerate>=1.12->foregrounds_diffusion==0.1.0) (3.24.0)\n",
      "Requirement already satisfied: fsspec>=2023.5.0 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from huggingface_hub>=0.21.0->accelerate>=1.12->foregrounds_diffusion==0.1.0) (2026.2.0)\n",
      "Requirement already satisfied: hf-xet<2.0.0,>=1.2.0 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from huggingface_hub>=0.21.0->accelerate>=1.12->foregrounds_diffusion==0.1.0) (1.2.0)\n",
      "Requirement already satisfied: httpx<1,>=0.23.0 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from huggingface_hub>=0.21.0->accelerate>=1.12->foregrounds_diffusion==0.1.0) (0.28.1)\n",
      "Requirement already satisfied: shellingham in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from huggingface_hub>=0.21.0->accelerate>=1.12->foregrounds_diffusion==0.1.0) (1.5.4)\n",
      "Requirement already satisfied: typer-slim in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from huggingface_hub>=0.21.0->accelerate>=1.12->foregrounds_diffusion==0.1.0) (0.23.1)\n",
      "Requirement already satisfied: typing-extensions>=4.1.0 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from huggingface_hub>=0.21.0->accelerate>=1.12->foregrounds_diffusion==0.1.0) (4.15.0)\n",
      "Requirement already satisfied: anyio in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from httpx<1,>=0.23.0->huggingface_hub>=0.21.0->accelerate>=1.12->foregrounds_diffusion==0.1.0) (4.12.1)\n",
      "Requirement already satisfied: certifi in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from httpx<1,>=0.23.0->huggingface_hub>=0.21.0->accelerate>=1.12->foregrounds_diffusion==0.1.0) (2026.1.4)\n",
      "Requirement already satisfied: httpcore==1.* in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from httpx<1,>=0.23.0->huggingface_hub>=0.21.0->accelerate>=1.12->foregrounds_diffusion==0.1.0) (1.0.9)\n",
      "Requirement already satisfied: idna in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from httpx<1,>=0.23.0->huggingface_hub>=0.21.0->accelerate>=1.12->foregrounds_diffusion==0.1.0) (3.11)\n",
      "Requirement already satisfied: h11>=0.16 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from httpcore==1.*->httpx<1,>=0.23.0->huggingface_hub>=0.21.0->accelerate>=1.12->foregrounds_diffusion==0.1.0) (0.16.0)\n",
      "Requirement already satisfied: sympy>=1.13.3 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from torch>=2.10->foregrounds_diffusion==0.1.0) (1.14.0)\n",
      "Requirement already satisfied: networkx>=2.5.1 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from torch>=2.10->foregrounds_diffusion==0.1.0) (3.6.1)\n",
      "Requirement already satisfied: jinja2 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from torch>=2.10->foregrounds_diffusion==0.1.0) (3.1.6)\n",
      "Requirement already satisfied: cuda-bindings==12.9.4 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from torch>=2.10->foregrounds_diffusion==0.1.0) (12.9.4)\n",
      "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.8.93 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from torch>=2.10->foregrounds_diffusion==0.1.0) (12.8.93)\n",
      "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.8.90 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from torch>=2.10->foregrounds_diffusion==0.1.0) (12.8.90)\n",
      "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.8.90 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from torch>=2.10->foregrounds_diffusion==0.1.0) (12.8.90)\n",
      "Requirement already satisfied: nvidia-cudnn-cu12==9.10.2.21 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from torch>=2.10->foregrounds_diffusion==0.1.0) (9.10.2.21)\n",
      "Requirement already satisfied: nvidia-cublas-cu12==12.8.4.1 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from torch>=2.10->foregrounds_diffusion==0.1.0) (12.8.4.1)\n",
      "Requirement already satisfied: nvidia-cufft-cu12==11.3.3.83 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from torch>=2.10->foregrounds_diffusion==0.1.0) (11.3.3.83)\n",
      "Requirement already satisfied: nvidia-curand-cu12==10.3.9.90 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from torch>=2.10->foregrounds_diffusion==0.1.0) (10.3.9.90)\n",
      "Requirement already satisfied: nvidia-cusolver-cu12==11.7.3.90 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from torch>=2.10->foregrounds_diffusion==0.1.0) (11.7.3.90)\n",
      "Requirement already satisfied: nvidia-cusparse-cu12==12.5.8.93 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from torch>=2.10->foregrounds_diffusion==0.1.0) (12.5.8.93)\n",
      "Requirement already satisfied: nvidia-cusparselt-cu12==0.7.1 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from torch>=2.10->foregrounds_diffusion==0.1.0) (0.7.1)\n",
      "Requirement already satisfied: nvidia-nccl-cu12==2.27.5 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from torch>=2.10->foregrounds_diffusion==0.1.0) (2.27.5)\n",
      "Requirement already satisfied: nvidia-nvshmem-cu12==3.4.5 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from torch>=2.10->foregrounds_diffusion==0.1.0) (3.4.5)\n",
      "Requirement already satisfied: nvidia-nvtx-cu12==12.8.90 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from torch>=2.10->foregrounds_diffusion==0.1.0) (12.8.90)\n",
      "Requirement already satisfied: nvidia-nvjitlink-cu12==12.8.93 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from torch>=2.10->foregrounds_diffusion==0.1.0) (12.8.93)\n",
      "Requirement already satisfied: nvidia-cufile-cu12==1.13.1.3 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from torch>=2.10->foregrounds_diffusion==0.1.0) (1.13.1.3)\n",
      "Requirement already satisfied: triton==3.6.0 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from torch>=2.10->foregrounds_diffusion==0.1.0) (3.6.0)\n",
      "Requirement already satisfied: cuda-pathfinder~=1.1 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from cuda-bindings==12.9.4->torch>=2.10->foregrounds_diffusion==0.1.0) (1.3.4)\n",
      "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from sympy>=1.13.3->torch>=2.10->foregrounds_diffusion==0.1.0) (1.3.0)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from jinja2->torch>=2.10->foregrounds_diffusion==0.1.0) (3.0.3)\n",
      "Requirement already satisfied: typer>=0.23.1 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from typer-slim->huggingface_hub>=0.21.0->accelerate>=1.12->foregrounds_diffusion==0.1.0) (0.23.1)\n",
      "Requirement already satisfied: click>=8.0.0 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from typer>=0.23.1->typer-slim->huggingface_hub>=0.21.0->accelerate>=1.12->foregrounds_diffusion==0.1.0) (8.3.1)\n",
      "Requirement already satisfied: rich>=10.11.0 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from typer>=0.23.1->typer-slim->huggingface_hub>=0.21.0->accelerate>=1.12->foregrounds_diffusion==0.1.0) (14.3.2)\n",
      "Requirement already satisfied: annotated-doc>=0.0.2 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from typer>=0.23.1->typer-slim->huggingface_hub>=0.21.0->accelerate>=1.12->foregrounds_diffusion==0.1.0) (0.0.4)\n",
      "Requirement already satisfied: markdown-it-py>=2.2.0 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from rich>=10.11.0->typer>=0.23.1->typer-slim->huggingface_hub>=0.21.0->accelerate>=1.12->foregrounds_diffusion==0.1.0) (4.0.0)\n",
      "Requirement already satisfied: pygments<3.0.0,>=2.13.0 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from rich>=10.11.0->typer>=0.23.1->typer-slim->huggingface_hub>=0.21.0->accelerate>=1.12->foregrounds_diffusion==0.1.0) (2.19.2)\n",
      "Requirement already satisfied: mdurl~=0.1 in /home/apb86/diffusion_project_env/lib64/python3.11/site-packages (from markdown-it-py>=2.2.0->rich>=10.11.0->typer>=0.23.1->typer-slim->huggingface_hub>=0.21.0->accelerate>=1.12->foregrounds_diffusion==0.1.0) (0.1.2)\n",
      "Building wheels for collected packages: foregrounds_diffusion\n",
      "  Building editable for foregrounds_diffusion (pyproject.toml) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for foregrounds_diffusion: filename=foregrounds_diffusion-0.1.0-0.editable-py3-none-any.whl size=4647 sha256=7e3715cbdae5292835f7559969a51225de89123449fcf771b6b276fe9364ec10\n",
      "  Stored in directory: /tmp/pip-ephem-wheel-cache-5_jou8cm/wheels/ad/05/47/d622ea03cc2c607997f65b7643dab9c4f27fa5f37d82097115\n",
      "Successfully built foregrounds_diffusion\n",
      "Installing collected packages: foregrounds_diffusion\n",
      "  Attempting uninstall: foregrounds_diffusion\n",
      "    Found existing installation: foregrounds_diffusion 0.1.0\n",
      "    Uninstalling foregrounds_diffusion-0.1.0:\n",
      "      Successfully uninstalled foregrounds_diffusion-0.1.0\n",
      "Successfully installed foregrounds_diffusion-0.1.0\n",
      "\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m26.0.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m26.1.1\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!pip install -e ~/cmb_foregrounds_diffusion/"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26f81c1b",
   "metadata": {},
   "source": [
    "## 1  Configuration\n",
    "\n",
    "Key physical parameters:\n",
    "- **NSIDE_IN = 8192** — native resolution of the AGORA CIB and tSZ FITS files.\n",
    "- **NSIDE_OUT = 2048** — output resolution for patch extraction.  Downgrading\n",
    "  8192 → 2048 averages over 16 pixels, suppressing small-scale noise without\n",
    "  losing the angular scales of interest (ℓ < 7000 → θ > 1.6′).\n",
    "- **PTSRC_THRESH_MJY = 2** — flux density threshold at 150 GHz.  Pixel values\n",
    "  above 2 mJy/beam are flagged and inpainted before training.  This matches the\n",
    "  cut applied to the training data filename suffix `_2mJy`.\n",
    "- **M500C_THRESHOLD = 3e14 M☉** — cluster mass threshold.  Only clusters above\n",
    "  this mass receive an apodised circular mask of radius 3θ₅₀₀."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2132dbc5-f8a7-4b72-bb53-8afc4d69b9bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get_apodised_mdpl2_cluster_mask loads the filtered halo catalogue, projects\n",
    "# each cluster to a HEALPix disc of angular radius howmanythetaforclusters × θ_500,\n",
    "# and applies a raised-cosine apodisation taper to avoid sharp mask edges.\n",
    "# Returns a float mask in [0, 1] (0 = fully masked, 1 = fully unmasked).\n",
    "# get_point_source_mask_in_healpix applies a spectral index correction\n",
    "# (S ∝ ν^α with α = 3 for CIB) to convert the threshold from freq0 to freq,\n",
    "# then identifies all pixels with corrected flux > threshold_mjy_freq0 [mJy/sr].\n",
    "# Returns a list of HEALPix pixel indices to mask.\n",
    "import numpy as np\n",
    "import healpy as hp\n",
    "\n",
    "from pathlib import Path\n",
    "\n",
    "from foregrounds_diffusion.masking import (\n",
    "    get_point_source_mask_in_healpix,\n",
    "    get_apodised_mdpl2_cluster_mask,\n",
    "    get_mdpl2_conversion_factors_K_to_MjyperSr,\n",
    "    get_peak_masks,\n",
    "    inpaint_masked_regions,\n",
    ")\n",
    "\n",
    "# Anchor to the project root regardless of where the notebook is launched from\n",
    "PROJECT_ROOT = Path(\"/home/apb86/cmb_foregrounds_diffusion\")\n",
    "HPC_WORK = Path(\"/home/apb86/rds/hpc-work\")\n",
    "\n",
    "CIB_FITS = PROJECT_ROOT / \"data\" / \"agora_len_mag_cibmap_act_150ghz.fits\"\n",
    "TSZ_FITS = PROJECT_ROOT / \"data\" / \"agora_ltszNG_bahamas80_bnd_unb_1.0e+12_1.0e+18_lensed.fits\"\n",
    "HALO_CAT = HPC_WORK / \"halo_catalogue\" / \"halo_catalogue_m500gt3e14.npy.npz\"\n",
    "\n",
    "FREQ              = 150.      # GHz\n",
    "PTSRC_THRESH_MJY  = 2.0       # mJy threshold at 150 GHz\n",
    "M500C_THRESHOLD   = 3e14      # M_sun\n",
    "NSIDE_IN = 8192\n",
    "NSIDE_OUT = 2048\n",
    "TSZ_SPECTRAL_FACTOR = -0.952 * 2.7255e6 # uK per unit Compton-y at 150 GHz\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a77cf11",
   "metadata": {},
   "source": [
    "## 2  Load raw maps\n",
    "\n",
    "Read both AGORA full-sky maps from the local FITS files.  The CIB map is the\n",
    "lensed, magnification-corrected 150 GHz intensity in Jy/sr.  The tSZ map is\n",
    "the lensed non-Gaussian Compton-*y* parameter from the BAHAMAS hydrodynamical\n",
    "simulation.  Loading at native NSIDE = 8192 ensures point-source detection uses\n",
    "the full angular resolution before degrading."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a328b666-d875-4c79-ac81-ad38691f6103",
   "metadata": {},
   "outputs": [],
   "source": [
    "cib_jy_sr = hp.read_map(CIB_FITS)\n",
    "tsz_y = hp.read_map(TSZ_FITS)\n",
    "\n",
    "print(f\"NSIDE = {NSIDE_IN},  Npix = {hp.nside2npix(NSIDE_IN):,}\")\n",
    "print(f\"CIB — min: {cib_jy_sr.min():.4f}  max: {cib_jy_sr.max():.4f}  std: {cib_jy_sr.std():.4f}\")\n",
    "print(f\"tSZ — min: {tsz_y.min():.4f}  max: {tsz_y.max():.4f}  std: {tsz_y.std():.4f}\")\n",
    "\n",
    "rng = np.random.default_rng(seed=42)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae7603e7",
   "metadata": {},
   "source": [
    "## 3  Convert CIB units\n",
    "\n",
    "The raw CIB FITS file is in Jy/sr.  `get_point_source_mask_in_healpix` expects\n",
    "MJy/sr (mega-Jansky per steradian), so we multiply by 10⁻⁶.  This unit choice\n",
    "is a historical convention in the AGORA pipeline; the absolute value does not\n",
    "affect the mask geometry, only the threshold comparison."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "8aa5310c-acfe-427c-b301-08f9061d9caf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -------------------------------------------------------------------------\n",
    "# 2. Convert CIB: Jy/sr → mJy/sr for masking\n",
    "# -------------------------------------------------------------------------\n",
    "cib_mjy_sr = cib_jy_sr * 1e-6"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "185121dc",
   "metadata": {},
   "source": [
    "## 4  Point-source mask\n",
    "\n",
    "`get_point_source_mask_in_healpix` identifies HEALPix pixels whose inferred\n",
    "flux density (after applying a spectral index correction from `freq0 = 150 GHz`)\n",
    "exceeds 2 mJy.  The returned array of pixel indices is converted to a binary\n",
    "mask (`0 = masked, 1 = unmasked`).  Inpainting with Gaussian noise replaces\n",
    "masked pixels with a locally consistent background, preventing the patch\n",
    "extractor from producing patches with sharp mask edges."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "f559ac50-b950-472c-84dc-d4bca8edc086",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Point-source mask: 26,154 pixels masked (0.00%)\n"
     ]
    }
   ],
   "source": [
    "# get_point_source_mask_in_healpix applies a spectral index correction\n",
    "# (S ∝ ν^α with α = 3 for CIB) to convert the threshold from freq0 to freq,\n",
    "# then identifies all pixels with corrected flux > threshold_mjy_freq0 [mJy/sr].\n",
    "# Returns a list of HEALPix pixel indices to mask.\n",
    "# -------------------------------------------------------------------------\n",
    "# 3. Point-source mask from CIB at full resolution\n",
    "# -------------------------------------------------------------------------\n",
    "ptsrc_pixels = get_point_source_mask_in_healpix(\n",
    "    freq=FREQ,\n",
    "    hmap_Mjy_per_sr=cib_mjy_sr,\n",
    "    threshold_mjy_freq0=PTSRC_THRESH_MJY,\n",
    "    freq0=150.,\n",
    ")\n",
    "ptsrc_mask = np.ones(hp.nside2npix(NSIDE_IN), dtype=np.float32)\n",
    "ptsrc_mask[ptsrc_pixels] = 0.\n",
    "print(f\"Point-source mask: {(ptsrc_mask == 0).sum():,} pixels masked \"\n",
    "      f\"({100 * (ptsrc_mask == 0).mean():.2f}%)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "180e11bf",
   "metadata": {},
   "source": [
    "## 5  Cluster mask\n",
    "\n",
    "`get_apodised_mdpl2_cluster_mask` places a smoothly apodised circular mask of\n",
    "angular radius `howmanythetaforclusters × θ₅₀₀` around each cluster in the\n",
    "filtered halo catalogue.  The apodisation (raised-cosine taper) prevents\n",
    "ringing artefacts when cutting patches that overlap a cluster boundary.  The\n",
    "mask is built directly at NSIDE_OUT = 2048 to avoid the computational cost of\n",
    "working at full resolution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "816a844c-be0f-4096-8b57-d98ebd612878",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\tTotal clusters to mask: 6686\n",
      "Deleted outdated persistence file, no further action needed (/home/apb86/.colossus/cache/cosmology/planck15_01025226215067a4148b3ddb0d531324, versions 1.2.10/1.0.0).\n",
      "0\n",
      "1000\n",
      "2000\n",
      "3000\n",
      "4000\n",
      "5000\n",
      "6000\n",
      "0\n",
      "5000\n",
      "Starting apodisation\n",
      "Cluster mask: sky fraction retained = 0.987\n"
     ]
    }
   ],
   "source": [
    "# get_apodised_mdpl2_cluster_mask loads the filtered halo catalogue, projects\n",
    "# each cluster to a HEALPix disc of angular radius howmanythetaforclusters × θ_500,\n",
    "# and applies a raised-cosine apodisation taper to avoid sharp mask edges.\n",
    "# Returns a float mask in [0, 1] (0 = fully masked, 1 = fully unmasked).\n",
    "# -------------------------------------------------------------------------\n",
    "# 4. Cluster mask at output resolution (cheaper to build at NSIDE=2048)\n",
    "# -------------------------------------------------------------------------\n",
    "cluster_mask = get_apodised_mdpl2_cluster_mask(\n",
    "    nside=NSIDE_OUT,\n",
    "    halo_cat_fname=HALO_CAT,\n",
    "    m500c_threshold=M500C_THRESHOLD,\n",
    "    howmanythetaforclusters=3,\n",
    "    apodise=True,\n",
    ")\n",
    "print(f\"Cluster mask: sky fraction retained = {cluster_mask.mean():.3f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac12dac1-6b83-4e2b-a0af-8ef735c23cc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inpaint the point-source mask at full resolution (NSIDE=8192) BEFORE degrading\n",
    "# to NSIDE=2048.  This order matters: degrading masked pixels first would spread\n",
    "# zero-padding artefacts over a 4× larger angular scale.\n",
    "# -------------------------------------------------------------------------\n",
    "# 5. Apply point-source mask at full resolution, then degrade\n",
    "# -------------------------------------------------------------------------\n",
    "ptsrc_mask_float = ptsrc_mask.astype(np.float32)\n",
    "cib_inpainted = inpaint_masked_regions(cib_mjy_sr, ptsrc_mask_float, rng=rng)\n",
    "tsz_inpainted = inpaint_masked_regions(tsz_y,      ptsrc_mask_float, rng=rng)\n",
    "\n",
    "cib_inpainted_2048 = hp.ud_grade(cib_inpainted, nside_out=NSIDE_OUT)\n",
    "tsz_inpainted_2048 = hp.ud_grade(tsz_inpainted, nside_out=NSIDE_OUT)\n",
    "\n",
    "# Also degrade point-source mask itself to check coverage at output resolution\n",
    "ptsrc_mask_2048 = hp.ud_grade(ptsrc_mask, nside_out=NSIDE_OUT)\n",
    "ptsrc_mask_2048 = np.clip(ptsrc_mask_2048, 0., 1.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1829e32e-7434-4406-98af-7fa28ef889f1",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "# -------------------------------------------------------------------------\n",
    "# 6. Apply cluster mask at NSIDE=2048\n",
    "# -------------------------------------------------------------------------\n",
    "cib_masked = inpaint_masked_regions(cib_inpainted_2048 * cluster_mask,\n",
    "                                    cluster_mask, rng=rng)\n",
    "tsz_masked = inpaint_masked_regions(tsz_inpainted_2048 * cluster_mask,\n",
    "                                    cluster_mask, rng=rng)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53ae6cad",
   "metadata": {},
   "source": [
    "## 6  Apply masks and degrade\n",
    "\n",
    "The two masking operations are applied in sequence at their native resolutions:\n",
    "first the point-source mask (inpainting at NSIDE = 8192 before downgrading\n",
    "to 2048), then the cluster mask (applied at NSIDE = 2048).  Inpainting is\n",
    "done after each mask step so that the two masked regions receive independent\n",
    "noise realisations drawn from the correct local statistics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "797e07e7-62a3-4496-b405-1f59e5394d2c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CIB — min: -9.8152e+00  max: 1.2193e+02 uK\n",
      "tSZ — min: -1.9938e+02  max: 8.3041e+00 uK\n"
     ]
    }
   ],
   "source": [
    "# -------------------------------------------------------------------------\n",
    "# 7. Convert to physical units\n",
    "#    CIB: MJy/sr → uK at 150 GHz\n",
    "#    tSZ: Compton-y → uK at 150 GHz\n",
    "# -------------------------------------------------------------------------\n",
    "conv_K_to_MJyperSr = 375.876\n",
    "cib_masked_uK = (cib_masked / conv_K_to_MJyperSr) * 1e6\n",
    "tsz_masked_uK = tsz_masked * TSZ_SPECTRAL_FACTOR\n",
    "\n",
    "print(f\"CIB — min: {cib_masked_uK.min():.4e}  max: {cib_masked_uK.max():.4e} uK\")\n",
    "print(f\"tSZ — min: {tsz_masked_uK.min():.4e}  max: {tsz_masked_uK.max():.4e} uK\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "0a944800-3d51-47b3-bbf7-96f514326036",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "setting the output map dtype to [dtype('float32')]\n",
      "setting the output map dtype to [dtype('float32')]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saved masked maps.\n"
     ]
    }
   ],
   "source": [
    "# -------------------------------------------------------------------------\n",
    "# 8. Save\n",
    "# -------------------------------------------------------------------------\n",
    "hp.write_map(PROJECT_ROOT / \"data\" / \"cib_150_masked.fits\", cib_masked_uK.astype(np.float32),\n",
    "             overwrite=True)\n",
    "hp.write_map(PROJECT_ROOT / \"data\" / \"tsz_150_masked.fits\", tsz_masked_uK.astype(np.float32),\n",
    "             overwrite=True)\n",
    "print(\"Saved masked maps.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f55c65b7-d714-4dc6-bebe-9b94bee041d3",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "95603c0a",
   "metadata": {},
   "source": [
    "## 7  Convert to physical units\n",
    "\n",
    "Convert both channels to µK (micro-Kelvin) so that patch normalisation\n",
    "parameters are physically interpretable:\n",
    "- **CIB** (MJy/sr → µK): divide by the conversion factor 375.876 K/(MJy/sr)\n",
    "  at 150 GHz, then multiply by 10⁶.\n",
    "- **tSZ** (Compton-y → µK): multiply by the spectral function\n",
    "  g(x) = x coth(x/2) − 4 evaluated at x = hν / kT_CMB for 150 GHz,\n",
    "  giving g(150 GHz) ≈ −4.031.  The negative sign reflects the tSZ\n",
    "  decrement below the CMB blackbody at frequencies < 218 GHz."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d1aaa55",
   "metadata": {},
   "source": [
    "## 8  Save masked maps\n",
    "\n",
    "Write both masked, inpainted, and unit-converted maps as FITS files for use\n",
    "in notebook 03 (`03_patch_extraction.ipynb`).  The CIB output is in µK and\n",
    "the tSZ output is in µK (Compton-y × g factor).  These files are the direct\n",
    "inputs to the flat-sky patch extractor."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}