From 68452bfbf6243b5b6a2e9d34a82e519a5aec4148 Mon Sep 17 00:00:00 2001
From: Georgios Zacharegkas <gzacharegkas@gmail.com>
Date: Wed, 8 Jul 2026 12:49:12 -0500
Subject: [PATCH 1/9] adding tutorials in docs and small changes

---
 .../lightcone_generators/mc_lightcone.py      |   3 +-
 .../mc_lightcone_subhalos.py                  |   2 +-
 diffhalos/mah/diffmahnet_utils.py             |   9 +-
 docs/source/diffhalos_quickstart.ipynb        | 466 +++++++++++++++++-
 4 files changed, 460 insertions(+), 20 deletions(-)

diff --git a/diffhalos/lightcone_generators/mc_lightcone.py b/diffhalos/lightcone_generators/mc_lightcone.py
index d7715cc..7271b23 100644
--- a/diffhalos/lightcone_generators/mc_lightcone.py
+++ b/diffhalos/lightcone_generators/mc_lightcone.py
@@ -450,12 +450,11 @@ def weighted_lc(
                 Base-10 log of z=0 age of the Universe for the input cosmology
 
             cen_weight: ndarray of shape (n_halos_tot, )
-
                 For centrals, cen_weight is determined by the halo mass function (HMF)
                 For satellites, cen_weight is HMF weight of the associated central
                 For satellites, cen_weight = halopop.cen_weight[halopop.halo_indx]
 
-            central : ndarray of shape (n_halos_tot, )
+            central: ndarray of shape (n_halos_tot, )
                 Integer equals 1 for central halos and 0 for subhalos
 
             sat_weight: ndarray of shape (n_halos_tot, )
diff --git a/diffhalos/lightcone_generators/mc_lightcone_subhalos.py b/diffhalos/lightcone_generators/mc_lightcone_subhalos.py
index 4000be6..9359c77 100644
--- a/diffhalos/lightcone_generators/mc_lightcone_subhalos.py
+++ b/diffhalos/lightcone_generators/mc_lightcone_subhalos.py
@@ -212,7 +212,7 @@ def weighted_lc_subhalos(
     subhalo_model_key=DEFAULT_DIFFMAHNET_SAT_MODEL,
 ):
     """
-    Generate a subhalo lightcone
+    Generate a weighted subhalo lightcone
 
     Parameters
     ----------
diff --git a/diffhalos/mah/diffmahnet_utils.py b/diffhalos/mah/diffmahnet_utils.py
index a3001a2..27b4581 100644
--- a/diffhalos/mah/diffmahnet_utils.py
+++ b/diffhalos/mah/diffmahnet_utils.py
@@ -23,7 +23,6 @@
 )
 
 from . import diffmahnet
-from .utils import rescale_mah_parameters
 
 DEFAULT_MAH_UPARAMS = get_unbounded_mah_params(DEFAULT_MAH_PARAMS)
 
@@ -210,9 +209,11 @@ def get_mah_from_unbounded_params(
     mah_params_uncorrected = DEFAULT_MAH_PARAMS._make(mah_params_bound)
     _, log_mah = mah_halopop(mah_params_uncorrected, t_grid, logt0)
 
-    mah_params_corrected = rescale_mah_parameters(
-        mah_params_uncorrected, m_vals, log_mah[:, -1]
-    )
+    # rescale mah parameters
+    delta_logm_obs = log_mah[:, -1] - m_vals
+    logm0_rescaled = mah_params_uncorrected.logm0 - delta_logm_obs
+    mah_params_corrected = mah_params_uncorrected._replace(logm0=logm0_rescaled)
+
     _, log_mah = mah_halopop(mah_params_corrected, t_grid, logt0)
 
     return log_mah
diff --git a/docs/source/diffhalos_quickstart.ipynb b/docs/source/diffhalos_quickstart.ipynb
index 47e2f66..f866951 100644
--- a/docs/source/diffhalos_quickstart.ipynb
+++ b/docs/source/diffhalos_quickstart.ipynb
@@ -13,6 +13,82 @@
     "The cell below shows how to generate a lightcone of halos and subhalos, and their mass assembly histories."
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7fd7d119-fa3a-47df-ad36-55a5416e35f0",
+   "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=20)\n",
+    "plt.rcParams.update({'xtick.direction': 'in', 'ytick.direction': 'in'})"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "19e4ea51-a236-4b34-8f62-48ddafd84f26",
+   "metadata": {},
+   "source": [
+    "### Monte carlo lightcone\n",
+    "\n",
+    "In this mode, the generated lightcone returns halos where each distinct halo has its own mass and redshift. Below we generate a Monte Carlo realization of the lightcone, including host halos and subhalos over a specified range of redshift, and within a given sky area."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "28c7fc3e-ed07-40c1-8006-00160b5fb9cf",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from diffhalos.lightcone_generators import mc_lc"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "258fab77-7d71-4550-b476-afcb9d392d9b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from jax import random as jran"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ed5a4c3c-432a-4977-9a05-56d27c9869a9",
+   "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)\n",
+    "print(halopop_mc._fields)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d24ad2fe-b244-4c8b-8ec6-bbc198c5760d",
+   "metadata": {},
+   "source": [
+    "### Mass-function weighting\n",
+    "\n",
+    "In the weighted version of the generator, every halo and subhalo in the returned lightcone has a _weight_ that encodes the multiplicity of the halo. For any input choice of the number of returned halos, the weights are determined according to the halo and subhalo mass functions. For large enough halo populations, the two versions of the lightcone agree almost perfectly, as demonstrated below where we plot the redshift distributions from the Monte-Carlo (MC) and quasi-Monte Carlo (QMC, i.e. weighted) generators."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -22,35 +98,399 @@
    "source": [
     "from diffhalos import weighted_lc\n",
     "\n",
-    "from jax import random as jran\n",
     "ran_key = jran.key(0)\n",
     "\n",
-    "n_host_halos = 200\n",
-    "z_min, z_max = 0.4, 2.5\n",
-    "lgmp_min, lgmp_max = 10.5, 15.5\n",
-    "sky_area_degsq = 500.0\n",
+    "n_host_halos = 70_000\n",
+    "z_min, z_max = 0.4, 3.5\n",
+    "lgmp_min, lgmp_max = 12.5, 15.5\n",
+    "sky_area_degsq = 10.0\n",
     "\n",
     "args = (ran_key, n_host_halos, z_min, z_max, lgmp_min, lgmp_max, sky_area_degsq)\n",
-    "halopop = weighted_lc(*args)\n",
-    "print(halopop._fields)"
+    "halopop_qmc = weighted_lc(*args)\n",
+    "print(halopop_qmc._fields)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "10e881ea-dee6-4ea3-a7a3-ca6184c8f1c6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, ax = plt.subplots(figsize=([7,5]))\n",
+    "\n",
+    "plt.hist(halopop_mc.z_obs, bins=30, color='darkslategray', alpha=0.8)\n",
+    "plt.hist(halopop_qmc.z_obs, weights=halopop_qmc.halo_weight, bins=30, color='darkorange', histtype='step')\n",
+    "\n",
+    "plt.xlabel(r'$z$')\n",
+    "plt.ylabel(r'$dn_{\\rm halo} / dz$')\n",
+    "\n",
+    "plt.yscale('log')\n",
+    "\n",
+    "custom_handles = [\n",
+    "    mpatches.Patch(edgecolor=\"darkslategray\", lw=1, color='darkslategray', label=r\"${\\rm MC}$\"),\n",
+    "    mpatches.Patch(edgecolor=\"darkorange\", fill=False, lw=1, color='darkorange', label=r\"${\\rm QMC}$\"),\n",
+    "]\n",
+    "\n",
+    "ax.legend(handles=custom_handles, bbox_to_anchor=(0.8, 1.15), frameon=False, ncols=2, fontsize=18)"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "d24ad2fe-b244-4c8b-8ec6-bbc198c5760d",
+   "id": "0ba30017-0836-48e2-a78c-9e9c23582703",
+   "metadata": {},
+   "source": [
+    "### Compare host and subhalo populations separately by identifying them in the lightcone output"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bd08f3ee-8633-46dc-8f55-8d30b8fc17c5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2b2ea5ff-fd7c-4567-bcfc-b4646de510f9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "is_cen_mc = np.where(halopop_mc.central==1)[0]\n",
+    "is_sat_mc = np.where(halopop_mc.central==0)[0]\n",
+    "is_cen_qmc = np.where(halopop_qmc.central==1)[0]\n",
+    "is_sat_qmc = np.where(halopop_qmc.central==0)[0]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "054d96cb-b0e9-4550-af73-ebe7d7e8a89b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, ax = plt.subplots(figsize=([7,5]))\n",
+    "\n",
+    "plt.hist(halopop_mc.z_obs[is_cen_mc], bins=30, color='darkslategray', alpha=0.8)\n",
+    "plt.hist(halopop_qmc.z_obs[is_cen_mc], weights=halopop_qmc.halo_weight[is_cen_mc], bins=30, color='darkorange', histtype='step')\n",
+    "\n",
+    "plt.xlabel(r'$z$')\n",
+    "plt.ylabel(r'$dn_{\\rm host} / dz$')\n",
+    "\n",
+    "plt.yscale('log')\n",
+    "\n",
+    "custom_handles = [\n",
+    "    mpatches.Patch(edgecolor=\"darkslategray\", lw=1, color='darkslategray', label=r\"${\\rm MC}$\"),\n",
+    "    mpatches.Patch(edgecolor=\"darkorange\", fill=False, lw=1, color='darkorange', label=r\"${\\rm QMC}$\"),\n",
+    "]\n",
+    "\n",
+    "ax.legend(handles=custom_handles, bbox_to_anchor=(0.8, 1.15), frameon=False, ncols=2, fontsize=18)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f1b3f088-c44a-489f-8d9c-67a5527be9dd",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, ax = plt.subplots(figsize=([7,5]))\n",
+    "\n",
+    "plt.hist(halopop_mc.z_obs[is_sat_mc], bins=30, color='darkslategray', alpha=0.8)\n",
+    "plt.hist(halopop_qmc.z_obs[is_sat_qmc], weights=halopop_qmc.halo_weight[is_sat_qmc], bins=30, color='darkorange', histtype='step')\n",
+    "\n",
+    "plt.xlabel(r'$z$')\n",
+    "plt.ylabel(r'$dn_{\\rm sub} / dz$')\n",
+    "\n",
+    "plt.yscale('log')\n",
+    "\n",
+    "custom_handles = [\n",
+    "    mpatches.Patch(edgecolor=\"darkslategray\", lw=1, color='darkslategray', label=r\"${\\rm MC}$\"),\n",
+    "    mpatches.Patch(edgecolor=\"darkorange\", fill=False, lw=1, color='darkorange', label=r\"${\\rm QMC}$\"),\n",
+    "]\n",
+    "\n",
+    "ax.legend(handles=custom_handles, bbox_to_anchor=(0.8, 1.15), frameon=False, ncols=2, fontsize=18)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "517bed1c-4375-47f0-b94e-f0a6f05391b4",
    "metadata": {},
    "source": [
-    "### mass-function weighting\n",
+    "### Host halo generator compared to exact theoretical prediction\n",
     "\n",
-    "Every halo and subhalo in the returned lightcone has a _weight_ that encodes the multiplicity of the halo. For any input choice of the number of returned halos, the weights are determined according to the halo and subhalo mass functions..."
+    "If we generate a host halo population at a specific redshift and bin in mass to infer their distribution, the result must agree with the exact theory prediction from the halo mass function (HMF) underlying model. We show that this is the case below, at various redshifts."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6336d498-3e5b-4bb8-b020-53e399f00c6f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from diffhalos.hmf.mc_hosts import mc_host_halos_singlez\n",
+    "from diffhalos.hmf.hmf_model import _compute_nhalos_tot, DEFAULT_HMF_PARAMS, predict_diff_hmf"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ee07b8b3-f7db-4860-af83-b98f7602b972",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ran_key = jran.key(0)\n",
+    "\n",
+    "lgmp_min = 11.0\n",
+    "lgmsub_min = lgmp_min - 0.01\n",
+    "Lbox = 1000.0\n",
+    "volume_com_mpc = Lbox**3\n",
+    "\n",
+    "z_grid = [0.01, 0.5, 1.5, 2.5, 4.0]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3bb8dbb4-9fcb-4e88-aaac-28838b11f2b3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "bins = 30"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fb21da77-2a51-437b-8f1b-d10abe3ff137",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.cm as cm\n",
+    "colors = cm.viridis(np.linspace(0,1,len(z_grid)))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bb0cf3f3-9360-448c-a80e-ded13691239c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tick_fontsize = 17\n",
+    "label_fontsize = 18\n",
+    "legend_fontsize = 16"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bfb0cd2d-14dc-4ed7-a963-1d9c4ef53307",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, ax = plt.subplots(figsize=([7,5]))\n",
+    "\n",
+    "for iz, zi in enumerate(z_grid):\n",
+    "    \n",
+    "    z_min = zi-0.005\n",
+    "    z_max = zi+0.005\n",
+    "    \n",
+    "    args = (ran_key, lgmp_min, lgmsub_min, z_min, z_max, sky_area_degsq)\n",
+    "\n",
+    "    # MC host halos\n",
+    "    mc_logmhalo = mc_host_halos_singlez(ran_key, lgmp_min, zi, volume_com_mpc)\n",
+    "\n",
+    "    # bin generated host masses\n",
+    "    hist_data = np.histogram(mc_logmhalo, bins=bins, density=False)\n",
+    "\n",
+    "    dlogm_bin_edges = hist_data[1]\n",
+    "    dlogm_bins = 0.5 * (dlogm_bin_edges[1:] + dlogm_bin_edges[:-1])\n",
+    "\n",
+    "    n_halo = mc_logmhalo.size\n",
+    "    n_tot = _compute_nhalos_tot(\n",
+    "        DEFAULT_HMF_PARAMS,\n",
+    "        lgmp_min,\n",
+    "        zi,\n",
+    "        volume_com_mpc,\n",
+    "    )\n",
+    "\n",
+    "    # normalize counts\n",
+    "    dnhalo_bins = hist_data[0] / np.diff(dlogm_bin_edges) / volume_com_mpc\n",
+    "    dnhalo_bins *= n_tot / n_halo\n",
+    "\n",
+    "    # theory prediction\n",
+    "    diff_hmf = predict_diff_hmf(DEFAULT_HMF_PARAMS, dlogm_bins, zi)\n",
+    "\n",
+    "    ax.plot(dlogm_bins, dnhalo_bins, color=colors[iz], lw=2)\n",
+    "    ax.plot(dlogm_bins, 10**diff_hmf, color=colors[iz], ls='--', lw=2)\n",
+    "\n",
+    "ax.tick_params(axis='both', which='major', labelsize=tick_fontsize)\n",
+    "ax.tick_params(axis='both', which='minor', labelsize=tick_fontsize)\n",
+    "\n",
+    "ax.set_xlabel(r\"$m_{\\rm halo}$\", fontsize=label_fontsize)\n",
+    "ax.set_ylabel(r\"$n_{\\rm halo}\\; [{\\rm Mpc}^{-3}]$\", fontsize=label_fontsize)\n",
+    "\n",
+    "ax.set_yscale('log')\n",
+    "\n",
+    "ax.set_xlim([11.0, 16.0])\n",
+    "ax.set_ylim([1e-8, 6e-2])\n",
+    "\n",
+    "custom_handles = [\n",
+    "    Line2D([0], [0], color='k', lw=2, label=r\"{\\rm MC\\; hosts}\"),\n",
+    "    Line2D([0], [0], color=colors[0], lw=2, label=r\"$z=%.1f$\"%z_grid[0]),\n",
+    "    Line2D([0], [0], color='k', lw=2, ls='--', label=r\"{\\rm Theory}\"),\n",
+    "    Line2D([0], [0], color=colors[-1], lw=2, label=r\"$z=%.1f$\"%z_grid[-1]),\n",
+    "]\n",
+    "\n",
+    "ax.legend(handles=custom_handles, bbox_to_anchor=(0.8, 1.2), frameon=False, ncols=2, fontsize=legend_fontsize)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "55b555d0-1397-48bf-8791-0c971e66c07f",
+   "metadata": {},
+   "source": [
+    "### Subhalo generator compared to exact theoretical prediction\n",
+    "\n",
+    "Similarly to hosts, if we generate a subhalo population, conditioned on the host mass at a given redshift, and bin in $\\mu \\equiv M_{\\rm sub}/M_{\\rm host}$ to infer their distribution, the result must agree with the exact theory prediction from the conditional subhalo mass function (CSHMF) underlying model. We show that this is the case below, for host halo masses inside which the subhalos reside. Since we have shown that the MC and QMC subhalo lightcones agree with each other, we will utilize the weighted lightcone generator below to speed up the computations and to reduce the memory footprint."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4f59e25f-f49b-43e3-a79f-686d932ee52c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from diffhalos.lightcone_generators import weighted_lc_subhalos\n",
+    "from diffhalos.lightcone_generators.utils import generate_mock_cenpop\n",
+    "from diffhalos.ccshmf.ccshmf_model import predict_diff_cshmf, DEFAULT_CCSHMF_PARAMS"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ec497544-658d-4f2c-be9a-70196c2420df",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "lgmhost_arr = np.linspace(9.0, 12.0, 4)\n",
+    "lgmsub_min = 6.0\n",
+    "\n",
+    "z_min = 0.99\n",
+    "z_max = 1.01"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3ea687df-e2d2-4a76-afc7-aeb01b752b8a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.cm as cm\n",
+    "colors = cm.viridis(np.linspace(0,1,len(lgmhost_arr)))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a1ae090f-593b-4b38-a5dd-f1928acaf988",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "bins = 30\n",
+    "n_mu_per_host = 500"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "880e5a05-93bc-40d9-865b-2112fb22ca2d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, ax = plt.subplots(figsize=([7,5]))\n",
+    "\n",
+    "for im, mi in enumerate(lgmhost_arr):\n",
+    "\n",
+    "    logmp_min = mi-0.01\n",
+    "    logmp_max = mi+0.01\n",
+    "    \n",
+    "    cenpop = generate_mock_cenpop(z_min, z_max, logmp_min, logmp_max, n_cens=1)\n",
+    "    \n",
+    "    subpop = weighted_lc_subhalos(ran_key, cenpop, lgmsub_min, n_mu_per_host=n_mu_per_host)\n",
+    "\n",
+    "    # bin in mu\n",
+    "    hist_data = np.histogram(subpop.logmu_obs, weights=subpop.sat_weight, bins=bins, density=False)\n",
+    "    dlogmu_bin_edges = hist_data[1]\n",
+    "    dlogmu_bins = 0.5 * (dlogmu_bin_edges[1:] + dlogmu_bin_edges[:-1])\n",
+    "    \n",
+    "    # normalize counts\n",
+    "    dnsub_bins = hist_data[0] / np.diff(dlogmu_bin_edges)\n",
+    "\n",
+    "    # theory prediction\n",
+    "    logmu_arr = np.linspace(subpop.logmu_obs.min(), subpop.logmu_obs.max(), 200)\n",
+    "    cshmf_theory = predict_diff_cshmf(DEFAULT_CCSHMF_PARAMS, mi, logmu_arr)\n",
+    "\n",
+    "    ax.plot(10**dlogmu_bins, dnsub_bins, color=colors[im], lw=2)\n",
+    "    ax.plot(10**logmu_arr, 10**cshmf_theory, color=colors[im], ls='--', lw=2)\n",
+    "\n",
+    "ax.tick_params(axis='both', which='major', labelsize=tick_fontsize)\n",
+    "ax.tick_params(axis='both', which='minor', labelsize=tick_fontsize)\n",
+    "\n",
+    "ax.set_xlabel(r\"$\\mu$\", fontsize=label_fontsize)\n",
+    "ax.set_ylabel(r\"$n_{\\rm sub} (\\mu | m_{\\rm host})$\", fontsize=label_fontsize)\n",
+    "\n",
+    "ax.set_xscale('log')\n",
+    "ax.set_yscale('log')\n",
+    "\n",
+    "custom_handles = [\n",
+    "    Line2D([0], [0], color='k', lw=2, label=r\"{\\rm MC\\; subhalos}\"),\n",
+    "    Line2D([0], [0], color=colors[0], lw=2, label=r\"$m_{\\rm host}=%.1f$\"%lgmhost_arr[0]),\n",
+    "    Line2D([0], [0], color='k', lw=2, ls='--', label=r\"{\\rm Theory}\"),\n",
+    "    Line2D([0], [0], color=colors[-1], lw=2, label=r\"$m_{\\rm host}=%.1f$\"%lgmhost_arr[-1]),\n",
+    "]\n",
+    "\n",
+    "ax.legend(handles=custom_handles, bbox_to_anchor=(0.9, 1.2), frameon=False, ncols=2, fontsize=legend_fontsize)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d9d911db-b149-49ef-9a6c-f1e16a00211e",
+   "metadata": {},
+   "source": [
+    "### Something related to MAHs?"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "53267d56-7659-48f0-8c61-fafacda552c8",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "82884cfa-1b30-438a-b644-d2459054c832",
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "diffstuff",
    "language": "python",
-   "name": "python3"
+   "name": "diffstuff"
   },
   "language_info": {
    "codemirror_mode": {
@@ -62,7 +502,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.13.11"
+   "version": "3.11.9"
   }
  },
  "nbformat": 4,

From e340a32e6a7c955119b6feda72c34ddc048ff431 Mon Sep 17 00:00:00 2001
From: Georgios Zacharegkas <gzacharegkas@gmail.com>
Date: Wed, 8 Jul 2026 13:02:47 -0500
Subject: [PATCH 2/9] updating docs

---
 README.rst       | 9 ++++++++-
 requirements.txt | 3 ++-
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/README.rst b/README.rst
index bfb4071..a25b742 100644
--- a/README.rst
+++ b/README.rst
@@ -24,4 +24,11 @@ To run the suite of unit tests::
 To build html of test coverage::
 
     $ pytest -v --cov --cov-report html
-    $ open htmlcov/index.html
\ No newline at end of file
+    $ open htmlcov/index.html
+
+Demo notebooks
+--------------
+The `diffhalos_quickstart.ipynb` notebook illustrates how to use some of the core functions of diffhalos. In particular,
+it demonstrates how to generate halo and subhalo lightcones, while it also shows that the output from the Monte-Carlo
+lightcone realizations agree with the exact theory predictions from the halo mass function and the conditional
+subhalo mas function implemented in diffhalos.
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index ec9ee1e..dad8b21 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -3,4 +3,5 @@ jax
 dsps
 diffmah
 dsps
-flowjax
\ No newline at end of file
+flowjax
+matplotlib
\ No newline at end of file

From 4268feb19d42fe1b825c2dc5af7122400f01e0af Mon Sep 17 00:00:00 2001
From: Georgios Zacharegkas <gzacharegkas@gmail.com>
Date: Thu, 9 Jul 2026 11:54:29 -0500
Subject: [PATCH 3/9] adding MAH notebook in docs

---
 diffhalos/mah/diffmahnet_utils.py      |   2 +-
 docs/source/MAHs_with_Diffhalos.ipynb  | 411 +++++++++++++++++++++++++
 docs/source/diffhalos_quickstart.ipynb |  26 +-
 3 files changed, 418 insertions(+), 21 deletions(-)
 create mode 100644 docs/source/MAHs_with_Diffhalos.ipynb

diff --git a/diffhalos/mah/diffmahnet_utils.py b/diffhalos/mah/diffmahnet_utils.py
index 27b4581..6c874bb 100644
--- a/diffhalos/mah/diffmahnet_utils.py
+++ b/diffhalos/mah/diffmahnet_utils.py
@@ -57,7 +57,7 @@ def mc_mah_cenpop(
     m_obs: ndarray of shape (n_cens, )
         grid of base-10 log of mass of the halos at observation, in Msun
 
-    t_obs: ndarray of shape (n_subs, )
+    t_obs: ndarray of shape (n_cens, )
         cosmic time at observation of each halo, in Gyr
 
     randkey: key
diff --git a/docs/source/MAHs_with_Diffhalos.ipynb b/docs/source/MAHs_with_Diffhalos.ipynb
new file mode 100644
index 0000000..6253471
--- /dev/null
+++ b/docs/source/MAHs_with_Diffhalos.ipynb
@@ -0,0 +1,411 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "cfda5b4a-adde-49df-88ce-a7a7c9c14d01",
+   "metadata": {},
+   "source": [
+    "# Calculating Mass Assembly Histories (MAHs) with DiffnahNet within Diffhalos\n",
+    "\n",
+    "In `Diffhalos` we use `DiffmahNet` in order to assign MAHs to the halo populations. `DiffmahNet` is a population-level model of the `Diffmah` model for predicting MAHs of individual halos. Once the halos are generated by sampling from the halo and subhalo mass function, for hosts ans subhalos respectively, we `DiffmahNet` pastes a MAH onto each of the (sub)halo in the lightcone, conditioned on the mass $M_{\\rm obs} \\equiv M_{\\rm halo} (t_{\\rm obs})$ of the (sub)halo 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 actually 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 set 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 halo lightcone outputs can be utilized to generate halo MAHs. Each halo in the generated MC lightcone in what follows is assigned a set of MAH parameters (already properly scaled), which can be used to produce the assembly history of each halo, some of which we plot 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)), 5))\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')\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": "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
+}
diff --git a/docs/source/diffhalos_quickstart.ipynb b/docs/source/diffhalos_quickstart.ipynb
index f866951..8e488cb 100644
--- a/docs/source/diffhalos_quickstart.ipynb
+++ b/docs/source/diffhalos_quickstart.ipynb
@@ -140,7 +140,9 @@
    "id": "0ba30017-0836-48e2-a78c-9e9c23582703",
    "metadata": {},
    "source": [
-    "### Compare host and subhalo populations separately by identifying them in the lightcone output"
+    "### Compare host and subhalo populations separately by identifying them in the lightcone output\n",
+    "\n",
+    "So far we have plotted halos and subhalos together. Below we demonstrate how one can distinguish between the two populations in the generated lightcones and use them separately. As we show in the plots, for each populations separately, MC and QMC are still in great agreement with each other."
    ]
   },
   {
@@ -223,7 +225,7 @@
    "source": [
     "### Host halo generator compared to exact theoretical prediction\n",
     "\n",
-    "If we generate a host halo population at a specific redshift and bin in mass to infer their distribution, the result must agree with the exact theory prediction from the halo mass function (HMF) underlying model. We show that this is the case below, at various redshifts."
+    "If we generate a halo population, from a lightcone generation, at a specific redshift and bin the halos in mass to infer their distribution, the result must agree with the exact theory prediction from the underlying halo mass function (HMF) model. We show that this is true below, at various redshifts."
    ]
   },
   {
@@ -358,7 +360,7 @@
    "source": [
     "### Subhalo generator compared to exact theoretical prediction\n",
     "\n",
-    "Similarly to hosts, if we generate a subhalo population, conditioned on the host mass at a given redshift, and bin in $\\mu \\equiv M_{\\rm sub}/M_{\\rm host}$ to infer their distribution, the result must agree with the exact theory prediction from the conditional subhalo mass function (CSHMF) underlying model. We show that this is the case below, for host halo masses inside which the subhalos reside. Since we have shown that the MC and QMC subhalo lightcones agree with each other, we will utilize the weighted lightcone generator below to speed up the computations and to reduce the memory footprint."
+    "Similarly to halos, if we generate a subhalo population, conditioned on the host mass at a given redshift, and bin in $\\mu \\equiv M_{\\rm sub}/M_{\\rm host}$ to infer their distribution, the result must agree with the exact theory prediction from the conditional subhalo mass function (CSHMF) model. We show that this is the case below, for various host halo masses. Since we have shown already that the MC and QMC subhalo lightcones agree with each other, below we utilize the weighted lightcone generator to speed up the computations and to reduce the memory footprint."
    ]
   },
   {
@@ -461,26 +463,10 @@
     "ax.legend(handles=custom_handles, bbox_to_anchor=(0.9, 1.2), frameon=False, ncols=2, fontsize=legend_fontsize)"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "id": "d9d911db-b149-49ef-9a6c-f1e16a00211e",
-   "metadata": {},
-   "source": [
-    "### Something related to MAHs?"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "53267d56-7659-48f0-8c61-fafacda552c8",
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "82884cfa-1b30-438a-b644-d2459054c832",
+   "id": "d7c4c261-ac40-4b46-9327-a2b4fa5bc3a9",
    "metadata": {},
    "outputs": [],
    "source": []

From aca11c88503eba8ab9fe6e71d069edef3f5f7ed7 Mon Sep 17 00:00:00 2001
From: Georgios Zacharegkas <gzacharegkas@gmail.com>
Date: Thu, 9 Jul 2026 12:01:07 -0500
Subject: [PATCH 4/9] adding new notebook index

---
 docs/source/index.rst | 1 +
 1 file changed, 1 insertion(+)

diff --git a/docs/source/index.rst b/docs/source/index.rst
index 651b91f..cb64533 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -16,6 +16,7 @@ User Guide
 
    installation.rst
    diffhalos_quickstart.ipynb
+   MAHs_with_Diffhalos.ipynb
    tutorials.rst
 
 See :ref:`Citation Information <cite_info>` for how to acknowledge diffhalos.

From e7f449a0e7488fdd8354cca5adce796f62dd9e3b Mon Sep 17 00:00:00 2001
From: Georgios Zacharegkas <gzacharegkas@gmail.com>
Date: Fri, 10 Jul 2026 10:29:42 -0500
Subject: [PATCH 5/9] updating github workflows

---
 .github/workflows/linting.yaml          | 2 +-
 .github/workflows/test_main_branch.yaml | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/linting.yaml b/.github/workflows/linting.yaml
index 09cb65c..aead1b4 100644
--- a/.github/workflows/linting.yaml
+++ b/.github/workflows/linting.yaml
@@ -12,7 +12,7 @@ jobs:
     runs-on: "ubuntu-latest"
 
     steps:
-      - uses: actions/checkout@v6
+      - uses: actions/checkout@v7
 
       - uses: conda-incubator/setup-miniconda@v4
         with:
diff --git a/.github/workflows/test_main_branch.yaml b/.github/workflows/test_main_branch.yaml
index fe9f6f8..de42dbc 100644
--- a/.github/workflows/test_main_branch.yaml
+++ b/.github/workflows/test_main_branch.yaml
@@ -16,7 +16,7 @@ jobs:
     runs-on: "ubuntu-latest"
 
     steps:
-      - uses: actions/checkout@v6
+      - uses: actions/checkout@v7
         with:
           fetch-depth: 0
 
@@ -55,7 +55,7 @@ jobs:
           pytest -v diffhalos --cov --cov-report=xml
 
       - name: Upload coverage reports to Codecov
-        uses: codecov/codecov-action@v6
+        uses: codecov/codecov-action@v7
         with:
           token: ${{ secrets.CODECOV_TOKEN }}
 

From ed6c7ebe3813e4c5a1491748eaac4e5e0fe596bb Mon Sep 17 00:00:00 2001
From: Georgios Zacharegkas <gzacharegkas@gmail.com>
Date: Fri, 10 Jul 2026 10:39:57 -0500
Subject: [PATCH 6/9] updating docs and renaming notebook

---
 README.rst                                            | 11 +++++++----
 ...{MAHs_with_Diffhalos.ipynb => diffhalos_mah.ipynb} |  0
 2 files changed, 7 insertions(+), 4 deletions(-)
 rename docs/source/{MAHs_with_Diffhalos.ipynb => diffhalos_mah.ipynb} (100%)

diff --git a/README.rst b/README.rst
index a25b742..455a9f3 100644
--- a/README.rst
+++ b/README.rst
@@ -28,7 +28,10 @@ To build html of test coverage::
 
 Demo notebooks
 --------------
-The `diffhalos_quickstart.ipynb` notebook illustrates how to use some of the core functions of diffhalos. In particular,
-it demonstrates how to generate halo and subhalo lightcones, while it also shows that the output from the Monte-Carlo
-lightcone realizations agree with the exact theory predictions from the halo mass function and the conditional
-subhalo mas function implemented in diffhalos.
\ No newline at end of file
+The `diffhalos_quickstart.ipynb` notebook illustrates how to use some of the core functions of diffhalos. In particular, it demonstrates 
+how to generate halo and subhalo lightcones using the Monte-Carlo and quasi-Monte Carlo versions of the generator. In addition, 
+it shows that the lightcone outputs agree well with the exact theory predictions from the halo and subhalo mass functions in diffhalos. 
+
+The `diffhalos_mah.ipynb` notebook goes over the basics of using the `DiffmahNet` model in `Diffhalos` to generate mass assembly 
+histories (MAHs) of populations of dark matter halos. It also demonstrates how the MAH of halos is part of the standard output 
+from lightcone generation.
\ No newline at end of file
diff --git a/docs/source/MAHs_with_Diffhalos.ipynb b/docs/source/diffhalos_mah.ipynb
similarity index 100%
rename from docs/source/MAHs_with_Diffhalos.ipynb
rename to docs/source/diffhalos_mah.ipynb

From 4d668363f9aaee1058cf91daf1eab7876b723afa Mon Sep 17 00:00:00 2001
From: Georgios Zacharegkas <gzacharegkas@gmail.com>
Date: Fri, 10 Jul 2026 10:42:05 -0500
Subject: [PATCH 7/9] updates to readme

---
 README.rst | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/README.rst b/README.rst
index 455a9f3..9500dbd 100644
--- a/README.rst
+++ b/README.rst
@@ -28,10 +28,10 @@ To build html of test coverage::
 
 Demo notebooks
 --------------
-The `diffhalos_quickstart.ipynb` notebook illustrates how to use some of the core functions of diffhalos. In particular, it demonstrates 
+The :code:`diffhalos_quickstart.ipynb` notebook illustrates how to use some of the core functions of diffhalos. In particular, it demonstrates 
 how to generate halo and subhalo lightcones using the Monte-Carlo and quasi-Monte Carlo versions of the generator. In addition, 
 it shows that the lightcone outputs agree well with the exact theory predictions from the halo and subhalo mass functions in diffhalos. 
 
-The `diffhalos_mah.ipynb` notebook goes over the basics of using the `DiffmahNet` model in `Diffhalos` to generate mass assembly 
+The :code:`diffhalos_mah.ipynb` notebook goes over the basics of using the :code:`DiffmahNet` model in :code:`Diffhalos` to generate mass assembly 
 histories (MAHs) of populations of dark matter halos. It also demonstrates how the MAH of halos is part of the standard output 
 from lightcone generation.
\ No newline at end of file

From fb839205a9619b4a3c90d08d3fbaadaaccb72904 Mon Sep 17 00:00:00 2001
From: Georgios Zacharegkas <gzacharegkas@gmail.com>
Date: Fri, 10 Jul 2026 15:23:35 -0500
Subject: [PATCH 8/9] changes to readme and adding new notebook

---
 README.rst                             |  26 ++-
 docs/source/diffhalos_mah.ipynb        |  21 +-
 docs/source/diffhalos_mc_vs_qmc.ipynb  | 279 +++++++++++++++++++++++++
 docs/source/diffhalos_quickstart.ipynb |  48 ++---
 docs/source/index.rst                  |   4 +-
 docs/source/tutorials.rst              |   5 -
 6 files changed, 328 insertions(+), 55 deletions(-)
 create mode 100644 docs/source/diffhalos_mc_vs_qmc.ipynb
 delete mode 100644 docs/source/tutorials.rst

diff --git a/README.rst b/README.rst
index 9500dbd..41379ce 100644
--- a/README.rst
+++ b/README.rst
@@ -1,6 +1,14 @@
 diffhalos
 ============
 
+Diffhalos is a generative model of cosmological lightcones of dark matter halos, subhalos, and their merger trees. 
+This library is written in JAX, and makes differentiable predictions for halo populations in a lightcone, 
+or at a single redshift.
+
+Documentation
+-------------
+Online documentation is available at `diffhalos.readthedocs.io <https://diffhalos.readthedocs.io/en/latest/>`_
+
 Installation
 ------------
 To install diffhalos into your environment from the source code::
@@ -8,12 +16,6 @@ To install diffhalos into your environment from the source code::
     $ cd /path/to/root/diffhalos
     $ pip install .
 
-
-Documentation
--------------
-Online documentation is available at `diffhalos.readthedocs.io <https://diffhalos.readthedocs.io/en/latest/>`_
-
-
 Testing
 -------
 To run the suite of unit tests::
@@ -30,8 +32,12 @@ Demo notebooks
 --------------
 The :code:`diffhalos_quickstart.ipynb` notebook illustrates how to use some of the core functions of diffhalos. In particular, it demonstrates 
 how to generate halo and subhalo lightcones using the Monte-Carlo and quasi-Monte Carlo versions of the generator. In addition, 
-it shows that the lightcone outputs agree well with the exact theory predictions from the halo and subhalo mass functions in diffhalos. 
+it shows that the lightcone outputs agree well with the exact theory predictions from the halo and subhalo mass functions in Diffhalos. 
+
+In the notebook :code:`diffhalos_mah.ipynb` we go over the basics of using the DiffmahNet model in Diffhalos to generate mass assembly 
+histories (MAHs) of populations of dark matter halos. It also demonstrates how the MAH of synthetic halos is part of the standard output 
+in the lightcones.
 
-The :code:`diffhalos_mah.ipynb` notebook goes over the basics of using the :code:`DiffmahNet` model in :code:`Diffhalos` to generate mass assembly 
-histories (MAHs) of populations of dark matter halos. It also demonstrates how the MAH of halos is part of the standard output 
-from lightcone generation.
\ No newline at end of file
+We demonstrate and explain in more detail the fundamentals of using the Monte Carlo and quasi-Monte Carlo generators, and discuss their differences, 
+in the notebook :code:`diffhalos_mc_vs_qmc.ipynb`. In there, we cover the basics of how to use each version for host halos, 
+and we talk about how they differ in their approch to lightcone generation.
\ No newline at end of file
diff --git a/docs/source/diffhalos_mah.ipynb b/docs/source/diffhalos_mah.ipynb
index 6253471..2fd943d 100644
--- a/docs/source/diffhalos_mah.ipynb
+++ b/docs/source/diffhalos_mah.ipynb
@@ -5,9 +5,9 @@
    "id": "cfda5b4a-adde-49df-88ce-a7a7c9c14d01",
    "metadata": {},
    "source": [
-    "# Calculating Mass Assembly Histories (MAHs) with DiffnahNet within Diffhalos\n",
+    "# Generating Mass Assembly Histories in Diffhalos\n",
     "\n",
-    "In `Diffhalos` we use `DiffmahNet` in order to assign MAHs to the halo populations. `DiffmahNet` is a population-level model of the `Diffmah` model for predicting MAHs of individual halos. Once the halos are generated by sampling from the halo and subhalo mass function, for hosts ans subhalos respectively, we `DiffmahNet` pastes a MAH onto each of the (sub)halo in the lightcone, conditioned on the mass $M_{\\rm obs} \\equiv M_{\\rm halo} (t_{\\rm obs})$ of the (sub)halo at the time of observation $t_{\\rm obs}$. "
+    "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}$. "
    ]
   },
   {
@@ -45,7 +45,7 @@
    "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 actually satisfied."
+    "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."
    ]
   },
   {
@@ -156,7 +156,7 @@
    "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 set 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}$."
+    "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}$."
    ]
   },
   {
@@ -280,7 +280,7 @@
     "           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)"
+    "_=ax.legend(handles=custom_handles, bbox_to_anchor=(0.9, 1.23), frameon=False, ncols=1, fontsize=17)"
    ]
   },
   {
@@ -290,7 +290,7 @@
    "source": [
     "### Generating MAH lightcones\n",
     "\n",
-    "Below we show that the halo lightcone outputs can be utilized to generate halo MAHs. Each halo in the generated MC lightcone in what follows is assigned a set of MAH parameters (already properly scaled), which can be used to produce the assembly history of each halo, some of which we plot for demonstration purposes."
+    "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."
    ]
   },
   {
@@ -363,19 +363,14 @@
    "source": [
     "fig, ax = plt.subplots(figsize=([7,5]))\n",
     "\n",
-    "indx = np.asarray(random.sample(range(0, len(halopop_mc.t_obs)), 5))\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')\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)"
+    "_=ax.set_yscale('log')"
    ]
   },
   {
diff --git a/docs/source/diffhalos_mc_vs_qmc.ipynb b/docs/source/diffhalos_mc_vs_qmc.ipynb
new file mode 100644
index 0000000..2b0c6a2
--- /dev/null
+++ b/docs/source/diffhalos_mc_vs_qmc.ipynb
@@ -0,0 +1,279 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "45059532",
+   "metadata": {},
+   "source": [
+    "# Differences between Monte Carlo and quasi-Monte Carlo\n",
+    "\n",
+    "This notebook presents the main differences between the two modes of lightcone generation with Diffhalos: the Monte Carlo (MC) and the quasi-Monte Carlo (QMC), also called _weighted_, lightcone.\n",
+    "\n",
+    "The basic difference in the philosophy behind the two versions in the case of host halos is the following: while in the MC version we randomly distribute each halo in redshift $z$ and mass $M_{\\rm halo}$ (using the halo mass function; HMF), in the case of QMC we contruct a _uniform_ grid of $(z,m_{\\rm halo})$, where $m_{\\rm halo} \\equiv \\log_{10} M_{\\rm halo}$. Then, in the QMC lightcone each point is assigned a _weight_ that represents the expected abundance of halos, according to an underlying HMF model. \n",
+    "\n",
+    "For subhalos the difference between the two versions of the lightcone is analogous to the host halo generation above. In the MC lightcone, for each _individual_ (host) halo, generated according to the above, we draw samples from the conditional subhalo mass fucntion (CSHMF; as a function of the parameter $\\mu \\equiv M_{\\rm sub}/M_{\\rm host}$), and thus populate each of them with the subhalos they host. In the QMC version, at each point in the preconstructed grid of redshift and mass for hosts from above, we first create a grid of $\\mu$ values (typically $3-5$) between a minimum and a maximum, and then we assign a _weight_ to each point that represents the abundance of subhalos."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7fd7d119-fa3a-47df-ad36-55a5416e35f0",
+   "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=17)\n",
+    "plt.rcParams.update({'xtick.direction': 'in', 'ytick.direction': 'in'})"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "19e4ea51-a236-4b34-8f62-48ddafd84f26",
+   "metadata": {},
+   "source": [
+    "### Visualizing a Monte Carlo lightcone\n",
+    "\n",
+    "In the cells below, we create a MC lightcone and we visualize how the points populate the $z-m_{\\rm halo}$ plane."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "28c7fc3e-ed07-40c1-8006-00160b5fb9cf",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from diffhalos.lightcone_generators.mc_lightcone_halos import mc_lc_halos"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "258fab77-7d71-4550-b476-afcb9d392d9b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from jax import random as jran"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ed5a4c3c-432a-4977-9a05-56d27c9869a9",
+   "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.1, 4.0\n",
+    "sky_area_degsq = 10.0\n",
+    "\n",
+    "args = (ran_key, lgmp_min, z_min, z_max, sky_area_degsq)\n",
+    "\n",
+    "halopop_mc = mc_lc_halos(*args)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0d7af999-2ebd-49f2-a467-e1ff0b7c9618",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plt.figure(figsize=([6,5]))\n",
+    "\n",
+    "plt.scatter(halopop_mc.z_obs, 10**halopop_mc.logmp_obs, s=1, color='cadetblue')\n",
+    "\n",
+    "plt.yscale('log')\n",
+    "\n",
+    "plt.xlabel(r\"$z$\")\n",
+    "_=plt.ylabel(r\"$M_{\\rm halo}\\; [{\\rm M_\\odot}]$\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "aeb105ba-832d-488b-8644-d751ee4f6da1",
+   "metadata": {},
+   "source": [
+    "The general trends we observe are as expected. From high to low redshift the maximum halo mass increases, as more massive halos form over cosmic time. Also, the abundance of low-mass halos is larger per redshift slice due to the shape of the HMF which increases as a steep power-law in $\\Lambda$CDM cosmologies at the low-mass regime.\n",
+    "\n",
+    "The total number of halos in the generated lightcone is not predetermined, but rather computed according to the input parameters. In particular, between redshifts $z_{\\min} < z < z_{\\max}$ and given a range of halo mass $M_{\\min} < M_{\\rm halo} < M_{\\max}$, we first compute a mean number of halos using the cumulative HMF $n_{\\rm halo}(>M_{\\rm halo},z)$. Once this is done for the whole lightcone, we generate Poisson realizations of the number of halos, so that each generated lightcone is slightly different."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e8a2e8e4-713e-41fe-b390-e547a01bc26a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "halopop_mc.z_obs.size"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d24ad2fe-b244-4c8b-8ec6-bbc198c5760d",
+   "metadata": {},
+   "source": [
+    "### Visualizing a quasi-Monte Carlo lightcone\n",
+    "\n",
+    "We now create a visualization of a weighted halo lightcone."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4581335b-aaaa-4d62-8292-60e4a76945bd",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from diffhalos.lightcone_generators.mc_lightcone_halos import weighted_lc_halos\n",
+    "\n",
+    "ran_key = jran.key(0)\n",
+    "\n",
+    "n_host_halos = 10000\n",
+    "z_min, z_max = 0.1, 4.0\n",
+    "lgmp_min, lgmp_max = 12.5, 15.5\n",
+    "sky_area_degsq = 10.0\n",
+    "\n",
+    "args = (ran_key, n_host_halos, z_min, z_max, lgmp_min, lgmp_max, sky_area_degsq)\n",
+    "halopop_qmc = weighted_lc_halos(*args)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f0e12a15-a193-4ad7-9c08-1778a9016d22",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "10e881ea-dee6-4ea3-a7a3-ca6184c8f1c6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plt.figure(figsize=([6,5]))\n",
+    "\n",
+    "scatter = plt.scatter(halopop_qmc.z_obs, 10**halopop_qmc.logmp_obs, s=1, c=np.log10(halopop_qmc.cen_weight), cmap='viridis')\n",
+    "\n",
+    "cbar = plt.colorbar(scatter)\n",
+    "cbar.set_label(r'$\\log_{10} n_{\\rm halo}$')\n",
+    "\n",
+    "plt.yscale('log')\n",
+    "\n",
+    "plt.xlabel(r\"$z$\")\n",
+    "_=plt.ylabel(r\"$M_{\\rm halo}\\; [{\\rm M_\\odot}]$\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ff90c5eb-e2dd-4af8-aa05-74d2e7481602",
+   "metadata": {},
+   "source": [
+    "As we see, the points are now uniformly distributed in redshist and halo mass. Additionally, we see a general trend where the weights (the plot is color-coded based on the base-10 log of the halo weights, in this case the differential HMF model prediction) increase from top-right to bottom-left, as lower-mass, lower-redshift halos are generally more abundant."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e271a49b-f60d-4153-9b95-4cff69391ed9",
+   "metadata": {},
+   "source": [
+    "### Using the weighted lightcones\n",
+    "\n",
+    "When using the QMC lightcone outputs for any kind of computation involving halo statistics, attention is needed to ensure that the weights are utilized properly. Since each point in the generated grid is simply the result of uniformly distributing halos in redshift and mass, the weights encode all the necessary information of the halo abundance when computing statistical quantities.\n",
+    "\n",
+    "For example, to plot the distribution of the generated halos as a function of redshift, we need to use the property `cen_weight` in the output to properly weight the histogram. Below we demonstrate this."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3f802cb5-5bd4-482a-adb5-b6b8082f056e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plt.figure(figsize=([5,4]))\n",
+    "\n",
+    "_=plt.hist(halopop_qmc.z_obs, weights=halopop_qmc.cen_weight, bins=20, color='darkred', alpha=0.6)\n",
+    "\n",
+    "plt.yscale('log')\n",
+    "\n",
+    "plt.xlabel(r\"$z$\")\n",
+    "_=plt.ylabel(r\"$N_{\\rm halos}$\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3995ed13-f1f2-4b27-83a1-6f64019a4f23",
+   "metadata": {},
+   "source": [
+    "When using the MC lightcone, there is not need to weight since the generated halos in it already have the proper abundance. In addition, in the limit where the number of generated halos is large enough, the MC and QMC lightcones should look almost identical. Below we show that this is roughly the case, even for the relatively low number of halos in the lightcones in this notebook. For a more careful comparison, see `diffhalos_quickstart.ipynb`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "94fa541f-cbec-4ecd-ba64-b4dd76b46890",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, ax = plt.subplots(figsize=([5,4]))\n",
+    "\n",
+    "_=plt.hist(halopop_qmc.z_obs, weights=halopop_qmc.cen_weight, bins=20, color='darkred', alpha=0.6)\n",
+    "_=plt.hist(halopop_mc.z_obs, bins=20, color='darkslateblue', histtype='step')\n",
+    "\n",
+    "plt.yscale('log')\n",
+    "\n",
+    "plt.xlabel(r\"$z$\")\n",
+    "plt.ylabel(r\"$N_{\\rm halos}$\")\n",
+    "\n",
+    "custom_handles = [\n",
+    "    mpatches.Patch(color=\"darkred\", lw=1, label=r\"${\\rm QMC}$\"),\n",
+    "    mpatches.Patch(color=\"darkslateblue\", fill=False, lw=1, label=r\"${\\rm MC}$\"),\n",
+    "]\n",
+    "\n",
+    "_=ax.legend(handles=custom_handles, bbox_to_anchor=(0.9, 1.15), frameon=False, ncols=2, fontsize=18)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2c1b4bba-cfbb-45db-adfc-b87bab8b155a",
+   "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
+}
diff --git a/docs/source/diffhalos_quickstart.ipynb b/docs/source/diffhalos_quickstart.ipynb
index 8e488cb..164a089 100644
--- a/docs/source/diffhalos_quickstart.ipynb
+++ b/docs/source/diffhalos_quickstart.ipynb
@@ -8,25 +8,7 @@
     "# Quickstart\n",
     "\n",
     "This quickstart guide gives an overview of the core features of Diffhalos,\n",
-    "and demonstrates how to use the primary functions of the library.\n",
-    "\n",
-    "The cell below shows how to generate a lightcone of halos and subhalos, and their mass assembly histories."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "7fd7d119-fa3a-47df-ad36-55a5416e35f0",
-   "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=20)\n",
-    "plt.rcParams.update({'xtick.direction': 'in', 'ytick.direction': 'in'})"
+    "and demonstrates how to use the primary functions of the library."
    ]
   },
   {
@@ -89,6 +71,22 @@
     "In the weighted version of the generator, every halo and subhalo in the returned lightcone has a _weight_ that encodes the multiplicity of the halo. For any input choice of the number of returned halos, the weights are determined according to the halo and subhalo mass functions. For large enough halo populations, the two versions of the lightcone agree almost perfectly, as demonstrated below where we plot the redshift distributions from the Monte-Carlo (MC) and quasi-Monte Carlo (QMC, i.e. weighted) generators."
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7fd7d119-fa3a-47df-ad36-55a5416e35f0",
+   "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=20)\n",
+    "plt.rcParams.update({'xtick.direction': 'in', 'ytick.direction': 'in'})"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -132,7 +130,7 @@
     "    mpatches.Patch(edgecolor=\"darkorange\", fill=False, lw=1, color='darkorange', label=r\"${\\rm QMC}$\"),\n",
     "]\n",
     "\n",
-    "ax.legend(handles=custom_handles, bbox_to_anchor=(0.8, 1.15), frameon=False, ncols=2, fontsize=18)"
+    "_=ax.legend(handles=custom_handles, bbox_to_anchor=(0.8, 1.15), frameon=False, ncols=2, fontsize=18)"
    ]
   },
   {
@@ -190,7 +188,7 @@
     "    mpatches.Patch(edgecolor=\"darkorange\", fill=False, lw=1, color='darkorange', label=r\"${\\rm QMC}$\"),\n",
     "]\n",
     "\n",
-    "ax.legend(handles=custom_handles, bbox_to_anchor=(0.8, 1.15), frameon=False, ncols=2, fontsize=18)"
+    "_=ax.legend(handles=custom_handles, bbox_to_anchor=(0.8, 1.15), frameon=False, ncols=2, fontsize=18)"
    ]
   },
   {
@@ -215,7 +213,7 @@
     "    mpatches.Patch(edgecolor=\"darkorange\", fill=False, lw=1, color='darkorange', label=r\"${\\rm QMC}$\"),\n",
     "]\n",
     "\n",
-    "ax.legend(handles=custom_handles, bbox_to_anchor=(0.8, 1.15), frameon=False, ncols=2, fontsize=18)"
+    "_=ax.legend(handles=custom_handles, bbox_to_anchor=(0.8, 1.15), frameon=False, ncols=2, fontsize=18)"
    ]
   },
   {
@@ -225,7 +223,7 @@
    "source": [
     "### Host halo generator compared to exact theoretical prediction\n",
     "\n",
-    "If we generate a halo population, from a lightcone generation, at a specific redshift and bin the halos in mass to infer their distribution, the result must agree with the exact theory prediction from the underlying halo mass function (HMF) model. We show that this is true below, at various redshifts."
+    "If we produce a halo population, from a lightcone generation, at a specific redshift and bin the halos in mass to infer their distribution, the result must agree with the exact theory prediction from the underlying halo mass function (HMF) model. We show that this is true below, at various redshifts. Note that we define $m_{\\rm halo} \\equiv \\log_{10} M_{\\rm halo}$."
    ]
   },
   {
@@ -350,7 +348,7 @@
     "    Line2D([0], [0], color=colors[-1], lw=2, label=r\"$z=%.1f$\"%z_grid[-1]),\n",
     "]\n",
     "\n",
-    "ax.legend(handles=custom_handles, bbox_to_anchor=(0.8, 1.2), frameon=False, ncols=2, fontsize=legend_fontsize)"
+    "_=ax.legend(handles=custom_handles, bbox_to_anchor=(0.8, 1.2), frameon=False, ncols=2, fontsize=legend_fontsize)"
    ]
   },
   {
@@ -460,7 +458,7 @@
     "    Line2D([0], [0], color=colors[-1], lw=2, label=r\"$m_{\\rm host}=%.1f$\"%lgmhost_arr[-1]),\n",
     "]\n",
     "\n",
-    "ax.legend(handles=custom_handles, bbox_to_anchor=(0.9, 1.2), frameon=False, ncols=2, fontsize=legend_fontsize)"
+    "_=ax.legend(handles=custom_handles, bbox_to_anchor=(0.9, 1.2), frameon=False, ncols=2, fontsize=legend_fontsize)"
    ]
   },
   {
diff --git a/docs/source/index.rst b/docs/source/index.rst
index cb64533..5fd5901 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -16,7 +16,7 @@ User Guide
 
    installation.rst
    diffhalos_quickstart.ipynb
-   MAHs_with_Diffhalos.ipynb
-   tutorials.rst
+   diffhalos_mah.ipynb
+   diffhalos_mc_vs_qmc.ipynb
 
 See :ref:`Citation Information <cite_info>` for how to acknowledge diffhalos.
diff --git a/docs/source/tutorials.rst b/docs/source/tutorials.rst
deleted file mode 100644
index 5aed4b2..0000000
--- a/docs/source/tutorials.rst
+++ /dev/null
@@ -1,5 +0,0 @@
-.. _tutorials:
-
-Tutorials
-===============
-This section of the docs provides some worked examples of using diffhalos in different applications...
\ No newline at end of file

From 66b13b4526a94c05f1c9ad9620b5882819b41dbc Mon Sep 17 00:00:00 2001
From: Georgios Zacharegkas <gzacharegkas@gmail.com>
Date: Sat, 11 Jul 2026 08:24:07 -0500
Subject: [PATCH 9/9] small text updates to the docs

---
 docs/source/diffhalos_mc_vs_qmc.ipynb | 6 ++----
 docs/source/index.rst                 | 5 ++---
 2 files changed, 4 insertions(+), 7 deletions(-)

diff --git a/docs/source/diffhalos_mc_vs_qmc.ipynb b/docs/source/diffhalos_mc_vs_qmc.ipynb
index 2b0c6a2..1be8cc1 100644
--- a/docs/source/diffhalos_mc_vs_qmc.ipynb
+++ b/docs/source/diffhalos_mc_vs_qmc.ipynb
@@ -7,9 +7,7 @@
    "source": [
     "# Differences between Monte Carlo and quasi-Monte Carlo\n",
     "\n",
-    "This notebook presents the main differences between the two modes of lightcone generation with Diffhalos: the Monte Carlo (MC) and the quasi-Monte Carlo (QMC), also called _weighted_, lightcone.\n",
-    "\n",
-    "The basic difference in the philosophy behind the two versions in the case of host halos is the following: while in the MC version we randomly distribute each halo in redshift $z$ and mass $M_{\\rm halo}$ (using the halo mass function; HMF), in the case of QMC we contruct a _uniform_ grid of $(z,m_{\\rm halo})$, where $m_{\\rm halo} \\equiv \\log_{10} M_{\\rm halo}$. Then, in the QMC lightcone each point is assigned a _weight_ that represents the expected abundance of halos, according to an underlying HMF model. \n",
+    "This notebook presents the main differences between the two modes of lightcone generation with Diffhalos: the Monte Carlo (MC) and the quasi-Monte Carlo (QMC), also called _weighted_, lightcone. While in the MC version we randomly distribute each halo in redshift $z$ and mass $M_{\\rm halo}$ (using the halo mass function; HMF), in the case of QMC we contruct a _uniform_ grid of $(z,m_{\\rm halo})$, where $m_{\\rm halo} \\equiv \\log_{10} M_{\\rm halo}$. Then, in the QMC lightcone each point is assigned a _weight_ that represents the expected abundance of halos, according to an underlying HMF model. \n",
     "\n",
     "For subhalos the difference between the two versions of the lightcone is analogous to the host halo generation above. In the MC lightcone, for each _individual_ (host) halo, generated according to the above, we draw samples from the conditional subhalo mass fucntion (CSHMF; as a function of the parameter $\\mu \\equiv M_{\\rm sub}/M_{\\rm host}$), and thus populate each of them with the subhalos they host. In the QMC version, at each point in the preconstructed grid of redshift and mass for hosts from above, we first create a grid of $\\mu$ values (typically $3-5$) between a minimum and a maximum, and then we assign a _weight_ to each point that represents the abundance of subhalos."
    ]
@@ -218,7 +216,7 @@
    "id": "3995ed13-f1f2-4b27-83a1-6f64019a4f23",
    "metadata": {},
    "source": [
-    "When using the MC lightcone, there is not need to weight since the generated halos in it already have the proper abundance. In addition, in the limit where the number of generated halos is large enough, the MC and QMC lightcones should look almost identical. Below we show that this is roughly the case, even for the relatively low number of halos in the lightcones in this notebook. For a more careful comparison, see `diffhalos_quickstart.ipynb`."
+    "When using the MC lightcone, there is no need to weight since the generated halos in it already have the proper abundance. In addition, in the limit where the number of generated halos is large enough, the MC and QMC lightcones should look almost identical. Below we show that this is roughly the case, even for the relatively low number of halos in the lightcones in this notebook. For a more careful comparison, see `diffhalos_quickstart.ipynb`."
    ]
   },
   {
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 5fd5901..3930756 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -1,11 +1,10 @@
 Diffhalos: Differentiable Populations of Dark Matter Halos
 ==========================================================
 
-Diffhalos is a generative model of cosmological volumes of dark matter halos, subhalos,
-and their merger trees, and these docs show you how to use it.
+Diffhalos is a generative model of cosmological volumes of dark matter halos, subhalos, and their merger trees.
 The Diffhalos library is written in `JAX <https://jax.readthedocs.io/>`__,
 and makes differentiable predictions for halo populations in a lightcone, or at a single redshift.
-The code is publicly available on
+The code is publicly available on 
 `GitHub <https://github.com/ArgonneCPAC/diffhalos/>`__.
 
 User Guide