{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cfda5b4a-adde-49df-88ce-a7a7c9c14d01",
   "metadata": {},
   "source": [
    "# Generating Mass Assembly Histories in Diffhalos\n",
    "\n",
    "In Diffhalos we use DiffmahNet in order to assign MAHs to the (sub)halo populations in the lightcones. DiffmahNet is a population-level model of the Diffmah model for predicting mass assembly histories (MAHs) of individual halos. Once the halos are generated by sampling from the halo and subhalo mass function for hosts and subhalos, respectively, DiffmahNet pastes a MAH onto each (sub)halo in the lightcone, conditioned on mass its $M_{\\rm obs} \\equiv M_{\\rm halo} (t_{\\rm obs})$ at the time of observation $t_{\\rm obs}$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f808f1d-78b3-476a-86ed-2277b43e1c61",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.lines import Line2D\n",
    "import matplotlib.patches as mpatches\n",
    "\n",
    "plt.rc(\"text\", usetex=True)\n",
    "plt.rc(\"font\", family=\"serif\", size=18)\n",
    "plt.rcParams.update({'xtick.direction': 'in', 'ytick.direction': 'in'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a384452-ba8a-448a-9b15-93e69b7d877e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from jax import random as jran\n",
    "from jax import numpy as jnp"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffa79a4c-db5b-4507-a3f1-0b5c3478c478",
   "metadata": {},
   "source": [
    "### Halo MAHs\n",
    "\n",
    "Below, we demonstrate how to produce MAHs for a population of halos, conditioned on a single set of $(t_{\\rm obs}, M_{\\rm obs})$. In what follows, one must be careful to _rescale_ the MAH parameters from DiffmahNet so that the conditions are exactly satisfied."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a35a373-b21a-4165-84d8-f0f514fe64e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffhalos.mah.diffmahnet import log_mah_kern\n",
    "from diffhalos.mah.diffmahnet_utils import mc_mah_cenpop\n",
    "from diffhalos.mah.utils import rescale_mah_parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c68d8f1c-e73d-4bc1-9aeb-004261acd885",
   "metadata": {},
   "outputs": [],
   "source": [
    "ran_key = jran.key(0)\n",
    "\n",
    "n_cens = 100\n",
    "\n",
    "# conditions for MAH generation\n",
    "m_obs = np.array([12.5])\n",
    "t_obs = np.array([13.5])\n",
    "\n",
    "n_sample = 100\n",
    "logt0 = np.log10(13.8)\n",
    "\n",
    "m_vals, t_vals = [\n",
    "    jnp.repeat(x.flatten(), n_sample)\n",
    "    for x in np.stack(\n",
    "        [m_obs, t_obs],\n",
    "        axis=-1,\n",
    "    ).T\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdc3d36c-9fea-47b8-a906-4ddba42c4ec1",
   "metadata": {},
   "outputs": [],
   "source": [
    "centrals_model_key = \"cenflow_v2_0_64bit.eqx\"\n",
    "\n",
    "# compute the MAH parameters\n",
    "mah_params_uncorr = mc_mah_cenpop(m_vals, t_vals, ran_key, centrals_model_key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35ec2766-df00-4e4f-9fe0-b4740a0445dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_t = 100\n",
    "t_min = 0.5\n",
    "t_grid = jnp.linspace(t_min, t_vals, n_t).T\n",
    "\n",
    "# compute observed mass with corrected parameters\n",
    "logmp_obs_uncorr = log_mah_kern(mah_params_uncorr, t_grid, logt0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff8a58f8-5565-4614-9ccf-10eef0ed413e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# apply correction to MAH parameters to ensure that the M_halo(t_obs) = M_obs is true for all halos\n",
    "mah_params = rescale_mah_parameters(mah_params_uncorr, m_vals, logmp_obs_uncorr[:,-1])\n",
    "logmp_obs = log_mah_kern(mah_params, t_grid, logt0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8824f722-cfac-4a05-bcd6-60c26016844d",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=([7,5]))\n",
    "\n",
    "_=ax.plot(t_grid.T, 10**logmp_obs.T)\n",
    "\n",
    "ax.set_xlabel(r\"$t\\; [{\\rm Gyr}]$\")\n",
    "ax.set_ylabel(r\"$M_{\\rm halo} (t)\\; [{\\rm M_\\odot}]$\")\n",
    "\n",
    "ax.set_yscale('log')\n",
    "\n",
    "ax.set_xlim([0.5, 13.6])\n",
    "ax.set_ylim([10**7, 10**12.8])\n",
    "\n",
    "_=ax.set_title(r\"$M_{\\rm obs}=10^{12.5}\\; {\\rm M_\\odot}$, \\; $t_{\\rm obs}=13.5\\; {\\rm Gyr}$\", fontsize=18)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7185409-9254-4339-ad8a-e07b329d4d18",
   "metadata": {},
   "source": [
    "### Different choice of $(t_{\\rm obs}, M_{\\rm obs})$\n",
    "\n",
    "There is nothing special to the conditions we chose above. Below we show that for different sets of $(t_{\\rm obs}, M_{\\rm obs})$, arbitrarily chosen, the halo population that is produced has a MAH that converges at $M_{\\rm obs}$ at time $t_{\\rm obs}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fff7604-0229-4d23-853c-79affe6896ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "ran_key = jran.key(0)\n",
    "\n",
    "n_cens = 100\n",
    "\n",
    "# conditions for MAH generation\n",
    "m_obs_2 = np.array([9.5])\n",
    "t_obs_2 = np.array([6.0])\n",
    "\n",
    "n_sample = 100\n",
    "logt0 = np.log10(13.8)\n",
    "\n",
    "m_vals, t_vals = [\n",
    "    jnp.repeat(x.flatten(), n_sample)\n",
    "    for x in np.stack(\n",
    "        [m_obs_2, t_obs_2],\n",
    "        axis=-1,\n",
    "    ).T\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfda1516-1b80-42f1-8c00-07f4fa8b67e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute the MAH parameters\n",
    "mah_params_uncorr = mc_mah_cenpop(m_vals, t_vals, ran_key, centrals_model_key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89d2ce97-8d55-426d-8ae6-bbf74ff58931",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_t = 100\n",
    "t_min = 0.5\n",
    "t_grid_2 = jnp.linspace(t_min, t_vals, n_t).T\n",
    "\n",
    "# compute observed mass with corrected parameters\n",
    "logmp_obs_uncorr = log_mah_kern(mah_params_uncorr, t_grid_2, logt0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d143111-cc89-47ba-b0c2-931146203c96",
   "metadata": {},
   "outputs": [],
   "source": [
    "# apply correction to MAH parameters to ensure that the M_halo(t_obs) = M_obs is true for all halos\n",
    "mah_params = rescale_mah_parameters(mah_params_uncorr, m_vals, logmp_obs_uncorr[:,-1])\n",
    "logmp_obs_2 = log_mah_kern(mah_params, t_grid_2, logt0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82239a73-967a-4f70-b5c2-6833ab007209",
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffhalos.mah.diffmahnet_utils import get_mean_and_std_of_mah"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2db62b9-61fe-4afd-8873-5698a97b6cfc",
   "metadata": {},
   "outputs": [],
   "source": [
    "mah_mean, mah_max, mah_min = get_mean_and_std_of_mah(10**logmp_obs)\n",
    "mah_mean_2, mah_max_2, mah_min_2 = get_mean_and_std_of_mah(10**logmp_obs_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35e992af-d9cf-4160-9b28-d39a50469d9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=([7,5]))\n",
    "\n",
    "plt.plot(t_grid[0], mah_mean, ls='-', color='k')\n",
    "ax.plot(t_grid[0], mah_max, ls='-', lw=1, color='k')\n",
    "ax.plot(t_grid[0], mah_min, ls='-', lw=1, color='k')\n",
    "ax.fill_between(t_grid[0], mah_min, mah_max, color='gray', alpha=0.4)\n",
    "\n",
    "plt.plot(t_grid_2[0], mah_mean_2, ls='-', color='crimson')\n",
    "ax.plot(t_grid_2[0], mah_max_2, ls='-', lw=1, color='crimson')\n",
    "ax.plot(t_grid_2[0], mah_min_2, ls='-', lw=1, color='crimson')\n",
    "ax.fill_between(t_grid_2[0], mah_min_2, mah_max_2, color='crimson', alpha=0.4)\n",
    "\n",
    "ax.set_xlabel(r\"$t\\; [{\\rm Gyr}]$\")\n",
    "ax.set_ylabel(r\"$M_{\\rm halo} (t)\\; [{\\rm M_\\odot}]$\")\n",
    "\n",
    "\n",
    "ax.set_yscale('log')\n",
    "\n",
    "_=ax.set_xlim([0.9, 13.6])\n",
    "_=ax.set_ylim([10**7, 10**12.8])\n",
    "\n",
    "custom_handles = [\n",
    "    Line2D([0], [0], color='k', lw=2, \n",
    "           label=r\"$M_{\\rm obs}=10^{%.1f}\\; {\\rm M_\\odot}$, \\; $t_{\\rm obs}=%.1f\\; {\\rm Gyr}$\"%(m_obs[0],t_obs[0])),\n",
    "    Line2D([0], [0], color='crimson', lw=2, ls='-', \n",
    "           label=r\"$M_{\\rm obs}=10^{%.1f}\\; {\\rm M_\\odot}$, \\; $t_{\\rm obs}=%.1f\\; {\\rm Gyr}$\"%(m_obs_2[0],t_obs_2[0])),\n",
    "]\n",
    "\n",
    "_=ax.legend(handles=custom_handles, bbox_to_anchor=(0.9, 1.23), frameon=False, ncols=1, fontsize=17)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07a23b00-8d0c-4a71-915d-818061257ad5",
   "metadata": {},
   "source": [
    "### Generating MAH lightcones\n",
    "\n",
    "Below we show that the lightcone outputs by default include the MAH of each (sub)halo. In particular, each halo in the generated lightcone is assigned a set of MAH parameters, which can then be used to produce the mass assembly history of each halo using the Diffmah model. Below we plot some of the generated MAHs, chosen randomly for a few halos, for demonstration purposes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d75b2fd5-fee3-492e-8d1e-695ea2084381",
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffhalos.lightcone_generators import mc_lc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bab73625-c86a-4b0e-911a-da64dbb346e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "ran_key = jran.key(0)\n",
    "\n",
    "lgmp_min = 12.5\n",
    "lgmsub_min = lgmp_min - 0.01\n",
    "z_min, z_max = 0.4, 3.5\n",
    "sky_area_degsq = 10.0\n",
    "\n",
    "args = (ran_key, lgmp_min, lgmsub_min, z_min, z_max, sky_area_degsq)\n",
    "\n",
    "halopop_mc = mc_lc(*args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1dcfa723-df8b-4bab-8047-3f7e33e4dffb",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_t = 100\n",
    "t_min = 0.5\n",
    "t_grid = jnp.linspace(t_min, halopop_mc.t_obs, n_t).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a6c4f81-84b0-4372-b7f9-b20a9dc75a42",
   "metadata": {},
   "outputs": [],
   "source": [
    "logmp_obs_mah = log_mah_kern(halopop_mc.mah_params, t_grid, halopop_mc.logt0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "439bfeba-26c9-423b-8646-670c4bbcb6d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c891b9d5-7e14-470d-a509-ea794f10ff01",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=([7,5]))\n",
    "\n",
    "indx = np.asarray(random.sample(range(0, len(halopop_mc.t_obs)), 10))\n",
    "\n",
    "_=ax.plot(t_grid[indx].T, 10**logmp_obs_mah[indx].T)\n",
    "\n",
    "ax.set_xlabel(r\"$t\\; [{\\rm Gyr}]$\")\n",
    "ax.set_ylabel(r\"$M_{\\rm halo} (t)\\; [{\\rm M_\\odot}]$\")\n",
    "\n",
    "_=ax.set_yscale('log')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51c11b6c-4fc0-49e3-8713-30453e029000",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "diffstuff",
   "language": "python",
   "name": "diffstuff"
  },
  "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}