{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from h5py import File\n",
    "\n",
    "import csiborgtools\n",
    "\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_key(key, stationary_observer=False):\n",
    "    nsims = range(51)\n",
    "    bfs = [None] * len(nsims)\n",
    "    for nsim in nsims:\n",
    "        if stationary_observer:\n",
    "            fname = f\"/mnt/extraspace/rstiskalek/quijote/BulkFlow_fiducial/sBF_nsim_{nsim}.hdf5\"\n",
    "        else:\n",
    "            fname = f\"/mnt/extraspace/rstiskalek/quijote/BulkFlow_fiducial/BF_nsim_{nsim}.hdf5\"\n",
    "\n",
    "        with File(fname, 'r') as f:\n",
    "            if nsim == 0:\n",
    "                print(\"Keys are: \", list(f.keys()))\n",
    "                distances = f['distances'][...]\n",
    "            bfs[nsim] = f[key][...]\n",
    "\n",
    "    bfs = np.asarray(bfs)\n",
    "\n",
    "    ndist = len(distances)\n",
    "    bx = np.asarray([bfs[..., i, 0].reshape(-1, ) for i in range(ndist)]).T\n",
    "    by = np.asarray([bfs[..., i, 1].reshape(-1, ) for i in range(ndist)]).T\n",
    "    bz = np.asarray([bfs[..., i, 2].reshape(-1, ) for i in range(ndist)]).T\n",
    "\n",
    "    return distances, np.stack([bx, by, bz], axis=-1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Keys are:  ['bf_volume_halo', 'bf_volume_halo_uniform', 'bf_volume_part', 'bf_vrad_weighted_halo', 'bf_vrad_weighted_halo_uniform', 'bf_vrad_weighted_part', 'distances', 'obs_0', 'obs_1', 'obs_10', 'obs_11', 'obs_12', 'obs_13', 'obs_14', 'obs_15', 'obs_16', 'obs_17', 'obs_18', 'obs_19', 'obs_2', 'obs_20', 'obs_21', 'obs_22', 'obs_23', 'obs_24', 'obs_25', 'obs_26', 'obs_3', 'obs_4', 'obs_5', 'obs_6', 'obs_7', 'obs_8', 'obs_9']\n"
     ]
    }
   ],
   "source": [
    "distances, bf_volume_part = read_key(\"bf_volume_part\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bulk flow in the cosmic (simulation) rest frame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Keys are:  ['bf_volume_halo', 'bf_volume_halo_uniform', 'bf_volume_part', 'bf_vrad_weighted_halo', 'bf_vrad_weighted_halo_uniform', 'bf_vrad_weighted_part', 'distances', 'obs_0', 'obs_1', 'obs_10', 'obs_11', 'obs_12', 'obs_13', 'obs_14', 'obs_15', 'obs_16', 'obs_17', 'obs_18', 'obs_19', 'obs_2', 'obs_20', 'obs_21', 'obs_22', 'obs_23', 'obs_24', 'obs_25', 'obs_26', 'obs_3', 'obs_4', 'obs_5', 'obs_6', 'obs_7', 'obs_8', 'obs_9']\n",
      "Keys are:  ['bf_volume_halo', 'bf_volume_halo_uniform', 'bf_volume_part', 'bf_vrad_weighted_halo', 'bf_vrad_weighted_halo_uniform', 'bf_vrad_weighted_part', 'distances', 'obs_0', 'obs_1', 'obs_10', 'obs_11', 'obs_12', 'obs_13', 'obs_14', 'obs_15', 'obs_16', 'obs_17', 'obs_18', 'obs_19', 'obs_2', 'obs_20', 'obs_21', 'obs_22', 'obs_23', 'obs_24', 'obs_25', 'obs_26', 'obs_3', 'obs_4', 'obs_5', 'obs_6', 'obs_7', 'obs_8', 'obs_9']\n",
      "Keys are:  ['bf_volume_halo', 'bf_volume_halo_uniform', 'bf_volume_part', 'bf_vrad_weighted_halo', 'bf_vrad_weighted_halo_uniform', 'bf_vrad_weighted_part', 'distances', 'obs_0', 'obs_1', 'obs_10', 'obs_11', 'obs_12', 'obs_13', 'obs_14', 'obs_15', 'obs_16', 'obs_17', 'obs_18', 'obs_19', 'obs_2', 'obs_20', 'obs_21', 'obs_22', 'obs_23', 'obs_24', 'obs_25', 'obs_26', 'obs_3', 'obs_4', 'obs_5', 'obs_6', 'obs_7', 'obs_8', 'obs_9']\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 83500 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "distances, bf_volume_part = read_key(\"bf_volume_part\")\n",
    "distances, bf_vrad_weighted_part = read_key(\"bf_vrad_weighted_part\")\n",
    "distances, bf_vrad_weighted_halo = read_key(\"bf_vrad_weighted_halo\")\n",
    "\n",
    "plt.figure()\n",
    "plt.title(\"Bulk flow magnitude comparison\")\n",
    "cols = plt.rcParams[\"axes.prop_cycle\"].by_key()[\"color\"]\n",
    "\n",
    "# LCDM expectation\n",
    "Rs,mean,std,mode,p05,p16,p84,p95 = np.load(\"/mnt/users/rstiskalek/csiborgtools/data/BulkFlowPlot.npy\")\n",
    "m = Rs < 200\n",
    "plt.plot(Rs[m], p05[m], color=\"black\", alpha=0.25)\n",
    "plt.plot(Rs[m], p95[m], color=\"black\", alpha=0.25)\n",
    "plt.plot(Rs[m], p16[m], color=\"black\", alpha=0.5)\n",
    "plt.plot(Rs[m], p84[m], color=\"black\", alpha=0.5)\n",
    "plt.plot(Rs[m], mode[m], color=\"black\", label=r\"$\\Lambda\\mathrm{CDM}$\")\n",
    "\n",
    "# Volume averaging\n",
    "# y = np.linalg.norm(bf_volume_part, axis=-1)\n",
    "# ylowlow, ylow, ymed, yhigh, yhighhigh = np.percentile(y, [5, 16, 50, 84, 95], axis=0)\n",
    "# plt.fill_between(distances, ylowlow, yhighhigh, alpha=0.25, color=cols[0])\n",
    "# plt.fill_between(distances, ylow, yhigh, alpha=0.5, color=cols[0])\n",
    "# plt.plot(distances, ymed, label=r\"$B_i = \\langle v_i \\rangle_R$\", color=cols[0])\n",
    "\n",
    "# # 1 / r^2\n",
    "# y = np.linalg.norm(bf_vrad_weighted_part, axis=-1)\n",
    "# ylowlow, ylow, ymed, yhigh, yhighhigh = np.percentile(y, [5, 16, 50, 84, 95], axis=0)\n",
    "# plt.fill_between(distances, ylowlow, yhighhigh, alpha=0.25, color=cols[1])\n",
    "# plt.fill_between(distances, ylow, yhigh, alpha=0.5, color=cols[1])\n",
    "# plt.plot(distances, ymed, label=r\"$B_i = R^2 \\langle v_r \\hat{n}_i / r^2 \\rangle_R$\", color=cols[1])\n",
    "\n",
    "\n",
    "# 1 / r^2\n",
    "y = np.linalg.norm(bf_vrad_weighted_halo, axis=-1)\n",
    "ylowlow, ylow, ymed, yhigh, yhighhigh = np.percentile(y, [5, 16, 50, 84, 95], axis=0)\n",
    "plt.fill_between(distances, ylowlow, yhighhigh, alpha=0.25, color=cols[2])\n",
    "plt.fill_between(distances, ylow, yhigh, alpha=0.5, color=cols[2])\n",
    "plt.plot(distances, ymed, label=\"Halo radial velocities\", color=cols[2])\n",
    "\n",
    "\n",
    "\n",
    "plt.legend()\n",
    "plt.xlim(distances[2], distances.max())\n",
    "plt.ylim(0, 1000)\n",
    "\n",
    "plt.xlabel(r\"$R ~ [\\mathrm{Mpc} / h]$\")\n",
    "plt.ylabel(r\"$\\left\\| \\mathbf{B} \\right\\| ~ [\\mathrm{km} / s]$\")\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"../../plots/bf_comparison.png\", dpi=450)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1377, 100, 3)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bf_volume_part.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "150.0\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 32764 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from scipy.stats import norm\n",
    "\n",
    "k = -1\n",
    "\n",
    "print(distances[k])\n",
    "\n",
    "plt.figure()\n",
    "\n",
    "plt.hist(np.linalg.norm(bf_volume_part[:, k, :], axis=-1), bins=\"auto\", density=1, histtype=\"step\")\n",
    "\n",
    "\n",
    "plt.hist(np.linalg.norm(bf_vrad_weighted_part[:, k, :], axis=-1), bins=\"auto\", density=1, histtype=\"step\")\n",
    "\n",
    "plt.hist(np.linalg.norm(bf_vrad_weighted_halo[:, k, :], axis=-1), bins=\"auto\", density=1, histtype=\"step\")\n",
    "\n",
    "\n",
    "d = norm(387, 28)\n",
    "t = np.linspace(0, 500, 1000)\n",
    "plt.plot(t, d.pdf(t) / 3)\n",
    "\n",
    "plt.xlim(0, 500)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Component differences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 69708 characters]",
      "text/plain": [
       "<Figure size 1200x400 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "dx = bf_volume_part - bf_vrad_weighted_part\n",
    "\n",
    "fig, axs = plt.subplots(1, 3, figsize=(12, 4), sharey=True)\n",
    "fig.suptitle(\"Bulk flow component difference\")\n",
    "cols = plt.rcParams[\"axes.prop_cycle\"].by_key()[\"color\"]\n",
    "for i in range(3):\n",
    "    y = dx[..., i]\n",
    "    ylowlow, ylow, ymed, yhigh, yhighhigh = np.percentile(y, [5, 16, 50, 84, 95], axis=0)\n",
    "    axs[i].fill_between(distances, ylowlow, yhighhigh, alpha=0.25, color=cols[0])\n",
    "    axs[i].fill_between(distances, ylow, yhigh, alpha=0.5, color=cols[0])\n",
    "    axs[i].plot(distances, ymed, color=cols[0])\n",
    "\n",
    "    axs[i].set_xlabel(r\"$R ~ [\\mathrm{Mpc} / h]$\")\n",
    "    axs[i].set_xlim(distances[2], distances.max())\n",
    "    axs[i].axhline(0, color=\"black\", linestyle=\"--\", zorder=0)\n",
    "\n",
    "axs[0].set_ylabel(r\"$\\Delta B_i ~ [\\mathrm{km} / s]$\")\n",
    "fig.tight_layout()\n",
    "fig.savefig(\"../../plots/bf_difference.png\", dpi=450)\n",
    "fig.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "###  Bulk flow in the observer rest frame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Keys are:  ['bf_volume_halo', 'bf_volume_halo_uniform', 'bf_volume_part', 'bf_vrad_weighted_halo', 'bf_vrad_weighted_halo_uniform', 'bf_vrad_weighted_part', 'distances', 'obs_0', 'obs_1', 'obs_10', 'obs_11', 'obs_12', 'obs_13', 'obs_14', 'obs_15', 'obs_16', 'obs_17', 'obs_18', 'obs_19', 'obs_2', 'obs_20', 'obs_21', 'obs_22', 'obs_23', 'obs_24', 'obs_25', 'obs_26', 'obs_3', 'obs_4', 'obs_5', 'obs_6', 'obs_7', 'obs_8', 'obs_9']\n",
      "Keys are:  ['bf_volume_halo', 'bf_volume_halo_uniform', 'bf_volume_part', 'bf_vrad_weighted_halo', 'bf_vrad_weighted_halo_uniform', 'bf_vrad_weighted_part', 'distances', 'obs_0', 'obs_1', 'obs_10', 'obs_11', 'obs_12', 'obs_13', 'obs_14', 'obs_15', 'obs_16', 'obs_17', 'obs_18', 'obs_19', 'obs_2', 'obs_20', 'obs_21', 'obs_22', 'obs_23', 'obs_24', 'obs_25', 'obs_26', 'obs_3', 'obs_4', 'obs_5', 'obs_6', 'obs_7', 'obs_8', 'obs_9']\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 76216 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "distances, bf_volume_part = read_key(\"bf_volume_part\", True)\n",
    "distances, bf_vrad_weighted_part = read_key(\"bf_vrad_weighted_part\", True)\n",
    "\n",
    "plt.figure()\n",
    "plt.title(\"Bulk flow magnitude comparison (observer rest frame)\")\n",
    "cols = plt.rcParams[\"axes.prop_cycle\"].by_key()[\"color\"]\n",
    "\n",
    "# Volume averaging\n",
    "y = np.linalg.norm(bf_volume_part, axis=-1)\n",
    "ylowlow, ylow, ymed, yhigh, yhighhigh = np.percentile(y, [5, 16, 50, 84, 95], axis=0)\n",
    "plt.fill_between(distances, ylowlow, yhighhigh, alpha=0.25, color=cols[0])\n",
    "plt.fill_between(distances, ylow, yhigh, alpha=0.5, color=cols[0])\n",
    "plt.plot(distances, ymed, label=r\"$B_i = \\langle v_i \\rangle_R$\", color=cols[0])\n",
    "\n",
    "# 1 / r^2\n",
    "y = np.linalg.norm(bf_vrad_weighted_part, axis=-1)\n",
    "ylowlow, ylow, ymed, yhigh, yhighhigh = np.percentile(y, [5, 16, 50, 84, 95], axis=0)\n",
    "plt.fill_between(distances, ylowlow, yhighhigh, alpha=0.25, color=cols[1])\n",
    "plt.fill_between(distances, ylow, yhigh, alpha=0.5, color=cols[1])\n",
    "plt.plot(distances, ymed, label=r\"$B_i = R^2 \\langle v_r \\hat{n}_i / r^2 \\rangle_R$\", color=cols[1])\n",
    "\n",
    "plt.legend()\n",
    "plt.xlim(distances[2], distances.max())\n",
    "plt.ylim(0, 1000)\n",
    "\n",
    "plt.xlabel(r\"$R ~ [\\mathrm{Mpc} / h]$\")\n",
    "plt.ylabel(r\"$\\left\\| \\mathbf{B} \\right\\| ~ [\\mathrm{km} / s]$\")\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"../../plots/bf_comparison_observer.png\", dpi=450)\n",
    "plt.show()"
   ]
  }
 ],
 "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
}
