{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "45059532",
   "metadata": {},
   "source": [
    "# Quickstart Guide to DSPS\n",
    "\n",
    "This demo notebook begins by downloading the DSPS default option for the SSP spectral library. These data are stored [at this URL](https://portal.nersc.gov/project/hacc/aphearin/DSPS_data/) in a flat hdf5 file with column names as expected by the `dsps.load_ssp_templates` function, which we will demonstrate below.\n",
    "\n",
    "When downloading and storing SSP libraries, you can optionally use the DSPS_DRN environment variable to specify the default location where DSPS will look for SSP libraries. But here we'll just save the downloaded data to `tempdata.h5`,  directly pass the filename to the data loader. The `load_ssp_templates` that we'll use to load these SSPs is just a convenience function - all of the DSPS functions that we'll demonstrate in this notebook accept plain arrays and floats as inputs, and so you can store your SSP data on disk in whatever format you like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7330b9d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "! curl https://portal.nersc.gov/project/hacc/aphearin/DSPS_data/ssp_data_fsps_v3.2_lgmet_age.h5 > tempdata.h5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92a4e8ef",
   "metadata": {},
   "source": [
    "### Inspect the SSP data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ffb47cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dsps import load_ssp_templates\n",
    "ssp_data = load_ssp_templates(fn='tempdata.h5')\n",
    "\n",
    "print(ssp_data._fields)\n",
    "\n",
    "print('\\nssp_lgmet.shape = {}'.format(ssp_data.ssp_lgmet.shape))\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": "16888a4b",
   "metadata": {},
   "source": [
    "The returned `ssp_data` is a namedtuple storing 4 ndarrays for the age-metallicity grid of the SSP spectra. Galaxy SEDs are calculated via probability-weighted sums of these spectral templates. For a galaxy observed at some $t_{\\rm obs}$, we'll calculate the restframe SED of two different models in the cells below:\n",
    "\n",
    "1. a galaxy with a tabulated star formation history (SFH), and metallicity Z distributed as a lognormal about some median Z, using the `calc_rest_sed_sfh_table_lognormal_mdf` function. \n",
    "2. a galaxy with SFH table and also tabulated history of metallicity (ZH), using the `calc_rest_sed_sfh_table_met_table` function.\n",
    "\n",
    "In the cells below, we'll randomly generate an SFH and ZH for a galaxy, and then plot the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d7fc23d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "gal_t_table = np.linspace(0.05, 13.8, 100) # age of the universe in Gyr\n",
    "gal_sfr_table = np.random.uniform(0, 10, gal_t_table.size) # SFR in Msun/yr\n",
    "\n",
    "gal_lgmet = -2.0 # log10(Z)\n",
    "gal_lgmet_scatter = 0.2 # lognormal scatter in the metallicity distribution function"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31daa8a6",
   "metadata": {},
   "source": [
    "The SED calculating functions require you specify the _time_ of the observation, `t_obs`, rather than the redshift, `z_obs`. We'll use the `age_at_z` function in `dsps.cosmology` to calculate the relationship between these two quantities, assuming the default redshift of DSPS. You could also use this same function to compute `gal_t_table` in case your input SFH is tabulated as a function of redshift."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14c9f1a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dsps.cosmology import age_at_z, DEFAULT_COSMOLOGY\n",
    "\n",
    "print(DEFAULT_COSMOLOGY)\n",
    "\n",
    "z_obs = 0.5\n",
    "t_obs = age_at_z(z_obs, *DEFAULT_COSMOLOGY) # age of the universe in Gyr at z_obs\n",
    "t_obs = t_obs[0] # age_at_z function returns an array, but SED functions accept a float for this argument"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "841d789c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dsps import calc_rest_sed_sfh_table_lognormal_mdf\n",
    "from dsps import calc_rest_sed_sfh_table_met_table\n",
    "\n",
    "sed_info = calc_rest_sed_sfh_table_lognormal_mdf(\n",
    "    gal_t_table, gal_sfr_table, gal_lgmet, gal_lgmet_scatter,\n",
    "    ssp_data.ssp_lgmet, ssp_data.ssp_lg_age_gyr, ssp_data.ssp_flux, t_obs)\n",
    "\n",
    "\n",
    "gal_lgmet_table = np.linspace(-3, -2, gal_t_table.size)\n",
    "\n",
    "sed_info2 = calc_rest_sed_sfh_table_met_table(\n",
    "    gal_t_table, gal_sfr_table, gal_lgmet_table, gal_lgmet_scatter,\n",
    "    ssp_data.ssp_lgmet, ssp_data.ssp_lg_age_gyr, ssp_data.ssp_flux, t_obs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "843af661",
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "\n",
    "fig, ax = plt.subplots(1, 1)\n",
    "__=ax.loglog()\n",
    "__=ax.plot(ssp_data.ssp_wave, sed_info.rest_sed)\n",
    "__=ax.plot(ssp_data.ssp_wave, sed_info2.rest_sed)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "228064ef",
   "metadata": {},
   "source": [
    "## Calculating photometry\n",
    "\n",
    "Now we'll use dsps.photometry to calculate the apparent and absolute magnitudes of an SED observed through a broadband filter. One additional ingredient we need for this calculation is the transmission curve of some filter. For this, we'll download one from the same public URL where we previously downloaded the SSP spectra, and just write the result to a temporary file. But the `load_transmission_curve` function also supports the use of the DSPS_DRN environment variable in case you want DSPS to remember your default data location. The only difference is that you should store your transmission curves in the `filters` subdirectory of DSPS_DRN. And as above, you can also ignore this data-loading convenience function and store your transmission curves wherever you like and in whatever format you prefer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db1aa5d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "! curl https://portal.nersc.gov/project/hacc/aphearin/DSPS_data/filters/lsst_g_transmission.h5 > tempfilter.h5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76509843",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dsps import load_transmission_curve\n",
    "lsst_g = load_transmission_curve(fn=\"tempfilter.h5\")\n",
    "\n",
    "print(lsst_g._fields)\n",
    "\n",
    "print('\\nlsst_g.wave.shape = {}'.format(lsst_g.wave.shape))\n",
    "print('lsst_g.transmission.shape = {}'.format(lsst_g.transmission.shape))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b972a67a",
   "metadata": {},
   "source": [
    "### Calculating absolute magnitude\n",
    "\n",
    "Since we have already calculated above the restframe SED at the time the galaxy is observed, then we can directly integrate the SED against the filter transmission curve to compute the absolute magnitude of the galaxy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed69f70c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dsps import calc_rest_mag\n",
    "rest_mag = calc_rest_mag(ssp_data.ssp_wave, sed_info.rest_sed, lsst_g.wave, lsst_g.transmission)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1bf9f0c3",
   "metadata": {},
   "source": [
    "### Calculating apparent magnitude\n",
    "\n",
    "To calculate the apparent magnitude, we need to redshift the SED and also apply the appropriate cosmological dimming factor to the restframe flux. These calculations are done under-the-hood with the `flat_wcdm.py` module in `dsps.cosmology`, and so we just need to pass in the same cosmological parameters used above to calculate `t_obs` from `z_obs`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b4d55b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dsps import calc_obs_mag\n",
    "\n",
    "obs_mag = calc_obs_mag(ssp_data.ssp_wave, sed_info.rest_sed, lsst_g.wave, lsst_g.transmission,\n",
    "                      z_obs, *DEFAULT_COSMOLOGY)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0486ed7",
   "metadata": {},
   "outputs": [],
   "source": [
    "! rm tempdata.h5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eda599c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "! rm tempfilter.h5"
   ]
  }
 ],
 "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.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}