{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5ce45be1-e5ce-417a-a6de-38b4752c54e8",
   "metadata": {},
   "source": [
    "# Generating a 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",
    "Before generating these docs, these data were already been downloaded, and so we will just load them directly. See the DSPS docs for download instructions."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87bf1831-76de-4bfa-aba5-4fab233aae32",
   "metadata": {},
   "source": [
    "#### The `load_ssp_templates` function in Diffsky loads an hdf5 file and packs it into a namedtuple with the expected field names.\n",
    "\n",
    "> **Note:** This tutorial uses unrealistically sparse SSP data to minimize memory demands of generating the documentation. You should use full-resolution SSP SEDs in real applications."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "842922d6-8fa1-406c-bea7-fc4437618007",
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffsky.data_loaders import load_ssp_templates\n",
    "ssp_data = load_ssp_templates(\n",
    "    drn=\"dsps_drn\",\n",
    "    bn=\"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.\n",
    "\n",
    "You can use any transmission curves to compute Diffsky photometry. You just need to create a namedtuple storing 'wave', 'transmission' of each transmission curve."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6524445f-a59f-4b0d-b302-8bd369fd3d92",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dsps.data_loaders import load_transmission_curve\n",
    "bn_base_pat = \"lsst_{}*\"\n",
    "filter_names = ('u', 'g', 'r', 'i', 'z', 'y')\n",
    "bn_pat_list = [bn_base_pat.format(x) for x in filter_names]\n",
    "tcurves = [load_transmission_curve(bn_pat=pat, drn=\"dsps_drn/filters\") for pat in bn_pat_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:\n",
    "* `ran_key`: Random number seed.\n",
    "* `z_min, z_max`: the redshift range\n",
    "* `lgmp_min` should be small enough to resolve the faintest galaxies in your population (values of 10-11 are typical for cosmological surveys).\n",
    "* `lgmsub_min` is the lowest mass of the subhalos. In most cases, this should typically be set equal to `lgmp_min`.\n",
    "* `sky_area_degsq` is the sky area in square degrees."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0e76a75-fec2-453e-8b5b-80b9452dc92e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffsky.experimental.mc_lightcone_generators import mc_lc_photdata\n",
    "\n",
    "z_min, z_max = 0.1, 0.15\n",
    "lgmp_min, lgmp_sub_min = 11.0, 11.0\n",
    "sky_area_degsq = 1.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "606041b1-7f08-45a9-b968-9aaa94d696cd",
   "metadata": {},
   "source": [
    "> **Note:** This tutorial demonstrates how to use the mc_lc_photdata and mc_lc_phot functions, which is currently imported from within `diffsky.experimental`. Code located within `diffsky.experimental` may not have a stable API, and will be relocated elsewhere in the library after development of this feature settles down. This tutorial is built dynamically with each PR, and so the syntax demonstrated here is always up-to-date with the main branch of the code."
   ]
  },
  {
   "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": [
    "ran_key, lc_halo_key = jran.split(ran_key, 2)\n",
    "args = (lc_halo_key,\n",
    "    z_min,\n",
    "    z_max,\n",
    "    lgmp_min,\n",
    "    lgmp_sub_min,\n",
    "    sky_area_degsq,\n",
    "    ssp_data,\n",
    "    tcurves,\n",
    "    z_phot_table,)\n",
    "\n",
    "lc_data = mc_lc_photdata(*args)\n",
    "print(lc_data._fields)\n",
    "print(f\"NUmber of halos = {lc_data.logmp_obs.size}\")"
   ]
  },
  {
   "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.\n",
    "\n",
    "For Monte Carlo realizations, we set `mc_merge = 1`: satellite galaxies that are either merged with their associated central (`p_merge=1`, or they continue to exist as they orbit within the host halo `p_merge=0`. The photometry and stellar mass of merged satellites and their centrals has been self-consistently calculated by the `mc_lc_phot` function, and so you can treat merged satellites as if they do not exist by simply masking them out of your analysis."
   ]
  },
  {
   "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 = 1  # select descrete realization of merging PDF for Monte Carlo lightcones\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": "code",
   "execution_count": null,
   "id": "0eec7b42-f9d4-47a6-a6de-196bc94c7b80",
   "metadata": {},
   "outputs": [],
   "source": [
    "phot_info.ssp_weights.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c171eea5-06f0-42d9-826c-4cf378179b04",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(np.unique(phot_info.p_merge))\n",
    "print(np.mean(phot_info.p_merge>0))"
   ]
  },
  {
   "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. Below we show how to check which parameters are available."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddb34de4-fb7b-4333-ae4c-664dc5f66195",
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffsky.param_utils import DIFFSKY_FIT_PARAMS\n",
    "print(DIFFSKY_FIT_PARAMS.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "426e0524-7307-4cd7-a4b2-e12a17971257",
   "metadata": {},
   "outputs": [],
   "source": [
    "cosmos_param_collection = DIFFSKY_FIT_PARAMS['c260710']\n",
    "feniks_param_collection = DIFFSKY_FIT_PARAMS['sdss_feniks_hizels_260710']"
   ]
  },
  {
   "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": "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": "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.p_merge==0], phot_info.logssfr_obs[phot_info.p_merge==0], 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[phot_info.p_merge==0], ri[phot_info.p_merge==0], 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",
    "ylim = ax.set_ylim(-0.8, 0.8)\n",
    "iz = phot_info.obs_mags[:, 3]-phot_info.obs_mags[:, 4]\n",
    "\n",
    "msk_cens = lc_data.is_central==1\n",
    "msk_sats = (lc_data.is_central==0) & (phot_info.p_merge==0)\n",
    "__=ax.scatter(lc_data.z_obs[msk_cens], iz[msk_cens], s=1, label=r'${\\rm centrals}$')\n",
    "__=ax.scatter(lc_data.z_obs[msk_sats], iz[msk_sats], s=1, label=r'${\\rm satellites}$')\n",
    "xlabel = ax.set_xlabel(r'${\\rm redshift}$')\n",
    "ylabel = ax.set_ylabel(r'${\\rm i-z}$')\n",
    "leg = ax.legend(markerscale=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71bb3160-09d1-474d-a006-add312491fe7",
   "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.13.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}