{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cba22801-e735-4d6a-b03f-177da05d3917",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax import random as jran\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "ran_key = jran.key(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc3a27e4-97d2-4f48-bce8-65f529b0ea03",
   "metadata": {},
   "source": [
    "## Populate 500 halos at z=0.5 with synthetic galaxies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce0447ba-9082-4dd0-944e-d5b7415209b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rgrspit_diffsky import mc_galpop\n",
    "from dsps.cosmology import DEFAULT_COSMOLOGY\n",
    "\n",
    "lgmp_min = 11.0\n",
    "n_halos = 500\n",
    "logmhost = np.linspace(lgmp_min, 15, n_halos)\n",
    "halo_radius = np.ones_like(logmhost)\n",
    "z_obs = 0.5\n",
    "mc_key, pos_key, vel_key = jran.split(ran_key, 3)\n",
    "halo_pos = jran.uniform(pos_key, shape=(n_halos, 3))\n",
    "halo_vel = jran.uniform(vel_key, shape=(n_halos, 3))\n",
    "Lbox = 2_000.0\n",
    "\n",
    "galcat = mc_galpop.mc_galpop_synthetic_subs(\n",
    "    mc_key,\n",
    "    logmhost,\n",
    "    halo_radius,\n",
    "    halo_pos,\n",
    "    halo_vel,\n",
    "    z_obs,\n",
    "    lgmp_min,\n",
    "    DEFAULT_COSMOLOGY,\n",
    "    Lbox,\n",
    ")\n",
    "print(galcat.keys())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "831f85ca-9fe3-4fa6-a759-c7b760dfb03e",
   "metadata": {},
   "source": [
    "### Plot some MAHs for halos of different mass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60c05bac-ae17-416b-8e00-7a7f859c0817",
   "metadata": {},
   "outputs": [],
   "source": [
    "mskm115 = np.abs(galcat['logmp_t_obs'] - 11.5) < 0.2\n",
    "mskm125 = np.abs(galcat['logmp_t_obs'] - 12.5) < 0.2\n",
    "mskm135 = np.abs(galcat['logmp_t_obs'] - 13.5) < 0.2\n",
    "mskm145 = np.abs(galcat['logmp_t_obs'] - 14.5) < 0.2\n",
    "\n",
    "fig, ax = plt.subplots(1, 1)\n",
    "xlim = ax.set_xlim(1.5, galcat['t_obs'])\n",
    "ylim = ax.set_ylim(5e9, 5e14)\n",
    "yscale = ax.set_yscale('log')\n",
    "\n",
    "mred = u'#d62728' \n",
    "morange = u'#ff7f0e'  \n",
    "mgreen = u'#2ca02c'\n",
    "mblue = u'#1f77b4' \n",
    "mpurple = u'#9467bd' \n",
    "for i in range(10):\n",
    "    __=ax.plot(galcat['t_table'], 10**galcat['log_mah_table'][mskm115][i], lw=0.5, color=mblue)\n",
    "    __=ax.plot(galcat['t_table'], 10**galcat['log_mah_table'][mskm125][i], lw=0.5, color=mgreen)\n",
    "    __=ax.plot(galcat['t_table'], 10**galcat['log_mah_table'][mskm135][i], lw=0.5, color=morange)\n",
    "    __=ax.plot(galcat['t_table'], 10**galcat['log_mah_table'][mskm145][i], lw=0.5, color=mred)\n",
    "\n",
    "xlabel = ax.set_xlabel(r'${\\rm cosmic\\ time\\ [Gyr]}$')\n",
    "ylabel = ax.set_ylabel(r'$M_{\\rm halo}\\ {\\rm [M_{\\odot}]}$')\n",
    "title = ax.set_title(r'${\\rm DiffmahPop}$-${\\rm generated\\ MAHs}$')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2d34a64-6e91-4dfe-b820-c575b1bd89b9",
   "metadata": {},
   "source": [
    "### Plot some SFHs for galaxies of different mass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ab1bd75-52a4-4336-9031-352b920e8f6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "mskm115 = np.abs(galcat['logsm_t_obs'] - 9.0) < 0.2\n",
    "mskm135 = np.abs(galcat['logsm_t_obs'] - 11.5) < 0.2\n",
    "\n",
    "fig, ax = plt.subplots(1, 1)\n",
    "xlim = ax.set_xlim(1.5, galcat['t_obs'])\n",
    "ylim = ax.set_ylim(1e-3, 2e3)\n",
    "yscale = ax.set_yscale('log')\n",
    "\n",
    "mred = u'#d62728' \n",
    "morange = u'#ff7f0e'  \n",
    "mgreen = u'#2ca02c'\n",
    "mblue = u'#1f77b4' \n",
    "mpurple = u'#9467bd' \n",
    "for i in range(10):\n",
    "    __=ax.plot(galcat['t_table'], galcat['sfh_table'][mskm115][i], lw=0.5, color=mblue)\n",
    "    __=ax.plot(galcat['t_table'], galcat['sfh_table'][mskm135][i], lw=0.5, color=mred)\n",
    "\n",
    "xlabel = ax.set_xlabel(r'${\\rm cosmic\\ time\\ [Gyr]}$')\n",
    "ylabel = ax.set_ylabel(r'${\\rm SFR}\\ {\\rm [M_{\\odot}/yr]}$')\n",
    "title = ax.set_title(r'${\\rm DiffstarPop}$-${\\rm generated\\ SFHs}$')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1647a842-2008-4f11-a895-f5321115f381",
   "metadata": {},
   "source": [
    "### Plot stellar-to-halo mass relation at $z_{\\rm obs}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9006158c-222b-44bf-aec4-c962563cdaa9",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1)\n",
    "__=ax.loglog()\n",
    "ylim = ax.set_ylim(10**7.1, 10**12.1)\n",
    "__=ax.scatter(10**galcat['logmp_t_obs'], 10**galcat['logsm_t_obs'], s=1)\n",
    "xlabel = ax.set_xlabel(r'$M_{\\rm halo}\\ {\\rm [M_{\\odot}]}$')\n",
    "ylabel = ax.set_ylabel(r'$M_{\\star}\\ {\\rm [M_{\\odot}]}$')\n",
    "title = ax.set_title(r'${\\rm stellar}$-${\\rm to}$-${\\rm halo\\ mass\\ relation\\ at\\ }z_{\\rm obs}$')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78b9590e-c3d8-4363-8b75-30e4446da49e",
   "metadata": {},
   "source": [
    "### Plot $F_{\\rm sat}(M_{\\star})$\n",
    "\n",
    "Calculate the satellite fraction as a function of stellar mass. Nearly all galaxies are satellites since merging has not been implemented yet in this repo."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e0096d7-7a85-48cf-a302-f8b8f89d563b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.stats import binned_statistic\n",
    "\n",
    "logmp_bins = np.linspace(11.25, 14.25, 40)\n",
    "logsm_bins = np.linspace(8, 11.75, 40)\n",
    "\n",
    "mean_fsat_vs_logsm, __, __ = binned_statistic(galcat['logsm_t_obs'], galcat['upid']!=-1, bins=logsm_bins)\n",
    "logsm_binmids = 0.5*(logsm_bins[:-1] + logsm_bins[1:])\n",
    "\n",
    "fig, ax = plt.subplots(1, 1)\n",
    "ylim = ax.set_ylim(-0.02, 1.2)\n",
    "xlim = ax.set_xlim(8, 11.75)\n",
    "__=ax.plot(np.linspace(0, 20, 100), np.zeros(100), ':', color='k')\n",
    "__=ax.plot(np.linspace(0, 20, 100), np.ones(100), ':', color='k')\n",
    "\n",
    "__=ax.plot(logsm_binmids, mean_fsat_vs_logsm)\n",
    "\n",
    "xlabel = ax.set_xlabel(r'$M_{\\rm halo}\\ {\\rm [M_{\\odot}]}$')\n",
    "ylabel = ax.set_ylabel(r'$\\langle F_{\\rm sat}\\ \\vert\\ M_{\\star}\\rangle$')\n",
    "title = ax.set_title(r'${\\rm satellite\\ fraction\\ at\\ }z_{\\rm obs}$')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71ef8695-edd1-4218-ae82-0e3789be5ab3",
   "metadata": {},
   "source": [
    "## Plot distributions of specific SFR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3e3379a-2efc-44c8-85e1-a764d8413e32",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1)\n",
    "\n",
    "msk_cens = galcat['upid']==-1\n",
    "lgssfr_bins = np.linspace(-14, -9, 50)\n",
    "\n",
    "msk_logsm = np.abs(galcat['logsm_t_obs']-11) < 0.2\n",
    "__=ax.hist(galcat['logssfr_t_obs'][msk_logsm], bins=lgssfr_bins, density=True, alpha=0.7, color=mred,\n",
    "          label=r'$M_{\\star}=10^{11}M_{\\odot}$')\n",
    "\n",
    "msk_logsm = np.abs(galcat['logsm_t_obs']-10) < 0.2\n",
    "__=ax.hist(galcat['logssfr_t_obs'][msk_logsm], bins=lgssfr_bins, density=True, alpha=0.7, color=mgreen,\n",
    "          label=r'$M_{\\star}=10^{10}M_{\\odot}$')\n",
    "\n",
    "msk_logsm = np.abs(galcat['logsm_t_obs']-9) < 0.2\n",
    "__=ax.hist(galcat['logssfr_t_obs'][msk_logsm], bins=lgssfr_bins, density=True, alpha=0.7, color=mpurple,\n",
    "          label=r'$M_{\\star}=10^{9}M_{\\odot}$')\n",
    "\n",
    "leg = ax.legend()\n",
    "xlabel = ax.set_xlabel(r'$\\log_{10}{\\rm sSFR\\ [yr^{-1}}]$')"
   ]
  }
 ],
 "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}