{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook is part of the SelfiSys project and is licensed under the GNU General Public License v3.0 or later (GPL-3.0-or-later). See the accompanying LICENSE file or visit [GNU GPL](https://www.gnu.org/licenses/gpl-3.0.en.html) for details.\n",
    "\n",
    "Tristan Hoellinger<br/>\n",
    "Institut d'Astrophysique de Paris</br>\n",
    "tristan.hoellinger@iap.fr\n",
    "\n",
    "# Compute the radial selection functions\n",
    "\n",
    "Define custom log-normal radial selection functions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up the environment and parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from selfisys.global_parameters import *\n",
    "\n",
    "wd = Path(OUTPUT_PATH) / \"expl_notebooks/selection/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "from pathlib import Path\n",
    "\n",
    "from scipy.integrate import simpson\n",
    "from scipy.interpolate import UnivariateSpline\n",
    "from astropy.cosmology import Planck18 as cosmo\n",
    "from classy import Class\n",
    "\n",
    "from selfisys.utils.tools import cosmo_vector_to_class_dict\n",
    "from selfisys.utils.plot_examples import *\n",
    "from selfisys.utils.plot_utils import *\n",
    "from selfisys.selection_functions import LognormalSelection\n",
    "\n",
    "Path(wd).mkdir(parents=True, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "L = 3.6\n",
    "Lcorner = L * np.sqrt(3)\n",
    "\n",
    "LogNorm = LognormalSelection()\n",
    "one_lognormal = LogNorm.one_lognormal\n",
    "\n",
    "def multiple_lognormal(x, ss, ll, rr):\n",
    "    return LogNorm.multiple_lognormal(x, None, ss, ll, rr)\n",
    "\n",
    "xx = np.linspace(1e-5, Lcorner, 1000)\n",
    "zz = np.linspace(0, 3.1, 1200)  # Redshift range"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Central parameters for the three populations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "means_com = np.array([0.6410, 1.9210, 2.9185])  # Arbitrary values for our mock populations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Redshift vs comoving distance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 248432 characters]",
      "text/plain": [
       "<Figure size 3600x1560 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 244872 characters]",
      "text/plain": [
       "<Figure size 3600x1500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_comoving_distance_redshift(zz, cosmo, means_com, L, Lcorner, wd)\n",
    "spline = redshift_distance_conversion(zz, cosmo, means_com, L, Lcorner, xx, wd)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Linear bias model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### For tiny $z$ (local, very late Universe)\n",
    "\n",
    "We rely on [\\[howlett2015clustering\\]](https://doi.org/10.48550/arXiv.1409.3238):\n",
    "\n",
    "> We measure Redshift-Space Distortions (RSD) in the two-point correlation function of a\n",
    "sample of 63,163 spectroscopically identified galaxies with z < 0.2, an epoch where there\n",
    "are currently only limited measurements, from the Sloan Digital Sky Survey (SDSS) Data\n",
    "Release 7 Main Galaxy Sample. Our sample, which we denote MGS, covers 6,813 deg2 with\n",
    "an effective redshift zeff = 0.15 and is described in our companion paper (Paper I), which\n",
    "concentrates on BAO measurements.\n",
    "\n",
    "Fitting their data with a $\\chi^2$ distribution, they measure \n",
    "$$\n",
    "f\\sigma_8(z=0.15)=0.43^{+0.15}_{-0.14}\n",
    "$$\n",
    "and\n",
    "$$\n",
    "b\\sigma_8(z=0.15)=1.20\\pm 0.15\n",
    "$$\n",
    "\n",
    "Further assuming $f=\\Omega_m^{0.55}$, we obtain:\n",
    "$$\n",
    "\\sigma_8(z=0.15)=0.43^{+0.15}_{-0.14}\\times\\Omega_m^{-0.55}=0.43^{+0.15}_{-0.14}\\times 0.31^{-0.55}\n",
    "$$\n",
    "$$\n",
    "\\sigma_8(z=0.15)\\simeq 0.82 \\implies b(z=0.15)\\simeq 1.47\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### For LRG bins\n",
    "\n",
    "We rely on [\\[gil2015power\\]](https://doi.org/10.48550/arXiv.1407.5668):\n",
    "\n",
    "> We analyse the anisotropic clustering of the Baryon Oscillation Spectroscopic Survey (BOSS) CMASS Data Release 11 sample, which consists of 690,827 galaxies in the redshift range 0.43 < z < 0.70 and has a sky coverage of 8498 deg2 corresponding to an effective volume of ∼ 6 Gpc3. We fit the Fourier space statistics, the power spectrum and bispectrum monopoles to measure the linear and quadratic bias parameters, b1 and b2, for a non-linear non-local bias model, the growth of structure parameter f and the amplitude of dark matter density fluctuations parametrised by σ8.\n",
    "\n",
    "They measure $b_{z_{\\textrm{eff}}}^{1.40}\\sigma_8(z_{\\textrm{eff}})=1.672\\pm 0.060$ and $\\sigma_8(z_{\\textrm{eff}})=0.597$, $z_{\\textrm{eff}}=0.57$, so we get $b_{z_{\\textrm{eff}}}^{1.40}=2.80\\pm 0.10$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2.0864048990680044"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "2.80**(1/1.40)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assuming $b(z)D(z)$ to be constant in the redshift range of interest, we obtain $b(z) \\simeq \\dfrac{b_{z_{eff}}\\times D(z_{\\text{eff}})}{D(z)} \\simeq \\dfrac{2.086\\times D(z_{\\text{eff}})}{D(z)}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute the linear growth factor with CLASS\n",
    "cosmo_dict = cosmo_vector_to_class_dict(planck_mean)\n",
    "Cosmology = Class()\n",
    "Cosmology.set(cosmo_dict)\n",
    "Cosmology.compute()\n",
    "Dz = Cosmology.get_background()[\"gr.fac. D\"]\n",
    "redshifts = Cosmology.get_background()[\"z\"]\n",
    "Cosmology.struct_cleanup()\n",
    "Cosmology.empty()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 178116 characters]",
      "text/plain": [
       "<Figure size 3600x1200 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure(figsize=(12, 4))\n",
    "plt.loglog(redshifts, Dz)\n",
    "plt.xlabel(r\"$z$\")\n",
    "plt.ylabel(\"normalised growth factor\")\n",
    "plt.grid(which=\"both\", axis=\"both\", linestyle=\"-\", linewidth=0.3, color=\"gray\", alpha=0.5)\n",
    "plt.title(r\"Growth factor using \\textsc{CLASS}\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.15   0.4925 0.8182]\n"
     ]
    }
   ],
   "source": [
    "z_means = np.round(spline(means_com), 4)\n",
    "print(z_means)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "growth_factor_z_eff = 0.7433515577264593\n",
      "b_zeff = 2.07\n",
      "numerator b(z): 1.5387377244937708\n",
      "b(0.4925)=1.99\n",
      "b(0.8182)=2.32\n"
     ]
    }
   ],
   "source": [
    "np.random.seed(1100)\n",
    "z_eff = 0.57\n",
    "z_eff_index = np.argmin(np.abs(redshifts - z_eff))\n",
    "growth_factor_z_eff = Dz[z_eff_index]\n",
    "print(\"growth_factor_z_eff =\", growth_factor_z_eff)\n",
    "b_zeff = np.round(2.80 ** (1 / 1.40) + 0.02 * np.random.uniform(-1, 1), 2)\n",
    "print(\"b_zeff =\", b_zeff)\n",
    "print(\"numerator b(z): {}\".format(b_zeff * growth_factor_z_eff))\n",
    "for z_mean in z_means[1:]:\n",
    "    bias = b_zeff * growth_factor_z_eff / Dz[np.argmin(np.abs(redshifts - z_mean))]\n",
    "    bias += 0.02 * np.random.uniform(-1, 1)\n",
    "    print(\n",
    "        \"b({})={:.2f}\".format(\n",
    "            z_mean,\n",
    "            bias,\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "b(z) \\simeq \\frac{2.086\\times D(z_{\\text{eff}})}{D(z)} \\simeq \\frac{1.54}{D(z)}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Toy model: lognormal distributions with respect to comoving distance"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Utilities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def integrate(x, res, res_mis=None):\n",
    "    \"\"\"\n",
    "    Compute the numerical integral of given data arrays using Simpson's\n",
    "    rule.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    x : array-like\n",
    "        The x-axis values for integration.\n",
    "    res : list of array-like\n",
    "        A list of y-axis values to integrate over the x-axis.\n",
    "    res_mis : list of array-like, optional\n",
    "        A list of y-axis values to integrate over the x-axis, for the\n",
    "        misspecified model. Default is None.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    tuple\n",
    "        - intes: list of float\n",
    "            The integrated values for res.\n",
    "        - intes_mis: list of float or None\n",
    "            The integrated values for res_mis if provided,\n",
    "            None otherwise.\n",
    "    \"\"\"\n",
    "\n",
    "    intes = []\n",
    "    for i, y in enumerate(res):\n",
    "        integral = simpson(y, x=x)\n",
    "        intes.append(integral)\n",
    "        print(f\"Integral of R_{i+1}(r): {integral:.4f}\")\n",
    "    total_res = np.sum(intes)\n",
    "    print(f\"Total: {total_res:.4f}\")\n",
    "\n",
    "    if res_mis is not None:\n",
    "        intes_mis = []\n",
    "        for i, y_mis in enumerate(res_mis):\n",
    "            integral_mis = simpson(y_mis, x=x)\n",
    "            intes_mis.append(integral_mis)\n",
    "            print(f\"Integral of R_{i+1}(r): {integral_mis:.4f}\")\n",
    "        total_mis = np.sum(intes_mis)\n",
    "        print(f\"Total: {total_mis:.4f}\")\n",
    "    else:\n",
    "        intes_mis = None\n",
    "\n",
    "    return intes, intes_mis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Means (comoving distance) of the 3 populations =  [0.641  1.921  2.9185]\n",
      "Means (redshift) of the 3 populations =  [0.15   0.4925 0.8182]\n"
     ]
    }
   ],
   "source": [
    "print(\"Means (comoving distance) of the 3 populations = \", means_com)\n",
    "print(\"Means (redshift) of the 3 populations = \", z_means)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Standard deviations for the three populations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Standard deviations for the 3 populations: [0.115   0.14925 0.18182]\n"
     ]
    }
   ],
   "source": [
    "cste = 0.1\n",
    "std1 = cste * (1+z_means[0])\n",
    "std2 = cste * (1+z_means[1])\n",
    "std3 = cste * (1+z_means[2])\n",
    "stds = np.array([std1, std2, std3])\n",
    "print(\"Standard deviations for the 3 populations:\", stds)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Corresponding lognormal distribution parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Modes (z) of the 3 populations: [0.142202 0.486072 0.811843]\n"
     ]
    }
   ],
   "source": [
    "mus = np.log(means_com**2 / np.sqrt(means_com**2 + stds**2))\n",
    "sig2s = np.log(1 + stds**2 / means_com**2)\n",
    "modes = np.exp(mus - sig2s)\n",
    "d_np = (cosmo.comoving_distance(zz) / 1e3).value\n",
    "z_modes = np.array([zz[np.argmin(np.abs(d_np - mode))] for mode in modes])\n",
    "print(\"Modes (z) of the 3 populations:\", np.round(z_modes, 6))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.11422018 0.14860717 0.18118432]\n"
     ]
    }
   ],
   "source": [
    "print(cste * (1 + z_modes))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Lognormal distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "params: [array([0.115  , 0.14925, 0.18182]), array([0.641 , 1.921 , 2.9185]), [1, 1, 1]]\n",
      "[20:43:31|\u001b[1;36mINFO      \u001b[00m]|\u001b[38;5;147m(selfisys.utils.plot_utils)\u001b[00m Plotting selection functions...\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 195668 characters]",
      "text/plain": [
       "<Figure size 3000x1500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Integral of R_1(r): 0.2778\n",
      "Integral of R_2(r): 0.2778\n",
      "Integral of R_3(r): 0.2778\n",
      "Total: 0.8333\n"
     ]
    }
   ],
   "source": [
    "max1 = np.max(one_lognormal(xx, std1, means_com[0]))\n",
    "max2 = np.max(one_lognormal(xx, std2, means_com[1]))\n",
    "max3 = np.max(one_lognormal(xx, std3, means_com[2]))\n",
    "rescale = np.max([max1, max2, max3])\n",
    "\n",
    "params = [stds, means_com, [1, 1, 1]]\n",
    "# Alternatively, one can rescale the selection functions using:\n",
    "# params = [stds, means_com, [1/rescale, 1/rescale, 1/rescale]]\n",
    "print(\"params:\", params)\n",
    "res = multiple_lognormal(xx, params[0], params[1], params[2])\n",
    "plot_selection_functions(xx, res, None, params, L, np.sqrt(3) * L)\n",
    "intes, _ = integrate(xx / L, res)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Actual selection functions\n",
    "\n",
    "To account for lightcone effects in the data model, we use three snapshots of the final overdensity fields, each corresponding to a central redshift associated with a specific selection function. For consistency, the selection functions are defined as lognormal distributions with respect to redshift, $z$, rather than comoving distance, $r$.\n",
    "\n",
    "We construct the mock selection functions as follows:\n",
    "1.\tDefine three rescaled lognormal distributions in redshift space, each characterised by a given mean (in $z$) and variance (in $z$).\n",
    "2.\tDefine three misspecified lognormal distributions by ensuring that the sum of their integrals matches that of the well-specified functions.\n",
    "3.\tConvert these lognormal distributions to functions of comoving distance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "d_np = cosmo.comoving_distance(zz).value / 1e3\n",
    "z_corner = zz[np.argmin(np.abs(d_np - Lcorner))]\n",
    "z_L = zz[np.argmin(np.abs(d_np - L))]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "zs = np.linspace(0.0001, z_corner, 1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Means (in redshift) of the 3 populations =  [0.15   0.4925 0.8182]\n",
      "standard deviations for the 3 populations =  [0.115   0.14925 0.18182]\n",
      "Modes (in redshift) of the 3 populations =  [0.075  0.4317 0.7611]\n"
     ]
    }
   ],
   "source": [
    "cste_z = 0.1\n",
    "std1_z = cste_z * (1 + z_means[0])\n",
    "std2_z = cste_z * (1 + z_means[1])\n",
    "std3_z = cste_z * (1 + z_means[2])\n",
    "stds_z = np.array([std1_z, std2_z, std3_z])\n",
    "print(\"Means (in redshift) of the 3 populations = \", z_means)\n",
    "print(\"standard deviations for the 3 populations = \", stds_z)\n",
    "\n",
    "mus_z = np.log(z_means**2 / np.sqrt(z_means**2 + stds_z**2))\n",
    "sig2s_z = np.log(1 + stds_z**2 / z_means**2)\n",
    "modes_z = np.exp(mus_z - sig2s_z)\n",
    "print(\"Modes (in redshift) of the 3 populations = \", np.round(modes_z, 4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Lognormals in redshift"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[20:43:32|\u001b[1;36mINFO      \u001b[00m]|\u001b[38;5;147m(selfisys.utils.plot_utils)\u001b[00m Plotting selection functions...\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 206012 characters]",
      "text/plain": [
       "<Figure size 3000x1500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "max1 = np.max(one_lognormal(zs, std1_z, z_means[0]))\n",
    "max2 = np.max(one_lognormal(zs, std2_z, z_means[1]))\n",
    "max3 = np.max(one_lognormal(zs, std3_z, z_means[2]))\n",
    "rescale = np.max([max1, max2, max3])\n",
    "\n",
    "params = [stds_z, z_means, [1 / rescale, 1 / rescale, 1 / rescale]]\n",
    "res = multiple_lognormal(zs, params[0], params[1], params[2])\n",
    "plot_selection_functions(zs, res, None, params, z_L, z_corner, axis=\"redshift\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Selection functions, defined as functions of redshifts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Integral of R_1(r): 0.9231\n",
      "Integral of R_2(r): 0.9231\n",
      "Integral of R_3(r): 0.9231\n",
      "Total: 2.7693\n",
      "> Well-specified parameters: [array([0.115  , 0.14925, 0.18182]), array([0.15  , 0.4925, 0.8182]), [np.float64(1.0), np.float64(1.0), np.float64(1.0)]]\n",
      "> Mis-specified parameters: [array([0.115  , 0.14925, 0.18182]), array([0.15  , 0.4925, 0.8182]), [np.float64(0.9), np.float64(1.199999810418996), np.float64(0.9)]]\n",
      "[20:43:32|\u001b[1;36mINFO      \u001b[00m]|\u001b[38;5;147m(selfisys.utils.plot_utils)\u001b[00m Plotting selection functions...\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 269568 characters]",
      "text/plain": [
       "<Figure size 3000x1500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Integral of R_1(r): 0.9231\n",
      "Integral of R_2(r): 0.9231\n",
      "Integral of R_3(r): 0.9231\n",
      "Total: 2.7693\n",
      "Integral of R_1(r): 0.8308\n",
      "Integral of R_2(r): 1.1077\n",
      "Integral of R_3(r): 0.8308\n",
      "Total: 2.7693\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 310364 characters]",
      "text/plain": [
       "<Figure size 3000x1500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Parameters for Model A\n",
    "params = [stds_z, z_means, [1.0 / max1, 1.0 / max1, 1.0 / max1]]\n",
    "res = multiple_lognormal(zs, params[0], params[1], params[2])\n",
    "\n",
    "# Integrate and print the results\n",
    "intes, _ = integrate(zs / z_L, res)\n",
    "\n",
    "# Define population transfer fractions\n",
    "diff_pop1 = 0.1  # Fraction of population 1 to transfer to population 2\n",
    "diff_pop3 = 0.1  # Fraction of population 3 to transfer to population 2\n",
    "\n",
    "# Compute total loss from populations 1 and 3 and adjust population 2\n",
    "tot_lost = diff_pop1 * intes[0] + diff_pop3 * intes[2]  # Total loss\n",
    "diff_pop2 = tot_lost / intes[1]  # Compensation for population 2\n",
    "\n",
    "# Parameters for Model B\n",
    "params_mis = [\n",
    "    stds_z,\n",
    "    z_means,\n",
    "    [\n",
    "        (1 - diff_pop1) / max1,  # Adjusted amplitude for population 1\n",
    "        (1 + diff_pop2) / max1,  # Adjusted amplitude for population 2\n",
    "        (1 - diff_pop3) / max1,  # Adjusted amplitude for population 3\n",
    "    ],\n",
    "]\n",
    "res_mis = multiple_lognormal(zs, params_mis[0], params_mis[1], params_mis[2])\n",
    "print(\"> Well-specified parameters:\", params)\n",
    "print(\"> Mis-specified parameters:\", params_mis)\n",
    "\n",
    "plot_selection_functions(\n",
    "    zs,\n",
    "    res,\n",
    "    res_mis,\n",
    "    params,\n",
    "    z_L,\n",
    "    z_corner,\n",
    "    axis=\"redshift\",\n",
    ")\n",
    "\n",
    "# Integrate and print the results for comparison\n",
    "intes, intes_mis = integrate(zs / z_L, res, res_mis)\n",
    "\n",
    "# Map back to comoving distance\n",
    "xx_of_zs = cosmo.comoving_distance(zs) / 1e3\n",
    "\n",
    "plot_selection_functions_def_in_z(\n",
    "    xx_of_zs=xx_of_zs,\n",
    "    res=res,\n",
    "    res_mis=res_mis,\n",
    "    z_means=z_means,\n",
    "    cosmo=cosmo,\n",
    "    L=L,\n",
    "    stds_z=stds_z,\n",
    "    wd=wd,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that these distributions are not lognormal with respect to comoving distance."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Lognormal distributions in $z$ (automated)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[20:43:33|\u001b[1;36mINFO      \u001b[00m]|\u001b[38;5;147m(selfisys.utils.plot_utils)\u001b[00m Plotting selection functions...\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 226924 characters]",
      "text/plain": [
       "<Figure size 3000x1500 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "cste_z = 0.1\n",
    "stds_z = np.array([cste_z * (1 + z) for z in z_means])\n",
    "params = [stds_z, z_means, [1, 1, 1]]\n",
    "\n",
    "zz, res = LogNorm.lognormals_z_to_x(xx, None, params, spline)\n",
    "\n",
    "plot_selection_functions(\n",
    "    xx,\n",
    "    res,\n",
    "    None,\n",
    "    params,\n",
    "    L,\n",
    "    np.sqrt(3) * L,\n",
    "    zz=zz,\n",
    "    zcorner=z_corner,\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "simbel",
   "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.13.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
