{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Stellar magnitude examples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the sub-directory `stars_extra` (within the `data` directory of an individual simulation), there is a set of files analogous to the snapshot files that contains pre-calculated star particle magnitudes in a wide range of bands. These can be used e.g. instead of stellar masses to make \"proper\" images, or to compute integrated galaxy magnitudes. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's start by importing the usual libraries (note that the last one may give a lot of warnings, ignore):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import hydrangea as hy\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For simplicity, let's work with simulation CE-0 here. Construct the (first) file of the snapshot, stellar magnitude, and subfind catalogues at snapshot 26 (z = 0.1):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim = hy.Simulation(0)\n",
    "snap_file = sim.get_snapshot_file(26)\n",
    "mag_file = sim.get_magnitude_file(26)\n",
    "sub_file = sim.get_subfind_file(26)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now set up the loader to read the (distributed) catalogues. The magnitude files contain two separate catalogues, one in which magnitudes are computed in the rest frame band (`restFrame`), and one in observer frame bands for the particular redshift of the snapshot (`redshifted`). Let's use the latter here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mags = hy.SplitFile(mag_file, 'redshifted')\n",
    "snap = hy.SplitFile(snap_file, part_type=4)\n",
    "sub = hy.SplitFile(sub_file, 'Subhalo')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Reading the actual data from the magnitude catalogue works exactly the same as for the \"basic\" snapshot catalogue. The only catch is that the names all contain a \"-\" sign, so accessing them directly as attributes of the `mags` object won't work (python would interpret the sign as a minus operator). Instead, we have to use the explicit `read_data` function; I store the data as an attribute of mags itself to keep things tidy, but it could also be stored in a separate variable. Let's look at the SDSS g and r bands here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mags.g = mags.read_data('g-Magnitude')\n",
    "mags.r = mags.read_data('r-Magnitude')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The stellar luminosities are stored as (absolute) magnitudes of individual star particles. Let's see whether they make sense by plotting them vs. stellar mass, colour-coded by the expansion factor at the stars' formation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.scatter(np.log10(snap.Mass[:10000]), mags.r[:10000], s=2,\n",
    "            c=snap.StellarFormationTime[:10000])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.scatter(snap.StellarFormationTime[:10000], mags.g[:10000] - mags.r[:10000], s=2,\n",
    "            c=np.log10(snap.Metallicity[:10000]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Both plots look vaguely reasonable. More massive star particles are in general brighter (more negative r-band magnitude), and at fixed mass younger stars (yellow -- formation at higher expansion factor) are brighter. Young stars show a strong reddening with age, but for stars older than a few Gyr (a_form <~ 0.6), metallicity is more important for determining the colour (lower = bluer)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compute galaxy magnitudes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is relatively straightforward to sum the individual stellar luminosities to get total luminosities per galaxy. The first thing we need for this is to convert from magnitudes to (linear) luminosities:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mags.lum_g = 10.0**(-0.4*mags.g)\n",
    "mags.lum_r = 10.0**(-0.4*mags.r)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Second, we have to associate star particles to subhaloes. For Subfind, this is straightforward because the Hydrangea library already provides functionality to compute the subhalo index for each particle in a snapshot. The only thing we need to do is specify which subfind catalogue we want to link against:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "snap.subfind_file = sub_file   # Here, we use the catalogue from the same snapshot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The magic `SubhaloIndex` attribute now contains (or rather will contain, once initialized on the fly) the subhalo index of each star particle. With this, we can easily sum up the lumimosities using numpy's `histogram`. Let's do a test first to make sure that this works -- (re-)computing the stellar masses (the first time this is run it may take a little while to compute the subhalo indices...)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mstar, edges = np.histogram(snap.SubhaloIndex, weights=snap.Mass.astype(np.float64),  # Casting as flt64 avoids rounding errors\n",
    "                            bins=np.arange(sub.MassType.shape[0]+1)-0.5)              # set up histogram bins so that each subhalo index sits in the centre of a separate bin"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compare the re-computed masses of the first ten subhaloes to those straight from the subfind catalogue (some very small differences may remain from rounding on the subfind side):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mstar[:10]/sub.MassType[:10, 4]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now analogously for the g and r band luminosities:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lum_g, edges = np.histogram(snap.SubhaloIndex, weights=mags.lum_g.astype(np.float64),\n",
    "                            bins=np.arange(sub.MassType.shape[0]+1)-0.5)\n",
    "lum_r, edges = np.histogram(snap.SubhaloIndex, weights=mags.lum_r.astype(np.float64),\n",
    "                            bins=np.arange(sub.MassType.shape[0]+1)-0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Convert back to (galaxy, absolute) magnitude and plot the colour-luminosity relation to see whether it looks reasonable (there may be some warning messages here due to galaxies with zero luminosity):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sub.mag_g = -5/2 * np.log10(lum_g)\n",
    "sub.mag_r = -5/2 * np.log10(lum_r)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ind_plot = np.nonzero(sub.mag_r < -17)[0]\n",
    "logssfr = np.log10(sub.StarFormationRate[ind_plot]/sub.MassType[ind_plot, 4])\n",
    "plt.scatter(sub.mag_r[ind_plot], sub.mag_g[ind_plot] - sub.mag_r[ind_plot],\n",
    "            s=15, c=np.clip(logssfr, -12, None))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Subhalo magnitudes for Cantor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Instead of subfind, the magnitudes can also be summed for Cantor subhaloes. The only difference is in finding the per-particle Cantor subhalo indices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cantor_file = sim.cantor_dir + 'Cantor_026.hdf5'\n",
    "cantor_id_file = sim.cantor_dir + 'Cantor_026_IDs.hdf5'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cantor_ids = hy.hdf5.read_data(cantor_id_file, 'IDs')\n",
    "cantor_offsets = hy.hdf5.read_data(cantor_file, 'Subhalo/Offset')\n",
    "cantor_lengths = hy.hdf5.read_data(cantor_file, 'Subhalo/Length')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Do the actual particle matching (the `SubhaloIndex` attribute of `SplitFile` does the same under the hood for Subfind)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "index_in_ids = hy.crossref.find_id_indices(snap.ParticleIDs, cantor_ids)[0]\n",
    "cantor_subhalo_index = hy.tools.ind_to_block(index_in_ids, cantor_offsets, cantor_lengths)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As a test, compute the Cantor r-band magnitudes and compare them to Subfind:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cantor_lum_r, edges = np.histogram(cantor_subhalo_index, weights=mags.lum_r.astype(np.float64),\n",
    "                                   bins=np.arange(cantor_lengths.shape[0]+1)-0.5)\n",
    "cantor_mag_r = -5/2 * np.log10(cantor_lum_r)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cantor_sf_index = hy.hdf5.read_data(cantor_file, 'Subhalo/IndexBySubfindID')\n",
    "plt.scatter(sub.mag_r[ind_plot], cantor_mag_r[cantor_sf_index[ind_plot]]-sub.mag_r[ind_plot], s=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This looks reasonable: most galaxies have (almost) the same magnitude in both catalogues, but a small number of them scatter to (mostly) lower, i.e. brighter, values. Also note that the two brightest galaxies have become (slightly) fainter; this is probabably because they are BCGs and Cantor assigns more particles to satellites, rather than to these centrals. "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}