{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2d6ac2eb-2554-4aab-9cce-95d8c7c374d3",
   "metadata": {},
   "source": [
    "# A tutorial to perform a cosmological analysis with BORG\n",
    "\n",
    "## Introduction"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29555fe6-a9cb-4d5a-b0e5-bd62a36601ff",
   "metadata": {},
   "source": [
    "For this tutorial, we will assume that you properly built BORG from source. We will consider that `ARES_SOURCE_PATH` points to the root of the ARES git that obtained by cloning from https://bitbucket.org/bayesian_lss_team/ares.git. `$ARES_BUILD_PATH` will be the build directory. If you executed `build.sh` without argument it will be located in `$ARES_SOURCE_PATH/build/`.\n",
    "\n",
    "After that we create a directory in `CHAIN_DIR` (which is to your convenience). For the purpose of this tutotrial, we will have the directory in your home folder in `borg_test_chain`.\n",
    "\n",
    "The first step is copying the example files in place:\n",
    "```bash\n",
    "mkdir $CHAIN_DIR\n",
    "cd $CHAIN_DIR\n",
    "# The BORG configuration file\n",
    "cp $ARES_SOURCE_PATH/examples/2mpp-chain.ini.txt .\n",
    "# Some data (actually not used but it is required to be there)\n",
    "cp $ARES_SOURCE_PATH/examples/2MPP.txt .\n",
    "# Mask/selection describing the data\n",
    "cp $ARES_SOURCE_PATH/examples/completeness_* .\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e206958f-f694-4b58-be31-dc65cd35844a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "CHAIN_DIR=$HOME/borg_test_chain\n",
    "ARES_SOURCE_PATH=$HOME/PROJECTS/ares_clean\n",
    "ARES_BUILD_PATH=$ARES_SOURCE_PATH/build\n",
    "cd $CHAIN_DIR\n",
    "rm -f $CHAIN_DIR/mcmc_*\n",
    "rm -f $CHAIN_DIR/restart.h5_0\n",
    "cp $ARES_SOURCE_PATH/examples/2mpp-chain.ini.txt .\n",
    "cp $ARES_SOURCE_PATH/examples/2MPP.txt .\n",
    "cp $ARES_SOURCE_PATH/examples/completeness_* ."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7c2d19c-2d76-4114-9adc-3532721a9583",
   "metadata": {},
   "source": [
    "The second step is to actually run BORG on the example configuration. We do it as follow\n",
    "```bash\n",
    "# Be sure to be in the chain directory\n",
    "cd $CHAIN_DIR\n",
    "# Run hades with INIT (to indicate you want to start from scratch) and the configuration file\n",
    "$ARES_BUILD_PATH/src/hades3 INIT 2mpp-chain.ini.txt\n",
    "```\n",
    "The code will run for some time. It is set to generate 1000 MCMC files before exiting cleanly. If you want to reduce that number for testing purpose on your laptop, it is possible by changing the configuration `number_to_generate` in `2mpp-chain.ini.txt`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d68f103a-bd00-4583-9440-be4a9e1aca47",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[  0/  1] [\u001b[35;1mWARNING\u001b[39;0m] | | No [likelihood] section in params tree\n",
      "[  0/  1] [\u001b[35;1mWARNING\u001b[39;0m] | | I used a default weight of 1\n",
      "[  0/  1] [\u001b[35;1mWARNING\u001b[39;0m] | | I used a default weight of 1\n",
      "[  0/  1] [\u001b[35;1mWARNING\u001b[39;0m] | | I used a default weight of 1\n",
      "[  0/  1] [\u001b[35;1mWARNING\u001b[39;0m] | | I used a default weight of 1\n",
      "[  0/  1] [\u001b[35;1mWARNING\u001b[39;0m] | | I used a default weight of 1\n",
      "[  0/  1] [\u001b[35;1mWARNING\u001b[39;0m] | | I used a default weight of 1\n",
      "[  0/  1] [\u001b[35;1mWARNING\u001b[39;0m] | | I used a default weight of 1\n",
      "[  0/  1] [\u001b[35;1mWARNING\u001b[39;0m] | | I used a default weight of 1\n",
      "[  0/  1] [\u001b[35;1mWARNING\u001b[39;0m] | | I used a default weight of 1\n",
      "[  0/  1] [\u001b[35;1mWARNING\u001b[39;0m] | | I used a default weight of 1\n",
      "[  0/  1] [\u001b[35;1mWARNING\u001b[39;0m] | | I used a default weight of 1\n",
      "[  0/  1] [\u001b[35;1mWARNING\u001b[39;0m] | | I used a default weight of 1\n",
      "[  0/  1] [\u001b[35;1mWARNING\u001b[39;0m] | | I used a default weight of 1\n",
      "[  0/  1] [\u001b[35;1mWARNING\u001b[39;0m] | | I used a default weight of 1\n",
      "[  0/  1] [\u001b[35;1mWARNING\u001b[39;0m] | | I used a default weight of 1\n",
      "[  0/  1] [\u001b[35;1mWARNING\u001b[39;0m] | | I used a default weight of 1\n",
      "[  0/  1] [\u001b[35;1mWARNING\u001b[39;0m] | | I will pick default bias parameters as the present array is too short.\n",
      "[  0/  1] [\u001b[35;1mWARNING\u001b[39;0m] | | I will pick default bias parameters as the present array is too short.\n",
      "[  0/  1] [\u001b[35;1mWARNING\u001b[39;0m] | | I will pick default bias parameters as the present array is too short.\n",
      "[  0/  1] [\u001b[35;1mWARNING\u001b[39;0m] | | I will pick default bias parameters as the present array is too short.\n",
      "[  0/  1] [\u001b[35;1mWARNING\u001b[39;0m] | | I will pick default bias parameters as the present array is too short.\n",
      "[  0/  1] [\u001b[35;1mWARNING\u001b[39;0m] | | I will pick default bias parameters as the present array is too short.\n",
      "[  0/  1] [\u001b[35;1mWARNING\u001b[39;0m] | | I will pick default bias parameters as the present array is too short.\n",
      "[  0/  1] [\u001b[35;1mWARNING\u001b[39;0m] | | I will pick default bias parameters as the present array is too short.\n",
      "[  0/  1] [\u001b[35;1mWARNING\u001b[39;0m] | | I will pick default bias parameters as the present array is too short.\n",
      "[  0/  1] [\u001b[35;1mWARNING\u001b[39;0m] | | I will pick default bias parameters as the present array is too short.\n",
      "[  0/  1] [\u001b[35;1mWARNING\u001b[39;0m] | | I will pick default bias parameters as the present array is too short.\n",
      "[  0/  1] [\u001b[35;1mWARNING\u001b[39;0m] | | I will pick default bias parameters as the present array is too short.\n",
      "[  0/  1] [\u001b[35;1mWARNING\u001b[39;0m] | | I will pick default bias parameters as the present array is too short.\n",
      "[  0/  1] [\u001b[35;1mWARNING\u001b[39;0m] | | I will pick default bias parameters as the present array is too short.\n",
      "[  0/  1] [\u001b[35;1mWARNING\u001b[39;0m] | | I will pick default bias parameters as the present array is too short.\n"
     ]
    }
   ],
   "source": [
    "%%bash\n",
    "CHAIN_DIR=$HOME/borg_test_chain\n",
    "ARES_SOURCE_PATH=$HOME/PROJECTS/ares_clean\n",
    "ARES_BUILD_PATH=$ARES_SOURCE_PATH/build\n",
    "cd $CHAIN_DIR\n",
    "$ARES_BUILD_PATH/src/hades3 INIT 2mpp-chain.ini.txt > /dev/null"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38655ab4-9ca7-42cc-9d20-8612b175975f",
   "metadata": {},
   "source": [
    "This should take about 20 mins on a 20-40 cores workstation. The 1000 MCMC files requires about 1 GB on hard drive. The default setup also has a short log of the actions taken by the code in the file `logares.txt_rank_0` . The log can be increased or reduced by setting `logfile_verbose_level` to a different number, though for debugging purpose in case of error, it is advised to put it back to debug level (level 5) and it is good practice to keep it above 3. You should now see a content like here below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "faa98dc6-eca3-4755-933f-f10ea8a406a4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mcmc_995.h5\n",
      "mcmc_996.h5\n",
      "mcmc_997.h5\n",
      "mcmc_998.h5\n",
      "mcmc_999.h5\n",
      "mcmc_99.h5\n",
      "mcmc_9.h5\n",
      "mock_data.h5\n",
      "restart.h5_0\n",
      "timing_stats_0.txt\n"
     ]
    }
   ],
   "source": [
    "!ls $HOME/borg_test_chain | tail -n 10"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ab87343-0dad-41e9-980c-e36faa36f6df",
   "metadata": {},
   "source": [
    " \n",
    "\n",
    "Now we will look at the out puts generated by the BORG run. Note, that you do not have to wait for the run to complete, but you can already investigate intermediate results while BORG still runs. BORG results are stored in two major HDF5 files, the restart and the mcmc files. The restart files contain all information on the state of the Markov Chain required to resume the Markov Chain if it has been interrupted. The restart file also contains static information, that will not change during the run, such as the data, selection functions and masks and other settings. The mcmc files contain the current state of the Markov Chain. They are indexed by the current step in the chain, and contain the current sampled values of density fields, power-spectra, galaxy bias and cosmological parameters etc."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93419148-982c-48f4-8e8a-9a5ace6646b2",
   "metadata": {},
   "source": [
    "## The required python preamble"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0557fc22-d062-4828-b793-1a7ebef820cc",
   "metadata": {},
   "source": [
    "The basic analysis tools are bundled in a python package called pyaquila (https://bitbucket.org/bayesian_lss_team/pyaquila.git). You may install it with:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "cc6f8368-36e6-48ea-8ef0-33cc54b47650",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting git+https://bitbucket.org/bayesian_lss_team/pyaquila.git\n",
      "  Cloning https://bitbucket.org/bayesian_lss_team/pyaquila.git to /tmp/pip-req-build-pmk55_co\n",
      "  Installing build dependencies ... \u001b[?25ldone\n",
      "\u001b[?25h  Getting requirements to build wheel ... \u001b[?25ldone\n",
      "\u001b[?25h    Preparing wheel metadata ... \u001b[?25ldone\n",
      "\u001b[?25hRequirement already satisfied (use --upgrade to upgrade): pyaquila==0.1.0 from git+https://bitbucket.org/bayesian_lss_team/pyaquila.git in /data/lavaux/venv/python39/lib/python3.9/site-packages\n",
      "Requirement already satisfied: fire<0.5,>=0.4 in /data/lavaux/venv/python39/lib/python3.9/site-packages (from pyaquila==0.1.0) (0.4.0)\n",
      "Requirement already satisfied: astropy<4.0,>=3.2 in /data/lavaux/venv/python39/lib/python3.9/site-packages (from pyaquila==0.1.0) (3.2.3)\n",
      "Requirement already satisfied: six in /data/lavaux/venv/python39/lib/python3.9/site-packages (from fire<0.5,>=0.4->pyaquila==0.1.0) (1.15.0)\n",
      "Requirement already satisfied: termcolor in /data/lavaux/venv/python39/lib/python3.9/site-packages (from fire<0.5,>=0.4->pyaquila==0.1.0) (1.1.0)\n",
      "Requirement already satisfied: numpy>=1.13 in /data/lavaux/venv/python39/lib/python3.9/site-packages (from astropy<4.0,>=3.2->pyaquila==0.1.0) (1.20.3)\n",
      "Building wheels for collected packages: pyaquila\n",
      "  Building wheel for pyaquila (PEP 517) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for pyaquila: filename=pyaquila-0.1.0-py3-none-any.whl size=34399 sha256=a173e3ed3cd23d9d5ca6f0435f723aa43e01d61343041fea8e4b9bb11b73788c\n",
      "  Stored in directory: /tmp/pip-ephem-wheel-cache-odgnvkhm/wheels/a6/64/82/4b145d8038ad90db0deb482757804b8b819c6591358c80d395\n",
      "Successfully built pyaquila\n",
      "\u001b[33mWARNING: You are using pip version 20.2.2; however, version 21.1.2 is available.\n",
      "You should consider upgrading via the '/data/lavaux/venv/python39/bin/python3.9 -m pip install --upgrade pip' command.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!pip install git+https://bitbucket.org/bayesian_lss_team/pyaquila.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4affe219-8bec-40bb-8f1b-8e2b5f76bd93",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import h5py as h5\n",
    "import matplotlib.pyplot as plt\n",
    "# This is pyaquila\n",
    "import aquila as aq\n",
    "import os\n",
    "%matplotlib inline\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8fdf4e2-5459-4ffc-a3bf-c1a9fc0e8522",
   "metadata": {},
   "source": [
    "Now please indicate the path where you stored your BORG run:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "91950528-add4-4a60-9f71-4ad35df86536",
   "metadata": {},
   "outputs": [],
   "source": [
    "chain_dir=f\"{os.environ['HOME']}/borg_test_chain\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6649b484-5734-4f85-a0b5-dfb8d8ce1b33",
   "metadata": {},
   "source": [
    "## Inspecting the restart file\n",
    "The restart file may be manually opened by"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "30042f60-2d1c-4b3c-9d06-0a01db72d3f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "hf = h5.File(f\"{chain_dir}/restart.h5_0\",mode=\"r\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cab81c9c-0e25-4552-912d-235be7b0c259",
   "metadata": {},
   "source": [
    "The content of the file can be investigated by listing the keys of the `scalars` section"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a4abd616-a96f-4b2a-9f71-1bb6c8bbb3d1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['ARES_version', 'BORG_final_density', 'BORG_version', 'K_MAX', 'K_MIN', 'L0', 'L1', 'L2', 'MCMC_STEP', 'N0', 'N1', 'N2', 'N2_HC', 'N2real', 'NCAT', 'NFOREGROUNDS', 'NUM_MODES', 'Ndata0', 'Ndata1', 'Ndata2', 'ares_heat', 'bias_sampler_blocked', 'borg_a_final', 'borg_a_initial', 'catalog_foreground_coefficient_0', 'catalog_foreground_coefficient_1', 'catalog_foreground_coefficient_10', 'catalog_foreground_coefficient_11', 'catalog_foreground_coefficient_12', 'catalog_foreground_coefficient_13', 'catalog_foreground_coefficient_14', 'catalog_foreground_coefficient_15', 'catalog_foreground_coefficient_2', 'catalog_foreground_coefficient_3', 'catalog_foreground_coefficient_4', 'catalog_foreground_coefficient_5', 'catalog_foreground_coefficient_6', 'catalog_foreground_coefficient_7', 'catalog_foreground_coefficient_8', 'catalog_foreground_coefficient_9', 'catalog_foreground_maps_0', 'catalog_foreground_maps_1', 'catalog_foreground_maps_10', 'catalog_foreground_maps_11', 'catalog_foreground_maps_12', 'catalog_foreground_maps_13', 'catalog_foreground_maps_14', 'catalog_foreground_maps_15', 'catalog_foreground_maps_2', 'catalog_foreground_maps_3', 'catalog_foreground_maps_4', 'catalog_foreground_maps_5', 'catalog_foreground_maps_6', 'catalog_foreground_maps_7', 'catalog_foreground_maps_8', 'catalog_foreground_maps_9', 'corner0', 'corner1', 'corner2', 'cosmology', 'fourierLocalSize', 'fourierLocalSize1', 'galaxy_bias_0', 'galaxy_bias_1', 'galaxy_bias_10', 'galaxy_bias_11', 'galaxy_bias_12', 'galaxy_bias_13', 'galaxy_bias_14', 'galaxy_bias_15', 'galaxy_bias_2', 'galaxy_bias_3', 'galaxy_bias_4', 'galaxy_bias_5', 'galaxy_bias_6', 'galaxy_bias_7', 'galaxy_bias_8', 'galaxy_bias_9', 'galaxy_bias_ref_0', 'galaxy_bias_ref_1', 'galaxy_bias_ref_10', 'galaxy_bias_ref_11', 'galaxy_bias_ref_12', 'galaxy_bias_ref_13', 'galaxy_bias_ref_14', 'galaxy_bias_ref_15', 'galaxy_bias_ref_2', 'galaxy_bias_ref_3', 'galaxy_bias_ref_4', 'galaxy_bias_ref_5', 'galaxy_bias_ref_6', 'galaxy_bias_ref_7', 'galaxy_bias_ref_8', 'galaxy_bias_ref_9', 'galaxy_data_0', 'galaxy_data_1', 'galaxy_data_10', 'galaxy_data_11', 'galaxy_data_12', 'galaxy_data_13', 'galaxy_data_14', 'galaxy_data_15', 'galaxy_data_2', 'galaxy_data_3', 'galaxy_data_4', 'galaxy_data_5', 'galaxy_data_6', 'galaxy_data_7', 'galaxy_data_8', 'galaxy_data_9', 'galaxy_nmean_0', 'galaxy_nmean_1', 'galaxy_nmean_10', 'galaxy_nmean_11', 'galaxy_nmean_12', 'galaxy_nmean_13', 'galaxy_nmean_14', 'galaxy_nmean_15', 'galaxy_nmean_2', 'galaxy_nmean_3', 'galaxy_nmean_4', 'galaxy_nmean_5', 'galaxy_nmean_6', 'galaxy_nmean_7', 'galaxy_nmean_8', 'galaxy_nmean_9', 'galaxy_schechter_0', 'galaxy_schechter_1', 'galaxy_schechter_10', 'galaxy_schechter_11', 'galaxy_schechter_12', 'galaxy_schechter_13', 'galaxy_schechter_14', 'galaxy_schechter_15', 'galaxy_schechter_2', 'galaxy_schechter_3', 'galaxy_schechter_4', 'galaxy_schechter_5', 'galaxy_schechter_6', 'galaxy_schechter_7', 'galaxy_schechter_8', 'galaxy_schechter_9', 'galaxy_sel_window_0', 'galaxy_sel_window_1', 'galaxy_sel_window_10', 'galaxy_sel_window_11', 'galaxy_sel_window_12', 'galaxy_sel_window_13', 'galaxy_sel_window_14', 'galaxy_sel_window_15', 'galaxy_sel_window_2', 'galaxy_sel_window_3', 'galaxy_sel_window_4', 'galaxy_sel_window_5', 'galaxy_sel_window_6', 'galaxy_sel_window_7', 'galaxy_sel_window_8', 'galaxy_sel_window_9', 'galaxy_selection_info_0', 'galaxy_selection_info_1', 'galaxy_selection_info_10', 'galaxy_selection_info_11', 'galaxy_selection_info_12', 'galaxy_selection_info_13', 'galaxy_selection_info_14', 'galaxy_selection_info_15', 'galaxy_selection_info_2', 'galaxy_selection_info_3', 'galaxy_selection_info_4', 'galaxy_selection_info_5', 'galaxy_selection_info_6', 'galaxy_selection_info_7', 'galaxy_selection_info_8', 'galaxy_selection_info_9', 'galaxy_selection_type_0', 'galaxy_selection_type_1', 'galaxy_selection_type_10', 'galaxy_selection_type_11', 'galaxy_selection_type_12', 'galaxy_selection_type_13', 'galaxy_selection_type_14', 'galaxy_selection_type_15', 'galaxy_selection_type_2', 'galaxy_selection_type_3', 'galaxy_selection_type_4', 'galaxy_selection_type_5', 'galaxy_selection_type_6', 'galaxy_selection_type_7', 'galaxy_selection_type_8', 'galaxy_selection_type_9', 'galaxy_synthetic_sel_window_0', 'galaxy_synthetic_sel_window_1', 'galaxy_synthetic_sel_window_10', 'galaxy_synthetic_sel_window_11', 'galaxy_synthetic_sel_window_12', 'galaxy_synthetic_sel_window_13', 'galaxy_synthetic_sel_window_14', 'galaxy_synthetic_sel_window_15', 'galaxy_synthetic_sel_window_2', 'galaxy_synthetic_sel_window_3', 'galaxy_synthetic_sel_window_4', 'galaxy_synthetic_sel_window_5', 'galaxy_synthetic_sel_window_6', 'galaxy_synthetic_sel_window_7', 'galaxy_synthetic_sel_window_8', 'galaxy_synthetic_sel_window_9', 'gravity.do_rsd', 'growth_factor', 'hades_accept_count', 'hades_attempt_count', 'hades_mass', 'hades_sampler_blocked', 'hmc_Elh', 'hmc_Eprior', 'hmc_bad_sample', 'hmc_force_save_final', 'localN0', 'localN1', 'localNdata0', 'localNdata1', 'localNdata2', 'localNdata3', 'localNdata4', 'localNdata5', 'momentum_field', 'nmean_sampler_blocked', 'projection_model', 's_field', 's_hat_field', 'sigma8_sampler_blocked', 'startN0', 'startN1', 'total_foreground_blocked']\n"
     ]
    }
   ],
   "source": [
    "print(list(hf['scalars'].keys()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e91c054e-810e-4981-9044-a6e95a633cda",
   "metadata": {},
   "source": [
    "And then do not forget to close the file to avoid later problems."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "6a22bf40-db45-49c6-b1ae-680d50443550",
   "metadata": {},
   "outputs": [],
   "source": [
    "hf.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce4fe68f-9ce3-4027-a28e-c3c90bef4318",
   "metadata": {},
   "source": [
    "It is better to use context to ensure HDF5 files are always opened/closed when not used. When leaving the context, for any reasons (including errors), the file will be closed cleanly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "96517eb1-1ec0-4314-a240-c9ba6231e28d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<HDF5 file \"restart.h5_0\" (mode r)>\n"
     ]
    }
   ],
   "source": [
    "with h5.File(f\"{chain_dir}/restart.h5_0\",mode=\"r\") as hf:\n",
    "    #Do something\n",
    "    print(hf)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82d446c9-5877-4131-85b0-6a1d1ebb5cfe",
   "metadata": {},
   "source": [
    "Now let's have a look at some fields. The input galaxy data projected on BORG mesh may be viewed with: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "ffa0be3c-8195-4ef4-a385-92433b4356b8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 13400 characters]\n",
      "text/plain": [
       "<Figure size 1152x576 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "with h5.File(f\"{chain_dir}/restart.h5_0\",mode=\"r\") as hf:\n",
    "    data = hf['scalars/galaxy_data_3']\n",
    "    #Plot data\n",
    "    fig, (ax1, ax2) = plt.subplots(1, 2,figsize=(16, 8))\n",
    "    ax1.set_title('A Slice through the y-z plane of the data cube')\n",
    "    im1=ax1.imshow(data[16,:,:])\n",
    "\n",
    "    ax2.set_title('A Slice through the x-z plane of the data cube')\n",
    "    im2=ax2.imshow(data[:,16,:])\n",
    "\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9177f5f-657d-4232-b1fb-052a02d6b794",
   "metadata": {},
   "source": [
    "As we have mask and selections we see that the data is only present in the center.\n",
    "\n",
    "## Investigating MCMC files\n",
    "\n",
    "MCMC files are indexed by the sample number $i_\\mathrm{samp}$. Each file can be opened separately. Suppose we want to open the th mcmc file, then:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "8392faa8-be21-4460-845e-990f3e1ec025",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<HDF5 file \"mcmc_10.h5\" (mode r)>\n"
     ]
    }
   ],
   "source": [
    "isamp=10 # sample number\n",
    "fname_mcmc=f\"{chain_dir}/mcmc_{isamp}.h5\"\n",
    "with h5.File(fname_mcmc, mode=\"r\") as hf:\n",
    "    # Do something\n",
    "    print(hf)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6c0ab8a-db52-44ce-8fa7-c602d2243f26",
   "metadata": {},
   "source": [
    "### Inspect the content of the mcmc files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "0891a9b0-884e-4ee8-9b6a-4db44a2e56df",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['BORG_final_density', 'catalog_foreground_coefficient_0', 'catalog_foreground_coefficient_1', 'catalog_foreground_coefficient_10', 'catalog_foreground_coefficient_11', 'catalog_foreground_coefficient_12', 'catalog_foreground_coefficient_13', 'catalog_foreground_coefficient_14', 'catalog_foreground_coefficient_15', 'catalog_foreground_coefficient_2', 'catalog_foreground_coefficient_3', 'catalog_foreground_coefficient_4', 'catalog_foreground_coefficient_5', 'catalog_foreground_coefficient_6', 'catalog_foreground_coefficient_7', 'catalog_foreground_coefficient_8', 'catalog_foreground_coefficient_9', 'cosmology', 'galaxy_bias_0', 'galaxy_bias_1', 'galaxy_bias_10', 'galaxy_bias_11', 'galaxy_bias_12', 'galaxy_bias_13', 'galaxy_bias_14', 'galaxy_bias_15', 'galaxy_bias_2', 'galaxy_bias_3', 'galaxy_bias_4', 'galaxy_bias_5', 'galaxy_bias_6', 'galaxy_bias_7', 'galaxy_bias_8', 'galaxy_bias_9', 'galaxy_nmean_0', 'galaxy_nmean_1', 'galaxy_nmean_10', 'galaxy_nmean_11', 'galaxy_nmean_12', 'galaxy_nmean_13', 'galaxy_nmean_14', 'galaxy_nmean_15', 'galaxy_nmean_2', 'galaxy_nmean_3', 'galaxy_nmean_4', 'galaxy_nmean_5', 'galaxy_nmean_6', 'galaxy_nmean_7', 'galaxy_nmean_8', 'galaxy_nmean_9', 'hades_accept_count', 'hades_attempt_count', 'hmc_Elh', 'hmc_Eprior', 's_field', 's_hat_field']\n"
     ]
    }
   ],
   "source": [
    "with h5.File(fname_mcmc, mode=\"r\") as hf:\n",
    "    # Do something\n",
    "    print(list(hf['scalars'].keys()))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f320f0bf-13c4-4b8f-a672-efe3e395c791",
   "metadata": {},
   "source": [
    "### Plotting density fields\n",
    "We can for instance be interested in plotting inferred initial and final density samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "763e3241-bf0d-49ce-8ffc-b27b581e7f17",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 24824 characters]\n",
      "text/plain": [
       "<Figure size 1152x576 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "with h5.File(fname_mcmc, mode=\"r\") as hf:\n",
    "\n",
    "    delta_in=np.array(hf['scalars/s_field'])\n",
    "    delta_fi=np.array(hf['scalars/BORG_final_density'])\n",
    "\n",
    "    fig, (ax1, ax2) = plt.subplots(1, 2,figsize=(16, 8))\n",
    "    ax1.set_title('initial density')\n",
    "    im1=ax1.imshow(delta_in[16,:,:])\n",
    "\n",
    "    ax2.set_title('final density')\n",
    "    im2=ax2.imshow(delta_fi[16,:,:])\n",
    "\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86848950-7c63-4584-b7b4-9aa7283a7541",
   "metadata": {},
   "source": [
    "We notice that no mask due to selection is visible there. That is because BORG provided a fair sample from the posterior distribution, which accounts for this mask while using the prior to fill the missing parts.\n",
    "\n",
    "### Plotting the power-spectrum\n",
    "\n",
    "The ARES repository provides some routines to analyse the BORG runs. A particularly useful routine calculates the posterior power-spectra of inferred initial density fields."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "b174d27f-c38a-42fb-829f-493238e8447d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 9160 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "ss = aq.analysis.Analysis(chain_dir,LSS_framework='BORG')\n",
    "\n",
    "# Indicate how many bins we want for the power spectrum.\n",
    "# This number is sticky internally as it involves precomputing some quantities.\n",
    "# Please keep it to a fixed number for each instance `ss`\n",
    "# `unit=True` is important for BORG >= 2.0 as shat is stored dimensionless\n",
    "opts=dict(Nbins=128,unit=True)\n",
    "\n",
    "#You can choose the sample number\n",
    "isamp=10\n",
    "#\n",
    "P_k,Pw,k_mode=ss.compute_power_shat_spectrum(isamp, **opts)\n",
    "\n",
    "# Here we compute the bin centered position for each P(kmode)\n",
    "k_mode_c = 0.5*(k_mode[1:]+k_mode[:-1])\n",
    "\n",
    "plt.axhline(1.0,color='r')\n",
    "plt.semilogx(k_mode_c,P_k)\n",
    "plt.xlabel('k ($h^{-1} Mpc)')\n",
    "plt.ylabel(r'$P(k)$ (dimensionless)')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c52e5e04-0f3f-4be5-8348-417a3065e877",
   "metadata": {},
   "source": [
    "The spectrum here-above is dimensionless as BORG works on variables prior to the convolution with primordial gravitational perturbations and primordial generation of density perturbations. If nothing bad happened and the cosmology in the ini file corresponds to the data, that spectrum should fluctuate around  `1`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54d8a59b-2fa4-43b9-903b-8fa96cfc676c",
   "metadata": {},
   "source": [
    "### Monitoring power-spectrum warm-up phase\n",
    "Rather than looking just at individual posterior sample power-spectra we can follow the evolution of power-spectra across the chain. Suppose you want to monitor the first 1000 samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "06b2c7dd-1ea3-4f46-9ef3-852ec3468809",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 54620 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "Nsamp=1000\n",
    "PPs=[]\n",
    "for isamp in np.arange(Nsamp):\n",
    "    # All samples have the same second and third tuple element\n",
    "    PPs.append(ss.compute_power_shat_spectrum(isamp, **opts)[0])\n",
    "\n",
    "    \n",
    "\n",
    "#plot power-spectra\n",
    "color_idx = np.linspace(0, 1, Nsamp)\n",
    "plt.axhline(1.0,color='r')\n",
    "idx=0\n",
    "for PP in PPs:\n",
    "    plt.semilogx(k_mode_c,PP,alpha=0.2,color=plt.cm.cool(color_idx[idx]), lw=1)\n",
    "    idx=idx+1\n",
    "plt.xlim([min(k_mode_c),max(k_mode_c)])\n",
    "plt.xlabel('k ($h^{-1} Mpc)')\n",
    "plt.ylabel(r'$P(k)$ (dimensionless)')\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89e82c86-2373-4a43-97bb-aa7c8a881275",
   "metadata": {},
   "source": [
    "Instead of looking at the power spectrum per mode we can look at the trace plot of the value for each mode. That allows to detect visually any oddity in the chain and observe the existence of a transition period. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "ae29d089-48a5-4793-a9bf-92a59e105dd4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 49684 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "PPs_trace = np.array(PPs).transpose()\n",
    "# We limit to the first 50 modes just for plotting purpose\n",
    "color_idx = np.linspace(0, 1, len(PPs_trace[:50]))\n",
    "plt.axhline(1.0,color='r')\n",
    "idx=0\n",
    "for i,PP in enumerate(PPs_trace[:50]):\n",
    "    plt.plot(PP,alpha=1.0,color=plt.cm.cool(color_idx[i]), lw=1)\n",
    "#plt.xlim([min(k_mode_c),max(k_mode_c)])\n",
    "#plt.xlabel('k ($h^{-1} Mpc)')\n",
    "plt.ylabel(r'$P(k)$ (dimensionless)')\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14566cdc-3e59-4cbf-9c4f-89e564db0600",
   "metadata": {},
   "source": [
    "### Exploring density field\n",
    "The above chain shows a transitionary phase (often called burn-in), which corresponds to a fast transition from a low probability to a high probability state. It is generally advised to truncate this part of the chain for short chains as it removes part of the samples that are highly unlikely. One may then compute interesting summaries of the chain. A typical interesting one is the mean and standard deviation per voxel of the density. Formally that corresponds to doing it the following operations:\n",
    "$$\n",
    "\\langle \\delta_{i,j,k}\\rangle = \\int_\\delta \\text{d}\\delta P(\\delta) \\delta_{i,j,k} \\simeq \\frac{1}{N_\\text{MCMC}} \\sum_{n} \\delta^{(n)}_{i,j,k}\n",
    "$$\n",
    "with $n$ running over the elements of the chain after burnin and $N_\\text{MCMC}$ the number of these elements. Here-below, we will compute this quantity and the standard deviation after truncating the first 200 samples which looks reasonable given the above trace plot of the power spectra."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "94d84978-b926-4525-92ab-40cbb0b6d4d9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1d4ae54377154073872bd6bce0c7a717",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/900 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# do_dev=True indicates to compute the standard deviation\n",
    "# start=500 indicates to start computing at the sample 500\n",
    "m=aq.utils.read_chain_avg_dev(chain_dir,['scalars.BORG_final_density'],do_dev=True,start=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49433544-d135-44f8-9441-f4e8ffa17902",
   "metadata": {},
   "source": [
    "Now we need to represent this field. To make a change, and actually uses physical coordinates (comoving coordinates), we will pcolormesh instead of imshow (which does a number of transpose to the field by default)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "c7056d9b-7001-40f9-b8d8-1ace2321a12f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 63120 characters]\n",
      "text/plain": [
       "<Figure size 1500x600 with 4 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "ix=np.linspace(-677.7/2,677.7/2,32)\n",
    "\n",
    "# Note the peculiar indexing that is required\n",
    "xm,ym,zm=np.meshgrid(ix,ix,ix,indexing='ij')\n",
    "\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2,figsize=(10, 4),dpi=150)\n",
    "\n",
    "ax1.set_title(\"Mean field\")\n",
    "mean_den=ax1.pcolormesh(xm[:,:,16],zm[:,16,:],m['scalars.BORG_final_density'][0][:,16,:],cmap=plt.cm.coolwarm,vmin=-1,vmax=1)\n",
    "ax1.set_aspect('equal')\n",
    "ax1.set_xlabel('X ($h^{-1}$ Mpc)')\n",
    "ax1.set_ylabel('Z ($h^{-1}$ Mpc)')\n",
    "plt.colorbar(mean_den,ax=ax1)\n",
    "\n",
    "ax2.set_title(\"Standard deviation\")\n",
    "std_den=ax2.pcolormesh(xm[:,:,16],zm[:,16,:],m['scalars.BORG_final_density'][1][:,16,:],cmap=plt.cm.gray_r)\n",
    "ax2.set_aspect('equal')\n",
    "ax2.set_xlabel('X ($h^{-1}$ Mpc)')\n",
    "ax2.set_ylabel('Z ($h^{-1}$ Mpc)')\n",
    "plt.colorbar(std_den,ax=ax2)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eea86fbc-29c6-4584-af30-c80946ce7660",
   "metadata": {},
   "source": [
    "In the above, we see that the standard deviation is varying quite a lot between different regions of the sky. That is expected given the selection that has been used in the model. The selection can be checked directly by looking at the field `scalars/galaxy_sel_window_0`  (and all the others depending on the number of catalogs) in the restart file."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.9.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
