{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5ce45be1-e5ce-417a-a6de-38b4752c54e8",
   "metadata": {},
   "source": [
    "# Generating a weighted Monte Carlo lightcone of Diffsky galaxies\n",
    "\n",
    "This notebook demonstrates how to generate a lightcone of diffsky galaxies with SEDs, star formation histories, dust, and other properties."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a7476aa-55fb-46f7-8a9d-a430b9568571",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from matplotlib import pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0920250d-f484-427c-9fda-323b53f6ba10",
   "metadata": {},
   "source": [
    "### Download stellar population synthesis data\n",
    "\n",
    "The very first thing we will do in this notebook is download some LSST-like transmission curves, and also some SEDs of simple stellar populations (SSPs) used in the SED modeling. You can skip this step of you have already cached previously-stored SSP data as described in the [DSPS Quickstart Guide](https://dsps.readthedocs.io/en/latest/quickstart.html).\n",
    "\n",
    "Note that for the demonstration purposes of these docs, we use a \"sparse\" set of SSP SEDs, but high-res SSPs should be used for more accurate SEDs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db896603-f5fe-467f-bc58-bfd7977dd7c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.makedirs('./drn_dsps_temp', exist_ok=True)\n",
    "\n",
    "filter_names = ('u', 'g', 'r', 'i', 'z', 'y')\n",
    "base_url = \"https://portal.nersc.gov/project/hacc/aphearin/DSPS_data/filters/lsst_{}_transmission.h5\"\n",
    "\n",
    "for filt in filter_names:\n",
    "    ! wget -q -P ./drn_dsps_temp {base_url.format(filt)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b11a7ecd-3577-4412-ac39-a52f99cfa5e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "! wget -q -P ./drn_dsps_temp https://portal.nersc.gov/project/hacc/aphearin/DSPS_data/fsps_v0.4.7_mist_c3k_a_kroupa_wNE_logGasU-2.0_logGasZ0.0.sparse.h5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87bf1831-76de-4bfa-aba5-4fab233aae32",
   "metadata": {},
   "source": [
    "#### The `load_ssp_templates` function in DSPS loads an hdf5 file and packs it into a namedtuple with the expected field names."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "842922d6-8fa1-406c-bea7-fc4437618007",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dsps.data_loaders import load_ssp_templates\n",
    "ssp_data = load_ssp_templates(\n",
    "    fn=\"drn_dsps_temp/fsps_v0.4.7_mist_c3k_a_kroupa_wNE_logGasU-2.0_logGasZ0.0.sparse.h5\")\n",
    "\n",
    "print(ssp_data._fields)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22aa0345-4e07-4048-98ea-ca75fbcad923",
   "metadata": {},
   "source": [
    "#### The `load_transmission_curve` function in DSPS loads an hdf5 file and packs it into a namedtuple with the expected field names."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6524445f-a59f-4b0d-b302-8bd369fd3d92",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dsps.data_loaders import load_transmission_curve\n",
    "fn_pat = \"drn_dsps_temp/lsst_{}_transmission.h5\"\n",
    "filter_names = ('u', 'g', 'r', 'i', 'z', 'y')\n",
    "fn_list = [fn_pat.format(x) for x in filter_names]\n",
    "tcurves = [load_transmission_curve(fn) for fn in fn_list]\n",
    "\n",
    "print(tcurves[0]._fields)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebbaba95-5f13-4784-b6fe-59b498c53e8c",
   "metadata": {},
   "source": [
    "#### Generate the halo lightcone and additional data needed to model/compute SEDs\n",
    "\n",
    "First specify halo lightcone specs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0e76a75-fec2-453e-8b5b-80b9452dc92e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffsky.experimental.lightcone_generators import weighted_lc_photdata\n",
    "\n",
    "num_halos = 500\n",
    "z_min, z_max = 0.1, 2.0\n",
    "lgmp_min, lgmp_max = 10.5, 15.0\n",
    "sky_area_degsq = 10.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "556f2d9a-d0c2-460c-ac87-07427d11a95c",
   "metadata": {},
   "source": [
    "#### Get a random number seed from JAX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58b03053-0182-4873-bcad-b78dd42b2dd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax import random as jran\n",
    "ran_key = jran.key(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cfc8ee62-9679-40ef-a94c-576cdbeb8715",
   "metadata": {},
   "source": [
    "#### Define a redshift table used for photometry interpolation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52892def-824e-4bcf-b347-2db6012bd1fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_z_phot_table = 15\n",
    "z_phot_table = np.linspace(z_min, z_max, n_z_phot_table)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f954c3a3-b9ef-4a4e-966f-4f30262990bd",
   "metadata": {},
   "source": [
    "#### Generate the lightcone of dark matter halos and additional SPS data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eab1fbb8-13f7-421e-9178-fd49128f3d1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "halo_lc_data = (num_halos, z_min, z_max, lgmp_min, lgmp_max, sky_area_degsq)\n",
    "phot_data = (ssp_data, tcurves, z_phot_table)\n",
    "\n",
    "ran_key, lc_halo_key = jran.split(ran_key, 2)\n",
    "args = (lc_halo_key, *halo_lc_data, *phot_data)\n",
    "\n",
    "lc_data = weighted_lc_photdata(*args)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51918c4d-b3a2-4211-b25c-18843f7755b0",
   "metadata": {},
   "source": [
    "#### Populate the lightcone with diffsky galaxies\n",
    "\n",
    "Now we will pass the `lc_data` to the diffsky SED kernels to populate the halo lightcone with galaxies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6004ccc-85c1-401b-bf01-31808233d4c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffsky.experimental.mc_phot import mc_lc_phot\n",
    "\n",
    "mc_merge = 0\n",
    "ran_key, sed_key = jran.split(ran_key, 2)\n",
    "phot_info, __, __ = mc_lc_phot(sed_key, lc_data, mc_merge)\n",
    "print(phot_info._fields)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98aa5ddd-be8d-4291-a4e9-9b61cb167848",
   "metadata": {},
   "source": [
    "### Choosing an alternative set of diffsky model parameters\n",
    "\n",
    "There are several choices for diffsky parameters that are based on ongoing work calibrating Diffsky to the COSMOS and FENIKS-UDS datasets: "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1b8e4a8-56b4-4611-aa30-c1bc9b6ecca1",
   "metadata": {},
   "source": [
    "#### COSMOS calibrated diffsky model parameters\n",
    "\n",
    "Below we show how to check which COSMOS parameters are available, and then select the `cosmos_260210` calibration to populate the lightcone."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddb34de4-fb7b-4333-ae4c-664dc5f66195",
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffsky.param_utils.cosmos_calibrations import COSMOS_PARAM_FITS\n",
    "print(COSMOS_PARAM_FITS.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "426e0524-7307-4cd7-a4b2-e12a17971257",
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffsky.param_utils.get_mock_params import get_param_collection_for_mock\n",
    "cosmos_param_collection = get_param_collection_for_mock(cosmos_fit='cosmos_260210')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80b1d1ce-8058-4527-a19a-7488290409a6",
   "metadata": {},
   "source": [
    "#### FENIKS-UDS calibrated diffsky model parameters\n",
    "\n",
    "Below we show how to load a FENIKS-UDS based calibration `feniks_260617`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73f258f7-f0a5-4b67-bb20-a934dfd6c7ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffsky.param_utils.get_calib_params import get_calib_params\n",
    "feniks_param_collection = get_calib_params(\n",
    "    calibration_dir=\"feniks_calibrations\", calibration_name=\"feniks_260617\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be2e23b4-8166-4d09-85f1-9d6c00e15a8d",
   "metadata": {},
   "source": [
    "#### We can generate the photometry using these alternative set of diffsky model parameters as below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06a1cdb6-0b51-4df5-8dd7-71bb515ff03d",
   "metadata": {},
   "outputs": [],
   "source": [
    "phot_info_feniks, _, _ = mc_lc_phot(sed_key, lc_data, mc_merge, param_collection = feniks_param_collection)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "437b9cf9-c091-46b3-bd20-a88ef4875826",
   "metadata": {},
   "source": [
    "#### For the _weighted_ lightcone, each halo has multiplicity according to its abundance in the volume\n",
    "\n",
    "This `cen_weight` and `sat_weight` columns need to be taken into account when predicting summary statistics from the weighted lightcone. The product of these two weights supplies the total weight of each galaxy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7a12968-c48a-4ea6-97ba-c349a73f7c1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "gal_weight = lc_data.cen_weight * lc_data.sat_weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5eda0e6a-dc9b-4a1d-a5d5-b0b3ccd2644b",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1)\n",
    "__=ax.loglog()\n",
    "__=ax.scatter(10**lc_data.logmp_obs, gal_weight, s=1)\n",
    "xlabel = ax.set_xlabel(r'$M_{\\rm halo}\\ [M_{\\odot}]$')\n",
    "ylabel = ax.set_ylabel(r'$N_{\\rm halos}$')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a6ec64f-2004-4177-8bfb-2859f3cdd57f",
   "metadata": {},
   "source": [
    "#### Calculate the halo mass function, accounting for halo weights\n",
    "\n",
    "The unweighted version uniformly spans $\\log_{10}M_{\\rm halo}$. The weighted version has the expected Schechter-type shape of the HMF."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e23d2748-b2d0-4f88-b722-e1f10a82d5f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1)\n",
    "yscale = ax.set_yscale('log')\n",
    "__=ax.hist(lc_data.logmp_obs, bins=100, alpha=0.7, label=r'${\\rm unweighted}$')\n",
    "__=ax.hist(lc_data.logmp_obs, bins=100, weights=gal_weight, alpha=0.7, label=r'${\\rm weighted}$')\n",
    "xlabel = ax.set_xlabel(r'$\\log_{10}M_{\\rm halo}/M_{\\odot}$')\n",
    "ylabel = ax.set_ylabel(r'$N_{\\rm halos}$')\n",
    "leg = ax.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5bfb9db-335e-4320-bd68-f7f041448221",
   "metadata": {},
   "source": [
    "#### Visually inspect the diversity of SFHs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f51a9e6-7b23-4d9a-b28a-ca25df550b08",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1)\n",
    "yscale = ax.set_yscale('log')\n",
    "ylim = ax.set_ylim(8e-4, 5e2)\n",
    "xscale = ax.set_xscale('log')\n",
    "xlim = ax.set_xlim(1, 15)\n",
    "\n",
    "n_plot = 5\n",
    "for i in range(n_plot):\n",
    "    __=ax.plot(lc_data.t_table, phot_info.sfh_table[i, :])\n",
    "\n",
    "xlabel = ax.set_xlabel(r'${\\rm cosmic\\ time\\ [Gyr]}$')\n",
    "ylabel = ax.set_ylabel(r'${\\rm SFR\\ [M_{\\odot}/yr]}$')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab680ef9-5c74-459f-8ab0-4ea48e9abf28",
   "metadata": {},
   "source": [
    "#### Visually inspect the sSFR PDF\n",
    "\n",
    "Note that the plot below shows the PDF for _all_ galaxies/halos in the lightcone, without accounting for the weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "302afa71-5a0c-4472-b72d-0d95689c9fc7",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1)\n",
    "xlim = ax.set_xlim(-15, -7)\n",
    "yscale = ax.set_yscale('log')\n",
    "__=ax.hist(phot_info.logssfr_obs, bins=70, alpha=0.7, weights=gal_weight, density=True)\n",
    "xlabel = ax.set_xlabel(r'${\\rm log_{10}(sSFR)}$')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a17a2e4a-7933-43c6-addc-f01d0b313c03",
   "metadata": {},
   "source": [
    "#### Visually inspect star-forming sequence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1c2de09-8245-44c0-ba91-0b192cd52428",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1)\n",
    "ylim = ax.set_ylim(-13.5, -7.5)\n",
    "xlim = ax.set_xlim(7.5, 13)\n",
    "__=ax.scatter(phot_info.logsm_obs, phot_info.logssfr_obs, s=1)\n",
    "xlabel = ax.set_xlabel(r'${\\rm log_{10}(M_{\\star})}$')\n",
    "ylabel = ax.set_ylabel(r'${\\rm log_{10}(sSFR)}$')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6eca1b85-c18a-4b91-a5e2-5687786e482c",
   "metadata": {},
   "source": [
    "### Plot color-color diagram\n",
    "\n",
    "Note that the plot below shows the PDF for _all_ galaxies/halos in the lightcone, without accounting for the weights, and without selecting a particular galaxy sample of interest."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32d2b400-394a-4b61-a5e9-423fbebb2d5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1)\n",
    "\n",
    "ri = phot_info.obs_mags[:, 2]-phot_info.obs_mags[:, 3]\n",
    "iz = phot_info.obs_mags[:, 3]-phot_info.obs_mags[:, 4]\n",
    "__=ax.scatter(iz, ri, s=1)\n",
    "xlabel = ax.set_xlabel(r'${\\rm i-z}$')\n",
    "ylabel = ax.set_ylabel(r'${\\rm r-i}$')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7929997-0377-4eee-88bf-050247dfa367",
   "metadata": {},
   "source": [
    "### Plot color--redshift relation\n",
    "\n",
    "Note that the plot below shows the PDF for _all_ galaxies/halos in the lightcone, without accounting for the weights, and without selecting a particular galaxy sample of interest."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf4348e5-0b51-4813-a7cf-9a4debe3ea65",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1)\n",
    "\n",
    "iz = phot_info.obs_mags[:, 3]-phot_info.obs_mags[:, 4]\n",
    "__=ax.scatter(lc_data.z_obs, iz, s=1)\n",
    "xlabel = ax.set_xlabel(r'${\\rm redshift}$')\n",
    "ylabel = ax.set_ylabel(r'${\\rm i-z}$')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67652be3-4861-4012-a14e-1ccb0c60ad16",
   "metadata": {},
   "source": [
    "#### Docs about accounting for the weights, computing high-res SEDs, and more coming soon..."
   ]
  }
 ],
 "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.13.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}