{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a000001",
   "metadata": {},
   "source": [
    "# DES Year 6 Key Project: 3×2pt Two-Point Statistics Pipeline\n",
    "\n",
    "This notebook illustrates how to compute the three two-point statistics used in the DES Y6 cosmological key project (KP) 3×2pt analysis and assemble them into a [two-point file](https://github.com/joezuntz/2point/tree/master):\n",
    "\n",
    "- **Cosmic shear** (ξ±): auto- and cross-correlations of the weak lensing shape catalog across source tomographic bin pairs ([DES Collaboration 2026](https://arxiv.org/abs/2602.10065)).\n",
    "- **Galaxy-galaxy lensing** (γt): cross-correlation between lens galaxies and source shapes ([Giannini et al. 2026](https://arxiv.org/abs/2601.15175)).\n",
    "- **Galaxy clustering** (w(θ)): auto-correlations of the lens galaxy sample ([Weaverdyck et al. 2026](https://arxiv.org/abs/2601.14484)).\n",
    "\n",
    "The final output is a FITS data vector in the `twopoint` format readable by CosmoSIS. \n",
    "\n",
    "**Required files to run this notebook** (update paths in Sections 1 and 3 as needed):\n",
    "- `desy6kp_cats_2024-11-07.hdf5`: index file linking all Y6KP catalogs.\n",
    "- `maglim_2024-11-07.hdf5`: MagLim lens sample described in [Weaverdyck et al. 2026](https://arxiv.org/abs/2601.14484).\n",
    "- `metadetect_2024-11-07.hdf5`: Metadetection shape catalog, masked from the original catalog described in [Yamamoto et al. 2026](https://arxiv.org/abs/2501.05665).\n",
    "- `maglimpp_wtheta_add_debias_enet_dv5.fits`: additive w(θ) debiasing correction for MagLim ([Weaverdyck et al. 2026](https://arxiv.org/abs/2601.14484)); provided for the fiducial selection.\n",
    "- `3x2pt_2025-07-28-20h_UNBLINDED_wmask_nob2.fits`: The released Y6 KP data vector. It includes the fiducial redshift distributions and covariance matrix.  Redshift distributions n(z) lens n(z) from [Giannini et al. 2025](https://arxiv.org/abs/2509.07964), source n(z) from [Yin et al. 2025](https://arxiv.org/abs/2510.23566). Covariance matrix computed with CosmoCov as described in [Sanchez-Cid et al. 2026](https://arxiv.org/abs/2601.14859). Note that if the selection is changed, the redshift distributions and covariance matrix will need to be recomputed accordingly.\n",
    "\n",
    "**Computational note**: The two-point measurements (Sections 5.1–5.3) are computationally expensive, typically taking several hours on a cluster with 30 CPU threads. Each bin pair is saved to a text file as soon as it is computed, so the notebook can be safely interrupted and restarted — bins with existing output files are skipped automatically. After the measurements are complete, continue from Section 5.4 to assemble the final data vector.\n",
    "\n",
    "If you have questions or comments about the notebook, email judit.prat.marti [at] upc.edu. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a000002",
   "metadata": {},
   "outputs": [],
   "source": [
    "import h5py\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import treecorr\n",
    "import os\n",
    "import datetime\n",
    "from astropy.io import fits\n",
    "import twopoint\n",
    "#print('TreeCorr version:', treecorr.__version__)\n",
    "# version used to run the oficial measurements was '5.1.1'\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a000003",
   "metadata": {},
   "source": [
    "## 1. Load the Catalog\n",
    "\n",
    "Set `folder_cats` to the directory containing the Y6KP HDF5 files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a000004",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Catalog loaded: /global/cfs/cdirs/des/y6kp-cats/final_version/desy6kp_cats_2024-11-07.hdf5\n",
      "Top-level groups: ['gold', 'maglim', 'mdet', 'ran_maglim', 'ran_redmagic', 'redmagic']\n"
     ]
    }
   ],
   "source": [
    "date = '2024-11-07'\n",
    "#folder_cats = '/path/to/catalogs/'  # update this to where the HDF5 files are located\n",
    "folder_cats = '/global/cfs/cdirs/des/y6kp-cats/final_version/' # in Perlmutter\n",
    "\n",
    "master_file_path = folder_cats + f'desy6kp_cats_{date}.hdf5'\n",
    "master = h5py.File(master_file_path, 'r')\n",
    "mdet = master['desy6kp/mdet/']\n",
    "\n",
    "print('Catalog loaded:', master_file_path)\n",
    "print('Top-level groups:', list(master['desy6kp'].keys()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a000005",
   "metadata": {},
   "source": [
    "## 2. Survey Area and Catalog Number Densities\n",
    "\n",
    "The effective survey area is pre-computed from the joint HEALSparse footprint mask and hardcoded below. See the commented code block below for how to re-derive it directly from the mask.\n",
    "\n",
    "We report the effective number density of the MagLim lens sample here. The Metadetect shape catalog statistics (number density and shape noise) are reported in Section 6, after the shear calibration responses R have been computed in Section 5.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a000006",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Survey area: 4031.036 sq. deg.  (14511730 arcmin²)\n"
     ]
    }
   ],
   "source": [
    "# Effective survey area pre-computed from the joint HEALSparse mask at nside=131072.\n",
    "area_deg    = 4031.036   # sq. deg.\n",
    "area_arcmin = area_deg * 60**2\n",
    "print(f'Survey area: {area_deg:.3f} sq. deg.  ({area_arcmin:.0f} arcmin²)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "181b7a86",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ----------------------------------------------------------------------------\n",
    "# Optional: compute the survey area directly from the HEALSparse mask and/or\n",
    "# apply the mask to an external catalog.\n",
    "#\n",
    "# The joint footprint mask (jointmask_DESY6_2024-06-04.hsp) is a binary\n",
    "# HEALSparse map at nside=131072. It is not required to run this notebook —\n",
    "# the area is hardcoded above — but it can be used to re-derive the area or\n",
    "# to apply the footprint cut to external catalogs.\n",
    "# Loading it requires ~450 GB of RAM, so it is only practical on a large node.\n",
    "#\n",
    "# import healpy as hp\n",
    "# import healsparse as hsp\n",
    "# Make sure version is '1.10.1' or above\n",
    "# print(hs.__version__)\n",
    "#\n",
    "# mask = hsp.HealSparseMap.read('jointmask_DESY6_2024-06-04.hsp')\n",
    "#\n",
    "# # Compute survey area from the mask\n",
    "# nside_sparse = 131072\n",
    "# pixarea     = hp.nside2pixarea(nside_sparse, degrees=True)\n",
    "# area_deg    = len(mask.valid_pixels) * pixarea\n",
    "# area_arcmin = area_deg * 60**2\n",
    "# print(f'Survey area from mask: {area_deg:.3f} sq. deg.')\n",
    "#\n",
    "# # Apply mask to an external catalog (ra, dec in degrees)\n",
    "# def create_mask(ra, dec):\n",
    "#     \"\"\"Return boolean array. True means the object is inside the footprint.\"\"\"\n",
    "#     mask = hsp.HealSparseMap.read('jointmask_DESY6_2024-06-04.hsp')\n",
    "#     mask_bool = mask.get_values_pos(ra, dec)\n",
    "#     return mask_bool\n",
    "#\n",
    "# # Example usage:\n",
    "# # in_footprint = create_mask(ra_array, dec_array)\n",
    "# # ra_cut, dec_cut = ra_array[in_footprint], dec_array[in_footprint]\n",
    "# ----------------------------------------------------------------------------\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a000007",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MagLim lens sample — object counts and effective number densities:\n",
      "  Bin 0: 1,852,538 objects,  neff = 0.1276 gal/arcmin²\n",
      "  Bin 1: 1,335,294 objects,  neff = 0.0920 gal/arcmin²\n",
      "  Bin 2: 1,413,738 objects,  neff = 0.0974 gal/arcmin²\n",
      "  Bin 3: 1,783,834 objects,  neff = 0.1228 gal/arcmin²\n",
      "  Bin 4: 1,391,521 objects,  neff = 0.0959 gal/arcmin²\n",
      "  Bin 5: 1,409,280 objects,  neff = 0.0970 gal/arcmin²\n"
     ]
    }
   ],
   "source": [
    "n_zlbins = 6  # number of lens (MagLim) tomographic bins\n",
    "n_zsbins = 4  # number of source (Metadetect) tomographic bins\n",
    "\n",
    "print('MagLim lens sample — object counts and effective number densities:')\n",
    "for i in range(n_zlbins):\n",
    "    w    = master[f'desy6kp/maglim/tomo_bin_{i}/weight'][:]\n",
    "    nobj = len(w)\n",
    "    neff = np.sum(w)**2 / np.sum(w**2) / area_arcmin\n",
    "    print(f'  Bin {i}: {nobj:>9,} objects,  neff = {neff:.4f} gal/arcmin²')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a000008",
   "metadata": {},
   "source": [
    "## 3. Measurement Configuration\n",
    "\n",
    "All three two-point statistics are measured in 26 logarithmically-spaced angular bins spanning 2.5–995 arcmin (in order to exactly match the DESY3 binning in the range 2.5-250 arcmin), using [TreeCorr](https://rmjarvis.github.io/TreeCorr/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a000009",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Angular range: 2.50 – 995.27 arcmin in 26 bins\n"
     ]
    }
   ],
   "source": [
    "output_folder = 'measurements'\n",
    "os.makedirs(output_folder, exist_ok=True)\n",
    "\n",
    "deltag     = 0.01   # shear step used in Metadetect image processing\n",
    "bin_slop   = 0.005\n",
    "angle_slop = 0.01\n",
    "\n",
    "# 26 log-spaced angular bins from 2.5 to ~995 arcmin, to match the exact binning from DES Y3 in the range from 2.5 to 250 arcmin.\n",
    "nbins      = 26\n",
    "bin_edges  = np.geomspace(2.5, 995.2679263837431, nbins + 1)\n",
    "min_sep    = bin_edges[0]\n",
    "max_sep    = bin_edges[-1]\n",
    "\n",
    "config = {\n",
    "    'nbins':       nbins,\n",
    "    'min_sep':     min_sep,\n",
    "    'max_sep':     max_sep,\n",
    "    'sep_units':   'arcmin',\n",
    "    'bin_slop':    bin_slop,\n",
    "    'num_threads': 30,\n",
    "    'angle_slop':  angle_slop,\n",
    "}\n",
    "\n",
    "print(f'Angular range: {min_sep:.2f} – {max_sep:.2f} arcmin in {nbins} bins')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a000010",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Path to the additive debiasing term for w(θ).\n",
    "# This corrects for a small bias in galaxy clustering introduced by the MagLim survey\n",
    "# property weights. The file contains columns:\n",
    "#   index_fulldv, index_wtheta, theta [arcmin], izbin, wtheta_add_debias\n",
    "# The wtheta_add_debias values are ADDED to the raw w(θ) auto-correlations.\n",
    "# Update this path to where you have stored the provided file.\n",
    "#path_wtheta_debias = 'maglimpp_wtheta_add_debias_enet_dv5.fits'\n",
    "path_wtheta_debias = '/global/cfs/projectdirs/des/nweaverd/des_y6/maglimpp_wtheta_add_debias_enet_dv5.fits' # in Perlmutter\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a000011",
   "metadata": {},
   "source": [
    "## 4. Helper Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a000012",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_response(mdet, zbin, deltag):\n",
    "    \"\"\"\n",
    "    Compute the weighted mean Metadetect shear response R for source bin `zbin`.\n",
    "\n",
    "    R is estimated from pairs of artificially sheared images (step ±deltag in each\n",
    "    ellipticity component). All two-point measurements involving source bin `zbin`\n",
    "    are divided by R to recover the calibrated shear signal.\n",
    "    \"\"\"\n",
    "    e1p = np.average(mdet[f'1p/tomo_bin_{zbin}/gauss_g_1'], weights=mdet[f'1p/tomo_bin_{zbin}/w'])\n",
    "    e1m = np.average(mdet[f'1m/tomo_bin_{zbin}/gauss_g_1'], weights=mdet[f'1m/tomo_bin_{zbin}/w'])\n",
    "    e2p = np.average(mdet[f'2p/tomo_bin_{zbin}/gauss_g_2'], weights=mdet[f'2p/tomo_bin_{zbin}/w'])\n",
    "    e2m = np.average(mdet[f'2m/tomo_bin_{zbin}/gauss_g_2'], weights=mdet[f'2m/tomo_bin_{zbin}/w'])\n",
    "\n",
    "    R11 = (e1p - e1m) / (2 * deltag)\n",
    "    R22 = (e2p - e2m) / (2 * deltag)\n",
    "    return (R11 + R22) / 2.0\n",
    "\n",
    "\n",
    "def build_lens_cat(zlbin):\n",
    "    \"\"\"Build a weighted TreeCorr catalog for MagLim lens bin `zlbin`.\"\"\"\n",
    "    maglim = master[f'desy6kp/maglim/tomo_bin_{zlbin}']\n",
    "    return treecorr.Catalog(\n",
    "        ra=maglim['RA'][:], dec=maglim['DEC'][:], w=maglim['weight'][:],\n",
    "        ra_units='deg', dec_units='deg'\n",
    "    )\n",
    "\n",
    "\n",
    "def build_ran_cat(zlbin):\n",
    "    \"\"\"Build a TreeCorr catalog for the MagLim random catalog for lens bin `zlbin`.\"\"\"\n",
    "    ran = master[f'desy6kp/ran_maglim/tomo_bin_{zlbin}']\n",
    "    return treecorr.Catalog(\n",
    "        ra=ran['ra'][:], dec=ran['dec'][:],\n",
    "        ra_units='deg', dec_units='deg'\n",
    "    )\n",
    "\n",
    "\n",
    "def build_source_cat(zsbin):\n",
    "    \"\"\"\n",
    "    Build a TreeCorr catalog for Metadetect source bin `zsbin`.\n",
    "\n",
    "    The inverse-variance weight `w` and the weighted-mean-subtracted ellipticity\n",
    "    components are passed to TreeCorr. The mean ellipticity subtraction removes\n",
    "    the additive shear bias.\n",
    "    \"\"\"\n",
    "    data = mdet[f'noshear/tomo_bin_{zsbin}']\n",
    "    ra   = data['ra'][:]\n",
    "    dec  = data['dec'][:]\n",
    "    e1   = data['gauss_g_1'][:]\n",
    "    e2   = data['gauss_g_2'][:]\n",
    "    w    = data['w'][:]\n",
    "\n",
    "    mean_e1 = np.average(e1, weights=w)\n",
    "    mean_e2 = np.average(e2, weights=w)\n",
    "\n",
    "    return treecorr.Catalog(\n",
    "        ra=ra, dec=dec,\n",
    "        g1=(e1 - mean_e1), g2=(e2 - mean_e2), w=w,\n",
    "        ra_units='deg', dec_units='deg'\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a000013",
   "metadata": {},
   "source": [
    "## 5. Two-Point Measurements\n",
    "\n",
    "### 5.1 Cosmic Shear (ξ±)\n",
    "\n",
    "We measure ξ+ and ξ− for all auto- and cross-correlations of the 4 source tomographic bins, following [DES Collaboration 2026](https://arxiv.org/abs/2602.10065).\n",
    "Each measurement is divided by the product of the shear calibration responses R for the two bins involved.\n",
    "\n",
    "Results are saved to text files in `measurements/`; bin pairs whose files already exist are skipped."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a000014",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Response R[0] = 0.856143\n",
      "  Source bin 0: 33,707,071 objects\n",
      "  Skipping cs 0×0 (file exists)\n",
      "  Skipping cs 0×1 (file exists)\n",
      "  Skipping cs 0×2 (file exists)\n",
      "  Skipping cs 0×3 (file exists)\n",
      "Response R[1] = 0.868506\n",
      "  Source bin 1: 34,580,888 objects\n",
      "  Skipping cs 1×1 (file exists)\n",
      "  Skipping cs 1×2 (file exists)\n",
      "  Skipping cs 1×3 (file exists)\n",
      "Response R[2] = 0.819427\n",
      "  Source bin 2: 34,646,416 objects\n",
      "  Skipping cs 2×2 (file exists)\n",
      "  Skipping cs 2×3 (file exists)\n",
      "Response R[3] = 0.691769\n",
      "  Source bin 3: 36,727,798 objects\n",
      "  Skipping cs 3×3 (file exists)\n",
      "\n",
      "All responses: {0: '0.8561', 1: '0.8685', 2: '0.8194', 3: '0.6918'}\n"
     ]
    }
   ],
   "source": [
    "R = {}   # shear responses; populated here, reused in GGL and catalog statistics\n",
    "m = {}\n",
    "m['cs']        = {}\n",
    "m['responses'] = {}\n",
    "\n",
    "th_edges = np.geomspace(min_sep, max_sep, nbins + 1)\n",
    "m['angular_binning'] = {'theta_min': th_edges[:-1], 'theta_max': th_edges[1:]}\n",
    "\n",
    "for zsbin1 in range(n_zsbins):\n",
    "    R[zsbin1] = compute_response(mdet, zsbin1, deltag)\n",
    "    m['responses'][zsbin1] = R[zsbin1]\n",
    "    print(f'Response R[{zsbin1}] = {R[zsbin1]:.6f}')\n",
    "\n",
    "    scat_1 = build_source_cat(zsbin1)\n",
    "    print(f'  Source bin {zsbin1}: {len(scat_1.ra):,} objects')\n",
    "\n",
    "    for zsbin2 in range(zsbin1, n_zsbins):\n",
    "        filename = os.path.join(output_folder, f'cs_zbin{zsbin1}_zbin{zsbin2}.txt')\n",
    "        if os.path.exists(filename):\n",
    "            print(f'  Skipping cs {zsbin1}×{zsbin2} (file exists)')\n",
    "            continue\n",
    "\n",
    "        R[zsbin2] = compute_response(mdet, zsbin2, deltag)\n",
    "        scat_2    = build_source_cat(zsbin2)\n",
    "\n",
    "        gg = treecorr.GGCorrelation(config)\n",
    "        if zsbin1 == zsbin2:\n",
    "            gg.process(scat_1, low_mem=True)\n",
    "        else:\n",
    "            gg.process(scat_1, scat_2)\n",
    "\n",
    "        theta = np.exp(gg.logr)\n",
    "        xip   = gg.xip / (R[zsbin1] * R[zsbin2])\n",
    "        xim   = gg.xim / (R[zsbin1] * R[zsbin2])\n",
    "\n",
    "        m['cs'][f'zbin{zsbin1}_zbin{zsbin2}'] = {'theta': theta, 'xip': xip, 'xim': xim}\n",
    "        np.savetxt(filename, np.c_[theta, xip, xim],\n",
    "                   header='theta [arcmin]  xip  xim')\n",
    "        print(f'  Saved cs {zsbin1}×{zsbin2}')\n",
    "\n",
    "print('\\nAll responses:', {k: f'{v:.4f}' for k, v in R.items()})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a000015",
   "metadata": {},
   "source": [
    "### 5.2 Galaxy-Galaxy Lensing (γt)\n",
    "\n",
    "We measure γt for all 24 combinations of 6 lens bins × 4 source bins, following [Giannini et al. 2026](https://arxiv.org/abs/2601.15175). Boost factors B are also saved alongside γt. gt_boosted corresponds to the fiducial measurement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a000016",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Lens bin 0: 1,852,538 galaxies, 92,626,900 randoms\n",
      "  Skipping ggl l0×s0 (file exists)\n",
      "  Skipping ggl l0×s1 (file exists)\n",
      "  Skipping ggl l0×s2 (file exists)\n",
      "  Skipping ggl l0×s3 (file exists)\n",
      "Lens bin 1: 1,335,294 galaxies, 66,764,700 randoms\n",
      "  Skipping ggl l1×s0 (file exists)\n",
      "  Skipping ggl l1×s1 (file exists)\n",
      "  Skipping ggl l1×s2 (file exists)\n",
      "  Skipping ggl l1×s3 (file exists)\n",
      "Lens bin 2: 1,413,738 galaxies, 70,686,900 randoms\n",
      "  Skipping ggl l2×s0 (file exists)\n",
      "  Skipping ggl l2×s1 (file exists)\n",
      "  Skipping ggl l2×s2 (file exists)\n",
      "  Skipping ggl l2×s3 (file exists)\n",
      "Lens bin 3: 1,783,834 galaxies, 89,191,700 randoms\n",
      "  Skipping ggl l3×s0 (file exists)\n",
      "  Skipping ggl l3×s1 (file exists)\n",
      "  Skipping ggl l3×s2 (file exists)\n",
      "  Skipping ggl l3×s3 (file exists)\n",
      "Lens bin 4: 1,391,521 galaxies, 69,576,050 randoms\n",
      "  Skipping ggl l4×s0 (file exists)\n",
      "  Skipping ggl l4×s1 (file exists)\n",
      "  Skipping ggl l4×s2 (file exists)\n",
      "  Skipping ggl l4×s3 (file exists)\n",
      "Lens bin 5: 1,409,280 galaxies, 70,464,000 randoms\n",
      "  Skipping ggl l5×s0 (file exists)\n",
      "  Skipping ggl l5×s1 (file exists)\n",
      "  Skipping ggl l5×s2 (file exists)\n",
      "  Skipping ggl l5×s3 (file exists)\n"
     ]
    }
   ],
   "source": [
    "m['ggl'] = {}\n",
    "\n",
    "for zlbin in range(n_zlbins):\n",
    "    lcat = build_lens_cat(zlbin)\n",
    "    rcat = build_ran_cat(zlbin)\n",
    "    print(f'Lens bin {zlbin}: {len(lcat.ra):,} galaxies, {len(rcat.ra):,} randoms')\n",
    "\n",
    "    for zsbin in range(n_zsbins):\n",
    "        filename = os.path.join(output_folder, f'ggl_zlbin{zlbin}_zsbin{zsbin}.txt')\n",
    "        if os.path.exists(filename):\n",
    "            print(f'  Skipping ggl l{zlbin}×s{zsbin} (file exists)')\n",
    "            continue\n",
    "\n",
    "        R[zsbin] = compute_response(mdet, zsbin, deltag)\n",
    "        scat     = build_source_cat(zsbin)\n",
    "\n",
    "        ng = treecorr.NGCorrelation(config)\n",
    "        ng.process(lcat, scat, low_mem=True)\n",
    "\n",
    "        nr = treecorr.NGCorrelation(config)\n",
    "        nr.process(rcat, scat, low_mem=True)\n",
    "\n",
    "        theta      = np.exp(ng.logr)\n",
    "        gt         = (ng.xi - nr.xi) / R[zsbin]\n",
    "\n",
    "        # Boost factor: ratio of random-to-lens weights, used to account for\n",
    "        # source-lens clustering \n",
    "        factor     = np.sum(rcat.w) / np.sum(lcat.w)\n",
    "        gt_boosted = (factor * ng.xi * ng.weight / nr.weight - nr.xi) / R[zsbin]\n",
    "        bf         = factor * ng.weight / nr.weight\n",
    "\n",
    "        m['ggl'][f'zlbin{zlbin}_zsbin{zsbin}'] = {\n",
    "            'theta': theta, 'gt': gt, 'gt_boosted': gt_boosted, 'bf': bf\n",
    "        }\n",
    "        np.savetxt(filename, np.c_[theta, gt, gt_boosted, bf],\n",
    "                   header='theta [arcmin]  gt  gt_boosted  boost_factor')\n",
    "        print(f'  Saved ggl l{zlbin}×s{zsbin}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a000017",
   "metadata": {},
   "source": [
    "### 5.3 Galaxy Clustering (w(θ))\n",
    "\n",
    "We measure w(θ) using the Landy-Szalay estimator for all 21 pairs of the 6 lens tomographic bins (6 auto + 15 cross), following [Weaverdyck et al. 2026](https://arxiv.org/abs/2601.14484). All pairs are included in the data vector file; the cross-correlations are not used in the cosmological inference (they are masked by scale cuts in CosmoSIS). The additive debiasing correction (Section 5.5) applies only to the auto-correlations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "a000018",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Skipping clu 0×0 (file exists)\n",
      "Skipping clu 0×1 (file exists)\n",
      "Skipping clu 0×2 (file exists)\n",
      "Skipping clu 0×3 (file exists)\n",
      "Skipping clu 0×4 (file exists)\n",
      "Skipping clu 0×5 (file exists)\n",
      "Skipping clu 1×1 (file exists)\n",
      "Skipping clu 1×2 (file exists)\n",
      "Skipping clu 1×3 (file exists)\n",
      "Skipping clu 1×4 (file exists)\n",
      "Skipping clu 1×5 (file exists)\n",
      "Skipping clu 2×2 (file exists)\n",
      "Skipping clu 2×3 (file exists)\n",
      "Skipping clu 2×4 (file exists)\n",
      "Skipping clu 2×5 (file exists)\n",
      "Skipping clu 3×3 (file exists)\n",
      "Skipping clu 3×4 (file exists)\n",
      "Skipping clu 3×5 (file exists)\n",
      "Skipping clu 4×4 (file exists)\n",
      "Skipping clu 4×5 (file exists)\n",
      "Skipping clu 5×5 (file exists)\n"
     ]
    }
   ],
   "source": [
    "m['clu'] = {}\n",
    "\n",
    "for zlbin1 in range(n_zlbins):\n",
    "    lcat_1 = build_lens_cat(zlbin1)\n",
    "    rcat_1 = build_ran_cat(zlbin1)\n",
    "\n",
    "    for zlbin2 in range(zlbin1, n_zlbins):\n",
    "        filename = os.path.join(output_folder, f'clu_zbin{zlbin1}_zbin{zlbin2}.txt')\n",
    "        if os.path.exists(filename):\n",
    "            print(f'Skipping clu {zlbin1}×{zlbin2} (file exists)')\n",
    "            continue\n",
    "\n",
    "        lcat_2 = build_lens_cat(zlbin2)\n",
    "        rcat_2 = build_ran_cat(zlbin2)\n",
    "\n",
    "        nn = treecorr.NNCorrelation(config)\n",
    "        rr = treecorr.NNCorrelation(config)\n",
    "        dr = treecorr.NNCorrelation(config)\n",
    "\n",
    "        if zlbin1 == zlbin2:\n",
    "            nn.process(lcat_1, low_mem=True)\n",
    "            rr.process(rcat_1, low_mem=True)\n",
    "            dr.process(lcat_1, rcat_1, low_mem=True)\n",
    "            rd = None\n",
    "        else:\n",
    "            rd = treecorr.NNCorrelation(config)\n",
    "            nn.process(lcat_1, lcat_2, low_mem=True)\n",
    "            rr.process(rcat_1, rcat_2, low_mem=True)\n",
    "            dr.process(lcat_1, rcat_2, low_mem=True)\n",
    "            rd.process(rcat_1, lcat_2, low_mem=True)\n",
    "\n",
    "        xi, varxi = nn.calculateXi(rr=rr, dr=dr, rd=rd)\n",
    "        theta = np.exp(nn.logr)\n",
    "\n",
    "        m['clu'][f'zbin{zlbin1}_zbin{zlbin2}'] = {'theta': theta, 'w_raw': xi}\n",
    "        np.savetxt(filename, np.c_[theta, xi],\n",
    "                   header='theta [arcmin]  w(theta) [raw, before debiasing]')\n",
    "        print(f'Saved clu {zlbin1}×{zlbin2}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a000019",
   "metadata": {},
   "source": [
    "### 5.4 Load Measurements from Files\n",
    "\n",
    "Run this cell after completing the measurement cells above, or when restarting from a previous cluster run. It loads all three two-point statistics from the saved text files. Continue to Section 5.5 to apply the w(θ) debiasing correction.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "a000020",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded:  10 cosmic shear pairs, 24 GGL pairs, 21 clustering pairs\n"
     ]
    }
   ],
   "source": [
    "m = {'cs': {}, 'ggl': {}, 'clu': {}}\n",
    "\n",
    "# Cosmic shear\n",
    "for zsbin1 in range(n_zsbins):\n",
    "    for zsbin2 in range(zsbin1, n_zsbins):\n",
    "        fname = os.path.join(output_folder, f'cs_zbin{zsbin1}_zbin{zsbin2}.txt')\n",
    "        if os.path.exists(fname):\n",
    "            data = np.loadtxt(fname)\n",
    "            m['cs'][f'zbin{zsbin1}_zbin{zsbin2}'] = {\n",
    "                'theta': data[:, 0], 'xip': data[:, 1], 'xim': data[:, 2]\n",
    "            }\n",
    "\n",
    "# Galaxy-galaxy lensing\n",
    "for zlbin in range(n_zlbins):\n",
    "    for zsbin in range(n_zsbins):\n",
    "        fname = os.path.join(output_folder, f'ggl_zlbin{zlbin}_zsbin{zsbin}.txt')\n",
    "        if os.path.exists(fname):\n",
    "            data = np.loadtxt(fname)\n",
    "            m['ggl'][f'zlbin{zlbin}_zsbin{zsbin}'] = {\n",
    "                'theta': data[:, 0], 'gt': data[:, 1],\n",
    "                'gt_boosted': data[:, 2], 'bf': data[:, 3]\n",
    "            }\n",
    "\n",
    "# Galaxy clustering (raw, before w(θ) debiasing)\n",
    "for zlbin1 in range(n_zlbins):\n",
    "    for zlbin2 in range(zlbin1, n_zlbins):\n",
    "        fname = os.path.join(output_folder, f'clu_zbin{zlbin1}_zbin{zlbin2}.txt')\n",
    "        if os.path.exists(fname):\n",
    "            data = np.loadtxt(fname)\n",
    "            m['clu'][f'zbin{zlbin1}_zbin{zlbin2}'] = {\n",
    "                'theta': data[:, 0], 'w_raw': data[:, 1]\n",
    "            }\n",
    "\n",
    "print(f\"Loaded:  {len(m['cs'])} cosmic shear pairs,\",\n",
    "      f\"{len(m['ggl'])} GGL pairs,\",\n",
    "      f\"{len(m['clu'])} clustering pairs\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a000021",
   "metadata": {},
   "source": [
    "### 5.5 Apply w(θ) Debiasing Correction\n",
    "\n",
    "An additive correction is applied to the raw w(θ) auto-correlations to remove the bias introduced by the MagLim survey property weights. The correction is pre-computed and provided as a FITS table. It applies only to the auto-correlations; cross-correlations are unaffected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "a000022",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "w(θ) debiasing correction applied.\n"
     ]
    }
   ],
   "source": [
    "assert nbins == 26, (\n",
    "    'The debiasing term was computed for 26 angular bins. '\n",
    "    'If you changed nbins you need to recompute the debiasing term.'\n",
    ")\n",
    "\n",
    "debias            = fits.open(path_wtheta_debias)[1].data\n",
    "wtheta_add_debias = debias['wtheta_add_debias']\n",
    "\n",
    "for zlbin in range(n_zlbins):\n",
    "    bin_key = f'zbin{zlbin}_zbin{zlbin}'\n",
    "    if bin_key not in m['clu']:\n",
    "        continue\n",
    "\n",
    "    start_idx    = zlbin * nbins\n",
    "    end_idx      = start_idx + nbins\n",
    "    debias_terms = wtheta_add_debias[start_idx:end_idx]\n",
    "\n",
    "    # Sanity check: angular scales in the debiasing table match our bins (within 2%)\n",
    "    rel_diff = ((debias['theta'][start_idx:end_idx] - m['clu'][bin_key]['theta'])\n",
    "                / m['clu'][bin_key]['theta'])\n",
    "    assert (np.abs(rel_diff) < 0.02).all(), (\n",
    "        f'Angular scale mismatch in lens bin {zlbin}: {rel_diff}'\n",
    "    )\n",
    "\n",
    "    m['clu'][bin_key]['w'] = m['clu'][bin_key]['w_raw'] + debias_terms\n",
    "\n",
    "# Cross-correlations require no debiasing correction\n",
    "for zlbin1 in range(n_zlbins):\n",
    "    for zlbin2 in range(zlbin1 + 1, n_zlbins):\n",
    "        bin_key = f'zbin{zlbin1}_zbin{zlbin2}'\n",
    "        if bin_key in m['clu']:\n",
    "            m['clu'][bin_key]['w'] = m['clu'][bin_key]['w_raw']\n",
    "\n",
    "print('w(θ) debiasing correction applied.')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a000023",
   "metadata": {},
   "source": [
    "## 6. Shape Catalog Statistics\n",
    "\n",
    "The shear calibration responses R computed above also allow us to report the effective number density and per-component shape noise of the Metadetect source catalog. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "a000024",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shear responses: {0: '0.8561', 1: '0.8685', 2: '0.8194', 3: '0.6918'}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Recompute responses in case this cell is run independently of Section 5.1\n",
    "R = {i: compute_response(mdet, i, deltag) for i in range(n_zsbins)}\n",
    "print('Shear responses:', {k: f'{v:.4f}' for k, v in R.items()})\n",
    "print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "a000025",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Metadetect source catalog — object counts and effective number densities:\n",
      "  Source bin 0: 33,707,071 objects,  neff = 2.047 gal/arcmin²\n",
      "  Source bin 1: 34,580,888 objects,  neff = 2.097 gal/arcmin²\n",
      "  Source bin 2: 34,646,416 objects,  neff = 2.138 gal/arcmin²\n",
      "  Source bin 3: 36,727,798 objects,  neff = 2.320 gal/arcmin²\n"
     ]
    }
   ],
   "source": [
    "print('Metadetect source catalog — object counts and effective number densities:')\n",
    "n_density = []\n",
    "for i in range(n_zsbins):\n",
    "    w    = master[f'desy6kp/mdet/noshear/tomo_bin_{i}/w'][:]\n",
    "    nobj = len(w)\n",
    "    neff = np.sum(w)**2 / np.sum(w**2) / area_arcmin\n",
    "    n_density.append(neff)\n",
    "    print(f'  Source bin {i}: {nobj:>9,} objects,  neff = {neff:.3f} gal/arcmin²')\n",
    "n_density = np.array(n_density)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "a000026",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Metadetect source catalog — per-component shape noise σe:\n",
      "  Source bin 0: σe = 0.2657\n",
      "  Source bin 1: σe = 0.2865\n",
      "  Source bin 2: σe = 0.2821\n",
      "  Source bin 3: σe = 0.3474\n",
      "sigma_e input to CosmoCov: 0.42\n",
      "\n",
      "Effective number densities that go into the covariance [gal/arcmin²]:\n",
      "  Source bin 0: 2.5377\n",
      "  Source bin 1: 2.2815\n",
      "  Source bin 2: 2.4444\n",
      "  Source bin 3: 1.7000\n"
     ]
    }
   ],
   "source": [
    "print('Metadetect source catalog — per-component shape noise σe:')\n",
    "sigma_e = []\n",
    "for i in range(n_zsbins):\n",
    "    data    = master[f'desy6kp/mdet/noshear/tomo_bin_{i}']\n",
    "    w       = data['w'][:]\n",
    "    e1      = data['gauss_g_1'][:]\n",
    "    e2      = data['gauss_g_2'][:]\n",
    "\n",
    "    mean_e1 = np.average(e1, weights=w)\n",
    "    mean_e2 = np.average(e2, weights=w)\n",
    "\n",
    "    sum_ws  = np.sum(w * R[i])\n",
    "    sum_w   = np.sum(w)\n",
    "    sum_w2  = np.sum(w**2)\n",
    "\n",
    "    var1    = np.sum(w**2 * (e1 - mean_e1)**2) / sum_ws**2\n",
    "    var2    = np.sum(w**2 * (e2 - mean_e2)**2) / sum_ws**2\n",
    "    se      = np.sqrt(0.5 * (var1 + var2) * sum_w**2 / sum_w2)\n",
    "\n",
    "    sigma_e.append(se)\n",
    "    print(f'  Source bin {i}: σe = {se:.4f}')\n",
    "sigma_e = np.array(sigma_e)\n",
    "\n",
    "# sigma_e_total convention from CosmoCov\n",
    "sigma_e_cosmocov = np.mean(sigma_e)*np.sqrt(2)\n",
    "print(f'sigma_e input to CosmoCov: {sigma_e_cosmocov:.2f}')\n",
    "\n",
    "# Effective number density for covariance (weighted combination across bins)\n",
    "s        = np.sum(n_density * sigma_e * np.sqrt(2)) / np.sum(n_density)\n",
    "\n",
    "\n",
    "# For the source number density that goes as input to the covariance we need to correct \n",
    "# to account for the mean values of the multiplicative shear bias, n^{i}source --> n^{i}_source * (1 + m^{i})^2\n",
    "# considering m^{i} values as in the Table III in the methods paper:\n",
    "mean_multiplicative_shear_bias_methods_paper = np.array([-3.4, 6.46, 15.94, 1.7])*10**(-3) \n",
    "neff_cov = 1.0 / ((sigma_e * np.sqrt(2))**2 / n_density / s**2) * (1+mean_multiplicative_shear_bias_methods_paper)**2\n",
    "print()\n",
    "print('Effective number densities that go into the covariance [gal/arcmin²]:')\n",
    "for i, ne in enumerate(neff_cov):\n",
    "    print(f'  Source bin {i}: {ne:.4f}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a000027",
   "metadata": {},
   "source": [
    "## 7. Assemble the 3×2pt two-point file"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a000028",
   "metadata": {},
   "source": [
    "### 7.1 Load Redshift Distributions\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "a000029",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the redshift distributions. In this case load them from the provided twopoint FITS file.\n",
    "# If using a different sample selection, provide your own n(z) file.\n",
    "nz_file = '3x2pt_2025-07-28-20h_UNBLINDED_wmask_nob2.fits'\n",
    "\n",
    "T_nz      = twopoint.TwoPointFile.from_fits(nz_file)\n",
    "nz_lens   = T_nz.kernels[0]\n",
    "nz_source = T_nz.kernels[1]\n",
    "assert nz_lens.name   == 'nz_lens',   f\"Expected nz_lens, got {nz_lens.name}\"\n",
    "assert nz_source.name == 'nz_source', f\"Expected nz_source, got {nz_source.name}\"\n",
    "\n",
    "assert n_zlbins==len(nz_lens.nzs), f'Lens n(z) bins number does not match.' \n",
    "assert n_zsbins==len(nz_source.nzs), f'Source n(z) bins number does not match.' "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a000030",
   "metadata": {},
   "source": [
    "### 7.2 Build the Data Vector\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "a000031",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data vector assembled.\n"
     ]
    }
   ],
   "source": [
    "from twopoint import Types\n",
    "\n",
    "theta = m['cs']['zbin0_zbin0']['theta']  # same angular grid for all bin pairs\n",
    "\n",
    "types = {\n",
    "    'xip':    (Types.galaxy_shear_plus_real,  Types.galaxy_shear_plus_real),\n",
    "    'xim':    (Types.galaxy_shear_minus_real, Types.galaxy_shear_minus_real),\n",
    "    'gammat': (Types.galaxy_position_real,    Types.galaxy_shear_plus_real),\n",
    "    'wtheta': (Types.galaxy_position_real,    Types.galaxy_position_real),\n",
    "}\n",
    "\n",
    "builder = twopoint.SpectrumCovarianceBuilder()\n",
    "\n",
    "# Cosmic shear: ξ+ and ξ−  (10 source bin pairs)\n",
    "for name in ['xip', 'xim']:\n",
    "    stype = types[name]\n",
    "    for bin1 in range(n_zsbins):\n",
    "        for bin2 in range(bin1, n_zsbins):\n",
    "            for angbin, ang in enumerate(theta):\n",
    "                value = m['cs'][f'zbin{bin1}_zbin{bin2}'][name][angbin]\n",
    "                builder.add_data_point('nz_source', 'nz_source',\n",
    "                                       stype[0], stype[1],\n",
    "                                       bin1 + 1, bin2 + 1, ang, angbin, value)\n",
    "\n",
    "# Galaxy-galaxy lensing: γt  (6 lens × 4 source bin pairs)\n",
    "stype = types['gammat']\n",
    "for lbin in range(n_zlbins):\n",
    "    for sbin in range(n_zsbins):\n",
    "        for angbin, ang in enumerate(theta):\n",
    "            value = m['ggl'][f'zlbin{lbin}_zsbin{sbin}']['gt_boosted'][angbin]\n",
    "            builder.add_data_point('nz_lens', 'nz_source',\n",
    "                                   stype[0], stype[1],\n",
    "                                   lbin + 1, sbin + 1, ang, angbin, value)\n",
    "\n",
    "# Galaxy clustering: w(θ)  (21 pairs: 6 auto + 15 cross)\n",
    "# Auto-correlations use w (debiased); cross-correlations use w (= w_raw, no debiasing applied).\n",
    "stype = types['wtheta']\n",
    "for zlbin1 in range(n_zlbins):\n",
    "    for zlbin2 in range(zlbin1, n_zlbins):\n",
    "        bin_key = f'zbin{zlbin1}_zbin{zlbin2}'\n",
    "        for angbin, ang in enumerate(theta):\n",
    "            value = m['clu'][bin_key]['w'][angbin]\n",
    "            builder.add_data_point('nz_lens', 'nz_lens',\n",
    "                                   stype[0], stype[1],\n",
    "                                   zlbin1 + 1, zlbin2 + 1, ang, angbin, value)\n",
    "\n",
    "print('Data vector assembled.')\n",
    "\n",
    "names = {builder.types[0]:\"xip\", builder.types[1]:\"xim\", builder.types[2]:\"gammat\", builder.types[3]:\"wtheta\"}\n",
    "builder.set_names(names)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a000032",
   "metadata": {},
   "source": [
    "### 7.3 Attach Covariance and Save to FITS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "a000033",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saved: desy6kp_3x2pt_2026-07-02-02h.fits\n"
     ]
    }
   ],
   "source": [
    "# The covariance matrix is computed externally with CosmoLike at the fiducial cosmology\n",
    "# ([Sanchez-Cid et al. 2026](https://arxiv.org/abs/2601.14859)).\n",
    "# Load it here before proceeding. For the fiducial setup, it has shape (N, N), where\n",
    "# N = 1690 is the size of the full uncut data vector, including w(θ) cross-correlations. \n",
    "# It should have the same ordering as the one used above to construct the data-vector.\n",
    "#\n",
    "#   Spectrum   Bin pairs   Angular bins   Elements   Start index\n",
    "#   ξ+         10          26             260         0\n",
    "#   ξ−         10          26             260         260\n",
    "#   γt         24          26             624         520\n",
    "#   w(θ)       21          26             546         1144\n",
    "#   ──────────────────────────────────────────────────────\n",
    "#   Total                                 1690\n",
    "#\n",
    "#covariance_file = 'path/to/cosmolike_covariance.fits' # in general\n",
    "covariance_file = '3x2pt_2025-07-28-20h_UNBLINDED_wmask_nob2.fits'\n",
    "covmat = twopoint.TwoPointFile.from_fits(covariance_file).covmat\n",
    "spectra, covmat_info = builder.generate(covmat, 'arcmin')\n",
    "\n",
    "T_out = twopoint.TwoPointFile(\n",
    "    spectra, [nz_lens, nz_source], windows=None, covmat_info=covmat_info\n",
    ")\n",
    "\n",
    "# Add ANGLEMIN / ANGLEMAX bin edges to each spectrum, required by CosmoSIS.\n",
    "log_lims = np.linspace(np.log(min_sep), np.log(max_sep), nbins + 1)\n",
    "lims     = np.exp(log_lims)\n",
    "\n",
    "for S in T_out.spectra:\n",
    "    S.angle_min = np.zeros(len(S))\n",
    "    S.angle_max = np.zeros(len(S))\n",
    "    for i, j in S.get_bin_pairs():\n",
    "        mask = S.get_pair_mask(i, j)\n",
    "        S.angle_min[mask] = lims[:-1]\n",
    "        S.angle_max[mask] = lims[1:]\n",
    "\n",
    "today   = datetime.datetime.now().strftime('%Y-%m-%d-%Hh')\n",
    "outfile = f'desy6kp_3x2pt_{today}.fits'\n",
    "T_out.to_fits(outfile, overwrite=True)\n",
    "print(f'Saved: {outfile}')\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "2pt_pipeline (py310)",
   "language": "python",
   "name": "2pt_pipeline"
  },
  "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.10.20"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
