{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "f235d0c0-d7fb-46fd-89df-3d66851f457f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/bartlett/fsigma8/borg_velocity/tests/tfr_inference.py:310: SyntaxWarning: invalid escape sequence '\\s'\n",
      "  sigma_m = np.median(df['e_Kmag'])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<module 'tfr_inference' from '/home/bartlett/fsigma8/borg_velocity/tests/tfr_inference.py'>"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import astropy.constants\n",
    "import astropy.units as apu\n",
    "import astropy.cosmology\n",
    "from astropy.coordinates import SkyCoord\n",
    "import borg_velocity.poisson_process as poisson_process\n",
    "from astropy.cosmology import LambdaCDM, z_at_value\n",
    "import jax.numpy as jnp\n",
    "\n",
    "import importlib\n",
    "import tfr_inference\n",
    "importlib.reload(tfr_inference)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "7811713f-d849-4372-90d4-c01aaab5e988",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.0028092475228007545 0.04950413092258361 0.011837616297618123 0.009863123740916042\n",
      "335 9990\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<>:3: SyntaxWarning: invalid escape sequence '\\s'\n",
      "<>:3: SyntaxWarning: invalid escape sequence '\\s'\n",
      "/tmp/ipykernel_13422/2464123570.py:3: SyntaxWarning: invalid escape sequence '\\s'\n",
      "  df = pd.read_csv(fname, sep='\\s+', names=columns)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x14ea40a974d0>]"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 23672 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 48340 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "columns = ['ID', 'RAdeg', 'DEdeg', 'cz2mrs', 'Kmag', 'Hmag', 'Jmag', 'e_Kmag', 'e_Hmah', 'e_Jmag', 'WHIc', 'e_WHIc']\n",
    "fname = '/data101/bartlett/fsigma8/PV_data/2MASS/table1.dat'\n",
    "df = pd.read_csv(fname, sep='\\s+', names=columns)\n",
    "    \n",
    "eta = np.log10(df['WHIc']) - 2.5\n",
    "\n",
    "bins = np.linspace(eta.min(), eta.max(), 30)\n",
    "plt.hist(eta, bins=bins, density=True)\n",
    "\n",
    "mu = np.median(eta)\n",
    "sigma = (np.percentile(eta, 84) - np.percentile(eta, 16)) / 2\n",
    "g = np.exp(- (bins - mu) ** 2 / 2 / sigma**2) / np.sqrt(2 * np.pi) / sigma\n",
    "plt.plot(bins, g)\n",
    "\n",
    "sigma_eta = df['e_WHIc'] / df['WHIc'] / np.log(10)\n",
    "print(sigma_eta.min(), sigma_eta.max(), sigma_eta.mean(), sigma_eta.median())\n",
    "\n",
    "print(df['cz2mrs'].min(),  df['cz2mrs'].max())\n",
    "z = df['cz2mrs'] / astropy.constants.c.to('km/s').value\n",
    "cosmo = astropy.cosmology.Planck18\n",
    "dL = cosmo.luminosity_distance(z).to(apu.Mpc).value # Mpc\n",
    "mu = 5 * np.log10(dL) + 25\n",
    "\n",
    "plt.figure()\n",
    "plt.plot(eta, mu, '.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "952f0354-76a0-4b76-a310-f47c7dbd48e5",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<>:20: SyntaxWarning: invalid escape sequence '\\s'\n",
      "<>:34: SyntaxWarning: invalid escape sequence '\\s'\n",
      "<>:20: SyntaxWarning: invalid escape sequence '\\s'\n",
      "<>:34: SyntaxWarning: invalid escape sequence '\\s'\n",
      "/tmp/ipykernel_13422/2032762086.py:20: SyntaxWarning: invalid escape sequence '\\s'\n",
      "  df = pd.read_csv(fname, sep='\\s+', names=columns)\n",
      "/tmp/ipykernel_13422/2032762086.py:34: SyntaxWarning: invalid escape sequence '\\s'\n",
      "  df = pd.read_csv(fname, sep='\\s+', names=columns)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "zcosmo 0.0039389587876126975 0.05462965177877669\n",
      "etatrue: -0.41705506563366274 0.5555464714961542\n",
      "muTFR 31.441390386741705 37.174221031390445\n",
      "mutrue 30.852170311886482 36.64413989655261\n",
      "sigma_TFR 0.3\n",
      "mtrue 5.590206617138733 16.490061642329344\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Text(0.5, 1.0, 'm')"
      ]
     },
     "execution_count": 64,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 15548 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 17092 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def estimate_data_parameters():\n",
    "    \n",
    "    \"\"\"\n",
    "    ID   2MASS XSC ID name (HHMMSSss+DDMMSSs)\n",
    "    RAdeg   Right ascension (J2000)\n",
    "    DEdeg   Declination (J2000)\n",
    "    cz2mrs  Heliocentric redshift from the 2MRS (km/s)\n",
    "    Kmag    NIR magnitudes in the K band from the 2MRS (mag)\n",
    "    Hmag    NIR magnitudes in the H band from the 2MRS (mag)\n",
    "    Jmag    NIR magnitudes in the J band from the 2MRS (mag)\n",
    "    e_Kmag    Error of the NIR magnitudes in K band from the (mag)\n",
    "    e_Hmag    Error of the NIR magnitudes in H band from the (mag)\n",
    "    e_Jmag    Error of the NIR magnitudes in J band from the (mag)\n",
    "    WHIc    Corrected HI width (km/s)\n",
    "    e_WHIc    Error of corrected HI width (km/s)\n",
    "    \"\"\"\n",
    "    \n",
    "    columns = ['ID', 'RAdeg', 'DEdeg', 'cz2mrs', 'Kmag', 'Hmag', 'Jmag', 'e_Kmag', 'e_Hmah', 'e_Jmag', 'WHIc', 'e_WHIc']\n",
    "    fname = '/data101/bartlett/fsigma8/PV_data/2MASS/table1.dat'\n",
    "    df = pd.read_csv(fname, sep='\\s+', names=columns)\n",
    "    \n",
    "    sigma_m = np.median(df['e_Kmag'])\n",
    "\n",
    "    eta = np.log10(df['WHIc']) - 2.5\n",
    "    sigma_eta = np.median(df['e_WHIc'] / df['WHIc'] / np.log(10))\n",
    "    \n",
    "    hyper_eta_mu = np.median(eta)\n",
    "    hyper_eta_sigma = (np.percentile(eta, 84) - np.percentile(eta, 16)) / 2\n",
    "    \n",
    "    return sigma_m, sigma_eta, hyper_eta_mu, hyper_eta_sigma\n",
    "\n",
    "columns = ['ID', 'RAdeg', 'DEdeg', 'cz2mrs', 'Kmag', 'Hmag', 'Jmag', 'e_Kmag', 'e_Hmah', 'e_Jmag', 'WHIc', 'e_WHIc']\n",
    "fname = '/data101/bartlett/fsigma8/PV_data/2MASS/table1.dat'\n",
    "df = pd.read_csv(fname, sep='\\s+', names=columns)\n",
    "\n",
    "# Get some parameters from the data\n",
    "sigma_m, sigma_eta, hyper_eta_mu, hyper_eta_sigma = estimate_data_parameters()\n",
    "\n",
    "# cosmo = astropy.cosmology.Planck18\n",
    "cosmo = LambdaCDM(H0 = 80, Om0 = 0.3, Ode0 = 0.7)\n",
    "\n",
    "# Other parameters to use\n",
    "L = 200.0\n",
    "N = 64\n",
    "xmin = -L/2\n",
    "Rmax = 100\n",
    "Nt = 2062\n",
    "alpha = 1.4\n",
    "mthresh = 11.25\n",
    "a_TFR = -23\n",
    "b_TFR = -8.2\n",
    "sigma_TFR = 0.3\n",
    "bias_epsilon = 1e-7\n",
    "\n",
    "dens = np.load('dens.npy')\n",
    "vel = np.load('vel.npy')\n",
    "N = dens.shape[0]\n",
    "\n",
    "phi = (1. + dens + bias_epsilon) ** alpha\n",
    "\n",
    " # Generate positions (comoving)\n",
    "xtrue = poisson_process.sample_3d(phi, Nt, L, (xmin, xmin, xmin))\n",
    "\n",
    "# Convert to RA, Dec, Distance (comoving)\n",
    "rtrue = np.sqrt(np.sum(xtrue** 2, axis=0))  # Mpc/h\n",
    "c = SkyCoord(x=xtrue[0], y=xtrue[1], z=xtrue[2], representation_type='cartesian')\n",
    "RA = c.spherical.lon.degree\n",
    "Dec = c.spherical.lat.degree\n",
    "r_hat = np.array(SkyCoord(ra=RA*apu.deg, dec=Dec*apu.deg).cartesian.xyz)\n",
    "\n",
    "# Compute cosmological redshift\n",
    "zcosmo = z_at_value(cosmo.comoving_distance, rtrue * apu.Mpc / cosmo.h).value\n",
    "\n",
    "print('zcosmo', zcosmo.min(), zcosmo.max())\n",
    "        \n",
    "# Compute luminosity distance\n",
    "# DO I NEED TO DO /h???\n",
    "# dL = (1 + zcosmo) * rtrue  # Mpc/h\n",
    "dL = cosmo.luminosity_distance(zcosmo).to(apu.Mpc).value # Mpc\n",
    "\n",
    "# Compute true distance modulus\n",
    "mutrue = 5 * np.log10(dL) + 25\n",
    "\n",
    " # Sample true linewidth (eta) from its prior\n",
    "etatrue = hyper_eta_mu + hyper_eta_sigma * np.random.randn(Nt)\n",
    "\n",
    "print('etatrue:', etatrue.min(), etatrue.max())\n",
    "\n",
    "# Obtain muTFR from mutrue using the intrinsic scatter\n",
    "muTFR = mutrue + sigma_TFR * np.random.randn(Nt)\n",
    "\n",
    "print('muTFR', muTFR.min(), muTFR.max())\n",
    "print('mutrue', mutrue.min(), mutrue.max())\n",
    "print('sigma_TFR', sigma_TFR)\n",
    "\n",
    "# Obtain apparent magnitude from the TFR\n",
    "mtrue = muTFR + (a_TFR + b_TFR * etatrue)\n",
    "\n",
    "print('mtrue', mtrue.min(), mtrue.max())\n",
    "\n",
    "# Scatter true observed apparent magnitudes and linewidths\n",
    "mobs = mtrue + sigma_m * np.random.randn(Nt)\n",
    "etaobs = etatrue + sigma_eta * np.random.randn(Nt)\n",
    "\n",
    "# Apply apparement magnitude cut\n",
    "m = mobs <= mthresh\n",
    "mobs = mobs[m]\n",
    "etaobs = etaobs[m]\n",
    "xtrue = xtrue[:,m]\n",
    "\n",
    "data_eta = np.log10(df['WHIc']) - 2.5\n",
    "plt.hist(etaobs, label='synthetic', histtype='step', density=True, bins=20)\n",
    "plt.hist(data_eta, label='data', histtype='step', density=True, bins=20)\n",
    "plt.legend()\n",
    "plt.title('eta')\n",
    "\n",
    "plt.figure()\n",
    "data_m = df['Kmag']\n",
    "plt.hist(mobs, label='synthetic', histtype='step', density=True, bins=20)\n",
    "plt.hist(data_m, label='data', histtype='step', density=True, bins=20)\n",
    "plt.legend()\n",
    "plt.title('m')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "id": "b7b1015b-9cc2-4f64-9da7-48601b0a159f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "203.125\n",
      "(26, 26, 26)\n",
      "-101.5625\n",
      "7.8125\n",
      "7.8125\n"
     ]
    }
   ],
   "source": [
    "L = 500\n",
    "xmin = - 250\n",
    "Rmax = 100\n",
    "N = 64\n",
    "x = np.linspace(xmin, xmin + L, N+1)\n",
    "i0 = np.argmin(np.abs(x + Rmax))\n",
    "i1 = np.argmin(np.abs(x - Rmax))\n",
    "\n",
    "Lsmall = x[i1] - x[i0]\n",
    "dens_small = dens[i0:i1, i0:i1, i0:i1]\n",
    "xmin_small = x[i0]\n",
    "print(Lsmall)\n",
    "print(dens_small.shape)\n",
    "print(xmin_small)\n",
    "\n",
    "print(x[1] - x[0])\n",
    "print(L / N)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0e2f630-355c-477c-8b35-ec570eb97382",
   "metadata": {},
   "source": [
    "# Check MB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "f6103c09-b314-4898-b615-33676f7fc66d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[STD    ] Building gravity model lpt\n",
      "[STD    ] | |     ___________                              \n",
      "[STD    ] | |    /-/_\"/-/_/-/|     __________________________ \n",
      "[STD    ] | |   /\"-/-_\"/-_//||                \u001b[34;1mBORG3\u001b[39;0m model\n",
      "[STD    ] | |  /__________/|/|     (c) Jens Jasche 2012 - 2019\n",
      "[STD    ] | |  |\"|_'='-]:+|/||        Guilhem Lavaux 2014 - 2019\n",
      "[STD    ] | |  |-+-|.|_'-\"||//     __________________________ \n",
      "[STD    ] | |  |[\".[:!+-'=|//     \n",
      "[STD    ] | |  |='!+|-:]|-|/       \n",
      "[STD    ] | |   ----------         \n",
      "[STD    ] | | \n",
      "[STD    ] | | Please acknowledge the following papers:\n",
      "[STD    ] | |   - Jasche & Lavaux (A&A, 2019, arXiv 1806.11117)\n",
      "[STD    ] | |   - Jasche & Wandelt (MNRAS, 2012, arXiv 1203.3639)\n",
      "[STD    ] | |   - Jasche & Kitaura (MNRAS, 2010, arXiv 0911.2496)\n",
      "[STD    ] | |   - Lavaux, Jasche & Leclercq (arXiV 1909.06396)\n",
      "[STD    ] | |   - And relevant papers depending on the used sub-module/contribution\n",
      "[STD    ] | | \n",
      "\n",
      "[STD    ] | | This is BORG version 6b1404bfd8011f03ee1b8635275f68d14d16bef3\n",
      "[\u001b[35;1mWARNING\u001b[39;0m] | | | Entering [/home/bartlett/borg/libLSS/physics/forwards/primordial_as.cpp]virtual void LibLSS::ForwardPrimordial_As::updateCosmo()\n",
      "[\u001b[35;1mWARNING\u001b[39;0m] | | | | Sigma8 is set, but sigma8 normalization is not supported for PRIMORDIAL_AS. Ignoring the supplied sigma8.\n",
      "[\u001b[35;1mWARNING\u001b[39;0m] | | | Done (in 2.1512e-05 seconds) (ctx='[/home/bartlett/borg/libLSS/physics/forwards/primordial_as.cpp]virtual void LibLSS::ForwardPrimordial_As::updateCosmo()')\n",
      "[STD    ] | Smoothing field with rsmooth=4\n",
      "[STD    ] | | Smoothing field with rsmooth=4\n",
      "[STD    ] | | Smoothing field with rsmooth=4\n",
      "[STD    ] | | Smoothing field with rsmooth=4\n",
      "[STD    ] \tMade 668 of 2000\n",
      "[STD    ] \tMade 1345 of 2000\n",
      "[STD    ] \tMade 2000 of 2000\n",
      "[STD    ] Obtaining peculiar velocities\n",
      "[STD    ] Radial projection\n",
      "[STD    ] Making MB data\n",
      "26.779654 162.9076 97.677765\n",
      "27.107544 164.64821 92.31098\n",
      "16.885567 111.88842 61.495857\n",
      "0\n",
      "0\n",
      "Low: 14.289446\n",
      "High 16.209558\n",
      "0.078125 250.0\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x146c1c6ac3e0>]"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 26312 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 14816 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 26312 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "importlib.reload(tfr_inference)\n",
    "\n",
    "# Get some parameters from the data\n",
    "sigma_m, sigma_eta, hyper_eta_mu, hyper_eta_sigma = tfr_inference.estimate_data_parameters()\n",
    "\n",
    "# Other parameters to use\n",
    "L = 500.0\n",
    "N = 64\n",
    "xmin = -L/2\n",
    "R_lim = L / 2\n",
    "Rmax = 100\n",
    "Nt = 2000\n",
    "alpha = 1.4\n",
    "mthresh = 11.25\n",
    "a_TFR = -23\n",
    "b_TFR = -8.2\n",
    "sigma_TFR = 0.3\n",
    "sigma_v = 150\n",
    "Nint_points = 201 \n",
    "Nsig = 10\n",
    "frac_sigma_r = 0.07  # WANT A BETTER WAY OF DOING THIS - ESTIMATE THROUGH SIGMAS FROM TFR\n",
    "interp_order = 1\n",
    "bias_epsilon = 1.e-7\n",
    "\n",
    "cpar, dens, vel = tfr_inference.get_fields(L, N, xmin)\n",
    "\n",
    "RA, Dec, czCMB, m_true, eta_true, m_obs, eta_obs, xtrue = tfr_inference.create_mock(\n",
    "        Nt, L, xmin, cpar, dens, vel, Rmax, alpha, mthresh,\n",
    "        a_TFR, b_TFR, sigma_TFR, sigma_m, sigma_eta, \n",
    "        hyper_eta_mu, hyper_eta_sigma, sigma_v, \n",
    "        interp_order=interp_order, bias_epsilon=bias_epsilon)\n",
    "\n",
    "\n",
    "plt.figure()\n",
    "rtrue = jnp.sqrt(jnp.sum(xtrue ** 2, axis=0))\n",
    "robs = czCMB / 100 \n",
    "plt.plot(rtrue, robs, '.')\n",
    "    \n",
    "MB_pos = tfr_inference.generateMBData(RA, Dec, czCMB, L, N, R_lim, Nsig, Nint_points, sigma_v, frac_sigma_r)\n",
    "r = jnp.sqrt(jnp.sum(MB_pos ** 2, axis=0))\n",
    "\n",
    "rtrue = jnp.sqrt(jnp.sum(xtrue ** 2, axis=0))\n",
    "\n",
    "for i in range(3):\n",
    "    print(r[i,0], r[i,-1], rtrue[i])\n",
    "    \n",
    "# Check truth always inside the range of the MB pos\n",
    "dr_low = rtrue - r[:,0]\n",
    "dr_high = r[:,-1] - rtrue\n",
    "m = dr_low < 0\n",
    "print(m.sum())\n",
    "m = dr_high < 0\n",
    "print(m.sum())\n",
    "print('Low:', dr_low.min())\n",
    "print('High', dr_high.min())\n",
    "print(r.min(), r.max())\n",
    "\n",
    "czCMB_new = ((1 + zcosmo) * (1 + vr_noised / tfr_inference.utils.speed_of_light) - 1) * tfr_inference. utils.speed_of_light\n",
    "\n",
    "plt.figure()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "67759f3d-0cb3-4771-bb07-f3d5bfc6848c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9c0b73c6-70d7-4ba7-9c6f-e98a53a1358a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([107., 113.,  94., 108., 100.,  93., 110.,  95.,  86.,  94.]),\n",
       " array([-1.99562025, -1.59607887, -1.19653749, -0.79699612, -0.39745474,\n",
       "         0.00208664,  0.40162802,  0.8011694 ,  1.20071077,  1.60025215,\n",
       "         1.99979353]),\n",
       " <BarContainer object of 10 artists>)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 10848 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "a = np.loadtxt('a.txt')\n",
    "plt.hist(a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aeb7135d-0fc7-43c1-81c9-1fd32e5f05d6",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "borg_new",
   "language": "python",
   "name": "borg_new"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
