{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "e6102874-229d-47d5-b908-95e6a78ea91b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from os.path import join as pjoin\n",
    "import h5py\n",
    "import scipy.stats\n",
    "\n",
    "import sys\n",
    "sys.path.insert(0, '..')\n",
    "import borg_velocity.poisson_process as poisson_process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "5a8a6016-9ce5-4e15-96b9-9f00374d8924",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.collections.PolyCollection at 0x7f395b4f2590>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 239184 characters]",
      "text/plain": [
       "<Figure size 1500x400 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "data_dir = '/data101/bartlett/fsigma8/borg_velocity/supranta_run_ic'\n",
    "\n",
    "with h5py.File(pjoin(data_dir, 'mock_data.h5')) as f:\n",
    "    delta = f['scalars/BORG_final_density'][:]\n",
    "\n",
    "N = delta.shape[0]\n",
    "fig, axs = plt.subplots(1, 3, figsize=(15,4))\n",
    "axs[0].pcolor(np.log10(2 + delta[N//2]))\n",
    "axs[1].pcolor(np.log10(2 + delta[:,N//2,:]))\n",
    "axs[2].pcolor(np.log10(2 + delta[...,N//2]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ddef2d7f-fc27-4a1c-b188-af987a180f8f",
   "metadata": {},
   "source": [
    "# Try a radial selection parameterisation\n",
    "\n",
    "$$\n",
    "P(r) = r^2 \\exp( - \\lambda r / R_{\\rm max})\n",
    "$$\n",
    "\n",
    "Need to choose an appropriate $\\lambda$. The peak of this distribution is at $r = 2 R / \\lambda$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "642d9533-b291-4d23-bf8c-41428f0b6ce9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.lines.Line2D at 0x7f3910d41350>"
      ]
     },
     "execution_count": 67,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 29776 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "Rmax = 1000\n",
    "lam = 5\n",
    "nsamp = 10000\n",
    "\n",
    "mydist = scipy.stats.gamma(a=3, scale=1/lam)\n",
    "xsamp = mydist.rvs(nsamp) * Rmax\n",
    "\n",
    "x = np.linspace(0, xsamp.max(), 100)\n",
    "ypred0 = lam ** 3 / 2 * (x / Rmax) ** 2 * np.exp(- lam * x / Rmax) / Rmax\n",
    "\n",
    "plt.hist(xsamp, bins=30, density=True)\n",
    "plt.plot(x, ypred0)\n",
    "plt.axvline(Rmax)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "800b3e7a-bfd1-4600-b46e-f76ef87d111b",
   "metadata": {},
   "source": [
    "Draw samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 140,
   "id": "16465202-da80-41c3-93fc-3cf33ef4a306",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f39108f9f90>]"
      ]
     },
     "execution_count": 140,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 28944 characters]",
      "text/plain": [
       "<Figure size 700x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 78160 characters]",
      "text/plain": [
       "<Figure size 1500x400 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "alpha = 1.4\n",
    "Nt = 5000\n",
    "bias_epsilon = 1e-7\n",
    "L = 500\n",
    "xmin = (-250, -250, -250)\n",
    "lam = 5\n",
    "Rmax = 200\n",
    "\n",
    "phi = (1. + delta + bias_epsilon) ** alpha\n",
    "\n",
    "# Create coordinates along each axis\n",
    "x = np.linspace(xmin[0], xmin[0] + L, N)\n",
    "y = np.linspace(xmin[1], xmin[1] + L, N)\n",
    "z = np.linspace(xmin[2], xmin[2] + L, N)\n",
    "\n",
    "# Create a 3D grid of coordinates\n",
    "X, Y, Z = np.meshgrid(x, y, z, indexing='ij')\n",
    "\n",
    "# Compute the radial distance from the origin (r = sqrt(x^2 + y^2 + z^2))\n",
    "r = np.sqrt(X**2 + Y**2 + Z**2)\n",
    "\n",
    "# Multiply phi by the corresponding radial selection\n",
    "# Do not include r^2 factor as we are sampling in Cartesian 3D,\n",
    "# so do not need the Jacobian term\n",
    "# phi *= lam ** 3 / 2 * (r / Rmax) ** 2 * np.exp(- lam * r / Rmax) / Rmax\n",
    "phi *= np.exp(- lam * r / Rmax)\n",
    "\n",
    "coord_true = np.zeros((3,Nt))\n",
    "coord_meas = np.zeros((3,Nt))\n",
    "\n",
    "# Make Nt samples\n",
    "ctrue = poisson_process.sample_3d(phi, Nt, L, xmin)\n",
    "# cmeas, sig_mu[i] = radially_scatter(ctrue, frac_sig_x)\n",
    "rtrue = np.sqrt(np.sum(ctrue**2,axis=0))\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(7,4), sharex=True)\n",
    "ax.hist(rtrue, density=True, bins=30)\n",
    "x = np.linspace(0, rtrue.max(), 100)\n",
    "ypred0 = lam ** 3 / 2 * (x / Rmax) ** 2 * np.exp(- lam * x / Rmax) / Rmax\n",
    "ax.plot(x, ypred0)\n",
    "\n",
    "fig, axs = plt.subplots(1, 3, figsize=(15,4))\n",
    "dx = L / 20\n",
    "m = ((ctrue[2] > -dx) & (ctrue[2] < dx))\n",
    "axs[0].plot(ctrue[0,m], ctrue[1,m], '.', ms=1)\n",
    "m = ((ctrue[1] > -dx) & (ctrue[1] < dx))\n",
    "axs[1].plot(ctrue[0,m], ctrue[2,m], '.', ms=1)\n",
    "m = ((ctrue[0] > -dx) & (ctrue[0] < dx))\n",
    "axs[2].plot(ctrue[1,m], ctrue[2,m], '.', ms=1)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b2a3e93-f643-4fc1-b27e-a8d18e64d9ce",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ba7fd1b-ad5f-4977-8aae-3672a94deb8f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "borg_env",
   "language": "python",
   "name": "borg_env"
  },
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
