{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6d1b82ae",
   "metadata": {},
   "source": [
    "# Calculating photometry for Roman/Rubin 2023 Diffsky galaxies\n",
    "\n",
    "This notebook illustrates how to calculate photometry predictions from the physical parameters of each diffsky galaxy. There's demo code showing how to compute exact photometry for a single galaxy at a time, and also a demo of calculating approximate photometry for large galaxy populations at once.\n",
    "\n",
    "#### Download mock galaxy data\n",
    "\n",
    "First we'll download a very small dataset that stores a downsampling of data from a single healpixel of the `roman_rubin_2023` mock."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "291c114f",
   "metadata": {},
   "outputs": [],
   "source": [
    "! curl https://portal.nersc.gov/project/hacc/aphearin/lsstdesc_diffsky_data/roman_rubin_2023_z_0_1_cutout_9043.testdata.hdf5 > diffsky.testdata.hdf5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "953dadfb",
   "metadata": {},
   "source": [
    "### Load the Diffsky data from the test healpixel\n",
    "\n",
    "The next cell directly reads the hdf5 file storing the mock data. This test file is formatted in the same way as the healpixels distributed on NERSC: the data is separated by the simulation snapshot, and there is a metadata column storing additional information. For demonstration purposes, we'll just load galaxies directly from the hdf5 file with the `load_diffsky_healpixel` convenience function. The returned mock stores the full collection of data from the snapshots, concatenated and stored as a flat ndarray for each column.\n",
    "\n",
    "DESC users working at NERSC may instead wish to use the [GCR](https://github.com/yymao/generic-catalog-reader)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "715551eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from lsstdesc_diffsky.io_utils import load_healpixel\n",
    "\n",
    "fn = \"diffsky.testdata.hdf5\"\n",
    "patlist = ('LSST', )\n",
    "mock, metadata = load_healpixel(fn, patlist)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0966525f",
   "metadata": {},
   "source": [
    "### Download and inspect template SEDs for SSPs\n",
    "\n",
    "In this next cell we'll download the template SEDs of the simple stellar populations used to compute the galaxy SEDs.\n",
    "\n",
    "**Note:** The demos below are the same single-metallicity kernels used to generate the `roman_rubin_2023` mock, and so in the following cell we download single-metallicitiy SSP SEDs. Multi-metallicity kernels and SSPs are also included in the `lsstdesc-diffsky` library, so take care to use the single-metallicity versions as shown below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "489a1cc6",
   "metadata": {},
   "outputs": [],
   "source": [
    "! curl https://portal.nersc.gov/project/hacc/aphearin/DSPS_data/ssp_data_fsps_v3.2_age.h5 > dsps_ssp_data_singlemet.h5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fca85a42",
   "metadata": {},
   "source": [
    "The `ssp_data` quantity stores all the information needed from the Simple Stellar Population templates to compute our galaxy SEDs. Note that these kernels are imported a custom version of dsps located within lsstdesc_diffsky:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12fa1b62",
   "metadata": {},
   "outputs": [],
   "source": [
    "from lsstdesc_diffsky.legacy.roman_rubin_2023.dsps.data_loaders.load_ssp_data import load_ssp_templates_singlemet\n",
    "ssp_data = load_ssp_templates_singlemet(fn='dsps_ssp_data_singlemet.h5')\n",
    "\n",
    "print(ssp_data._fields)\n",
    "\n",
    "print('ssp_lg_age_gyr.shape = {}'.format(ssp_data.ssp_lg_age_gyr.shape))\n",
    "print('ssp_wave.shape = {}'.format(ssp_data.ssp_wave.shape))\n",
    "print('ssp_flux.shape = {}'.format(ssp_data.ssp_flux.shape))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bb7b954",
   "metadata": {},
   "source": [
    "### Load transmission curves\n",
    "\n",
    "The dsps library ships with a few transmission curves as a convenience for getting started. Up-to-date transmission curves for a wide range of instruments can be found from publicly available sources such as the [kcorrect](https://github.com/blanton144/kcorrect/) library and elsewhere.\n",
    "\n",
    "**Note on units:** Wherever you get your transmission curves, double-check that your wavelengths are in angstroms, as these are the $\\lambda$ units used in the dsps library we'll use to calculate photometry.\n",
    "\n",
    "In this next cell, we'll just use dsps to generate on-the-fly a few transmission curves that roughly approximate LSST-band photometry in _ugrizy_."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e309508",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dsps.data_loaders.retrieve_fake_fsps_data import load_fake_filter_transmission_curves\n",
    "wave, *trans_curves = load_fake_filter_transmission_curves()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8df68d26",
   "metadata": {},
   "source": [
    "#### Interpolate transmission curves to a common wavelength\n",
    "\n",
    "We will vectorize our photometry computations across the set of filters, so with these calculations it's always necessary to interpolate the collection of filters to be defined by arrays of the same length. In this demo this is trivial since the fake transmission curve generator already returns all transmission curves to be defined on the same wavelength grid, but a real transmission curve is typically defined on its own specialized grid in wavelength. The code below handles this with the `interpolate_filter_trans_curves` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3dc95ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "from lsstdesc_diffsky.photometry.precompute_ssp_tables import interpolate_filter_trans_curves\n",
    "\n",
    "wave_filters = [wave for x in trans_curves]\n",
    "trans_filters = [x for x in trans_curves]\n",
    "rest_filter_waves, rest_filter_trans = interpolate_filter_trans_curves(wave_filters, trans_filters)\n",
    "obs_filter_waves, obs_filter_trans = interpolate_filter_trans_curves(wave_filters, trans_filters)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8dd72043",
   "metadata": {},
   "source": [
    "### Retrieve Diffsky parameters of each galaxy\n",
    "\n",
    "Each individual diffsky galaxy has its own parameters controlling its assembly history and SED. The `load_diffsky_params` function interprets the columns of the mock that store these parameters, and returns a collection of arrays that are formatted and shaped in the form expected by the function used to compute the SED of the disk, bulge, and knots."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40e13ca2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from lsstdesc_diffsky.io_utils import load_diffsky_params\n",
    "diffsky_params = load_diffsky_params(mock)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "614bfc77",
   "metadata": {},
   "source": [
    "### Retrieve DiffskyPop parameters for the `roman_rubin_2023` mock\n",
    "\n",
    "\n",
    "The DiffskyPop model has a number of parameters controlling the probabilistic relationships of the galaxy--halo connection. The next cell retrieves the values of these parameters used to generate the `roman_rubin_2023` mock. The returned quantity is a NamedTuple with a field name for the parameters of each DiffskyPop model ingredient."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41f52353",
   "metadata": {},
   "outputs": [],
   "source": [
    "from lsstdesc_diffsky.param_data import read_diffskypop_params\n",
    "diffskypop_params = read_diffskypop_params('roman_rubin_2023')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20e8c083",
   "metadata": {},
   "source": [
    "## Calculating exact photometry for an individual galaxy\n",
    "\n",
    "The code below shows how to compute the photometry of an individual diffsky galaxy. We'll see the corresponding kernel for calculating photometry of populations of diffsky galaxies. After that, we'll demo a much faster and more memory efficient kernel for calculating approximate photometry.\n",
    "\n",
    "**Note:** The demos below are the same single-metallicity kernels used to generate the `roman_rubin_2023` mock. Multi-metallicity kernels are also included in the `lsstdesc-diffsky` library, so take care to import from modules with `_singlemet` in the name as shown below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "defae7d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from lsstdesc_diffsky.defaults import OUTER_RIM_COSMO_PARAMS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfc9623b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from lsstdesc_diffsky.photometry.photometry_kernels_singlemet import calc_photometry_singlegal\n",
    "\n",
    "igal = 0\n",
    "\n",
    "args = (\n",
    "    mock['redshift'][igal],\n",
    "    diffsky_params.mah_params[igal],\n",
    "    diffsky_params.ms_params[igal],\n",
    "    diffsky_params.q_params[igal],\n",
    "    ssp_data,\n",
    "    diffskypop_params,\n",
    "    rest_filter_waves,\n",
    "    rest_filter_trans,\n",
    "    obs_filter_waves,\n",
    "    obs_filter_trans,\n",
    "    OUTER_RIM_COSMO_PARAMS)\n",
    "\n",
    "_res = calc_photometry_singlegal(*args)\n",
    "rest_mags, obs_mags, rest_mags_nodust, obs_mags_nodust = _res"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95948497",
   "metadata": {},
   "source": [
    "## Calculating exact photometry for a galaxy population"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7d11b24",
   "metadata": {},
   "outputs": [],
   "source": [
    "from lsstdesc_diffsky.photometry.photometry_kernels_singlemet import calc_photometry_galpop\n",
    "\n",
    "args = (\n",
    "    mock['redshift'],\n",
    "    diffsky_params.mah_params,\n",
    "    diffsky_params.ms_params,\n",
    "    diffsky_params.q_params,\n",
    "    ssp_data,\n",
    "    diffskypop_params,\n",
    "    rest_filter_waves,\n",
    "    rest_filter_trans,\n",
    "    obs_filter_waves,\n",
    "    obs_filter_trans,\n",
    "    OUTER_RIM_COSMO_PARAMS)\n",
    "\n",
    "_res = calc_photometry_galpop(*args)\n",
    "rest_mags, obs_mags, rest_mags_nodust, obs_mags_nodust = _res"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40ecfd3d",
   "metadata": {},
   "source": [
    "## Calculating approximate photometry for a galaxy population\n",
    "\n",
    "There are a couple of approximations we can make in calculating the photometry that offer orders-of-magnitude improvement in memory efficiency and runtime. These approximations give a statistically unbiased magnitude with typical scatter of $\\sigma(\\delta{\\rm mag})\\approx0.05$, and so calculating approximate photometry can be much more practical when studying large galaxy populations. The details behind the approximations are described in [these notes](https://www.overleaf.com/read/pqptngggwrhw#c4af95).\n",
    "\n",
    "\n",
    "#### Precompute photometry of SSP template SEDs\n",
    "\n",
    "In the next cell we compute the restframe magnitudes of our SSPs through each filter, and in the following cell we do the same calculation for apparent magnitudes of our SSPs on a grid of redshift."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3c09a3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from lsstdesc_diffsky.photometry.precompute_ssp_tables import precompute_ssp_obsmags_on_z_table_singlemet\n",
    "import numpy as np\n",
    "\n",
    "z_table = np.linspace(0.01, mock['redshift'].max()+0.05, 51)\n",
    "\n",
    "ssp_obsmag_table = precompute_ssp_obsmags_on_z_table_singlemet(\n",
    "    ssp_data.ssp_wave, ssp_data.ssp_flux,\n",
    "    obs_filter_waves, obs_filter_trans, z_table,\n",
    "    *OUTER_RIM_COSMO_PARAMS[:-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e946e7e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from lsstdesc_diffsky.photometry.precompute_ssp_tables import precompute_ssp_restmags_singlemet\n",
    "\n",
    "ssp_restmag_table = precompute_ssp_restmags_singlemet(\n",
    "    ssp_data.ssp_wave, ssp_data.ssp_flux, rest_filter_waves, rest_filter_trans)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a4345aa",
   "metadata": {},
   "source": [
    "### Compute galaxy photometry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee22b9e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from lsstdesc_diffsky.photometry.photometry_lc_interp_singlemet import get_diffsky_sed_info_singlemet\n",
    "from dsps.cosmology.flat_wcdm import age_at_z0\n",
    "from jax import random as jran\n",
    "\n",
    "ran_key = jran.PRNGKey(0)\n",
    "\n",
    "t0 = age_at_z0(*OUTER_RIM_COSMO_PARAMS[:-1])\n",
    "gal_t_table = np.linspace(0.1, t0, 100)\n",
    "\n",
    "ssp_z_table = np.linspace(mock['redshift'].min()/2, mock['redshift'].max()+0.1, 51)\n",
    "\n",
    "args = (ran_key,\n",
    "    mock['redshift'],\n",
    "    diffsky_params.mah_params,\n",
    "    diffsky_params.ms_params,\n",
    "    diffsky_params.q_params,\n",
    "    ssp_z_table,\n",
    "    ssp_restmag_table,\n",
    "    ssp_obsmag_table,\n",
    "    ssp_data,\n",
    "    gal_t_table,\n",
    "    rest_filter_waves,\n",
    "    rest_filter_trans,\n",
    "    obs_filter_waves,\n",
    "    obs_filter_trans,\n",
    "    diffskypop_params,\n",
    "    OUTER_RIM_COSMO_PARAMS)\n",
    "\n",
    "sed_info = get_diffsky_sed_info_singlemet(*args)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89ae3137",
   "metadata": {},
   "source": [
    "#### Interpreting the results\n",
    "\n",
    "The returned `sed_info` is a namedtuple that stores the calculated photometry along with numerous other quantities computed along the way.\n",
    "* Columns `gal_obsmags_dust` stores the apparent magnitudes of each galaxy on each input observer-frame filter\n",
    "* Columns `gal_restmags_dust` stores the absolute magnitude in the restframe of the galaxy\n",
    "* The `_dust` and `_nodust` quantities refer to whether or not dust attenuation in the emitting galaxy is included in the magnitude computation (in particular, neither quantity takes any account of dust in the Milky Way)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95d0d589",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(sed_info._fields)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd8ae963",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(sed_info.gal_obsmags_dust.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76d60f9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "\n",
    "fig, ax = plt.subplots(1, 1)\n",
    "ylim = ax.set_ylim(-0.2, 1.25)\n",
    "__=ax.scatter(\n",
    "    mock['redshift'], \n",
    "    sed_info.gal_restmags_dust[:, 1]-sed_info.gal_restmags_dust[:, 2], s=2)\n",
    "\n",
    "xlabel = ax.set_xlabel('redshift')\n",
    "ylabel = ax.set_ylabel('g-r')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a2c4ecf",
   "metadata": {},
   "outputs": [],
   "source": [
    "iband = 1\n",
    "\n",
    "fig, ax = plt.subplots(1, 1)\n",
    "xlim = ax.set_xlim(-0.05, 0.05)\n",
    "__=ax.hist(rest_mags[:, iband] - sed_info.gal_restmags_dust[:, iband])\n",
    "\n",
    "xlabel = ax.set_xlabel(r'$\\delta m_{g}$')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72c74ea8",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert np.allclose(rest_mags, sed_info.gal_restmags_dust, atol=0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7c3bafa",
   "metadata": {},
   "source": [
    "### Now clean up the temporary files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1720e86c",
   "metadata": {},
   "outputs": [],
   "source": [
    "! rm dsps_ssp_data_singlemet.h5\n",
    "! rm diffsky.testdata.hdf5"
   ]
  }
 ],
 "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.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}