{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load supported SSP templates\n",
    "\n",
    "This notebook shows how to load and use the supported SSP templates. Currently we have support for custom build SSP templates stored in hdf5 format for which we provide a template based on Bruzual&Charlot2003 models. Additionally we support all SSP templates that the pyPipe3D project uses. Those templates come in astronomy friendly fits file format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "from rubix.spectra.ssp.templates import BruzualCharlot2003\n",
    "\n",
    "BruzualCharlot2003"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "print(BruzualCharlot2003.age)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load SSP template via custom config\n",
    "\n",
    "This shows how to use a custom configuration to load an SSP template that is stored under some file location on your disk."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "config = {\n",
    "        \"name\": \"Bruzual & Charlot (2003)\",\n",
    "        \"format\": \"HDF5\",\n",
    "        \"source\": \"https://www.bruzual.org/bc03/\",\n",
    "        \"file_name\": \"BC03lr.h5\",\n",
    "        \"fields\": {\n",
    "            \"age\": {\n",
    "                \"name\": \"age\",\n",
    "                \"units\": \"Gyr\",\n",
    "                \"in_log\": False\n",
    "            },\n",
    "            \"metallicity\": {\n",
    "                \"name\": \"metallicity\",\n",
    "                \"units\": \"\",\n",
    "                \"in_log\": False\n",
    "            },\n",
    "            \"wavelength\": {\n",
    "                \"name\": \"wavelength\",\n",
    "                \"units\": \"Angstrom\",\n",
    "                \"in_log\": False\n",
    "            },\n",
    "            \"flux\": {\n",
    "                \"name\": \"flux\",\n",
    "                \"units\": \"Lsun/Angstrom\",\n",
    "                \"in_log\": False\n",
    "            }\n",
    "        }\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "from rubix.spectra.ssp.grid import HDF5SSPGrid\n",
    "ssp = HDF5SSPGrid.from_file(config, file_location=\"../rubix/spectra/ssp/templates\")\n",
    "ssp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "ssp.age.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "ssp.metallicity.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "ssp.wavelength.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "ssp.flux.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Let's plot some example spectra"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "plt.plot(ssp.wavelength,ssp.flux[0][0])\n",
    "plt.xlabel(r'$\\lambda$ [%s]'%config[\"fields\"][\"wavelength\"][\"units\"])\n",
    "plt.ylabel(r'Flux [%s]'%config[\"fields\"][\"flux\"][\"units\"])\n",
    "#plt.yscale(\"log\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "plt.plot(ssp.wavelength,ssp.flux[-1][-1])\n",
    "plt.xlabel(r'$\\lambda$ [%s]'%config[\"fields\"][\"wavelength\"][\"units\"])\n",
    "plt.ylabel(r'Flux [%s]'%config[\"fields\"][\"flux\"][\"units\"])\n",
    "#plt.yscale(\"log\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "for i in range(len(ssp.metallicity)):\n",
    "    plt.plot(ssp.wavelength,ssp.flux[i][0], label=r'Z=%0.3f'%ssp.metallicity[i])\n",
    "plt.xlabel(r'$\\lambda$ [%s]'%config[\"fields\"][\"wavelength\"][\"units\"])\n",
    "plt.ylabel(r'Flux [%s]'%config[\"fields\"][\"flux\"][\"units\"])\n",
    "#plt.yscale(\"log\")\n",
    "plt.xlim(0,10000)\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "ages = np.linspace(0,len(ssp.age),10)\n",
    "for age in ages:\n",
    "    plt.plot(ssp.wavelength,ssp.flux[0][int(age)], label='%.2f %s'%(ssp.age[int(age)], config[\"fields\"][\"age\"][\"units\"]))\n",
    "plt.xlabel(r'$\\lambda$ [%s]'%config[\"fields\"][\"wavelength\"][\"units\"])\n",
    "plt.ylabel(r'Flux [%s]'%config[\"fields\"][\"flux\"][\"units\"])\n",
    "#plt.yscale(\"log\")\n",
    "plt.xlim(0,5000)\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Automatic download supported SSP template\n",
    "Rubix supports automatic download of a supported SSP template from a specified url in case the template can't be found on disk under the file_location specified."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "config = {\n",
    "        \"name\": \"Mastar Charlot & Bruzual (2019)\",\n",
    "        \"format\": \"pyPipe3D\",\n",
    "        \"source\": \"https://ifs.astroscu.unam.mx/pyPipe3D/templates/\",\n",
    "        \"file_name\": \"MaStar_CB19.slog_1_5.fits.gz\",\n",
    "        \"fields\": {\n",
    "            \"age\": {\n",
    "                \"name\": \"age\",\n",
    "                \"units\": \"Gyr\",\n",
    "                \"in_log\": False\n",
    "            },\n",
    "            \"metallicity\": {\n",
    "                \"name\": \"metallicity\",\n",
    "                \"units\": \"\",\n",
    "                \"in_log\": False\n",
    "            },\n",
    "            \"wavelength\": {\n",
    "                \"name\": \"wavelength\",\n",
    "                \"units\": \"Angstrom\",\n",
    "                \"in_log\": False\n",
    "            },\n",
    "            \"flux\": {\n",
    "                \"name\": \"flux\",\n",
    "                \"units\": \"Lsun/Angstrom\",\n",
    "                \"in_log\": False\n",
    "            }\n",
    "        }\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "from rubix.spectra.ssp.grid import pyPipe3DSSPGrid\n",
    "ssp = pyPipe3DSSPGrid.from_file(config, file_location=\"../rubix/spectra/ssp/templates\")\n",
    "ssp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "ssp.age.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "ssp.metallicity.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "ssp.wavelength.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "ssp.flux.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lets plot some example spectra\n",
    "Example for Mastar templates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "plt.plot(ssp.wavelength,ssp.flux[0][0])\n",
    "plt.xlabel(r'$\\lambda$ [%s]'%config[\"fields\"][\"wavelength\"][\"units\"])\n",
    "plt.ylabel(r'Flux [%s]'%config[\"fields\"][\"flux\"][\"units\"])\n",
    "#plt.yscale(\"log\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "plt.plot(ssp.wavelength,ssp.flux[-1][-1])\n",
    "plt.xlabel(r'$\\lambda$ [%s]'%config[\"fields\"][\"wavelength\"][\"units\"])\n",
    "plt.ylabel(r'Flux [%s]'%config[\"fields\"][\"flux\"][\"units\"])\n",
    "#plt.yscale(\"log\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "for i in range(len(ssp.metallicity)):\n",
    "    plt.plot(ssp.wavelength,ssp.flux[i][0], label=r'Z=%0.3f'%ssp.metallicity[i])\n",
    "plt.xlabel(r'$\\lambda$ [%s]'%config[\"fields\"][\"wavelength\"][\"units\"])\n",
    "plt.ylabel(r'Flux [%s]'%config[\"fields\"][\"flux\"][\"units\"])\n",
    "#plt.yscale(\"log\")\n",
    "plt.xlim(2000,10000)\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NBVAL_SKIP\n",
    "ages = np.linspace(0,len(ssp.age),10)\n",
    "for age in ages:\n",
    "    plt.plot(ssp.wavelength,ssp.flux[0][int(age)], label='%.2f %s'%(ssp.age[int(age)], config[\"fields\"][\"age\"][\"units\"]))\n",
    "plt.xlabel(r'$\\lambda$ [%s]'%config[\"fields\"][\"wavelength\"][\"units\"])\n",
    "plt.ylabel(r'Flux [%s]'%config[\"fields\"][\"flux\"][\"units\"])\n",
    "#plt.yscale(\"log\")\n",
    "plt.xlim(2000,5000)\n",
    "plt.legend()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "rubix",
   "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}