{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "recorded-display",
   "metadata": {},
   "source": [
    "# Overview of Diffmah and DiffmahPop\n",
    "\n",
    "Starting from a collection of best-fitting approximations to halo merger trees, you can use the `calc_halo_history` function to compute the assembly history for every halo in the sample. Here we'll just demonstrate a few simple cases.\n",
    "\n",
    "Note that in these examples, we pass in arbitrary values for the early- and late-time indices. However, for real halos (and also for the results returned by the diffmah-provided MAH fitting script), $0 < \\alpha_{\\rm late} < \\alpha_{\\rm early}.$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8dc6f535",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "import matplotlib.cm as cm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "plain-arbor",
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffmah import mah_halopop, DEFAULT_MAH_PARAMS\n",
    "\n",
    "n_halos, n_times = 50, 100\n",
    "tarr = np.linspace(0.5, 13.8, n_times)\n",
    "logt0 = np.log10(tarr[-1])\n",
    "colors=cm.coolwarm(np.linspace(1, 0, n_halos)) # red first\n",
    "\n",
    "ZZ = np.zeros(n_halos)\n",
    "logtc = np.log10(np.linspace(1, 5, n_halos))\n",
    "logm0 = 12 + ZZ\n",
    "early, late = 2 + ZZ, 1 + ZZ\n",
    "t_peak = 14.0 + ZZ\n",
    "mah_params = DEFAULT_MAH_PARAMS._make((logm0, logtc, early, late, t_peak))\n",
    "dmhdt, log_mah = mah_halopop(mah_params, tarr, logt0)\n",
    "\n",
    "fig, ax = plt.subplots(1, 1)\n",
    "__=ax.loglog()\n",
    "for ih in range(n_halos):\n",
    "    __=ax.plot(tarr, 10**log_mah[ih, :], color=colors[ih])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "completed-macro",
   "metadata": {},
   "outputs": [],
   "source": [
    "logtc = 0.5 + ZZ\n",
    "early = np.linspace(1, 3, n_halos)\n",
    "late = 1 + ZZ\n",
    "mah_params = DEFAULT_MAH_PARAMS._make((logm0, logtc, early, late, t_peak))\n",
    "dmhdt, log_mah = mah_halopop(mah_params, tarr, logt0)\n",
    "\n",
    "\n",
    "fig, ax = plt.subplots(1, 1)\n",
    "__=ax.loglog()\n",
    "for ih in range(n_halos):\n",
    "    __=ax.plot(tarr, 10**log_mah[ih, :], color=colors[ih])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "constant-thunder",
   "metadata": {},
   "outputs": [],
   "source": [
    "tauc = 2.0 \n",
    "early = 3 + ZZ\n",
    "late = np.linspace(0.01, 3, n_halos)\n",
    "mah_params = DEFAULT_MAH_PARAMS._make((logm0, logtc, early, late, t_peak))\n",
    "dmhdt, log_mah = mah_halopop(mah_params, tarr, logt0)\n",
    "\n",
    "fig, ax = plt.subplots(1, 1)\n",
    "__=ax.loglog()\n",
    "for ih in range(n_halos):\n",
    "    __=ax.plot(tarr, 10**log_mah[ih, :], color=colors[ih])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acknowledged-exhibition",
   "metadata": {},
   "source": [
    "## Generating Monte Carlo realizations of halo MAHs with DiffmahPop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "066ccdaa-0b32-4b26-a319-2e525ae75a2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax import random as jran\n",
    "ran_key = jran.key(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "imported-jungle",
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffmah.diffmahpop_kernels import mc_cenpop, DEFAULT_DIFFMAHPOP_PARAMS\n",
    "\n",
    "cosmic_time = np.linspace(1.5, 13.8, 100)\n",
    "\n",
    "n_halos = 10\n",
    "lgm_obs = np.zeros(n_halos) + 13\n",
    "t_obs = np.zeros(n_halos) + 13.5\n",
    "\n",
    "halopop = mc_cenpop(DEFAULT_DIFFMAHPOP_PARAMS, cosmic_time, lgm_obs, t_obs, ran_key, logt0)\n",
    "\n",
    "fig, ax = plt.subplots(1, 1)\n",
    "__=ax.loglog()\n",
    "for ih in range(n_halos):\n",
    "    __=ax.plot(cosmic_time, 10**halopop.log_mah[ih, :])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8efb29d1-f748-4d52-b544-8b56bdfff114",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.12.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}