{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "76e775b8",
   "metadata": {},
   "source": [
    "# 09 — tSZ Cluster Stacking\n",
    "\n",
    "**Purpose:** Compare tSZ cluster profiles between Agora and DDPM samples using a\n",
    "signal-to-noise–based stacking approach.\n",
    "\n",
    "This notebook identifies cluster locations by SNR thresholding, stacks cutouts from\n",
    "both Agora and DDPM maps, and computes radial profiles to quantify agreement:\n",
    "\n",
    "1. **SNR selection** — pixels satisfying A·σ < |T − T̄| ≤ B·σ are selected in three\n",
    "   bins: (a) 5–10σ (~263k Agora clusters), (b) 10–20σ (~60k), (c) ≥20σ (~3.9k).\n",
    "\n",
    "2. **Cutout extraction** — extracts 22×22 pixel (≈31') cutouts centred on each selected\n",
    "   pixel from both the Agora and DDPM tSZ maps.\n",
    "\n",
    "3. **Stacking** — uniformly weighted average of all cutouts per SNR bin, producing a\n",
    "   mean cluster profile image Ŝ(θx, θy).\n",
    "\n",
    "4. **Radial profiles** — converts the 2D stacked images to 1D profiles using\n",
    "   `radial_profile`, with 10 linearly spaced bins of Δθ = 1' out to θ_max = 10'.\n",
    "   Plots the profile, Agora–DDPM difference, and their ratio.\n",
    "\n",
    "**Inputs:**\n",
    "- Test maps and DDPM samples: `data/low_pass/2mJy/*.npy`\n",
    "\n",
    "**Outputs:**\n",
    "- Stacked cutout arrays: `tsz_extracts/agora_tsz_stacks.npy`,\n",
    "  `tsz_extracts/ddpm_tsz_stacks.npy`\n",
    "- Radial profile plots: `plots/tsz_stacks.pdf`,\n",
    "  `plots/tsz_stacks_radial_profile.pdf` (Figure 3)\n",
    "\n",
    "**Key module functions:**\n",
    "- `foregrounds_diffusion.flatmaps.radial_profile`\n",
    "\n",
    "**Paper reference:** §4.2 (Figure 3 — stacked tSZ radial profiles in three SNR bins)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "402f30bc",
   "metadata": {},
   "source": [
    "## 1  Configuration\n",
    "\n",
    "Stacking parameters:\n",
    "- **SNR_BINS**: three S/N intervals `[3,6]`, `[6,9]`, `[9,∞]` used to\n",
    "  separate clusters by detection significance.  Higher-S/N bins probe more\n",
    "  massive clusters and produce cleaner stacked profiles.\n",
    "- **CUTOUT_PIX = 40**: cutout width of 40 pixels = 56′ centred on each peak.\n",
    "  This captures the full tSZ profile (typically ≈ 5–10′ radius for the cluster\n",
    "  population at z ~ 0.5) with enough padding for the radial profile estimator.\n",
    "- **PIXEL_RES = 1.40625′**: physical scale of the cutout pixels, used to\n",
    "  convert pixel offsets to angular separations for the radial profile."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d894504",
   "metadata": {},
   "outputs": [],
   "source": [
    "# select_snr_pixels(maps_nhw, snr_min, snr_max):\n",
    "#   Smooths each map with a matched filter and returns the (row, col) coordinates\n",
    "#   of pixels where signal-to-noise falls in [snr_min, snr_max].\n",
    "#   Uses only the first map in the stack as the detection template — the same\n",
    "#   coordinates are then used to extract cutouts from all maps.\n",
    "#\n",
    "# extract_cutouts(maps_nhw, pixel_coords, cutout_size, max_cutouts):\n",
    "#   Returns (N_cutouts, cutout_size, cutout_size) array; coordinates that\n",
    "#   would fall outside the map boundary are discarded.\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from pathlib import Path\n",
    "\n",
    "from foregrounds_diffusion.flatmaps import radial_profile\n",
    "from foregrounds_diffusion.preprocessing import renormalize_dm_maps, denormalize_dm_maps\n",
    "from foregrounds_diffusion.stacking import select_snr_pixels, extract_cutouts\n",
    "\n",
    "PTSRC      = 2\n",
    "CUTOUT_PIX = 22       # 22 px × 1.41 arcmin/px ≈ 31 arcmin\n",
    "PIXEL_RES  = 1.41     # arcmin / pixel\n",
    "\n",
    "SNR_BINS = [\n",
    "    (5,  10,  \"5–10σ\"),\n",
    "    (10, 20,  \"10–20σ\"),\n",
    "    (20, None, \"≥20σ\"),\n",
    "]\n",
    "\n",
    "PROJECT_ROOT = Path(\"/home/apb86/cmb_foregrounds_diffusion\")\n",
    "PATCHES_DIR = Path(f\"data/low_pass/{PTSRC}mJy\")\n",
    "\n",
    "OUT_DIR = Path(\"tsz_extracts\")\n",
    "OUT_DIR.mkdir(exist_ok=True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aabee446",
   "metadata": {},
   "source": [
    "## 2  Load and denormalise maps\n",
    "\n",
    "The DDPM samples are already in normalised space; denormalise to µK before\n",
    "stacking so that profile amplitudes are physically meaningful.  Agora maps\n",
    "are also denormalised from z-score to µK using the training-set mean and\n",
    "standard deviation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ddf7187",
   "metadata": {},
   "outputs": [],
   "source": [
    "cib_maps    = np.load(PATCHES_DIR / f\"CIB_map_150GHz_256_st6_minmax_{PTSRC}mJy_zero_lp.npy\")\n",
    "tsz_maps    = np.load(PATCHES_DIR / f\"tSZ3_map_150GHz_256_st6_minmax_{PTSRC}mJy_norm_lp.npy\")\n",
    "ddpm_raw    = np.load(PROJECT_ROOT / \"data\" / \"low_pass\" / f\"{PTSRC}mJy\" / f\"new_samples_cib_tsz_{PTSRC}mJy_zero_norm_6x6_w_au_lp.npy\")\n",
    "gauss_maps  = np.load(PATCHES_DIR / f\"gaussian_cib_tsz_{PTSRC}mJy_lp.npy\")\n",
    "\n",
    "# Collect pixel values from N_SAMPLES maps per source\n",
    "agora_cib_px  = cib_maps[:, :, :, 0].ravel()\n",
    "agora_tsz_px  = tsz_maps[:, :, :, 0].ravel()\n",
    "\n",
    "train_maps = np.concatenate([cib_maps, tsz_maps], axis=-1)  # (N, H, W, 2)\n",
    "print(f\"train_maps shape: {train_maps.shape}\")\n",
    "print(f\"ddpm_raw shape:   {ddpm_raw.shape}\")\n",
    "\n",
    "# Load saved normalisation parameters from patch extraction\n",
    "norm_params = np.load(PATCHES_DIR / f\"norm_params_{PTSRC}mJy.npy\")\n",
    "cib_mean, cib_std, tsz_mean, tsz_std = norm_params\n",
    "print(f\"CIB — mean: {cib_mean:.4f}  std: {cib_std:.4f}\")\n",
    "print(f\"tSZ — mean: {tsz_mean:.4f}  std: {tsz_std:.4f}\")\n",
    "\n",
    "ddpm_renorm = denormalize_dm_maps(ddpm_raw, cib_mean, cib_std, tsz_mean, tsz_std)\n",
    "print(f\"Denormalised DDPM — CIB range: [{ddpm_renorm[:,0].min():.4f}, {ddpm_renorm[:,0].max():.4f}]\")\n",
    "print(f\"Denormalised DDPM — tSZ range: [{ddpm_renorm[:,1].min():.4f}, {ddpm_renorm[:,1].max():.4f}]\")\n",
    "\n",
    "agora_tsz = tsz_maps[:, :, :, 0]   # (N, H, W)\n",
    "ddpm_tsz  = ddpm_renorm[:, 1]         # (N, H, W)\n",
    "print(f\"Agora tSZ maps: {agora_tsz.shape},  DDPM tSZ maps: {ddpm_tsz.shape}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "690c5295",
   "metadata": {},
   "source": [
    "## 3  Identify SNR pixels and extract cutouts\n",
    "\n",
    "`select_snr_pixels` smooths the tSZ map with a matched filter (or a simple\n",
    "Gaussian of width θ_beam) and returns pixel coordinates where the S/N\n",
    "exceeds the given thresholds.  `extract_cutouts` cuts out `CUTOUT_PIX ×\n",
    "CUTOUT_PIX` patches centred on each selected pixel from every map in the\n",
    "stack.  Cutouts that would extend beyond the patch boundary are discarded.\n",
    "\n",
    "The expected tSZ profile is a negative decrement (in µK at 150 GHz) with\n",
    "a central minimum and a positive ring from relativistic corrections — the\n",
    "DDPM's ability to reproduce this profile shape is a non-trivial test of\n",
    "whether it has learnt realistic cluster morphology."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a624a0dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# select_snr_pixels(maps_nhw, snr_min, snr_max):\n",
    "#   Smooths each map with a matched filter and returns the (row, col) coordinates\n",
    "#   of pixels where signal-to-noise falls in [snr_min, snr_max].\n",
    "#   Uses only the first map in the stack as the detection template — the same\n",
    "#   coordinates are then used to extract cutouts from all maps.\n",
    "#\n",
    "# extract_cutouts(maps_nhw, pixel_coords, cutout_size, max_cutouts):\n",
    "#   Returns (N_cutouts, cutout_size, cutout_size) array; coordinates that\n",
    "#   would fall outside the map boundary are discarded.\n",
    "agora_stacks, ddpm_stacks = {}, {}\n",
    "\n",
    "for smin, smax, label in SNR_BINS:\n",
    "    coords = select_snr_pixels(agora_tsz, smin, smax)\n",
    "    a_cuts = extract_cutouts(agora_tsz, coords, CUTOUT_PIX)\n",
    "    d_cuts = extract_cutouts(ddpm_tsz,  coords, CUTOUT_PIX)\n",
    "    agora_stacks[label] = a_cuts.mean(axis=0) if a_cuts is not None else None\n",
    "    ddpm_stacks[label]  = d_cuts.mean(axis=0) if d_cuts is not None else None\n",
    "    print(f\"{label}: stacked {len(a_cuts) if a_cuts is not None else 0} Agora / \"\n",
    "          f\"{len(d_cuts) if d_cuts is not None else 0} DDPM cutouts\")\n",
    "\n",
    "np.save(OUT_DIR / \"agora_tsz_stacks.npy\", {k: v for k, v in agora_stacks.items()}, allow_pickle=True)\n",
    "np.save(OUT_DIR / \"ddpm_tsz_stacks.npy\",  {k: v for k, v in ddpm_stacks.items()},  allow_pickle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cfa20a09",
   "metadata": {},
   "source": [
    "## 4  Compute mean stacked profiles and radial averages\n",
    "\n",
    "Average the cutout stack over all cutouts to produce the mean stacked image.\n",
    "Then compute the radial profile using `radial_profile` (from `flatmaps`)\n",
    "binned in steps of one pixel width.  The profile is compared between Agora\n",
    "and DDPM to assess whether the DDPM generates clusters of the correct\n",
    "amplitude and angular extent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "122fe0bf",
   "metadata": {},
   "outputs": [
    {
     "ename": "SyntaxError",
     "evalue": "invalid syntax (3947676670.py, line 3)",
     "output_type": "error",
     "traceback": [
      "  \u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[7]\u001b[39m\u001b[32m, line 3\u001b[39m\n\u001b[31m    \u001b[39m\u001b[31m\"\"\"Extract cutouts around pixel_coords from all maps.\"\"\"\u001b[39m\n       ^\n\u001b[31mSyntaxError\u001b[39m\u001b[31m:\u001b[39m invalid syntax\n"
     ]
    }
   ],
   "source": [
    "# select_snr_pixels(maps_nhw, snr_min, snr_max):\n",
    "#   Smooths each map with a matched filter and returns the (row, col) coordinates\n",
    "#   of pixels where signal-to-noise falls in [snr_min, snr_max].\n",
    "#   Uses only the first map in the stack as the detection template — the same\n",
    "#   coordinates are then used to extract cutouts from all maps.\n",
    "#\n",
    "# extract_cutouts(maps_nhw, pixel_coords, cutout_size, max_cutouts):\n",
    "#   Returns (N_cutouts, cutout_size, cutout_size) array; coordinates that\n",
    "#   would fall outside the map boundary are discarded.\n",
    "\"\"\"\n",
    "def extract_cutouts(maps_nhw, pixel_coords, cutout_size, max_cutouts=500):\n",
    "    \"\"\"Extract cutouts around pixel_coords from all maps.\"\"\"\n",
    "    half = cutout_size // 2\n",
    "    cutouts = []\n",
    "    coords = pixel_coords[:max_cutouts]  # cap to avoid memory issues\n",
    "    for i, m in enumerate(maps_nhw):\n",
    "        for (ri, rj) in coords:\n",
    "            ri0, ri1 = ri - half, ri + half\n",
    "            rj0, rj1 = rj - half, rj + half\n",
    "            if ri0 < 0 or ri1 > m.shape[0] or rj0 < 0 or rj1 > m.shape[1]:\n",
    "                continue\n",
    "            cutouts.append(m[ri0:ri1, rj0:rj1])\n",
    "    return np.array(cutouts, dtype=np.float32) if cutouts else None\n",
    "\n",
    "agora_stacks, ddpm_stacks = {}, {}\n",
    "for smin, smax, label in SNR_BINS:\n",
    "    coords = select_snr_pixels(agora_tsz, smin, smax)\n",
    "    a_cuts = extract_cutouts(agora_tsz, coords, CUTOUT_PIX)\n",
    "    d_cuts = extract_cutouts(ddpm_tsz,  coords, CUTOUT_PIX)\n",
    "    agora_stacks[label] = a_cuts.mean(axis=0) if a_cuts is not None else None\n",
    "    ddpm_stacks[label]  = d_cuts.mean(axis=0) if d_cuts is not None else None\n",
    "    n_a = len(a_cuts) if a_cuts is not None else 0\n",
    "    n_d = len(d_cuts) if d_cuts is not None else 0\n",
    "    print(f\"{label}: stacked {n_a} Agora / {n_d} DDPM cutouts\")\n",
    "\n",
    "np.save(OUT_DIR / \"agora_tsz_stacks.npy\", {k: v for k, v in agora_stacks.items()}, allow_pickle=True)\n",
    "np.save(OUT_DIR / \"ddpm_tsz_stacks.npy\",  {k: v for k, v in ddpm_stacks.items()},  allow_pickle=True)\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61360794",
   "metadata": {},
   "source": [
    "## 5  Plot stacked images and radial profiles\n",
    "\n",
    "Display the mean 2D stacked image for each S/N bin (left column) alongside\n",
    "the 1D radial profile (right column) for Agora and DDPM side by side.\n",
    "Error bars show the standard error on the mean across the cutout ensemble."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33b92ea3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build coordinate grids (degrees) for the cutout\n",
    "HALF    = CUTOUT_PIX // 2\n",
    "pix_deg = PIXEL_RES / 60.\n",
    "xi      = (np.arange(CUTOUT_PIX) - HALF) * pix_deg\n",
    "xgrid, ygrid = np.meshgrid(xi, xi)\n",
    "\n",
    "fig, axes = plt.subplots(len(SNR_BINS), 3, figsize=(14, 4 * len(SNR_BINS)))\n",
    "\n",
    "for row, (smin, smax, label) in enumerate(SNR_BINS):\n",
    "    a_stack = agora_stacks[label]\n",
    "    d_stack = ddpm_stacks[label]\n",
    "    if a_stack is None or d_stack is None:\n",
    "        continue\n",
    "\n",
    "    # Radial profiles: bin_size=1', maxbin=10', to_arcmins converts degrees→arcmin\n",
    "    prof_a = radial_profile(a_stack, xy=(xgrid, ygrid), bin_size=1., minbin=0., maxbin=10., to_arcmins=1)\n",
    "    prof_d = radial_profile(d_stack, xy=(xgrid, ygrid), bin_size=1., minbin=0., maxbin=10., to_arcmins=1)\n",
    "    theta  = prof_a[:, 0]   # arcmin\n",
    "\n",
    "    ax_img, ax_prof, ax_res = axes[row]\n",
    "\n",
    "    im = ax_img.imshow(a_stack - d_stack, origin=\"lower\", cmap=\"RdBu_r\")\n",
    "    ax_img.set_title(f\"{label} — Agora − DDPM stack\")\n",
    "    plt.colorbar(im, ax=ax_img)\n",
    "\n",
    "    ax_prof.plot(theta, prof_a[:, 1], label=\"Agora\", color=\"C0\")\n",
    "    ax_prof.fill_between(theta, prof_a[:,1]-prof_a[:,2], prof_a[:,1]+prof_a[:,2], alpha=0.3, color=\"C0\")\n",
    "    ax_prof.plot(theta, prof_d[:, 1], label=\"DDPM\",  color=\"C1\")\n",
    "    ax_prof.fill_between(theta, prof_d[:,1]-prof_d[:,2], prof_d[:,1]+prof_d[:,2], alpha=0.3, color=\"C1\")\n",
    "    ax_prof.set_xlabel(r\"$\\theta$ (arcmin)\");  ax_prof.set_ylabel(r\"$\\langle T \\rangle$\")\n",
    "    ax_prof.set_title(f\"{label} — radial profile\");  ax_prof.legend()\n",
    "\n",
    "    ratio = prof_d[:, 1] / (prof_a[:, 1] + 1e-30)\n",
    "    ax_res.plot(theta, ratio)\n",
    "    ax_res.axhline(1., color=\"k\", ls=\"--\", lw=0.8)\n",
    "    ax_res.set_xlabel(r\"$\\theta$ (arcmin)\");  ax_res.set_ylabel(\"DDPM / Agora\")\n",
    "    ax_res.set_title(f\"{label} — ratio\")\n",
    "\n",
    "plt.tight_layout()\n",
    "Path(\"plots\").mkdir(exist_ok=True)\n",
    "plt.savefig(\"plots/tsz_stacks_radial_profile.pdf\", bbox_inches=\"tight\")\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7534265-b861-4f18-a2c1-1903d7712984",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}