csiborgtools/notebooks/flow/flow_los.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Density & velocity fields alond a LOS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 128,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "# Copyright (C) 2024 Richard Stiskalek\n",
    "# This program is free software; you can redistribute it and/or modify it\n",
    "# under the terms of the GNU General Public License as published by the\n",
    "# Free Software Foundation; either version 3 of the License, or (at your\n",
    "# option) any later version.\n",
    "#\n",
    "# This program is distributed in the hope that it will be useful, but\n",
    "# WITHOUT ANY WARRANTY; without even the implied warranty of\n",
    "# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General\n",
    "# Public License for more details.\n",
    "#\n",
    "# You should have received a copy of the GNU General Public License along\n",
    "# with this program; if not, write to the Free Software Foundation, Inc.,\n",
    "# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import jax\n",
    "from jax import numpy as jnp\n",
    "from numpyro.infer import MCMC, NUTS, init_to_median\n",
    "\n",
    "import csiborgtools\n",
    "\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline\n",
    "\n",
    "paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LOS density & radial velocity plots "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 129,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-07-17 10:04:21.726937:   reading the catalogue,\n",
      "2024-07-17 10:04:21.744932:   reading the interpolated field,\n",
      "2024-07-17 10:04:21.770393:   calculating the radial velocity.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/mnt/users/rstiskalek/csiborgtools/csiborgtools/flow/flow_model.py:91: UserWarning: The number of radial steps is even. Skipping the first step at 0.0 because Simpson's rule requires an odd number of steps.\n",
      "  warn(f\"The number of radial steps is even. Skipping the first \"\n"
     ]
    }
   ],
   "source": [
    "fpath = \"/mnt/extraspace/rstiskalek/catalogs/PV_compilation.hdf5\"\n",
    "\n",
    "loader_carrick = csiborgtools.flow.DataLoader(\"Carrick2015\", [0], \"Pantheon+\", fpath, paths, ksmooth=0, )\n",
    "# loaders_csiborg2X = [csiborgtools.flow.DataLoader(\"csiborg2X\", i, \"LOSS\", fpath, paths, ksmooth=1, verbose=False) for i in range(20)]\n",
    "# loaders_csiborg2 = [csiborgtools.flow.DataLoader(\"csiborg2_main\", i, \"LOSS\", fpath, paths, ksmooth=1, verbose=False) for i in range(20)]\n",
    "\n",
    "# loader_CF4 = csiborgtools.flow.DataLoader(\"CF4gp\", [0], \"LOSS\", fpath, paths, ksmooth=0, )\n",
    "# loader_lilow = csiborgtools.flow.DataLoader(\"Lilow2024\", [0], \"LOSS\", fpath, paths, ksmooth=0, )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 130,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAArIAAAKwCAYAAABksZyLAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/H5lhTAAAACXBIWXMAAA9hAAAPYQGoP6dpAAC4D0lEQVR4nOzdd3QUZRsF8LvpCYRAKIEAISA19BYITTpSBWmKSFHBElQIgmIB/SxYABUJRZAiShEUKdJDR3qRJqFIxyTUhCSQtvv98TDZhBSyye7O7O79nZMzk93ZmScBws27z7yvzmAwGEBEREREZGOc1C6AiIiIiCg/GGSJiIiIyCYxyBIRERGRTWKQJSIiIiKbxCBLRERERDaJQZaIiIiIbBKDLBERERHZJAZZIiIiIrJJLmoXQEZ6vR7Xr1+Ht7c3dDqd2uUQERERWZ3BYMC9e/fg7+8PJ6fcx1wZZDXk+vXrKF++vNplEBEREanuypUrKFeuXK7HMMhqiLe3NwD5gytSpIjK1RARERFZX1xcHMqXL5+ei3LDIGsha9aswejRo6HX6/HOO+/g5ZdffuxrlHaCIkWKMMgSERGRQ8tLmyWDrAWkpqYiLCwMW7duhY+PDxo2bIhevXqhePHiapdGREREZDc4a4EF7N+/HzVr1kTZsmVRuHBhdO7cGRs3blS7LCIiIiK7YlNBdsaMGahTp076W+8hISFYt26dWa+xY8cOdO/eHf7+/tDpdPjjjz+yPS48PByBgYHw8PBAkyZNsH///vTnrl+/jrJly6Z/XrZsWVy7ds2sdRIRERE5OptqLShXrhy++OILVKlSBQaDAQsWLMDTTz+NI0eOoGbNmlmO3717N4KDg+Hq6prp8VOnTqF48eLw8/PL8pqEhATUrVsXL774Ip555pls61i6dCnCwsIwc+ZMNGnSBN9++y06deqEyMhIlCpVyjxfLBERkYYZDAakpqYiLS1N7VLIxjg7O8PFxcUsU43aVJDt3r17ps8/++wzzJgxA3v37s0SZPV6PUJDQ1GlShUsWbIEzs7OAIDIyEi0bdsWYWFhGDt2bJZrdO7cGZ07d861jilTpmDYsGEYOnQoAGDmzJn4888/MXfuXLz77rvw9/fPNAJ77do1BAcH53i+8PBwhIeH84cBERHZhOTkZPz3339ITExUuxSyUV5eXihTpgzc3NwKdB6dwWAwmKkmq0pLS8OyZcswePBgHDlyBEFBQVmOuX79Olq1aoUmTZpg4cKFuHDhAlq1aoXu3btj5syZj72GTqfDihUr0LNnz/THkpOT4eXlheXLl2d6fPDgwbh79y5WrlyJ1NRU1KhRA9u2bUu/2euvv/567M1ecXFx8PHxQWxsLGctICIiTdLr9Th79iycnZ1RsmRJuLm5cREfyjODwYDk5GTcuHEDaWlpqFKlSpZFD0zJQzY1IgsAx48fR0hICB48eIDChQtjxYoV2YZYAPD398eWLVvQsmVLDBgwAHv27EH79u0xY8aMfF//5s2bSEtLy9KW4Ofnh9OnTwMAXFxcMHnyZLRp0wZ6vR5jx47ljAVERGQXkpOTodfrUb58eXh5ealdDtkgT09PuLq64tKlS0hOToaHh0e+z2VzQbZatWo4evQoYmNjsXz5cgwePBjbt2/PMcwGBARg4cKFePLJJ1GpUiX8+OOPVvnNsUePHujRo4fFr0NERKSGxy0dSpQbc/39sbm/hW5ubqhcuTIaNmyIiRMnom7duvjuu+9yPD46OhrDhw9H9+7dkZiYiFGjRhXo+iVKlICzszOio6OzXKd06dIFOjcRERER5Z3NBdlH6fV6JCUlZfvczZs30a5dO9SoUQO///47IiIisHTpUrz99tv5vp6bmxsaNmyIiIiITDVEREQgJCQk3+clG5CcDPz5J7BsGRATo3Y1REREDs+mWgvGjRuHzp07IyAgAPfu3cOiRYuwbds2bNiwIcuxer0enTt3RoUKFbB06VK4uLggKCgImzZtQtu2bVG2bNlsR2fj4+Nx7ty59M8vXLiAo0ePwtfXFwEBAQCAsLAwDB48GI0aNUJwcDC+/fZbJCQkpM9iQHbo+nWgbVsgMlI+9/QEpk0DXnxR3bqIiEizLl68iIoVK+LIkSOoV6/eY48fMmQI7t69m+Mc9pSVTY3IxsTEYNCgQahWrRratWuHAwcOYMOGDejQoUOWY52cnPD555/jt99+yzS1Q926dbF582b07ds322scPHgQ9evXR/369QFIaK1fvz7Gjx+ffkz//v0xadIkjB8/HvXq1cPRo0exfv36bOelJTuQlgb06CEhtmhRoFo14P59YPhwYMcOtasjIiITREVF4Y033kClSpXg7u6O8uXLo3v37pneaTWX8uXL47///kOtWrXMfu7HSUlJwTvvvIPatWujUKFC8Pf3x6BBg3D9+vVMx92+fRvPP/88ihQpgqJFi+Kll15CfHx8+vMPHjzAkCFDULt2bbi4uGSasUmxbds26HS6LB9RUVGW/jJta0T2xx9/NOn47AIugPSQmp3WrVsjLzOSjRgxAiNGjDCpHrJRCxYAhw4BPj7A4cNAYCAwcCCwaBEwZgywdy/AqWeIiDTv4sWLaN68OYoWLYqvv/4atWvXRkpKCjZs2IDQ0ND02YdMkZaWBp1Ol+XmpeTkZLi5uVn0/pmPPvoIFy9exPz587M8l5iYiMOHD+PDDz9E3bp1cefOHbz11lvo0aMHDh48mH7c888/j//++w+bNm1CSkoKhg4diuHDh2PRokXpX5+npyfefPNN/Pbbb7nWExkZmWm6LGssEmVTI7JEVhcfD7z/vuyPHw9UrCih9ZtvAA8PYP9+YOdOdWskIlKbwQAkJFj/w8Sp8F9//XXodDrs378fvXv3RtWqVVGzZk2EhYVh7969AGTRI2UUs3z58nj99dczjVDOnz8fRYsWxapVqxAUFAR3d3dcvnwZgYGB+OSTTzBo0CAUKVIEw4cPx8WLF6HT6XD06NH01588eRLdunVDkSJF4O3tjZYtW+L8+fPZ1nvgwAGULFkSX375pcl/JD4+Pti0aRP69euHatWqoWnTppg2bRoOHTqEy5cvAwD++ecfrF+/HnPmzEGTJk3QokULfP/991iyZEn6yG2hQoUwY8YMDBs27LGhvFSpUihdunT6hzVmtmCQJcrN5MlAVBTwxBNAaKjx8VKlgBdekP25c9WpjYhIKxITgcKFrf9hwspit2/fxvr16xEaGopChQpleb5o0aIApDVx6tSpOHnyJBYsWIAtW7ZkWQk0MTERX375JebMmYOTJ0+mjzxOmjQJdevWxZEjR/Dhhx9muca1a9fQqlUruLu7Y8uWLTh06BBefPFFpKamZjl2y5Yt6NChAz777DO88847ef46cxMbGwudTpf+te7ZswdFixZFo0aN0o9p3749nJycsG/fPpPPX69ePZQpUwYdOnTA7t27zVLz49hUawGRVSUlAeHhsv/ZZ4C7e+bnX3gBmD0bWLECmDUr6/NERKQZ586dg8FgQPXq1XM9buTIken7gYGB+PTTT/Hqq69i+vTp6Y+npKRg+vTpqFu3bqbXtm3bFqNHj07//OLFi5meDw8Ph4+PD5YsWQJXV1cAQNWqVbPUsGLFCgwaNAhz5sxB//798/ol5urBgwd455138Nxzz6W//R8VFZXl7X8XFxf4+vqa1N9apkwZzJw5E40aNUJSUhLmzJmD1q1bY9++fWjQoIFZ6s8JgyxRTpYtA27cAMqVA3r3zvp88+ZA2bLAtWvAxo1A9+7Wr5GISAu8vKQVS43r5lFe7n8BgM2bN2PixIk4ffo04uLikJqaigcPHiAxMTF9JTM3NzfUqVMny2szjmxm5+jRo2jZsmV6iM3Ovn37sGbNGixfvjzLjVU7d+5E586d0z9PTk6GwWDA8uXL0x+bNWsWnn/++UyvS0lJQb9+/WAwGAq0umlOqlWrhmrVqqV/3qxZM5w/fx7ffPMNFi5caPbrZcQgS5QTZTT2lVcAl2z+qTg5Ad26yWjs1q0MskTkuHQ6IJu367WkSpUq0Ol0ud7QdfHiRXTr1g2vvfY
      "text/plain": [
       "<Figure size 700x700 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAArIAAAKxCAYAAACv7U8uAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/H5lhTAAAACXBIWXMAAA9hAAAPYQGoP6dpAACu9klEQVR4nOzdd3gUVRcG8HcT0iEJJST0Ir13CFIF6SIgxQ6IIhpUBGlKEVH5BMVGBDsWUIoCCkqvSu9SpfdQTUISUne+Pw6zm0ASstmZnS3v73n2mcnu7MxJCJuzd88916QoigIiIiIiIhfjZXQARERERET5wUSWiIiIiFwSE1kiIiIicklMZImIiIjIJTGRJSIiIiKXxESWiIiIiFwSE1kiIiIickkFjA7AlZnNZly8eBGFChWCyWQyOhwiIiIil6coCm7evImSJUvCyyv3MVcmsna4ePEiypQpY3QYRERERG7n3LlzKF26dK7HMJG1Q6FChQDIDzo4ONjgaIiIiIhcX3x8PMqUKWPJs3LDRNYOajlBcHAwE1kiIiIiDeWlbJOTvYiIiIjIJTGRJSIiIiKXxESWiIiIiFwSa2SJiIgoXzIyMpCWlmZ0GORifHx84O3trcm5mMgSERGRTRRFQUxMDGJjY40OhVxUaGgoIiIi7O7Dz0SWiIiIbKImscWLF0dgYCAXBaI8UxQFSUlJuHLlCgCgRIkSdp2PiSwRERHlWUZGhiWJLVq0qNHhkAsKCAgAAFy5cgXFixe3q8yAk72IiIgoz9Sa2MDAQIMjIVem/v7YW2PNRJaIiIhsxnICsodWvz9MZImIiIjIJTGR9XT79wNt2wIlSwIvvQQkJRkdEREREVGeMJH1ZLNmAQ0aAOvXA5cuATNmAO3bA8nJRkdGRETk8k6fPg2TyYS9e/fm6fgBAwagR48eusbkbpjIeqpDh4BXXgEyMoCePYGZM4HChYEtW4ARI4yOjoiISBcxMTF46aWXULFiRfj5+aFMmTJ46KGHsGbNGs2vVaZMGVy6dAm1atXS/Nz3kpaWhtGjR6N27doICgpCyZIl8fTTT+PixYtZjrtx4waeeOIJBAcHIzQ0FIMGDUJCQoLl8eTkZAwYMAC1a9dGgQIFsk20169fD5PJdNctJiZG72+T7bc8kqIAL74IpKYCXbsCv/wCmExA+fJA587AZ58BrVsDffsaHSkREZFmTp8+jfvvvx+hoaGYNm0aateujbS0NKxYsQJRUVE4cuSIzefMyMiAyWSCl1fWscHU1FT4+voiIiJCq/Dv8uabb+L06dOYPXv2XY8lJSVh9+7dGD9+POrWrYv//vsPr7zyCrp3746dO3dajnviiSdw6dIlrFq1CmlpaRg4cCAGDx6MuXPnWr6/gIAAvPzyy/jll19yjefo0aMIDg62fF28eHFtvtFccETWEy1eDGzYAPj7A9HRksQCQKdOwNixsj9ggIzOEhER3YuiAImJjr8pik1hvvjiizCZTNi+fTseeeQRVKlSBTVr1sTw4cOxdetWAMD06dMto5hlypTBiy++mGWEcvbs2QgNDcVvv/2GGjVqwM/PD2fPnkX58uUxefJkPP300wgODsbgwYOzLS04ePAgunXrhuDgYBQqVAgtW7bEiRMnso13x44dCAsLw3vvvWfzP0lISAhWrVqFvn37omrVqmjWrBlmzJiBXbt24ezZswCAw4cPY/ny5fjqq6/QtGlTtGjRAp9++il+/vlny8htUFAQZs6cieeee+6eSXnx4sURERFhud2Z3OuBiaynSUkBRo6U/eHDgXLlsj7+1ltAly7ArVvAY49x8hcREd1bUhJQsKDjbzb8jbpx4waWL1+OqKgoBAUF3fV4aGgoAMDLywuffPIJDh48iO+++w5r167FqFGj7vh2k/Dee+/hq6++wsGDBy0jj++//z7q1q2LPXv2YPz48Xdd48KFC2jVqhX8/Pywdu1a7Nq1C8888wzS09PvOnbt2rV48MEH8c4772D06NF5/j5zExcXB5PJZPlet2zZgtDQUDRq1MhyTPv27eHl5YVt27bZfP569eqhRIkSePDBB/H3339rEvO9sLTA03z/PXDiBBARAYwZc/fjBQoA8+YBNWsCZ84A770HTJrk+DiJiIg0dPz4cSiKgmrVquV63LBhwyz75cuXx9tvv40hQ4bgs88+s9yflpaGzz77DHXr1s3y3AceeAAjMs0zOX36dJbHo6OjERISgp9//hk+Pj4AgCpVqtwVw6JFi/D000/jq6++Qr9+/fL6LeYqOTkZo0ePxmOPPWb5+D8mJuauj/8LFCiAIkWK2FTfWqJECcyaNQuNGjVCSkoKvvrqK7Rp0wbbtm1DgwYNNIk/J0xk8yE6OhrR0dHIyMgwOhTbLVgg22HDgEKFsj+mYEFg6lTg0UeBL78Exo+XBJeIiCg7gYFApo/fHXrdPFLyWIawevVqTJkyBUeOHEF8fDzS09ORnJyMpKQky2pUvr6+qFOnzl3PzTyymZ29e/eiZcuWliQ2O9u2bcPSpUuxcOHCuyZWbdq0CZ07d7Z8nZqaCkVRsHDhQst9n3/+OZ544oksz0tLS0Pfvn2hKApmzpyZa4z5UbVqVVStWtXydfPmzXHixAl8+OGH+OGHHzS/XmbMTvIhKioKUVFRiI+PR0hIiNHh5F1sLLBunez37Jn7sT17AkWLSluuNWuAjh11D4+IiFyUyQRk83G9M6lcuTJMJlOuE7pOnz6Nbt264YUXXsA777yDIkWK4K+//sKgQYOQmppqSWQDAgKyXZkqu5KFzAICAu4Z53333YeiRYvim2++QdeuXbMkvY0aNcpSb/vJJ5/gwoULWWpow8PDs5xPTWLPnDmDtWvXZpmMFRERgStXrmQ5Pj09HTdu3LB7klqTJk3w119/2XWOvGCNrCdZuRJITweqVwey+SgjC19fqZEFgO++0z82IiIiHRUpUgQdO3ZEdHQ0EhMT73o8NjYWu3btgtlsxgcffIBmzZqhSpUqd7WrskedOnWwadMmpKWl5XhMsWLFsHbtWhw/fhx9+/bNcmxAQAAqVapkuRUpUgSFChXKcl+hTJ+2qknssWPHsHr1ahQtWjTLtSIjIy3ft2rt2rUwm81o2rSpXd/r3r17UaJECbvOkRdMZD3Jhg2y7dAhb8f37y/bRYuAuDh9YiIiInIQtSywSZMm+OWXX3Ds2DEcPnwYn3zyCSIjI1GpUiWkpaXh008/xcmTJ/HDDz9g1qxZml1/6NChiI+Px6OPPoqdO3fi2LFj+OGHH3D06NEsxxUvXhxr167FkSNH8Nhjj2U7Gexe0tLS0Lt3b+zcuRNz5sxBRkYGYmJiEBMTg9TUVABA9erV0alTJzz33HPYvn07/v77bwwdOhSPPvooSpYsaTnXoUOHsHfvXty4cQNxcXHYu3dvlpHhjz76CEuWLMHx48dx4MABDBs2DGvXrkVUVFT+flA2YCLrSdQh/pYt83Z8w4YyepucDMyfr19cREREDlCxYkXs3r0bbdu2xYgRI1CrVi08+OCDWLNmDWbOnIm6deti+vTpeO+991CrVi3MmTMHU6ZM0ez6RYsWxdq1a5GQkIDWrVujYcOG+PLLL7OtmY2IiMDatWvxzz//4IknnrB5Xs6FCxfw22+/4fz585ZuAupt8+bNluPmzJmDatWqoV27dujSpQtatGiBL774Isu5unTpgvr16+P333/H+vXrUb9+fdSvX9/yeGpqKkaMGIHatWujdevW2LdvH1avXo127drZ+BOynUnJa/Uz3UWtkY2Li8tSc+KUYmOBIkWk515MDHBHDU2Opk0DRo0CGjcGtm/XNUQiInJ+ycnJOHXqFCpUqAB/f3+jwyEXldvvkS35FUdkPcWWLZLEVqqU9yQWkIURfH2BHTuAPXt0C4+IiIjIVkxkPcXtFUvQvLltzwsLA3r1kv3PP9c2JiIiIiI7MJH1FOpys82a2f7c55+X7Zw5wM2b2sVEREREZAcmsp7AbAbUpeYiI21/fuvWQNWq0ux60SJtYyMiIpfEKTZkD61+f5jIeoIjR4D4eFkBpVYt259vMgF9+8r+779
      "text/plain": [
       "<Figure size 700x700 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# ks = [i for i in range(50)]\n",
    "# ks = [30, 50,]\n",
    "ks = np.random.choice(50, 2, replace=False)\n",
    "\n",
    "for k in ks:\n",
    "    fig, axs = plt.subplots(2, 1, figsize=(7, 7), sharex=True)\n",
    "    fig.subplots_adjust(wspace=0)\n",
    "    cols = plt.rcParams['axes.prop_cycle'].by_key()['color']\n",
    "\n",
    "    # # CSiBORG2\n",
    "    # x = loaders_csiborg2X[0].rdist\n",
    "    # y = np.asarray([loaders_csiborg2[i].los_density[k, :] for i in range(len(loaders_csiborg2X))])\n",
    "    # ylow, ymed, yhigh = np.percentile(y, [16, 50, 84], axis=0)\n",
    "    # axs[0].fill_between(x, ylow, yhigh, color=cols[0], alpha=0.25)\n",
    "    # axs[0].plot(x, ymed, color=cols[0], label=\"CSiBORG2\")\n",
    "\n",
    "    # y = np.asarray([loaders_csiborg2[i].los_radial_velocity[k, :] for i in range(len(loaders_csiborg2X))])\n",
    "    # ylow, ymed, yhigh = np.percentile(y, [16, 50, 84], axis=0)\n",
    "    # axs[1].fill_between(x, ylow, yhigh, color=cols[0], alpha=0.25)\n",
    "    # axs[1].plot(x, ymed, color=cols[0], label=\"CSiBORG2\")\n",
    "\n",
    "    # # CSiBORG2X\n",
    "    # x = loaders_csiborg2X[0].rdist\n",
    "    # y = np.asarray([loaders_csiborg2X[i].los_density[k, :] for i in range(len(loaders_csiborg2X))])\n",
    "    # ylow, ymed, yhigh = np.percentile(y, [16, 50, 84], axis=0)\n",
    "    # axs[0].fill_between(x, ylow, yhigh, color=cols[1], alpha=0.25)\n",
    "    # axs[0].plot(x, ymed, color=cols[1], label=\"CSiBORG2X\")\n",
    "\n",
    "    # y = np.asarray([loaders_csiborg2X[i].los_radial_velocity[k, :] for i in range(len(loaders_csiborg2X))])\n",
    "    # ylow, ymed, yhigh = np.percentile(y, [16, 50, 84], axis=0)\n",
    "    # axs[1].fill_between(x, ylow, yhigh, color=cols[1], alpha=0.25)\n",
    "    # axs[1].plot(x, ymed, color=cols[1], label=\"CSiBORG2X\")\n",
    "\n",
    "    # Plot Carrick+2015\n",
    "    axs[0].plot(loader_carrick.rdist, loader_carrick.los_density[0, k, :], color=\"red\", label=\"Carrick+2015\")\n",
    "    axs[1].plot(loader_carrick.rdist, loader_carrick.los_radial_velocity[0, k, :] * 0.43, color=\"red\")\n",
    "    axs[1].axvline(loader_carrick._rmax[0, k], c=\"black\", linestyle=\"--\")\n",
    "\n",
    "    # # Plot CF4\n",
    "    # c = cols[4]\n",
    "    # axs[0].plot(loader_CF4.rdist, loader_CF4.los_density[0, k, :], color=c, label=\"CF4\")\n",
    "    # axs[1].plot(loader_CF4.rdist, loader_CF4.los_radial_velocity[0, k, :], color=c)\n",
    "\n",
    "    # # Plot Lilow2024\n",
    "    # c = cols[5]\n",
    "    # axs[0].plot(loader_lilow.rdist, loader_lilow.los_density[0, k, :], color=c, label=\"Lilow+2024\")\n",
    "    # axs[1].plot(loader_lilow.rdist, loader_lilow.los_radial_velocity[0, k, :], color=c)\n",
    "    axs[1].axhline(0, color=\"black\", linestyle=\"--\")\n",
    "\n",
    "\n",
    "    axs[1].set_xlabel(r\"$r ~ [\\mathrm{Mpc} / h]$\")\n",
    "    axs[0].set_ylabel(r\"$\\rho_{\\rm LOS} / \\langle \\rho_{\\rm matter} \\rangle$\")\n",
    "    axs[1].set_ylabel(r\"$v_{\\rm LOS} ~ [\\mathrm{km/s}]$\")\n",
    "    axs[0].set_yscale(\"log\")\n",
    "\n",
    "    axs[0].legend(loc=\"upper right\")\n",
    "    # axs[0].set_xlim(0, 200)\n",
    "\n",
    "    fig.tight_layout(w_pad=0, h_pad=0)\n",
    "    fig.savefig(f\"../../plots/LOSS_los_{k}.png\", dpi=500, bbox_inches=\"tight\")\n",
    "\n",
    "    fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test running a model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-06-29 19:40:25.229961:   reading the catalogue,\n",
      "2024-06-29 19:40:25.243502:   reading the interpolated field,\n",
      "2024-06-29 19:40:25.261423:   calculating the radial velocity.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/mnt/users/rstiskalek/csiborgtools/csiborgtools/flow/flow_model.py:91: UserWarning: The number of radial steps is even. Skipping the first step at 0.0 because Simpson's rule requires an odd number of steps.\n",
      "  warn(f\"The number of radial steps is even. Skipping the first \"\n"
     ]
    }
   ],
   "source": [
    "fpath_data = \"/mnt/extraspace/rstiskalek/catalogs/PV_compilation.hdf5\"\n",
    "\n",
    "simname = \"Carrick2015\"\n",
    "catalogue = \"LOSS\"\n",
    "loader = csiborgtools.flow.DataLoader(simname, [0, 0], catalogue, fpath_data, paths, ksmooth=0, )\n",
    "\n",
    "SN_hyperparams = {\"e_mu_mean\": 0.1, \"e_mu_std\": 0.05,\n",
    "                  \"mag_cal_mean\": -18.25, \"mag_cal_std\": 0.5,\n",
    "                  \"alpha_cal_mean\": 0.148, \"alpha_cal_std\": 0.05,\n",
    "                  \"beta_cal_mean\": 3.112, \"beta_cal_std\": 1.0,\n",
    "                  }\n",
    "calibration_hyperparams = {\"Vext_std\": 250,\n",
    "                           \"alpha_mean\": 1.0, \"alpha_std\": 0.5,\n",
    "                           \"beta_mean\": 1.0, \"beta_std\": 0.5,\n",
    "                           \"sigma_v_mean\": 150., \"sigma_v_std\": 100.,\n",
    "                           \"sample_alpha\": True, \"sample_beta\": True,\n",
    "                           }\n",
    "get_model_kwargs = {\"zcmb_max\": 0.05}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Running HMC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Selected 50/50 galaxies.\n"
     ]
    }
   ],
   "source": [
    "model = csiborgtools.flow.get_model(loader, **get_model_kwargs)\n",
    "model_kwargs = {\"distmod_hyperparams\": SN_hyperparams, \"calibration_hyperparams\": calibration_hyperparams,}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "sample: 100%|██████████| 1000/1000 [02:10<00:00,  7.68it/s, 7 steps of size 4.49e-01. acc. prob=0.90]  \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "                 mean       std    median      5.0%     95.0%     n_eff     r_hat\n",
      "    Vext[0]     -3.71     69.92     -3.04   -123.73    103.87    469.72      1.00\n",
      "    Vext[1]    -27.47     95.52    -30.48   -151.20    172.63    308.02      1.00\n",
      "    Vext[2]    -59.27    131.26    -57.79   -273.64    137.55    456.29      1.00\n",
      "      alpha      1.09      0.38      1.10      0.50      1.69    400.05      1.00\n",
      "  alpha_cal      0.13      0.03      0.13      0.09      0.17    558.81      1.00\n",
      "       beta      0.43      0.11      0.44      0.27      0.61    341.86      1.00\n",
      "   beta_cal      3.54      0.18      3.54      3.23      3.81    606.77      1.00\n",
      "       e_mu      0.08      0.03      0.08      0.04      0.12    330.71      1.00\n",
      "    mag_cal    -18.19      0.04    -18.19    -18.25    -18.13    389.94      1.00\n",
      "    sigma_v    176.93     52.05    169.93    102.74    267.56    315.30      1.00\n",
      "\n",
      "Number of divergences: 0\n"
     ]
    }
   ],
   "source": [
    "kernel = NUTS(model, init_strategy=init_to_median(num_samples=100))\n",
    "mcmc = MCMC(kernel, num_warmup=500, num_samples=500)\n",
    "\n",
    "rng_key = jax.random.PRNGKey(5)\n",
    "mcmc.run(rng_key, extra_fields=(\"potential_energy\",), **model_kwargs)\n",
    "mcmc.print_summary()\n",
    "samples = mcmc.get_samples()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv_csiborg",
   "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}