{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "87fef6e1",
   "metadata": {},
   "source": [
    "# Loading waveforms.\n",
    "## Authors: Vaishak Prasad\n",
    "\n",
    "[comment]: <> (Abstract: TODO)\n",
    "\n",
    "**Notebook Status:** <font color='green'><b> Production </b></font> \n",
    "\n",
    "**Notes:** \n",
    "Last checked on 14 May 2023\n",
    "\n",
    "## Introduction:\n",
    "This notebook describes the routines for \n",
    "\n",
    "1. Loading RIT waveform data (Strain $rh/M$ and $\\Psi_4$) into mode arrays.\n",
    "1. Loading GT waveform data (Strain $rh/M$ and $\\Psi_4$) into mode arrays.\n",
    "1. Loading SpEC waveform data (Strain $rh/M$ and $\\Psi_4$) into mode arrays.\n",
    "1. Loading SpECTRE CCE waveform data (Strain $rh/M$ and $\\Psi_4$) into mode arrays.\n",
    "1. Trasformations between harmonic and angular basis."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be5a4654",
   "metadata": {},
   "source": [
    "<a id='toc'></a>\n",
    "\n",
    "# Table of Contents\n",
    "$$\\label{toc}$$\n",
    "\n",
    "This notebook is organized as follows:\n",
    "\n",
    "1. [Sec 1](#imports): Initialize the required Python libraries.\n",
    "1. [Sec 2](#rit)  : Load RIT data. \n",
    "    1. [Sec 2.1](#rit_psi4) : Weyl scalar $\\Psi_4$\n",
    "    1. [Sec 2.2](#rit_strain) : Strain $rh/M$\n",
    "1. [Sec 3](#gt) : Load GeorgiaTech (MAYA) Strain.\n",
    "1. [Sec 4](#spec)  : Load SpEC data. \n",
    "    1. [Sec 4.1](#sxs_psi4) : Weyl scalar $\\Psi_4$\n",
    "    1. [Sec 4.2](#sxs_strain) : Strain $rh/M$\n",
    "1. [Sec 5](#spectre)  : Load SpECTRE data. \n",
    "    1. [Sec 5.1](#spectre_psi4) : Weyl scalar $\\Psi_4$\n",
    "    1. [Sec 5.2](#sprectre_strain) : Strain $rh/M$\n",
    "1. [Sec 6](#transforms) : Transformations to and from SWSH basis."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ddf18b09",
   "metadata": {},
   "source": [
    "<a id='imports'></a>\n",
    "\n",
    "# Sec 1: Initialize the required Python libraries and functions. [Back to [top](#toc)\\]\n",
    "$$\\label{imports}$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "760f61c3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/vaishakprasad/Documents/Codes/waveformtools/waveformtools\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import waveformtools\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "#####################\n",
    "# matplotlib presets\n",
    "#########################################\n",
    "\n",
    "labelsize=16\n",
    "labelpad=14\n",
    "\n",
    "plt.rcParams.update({'font.size': 12})\n",
    "plt.rcParams.update({\"figure.figsize\" : (8,6)})\n",
    "plt.rcParams.update({\"axes.grid\" : True})\n",
    "plt.rcParams.update({\"axes.labelpad\" : labelpad})\n",
    "plt.rcParams.update({\"axes.labelsize\" : labelsize})\n",
    "plt.rcParams.update({\"figure.autolayout\" : True})\n",
    "plt.rcParams.update({\"grid.alpha\" : 0.3})\n",
    "plt.rcParams.update({\"grid.alpha\" : 0.3})\n",
    "plt.rcParams.update({'legend.markerscale': 12})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e42d17e8",
   "metadata": {},
   "source": [
    "<a id='rit'></a>\n",
    "\n",
    "## Sec 2: Loding of RIT waveform data \\[Back to [top](#toc)\\]\n",
    "\n",
    "\n",
    "\n",
    "Here the $\\Psi_4$ and strain $rh/M$ data from a simulation from the RIT catalogue is loaded info wavefoms modes_array.\n",
    "\n",
    "$$\\label{rit}$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "827921e7",
   "metadata": {},
   "source": [
    "<a id='rit_psi4'></a>\n",
    "## Sec 2.1: Load $\\Psi_4$ data of all available modes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d303daf5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from waveformtools import waveforms\n",
    "from waveformtools.waveforms import modes_array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2c963bcf",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Paths\n",
    "from pathlib import Path\n",
    "home = str(Path.home())\n",
    "# This path contains all tShe\n",
    "# psi4 data correcponding to the\n",
    "# simulation.\n",
    "psi4_data_dir = f\"{home}/NR/ExtrapPsi4_RIT-BBH-0001-n100-id3\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "61f7c7b3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Passing /home/vaishakprasad/NR/ExtrapPsi4_RIT-BBH-0001-n100-id3/None\n",
      "Loading RIT Psi4 type data.\n",
      "All files found ['rPsi4_l3_m3_rInf.asc', 'rPsi4_l4_m-3_rInf.asc', 'rPsi4_l4_m4_rInf.asc', 'rPsi4_l3_m-2_rInf.asc', 'rPsi4_l3_m1_rInf.asc', 'rPsi4_l4_m-1_rInf.asc', 'rPsi4_l2_m2_rInf.asc', 'rPsi4_l4_m1_rInf.asc', 'rPsi4_l2_m1_rInf.asc', 'rPsi4_l2_m0_rInf.asc', 'rPsi4_l4_m2_rInf.asc', 'rPsi4_l3_m-3_rInf.asc', 'rPsi4_l4_m-4_rInf.asc', 'rPsi4_l3_m0_rInf.asc', 'rPsi4_l4_m3_rInf.asc', 'rPsi4_l3_m2_rInf.asc', 'rPsi4_l2_m-2_rInf.asc', 'rPsi4_l2_m-1_rInf.asc', 'rPsi4_l4_m0_rInf.asc', 'rPsi4_l3_m-1_rInf.asc', 'rPsi4_l4_m-2_rInf.asc']\n",
      "The modes list is [[2, [-2, -1, 0, 1, 2]], [3, [-3, -2, -1, 0, 1, 2, 3]], [4, [-4, -3, -2, -1, 0, 1, 2, 3, 4]]]\n",
      "Reading in modes...\n",
      "Loading 2 -2\n",
      "Creating modes data\n",
      "Min dt 0.25 and Max dt 0.25\n",
      "Resampling at the finest timestep 0.25\n",
      "Loading 2 -1\n",
      "Loading 2 0\n",
      "Loading 2 1\n",
      "Loading 2 2\n",
      "Loading 3 -3\n",
      "Loading 3 -2\n",
      "Loading 3 -1\n",
      "Loading 3 0\n",
      "Loading 3 1\n",
      "Loading 3 2\n",
      "Loading 3 3\n",
      "Loading 4 -4\n",
      "Loading 4 -3\n",
      "Loading 4 -2\n",
      "Loading 4 -1\n",
      "Loading 4 0\n",
      "Loading 4 1\n",
      "Loading 4 2\n",
      "Loading 4 3\n",
      "Loading 4 4\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Name of the run\n",
    "label = f'q1a0_a RIT_rPsi4inf'\n",
    "\n",
    "# Create a modes array\n",
    "wf_psi = modes_array(label=label, data_dir=psi4_data_dir)\n",
    "\n",
    "wf_psi.load_modes(ftype='RIT', var_type='Psi4', centre=True)\n",
    "\n",
    "# Specify the sampling step\n",
    "#wf_psi.load_modes(ftype='RIT', var_type='Psi4', centre=True, resam_type=0.1)\n",
    "\n",
    "# Specify the resampling method\n",
    "# wf_psi.load_modes(ftype='RIT', var_type='Psi4', centre=True, resam_type=0.1, resam_kind='cubic')\n",
    "\n",
    "# Uncomment the below line to save \n",
    "# the modes array as waveformtools hdf file\n",
    "#wf_psi.save_modes(out_file_name='RIT_{label}_mp_psi4_resam_inf_new.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a70eabed",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'label': 'q1a0_a RIT_rPsi4inf',\n",
       " 'data_dir': '/home/vaishakprasad/NR/ExtrapPsi4_RIT-BBH-0001-n100-id3',\n",
       " 'file_name': None,\n",
       " 'key_format': None,\n",
       " 'ell_max': 4,\n",
       " 'modes_list': [[2, [-2, -1, 0, 1, 2]],\n",
       "  [3, [-3, -2, -1, 0, 1, 2, 3]],\n",
       "  [4, [-4, -3, -2, -1, 0, 1, 2, 3, 4]]],\n",
       " 'r_ext': inf,\n",
       " 'frequency_axis': None,\n",
       " 'out_file_name': None,\n",
       " 'maxtime': None,\n",
       " 'date': '2023-05-14',\n",
       " 'time': '18:34:16',\n",
       " 'key_ex': None,\n",
       " 'spin_weight': -2,\n",
       " 'actions': 'empty->load_modes->recenter',\n",
       " 'centered': True}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wf_psi.get_metadata()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1963997e",
   "metadata": {},
   "source": [
    "## Plot a few modes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "74ac105a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from waveformtools.compare import plot_modes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d5af9728",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Help on function plot_modes in module waveformtools.compare:\n",
      "\n",
      "plot_modes(wf1, nmodes=3, save_fig=False, xlim=[-1200, 100], ylim='auto', nstop=1, plot22=False)\n",
      "    Plot the first `nmodes` dominant modes of\n",
      "    the input waveforms\n",
      "    \n",
      "    Parameters\n",
      "    ----------\n",
      "    wf:           modes_array\n",
      "                              The list of `modes_array` waveforms\n",
      "                              to plot.\n",
      "    nmodes:     int\n",
      "                            The number of modes to plot.\n",
      "    xlim : list\n",
      "           [xmin, xmax] limits to plot.\n",
      "    tol:    float\n",
      "                    The tolerance to detect the modes.\n",
      "    \n",
      "    Returns\n",
      "    -------\n",
      "    Plots.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "help(plot_modes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "f588b929",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 45896 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 22596 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_modes(wf_psi, nmodes=3, xlim=[-800, 200])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f50c415",
   "metadata": {},
   "source": [
    "<a id='rit_strain'></a>\n",
    "\n",
    "## Sec 2.2: Load RIT Strain data $r h /M$ [Back to [top](#toc)]\n",
    "$$\\label{rit_strain}$$\n",
    "\n",
    "Here the strain $r h/M$ data from a simulation from the RIT catalogue is loaded into a `modes_array`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d829c349",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set path variables\n",
    "\n",
    "strain_data_dir = f'{home}/NR'\n",
    "strain_file_name = 'ExtrapStrain_RIT-BBH-0001-n100.h5'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "a342f87c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Passing /home/vaishakprasad/NR/ExtrapStrain_RIT-BBH-0001-n100.h5\n",
      "Loading RIT strain data.\n",
      "Interpolating using interp1d\n",
      "Fetching all keys from H5 file\n",
      "['NRTimes', 'amp_l2_m-1', 'amp_l2_m-2', 'amp_l2_m0', 'amp_l2_m1', 'amp_l2_m2', 'amp_l3_m-1', 'amp_l3_m-2', 'amp_l3_m-3', 'amp_l3_m0', 'amp_l3_m1', 'amp_l3_m2', 'amp_l3_m3', 'amp_l4_m-1', 'amp_l4_m-2', 'amp_l4_m-3', 'amp_l4_m-4', 'amp_l4_m0', 'amp_l4_m1', 'amp_l4_m2', 'amp_l4_m3', 'amp_l4_m4', 'auxiliary-info', 'phase_l2_m-1', 'phase_l2_m-2', 'phase_l2_m0', 'phase_l2_m1', 'phase_l2_m2', 'phase_l3_m-1', 'phase_l3_m-2', 'phase_l3_m-3', 'phase_l3_m0', 'phase_l3_m1', 'phase_l3_m2', 'phase_l3_m3', 'phase_l4_m-1', 'phase_l4_m-2', 'phase_l4_m-3', 'phase_l4_m-4', 'phase_l4_m0', 'phase_l4_m1', 'phase_l4_m2', 'phase_l4_m3', 'phase_l4_m4']\n",
      "All files found ['amp_l2_m-1', 'amp_l2_m-2', 'amp_l2_m0', 'amp_l2_m1', 'amp_l2_m2', 'amp_l3_m-1', 'amp_l3_m-2', 'amp_l3_m-3', 'amp_l3_m0', 'amp_l3_m1', 'amp_l3_m2', 'amp_l3_m3', 'amp_l4_m-1', 'amp_l4_m-2', 'amp_l4_m-3', 'amp_l4_m-4', 'amp_l4_m0', 'amp_l4_m1', 'amp_l4_m2', 'amp_l4_m3', 'amp_l4_m4', 'phase_l2_m-1', 'phase_l2_m-2', 'phase_l2_m0', 'phase_l2_m1', 'phase_l2_m2', 'phase_l3_m-1', 'phase_l3_m-2', 'phase_l3_m-3', 'phase_l3_m0', 'phase_l3_m1', 'phase_l3_m2', 'phase_l3_m3', 'phase_l4_m-1', 'phase_l4_m-2', 'phase_l4_m-3', 'phase_l4_m-4', 'phase_l4_m0', 'phase_l4_m1', 'phase_l4_m2', 'phase_l4_m3', 'phase_l4_m4']\n",
      "ell_max not provided.\n",
      "modes array not provided. Setting ell_max from file...\n",
      "Chosen ell max 4 Available ell_max 4\n",
      "Constructing the modes list\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/vaishakprasad/Documents/Codes/waveformtools/waveformtools/dataIO.py:675: FutureWarning: Unlike other reduction functions (e.g. `skew`, `kurtosis`), the default behavior of `mode` typically preserves the axis it acts along. In SciPy 1.11.0, this behavior will change: the default value of `keepdims` will become False, the `axis` over which the statistic is taken will be eliminated, and the value None will no longer be accepted. Set `keepdims` to True or False to avoid this warning.\n",
      "  dt_auto = mode(np.diff(time_axis))[0][0]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reading in modes...\n",
      "Loading 2 -2\n",
      "Creating modes data\n",
      "Default dt is 0.25\n",
      "Loading 2 -1\n",
      "Loading 2 0\n",
      "Loading 2 1\n",
      "Loading 2 2\n",
      "Loading 3 -3\n",
      "Loading 3 -2\n",
      "Loading 3 -1\n",
      "Loading 3 0\n",
      "Loading 3 1\n",
      "Loading 3 2\n",
      "Loading 3 3\n",
      "Loading 4 -4\n",
      "Loading 4 -3\n",
      "Loading 4 -2\n",
      "Loading 4 -1\n",
      "Loading 4 0\n",
      "Loading 4 1\n",
      "Loading 4 2\n",
      "Loading 4 3\n",
      "Loading 4 4\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Name of the run\n",
    "label = 'RIT q1a0_a strain'\n",
    "\n",
    "# Create a modes array\n",
    "wf_strain = modes_array(label=label, data_dir=strain_data_dir, file_name=strain_file_name)\n",
    "\n",
    "wf_strain.load_modes(ftype='RIT', var_type='Strain', centre=True)\n",
    "\n",
    "\n",
    "# Uncomment the below line to save \n",
    "# the modes array as waveformtools hdf file\n",
    "# wf_psi.save_modes(out_file_name='RIT_{label}_mp_psi4_resam_inf_new.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "d3d703ff",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 16892 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 18400 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_modes(wf_strain, nmodes=3, xlim=[-600, 100])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "917f2db4",
   "metadata": {},
   "source": [
    "<a id='gt'></a>\n",
    "\n",
    "## Sec 3: Loding of GT waveform data \\[Back to [top](#toc)\\]\n",
    "\n",
    "\n",
    "\n",
    "Here the strain $rh/M$ data from a simulation from the GT catalogue is loaded info wavefoms modes_array.\n",
    "\n",
    "$$\\label{gt}$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "3fd733de",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set path variables\n",
    "\n",
    "strain_data_dir = f'{home}/.cache/MAYA/data'\n",
    "strain_file_name = 'GT0001.h5'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "13be4ffe",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Passing /home/vaishakprasad/.cache/MAYA/data/GT0001.h5\n",
      "Loading RIT strain data.\n",
      "Interpolating using interp1d\n",
      "Fetching all keys from H5 file\n",
      "['amp_l2_m-1', 'amp_l2_m-2', 'amp_l2_m0', 'amp_l2_m1', 'amp_l2_m2', 'amp_l3_m-1', 'amp_l3_m-2', 'amp_l3_m-3', 'amp_l3_m0', 'amp_l3_m1', 'amp_l3_m2', 'amp_l3_m3', 'amp_l4_m-1', 'amp_l4_m-2', 'amp_l4_m-3', 'amp_l4_m-4', 'amp_l4_m0', 'amp_l4_m1', 'amp_l4_m2', 'amp_l4_m3', 'amp_l4_m4', 'amp_l5_m-1', 'amp_l5_m-2', 'amp_l5_m-3', 'amp_l5_m-4', 'amp_l5_m-5', 'amp_l5_m0', 'amp_l5_m1', 'amp_l5_m2', 'amp_l5_m3', 'amp_l5_m4', 'amp_l5_m5', 'amp_l6_m-1', 'amp_l6_m-2', 'amp_l6_m-3', 'amp_l6_m-4', 'amp_l6_m-5', 'amp_l6_m-6', 'amp_l6_m0', 'amp_l6_m1', 'amp_l6_m2', 'amp_l6_m3', 'amp_l6_m4', 'amp_l6_m5', 'amp_l6_m6', 'auxiliary-info', 'phase_l2_m-1', 'phase_l2_m-2', 'phase_l2_m0', 'phase_l2_m1', 'phase_l2_m2', 'phase_l3_m-1', 'phase_l3_m-2', 'phase_l3_m-3', 'phase_l3_m0', 'phase_l3_m1', 'phase_l3_m2', 'phase_l3_m3', 'phase_l4_m-1', 'phase_l4_m-2', 'phase_l4_m-3', 'phase_l4_m-4', 'phase_l4_m0', 'phase_l4_m1', 'phase_l4_m2', 'phase_l4_m3', 'phase_l4_m4', 'phase_l5_m-1', 'phase_l5_m-2', 'phase_l5_m-3', 'phase_l5_m-4', 'phase_l5_m-5', 'phase_l5_m0', 'phase_l5_m1', 'phase_l5_m2', 'phase_l5_m3', 'phase_l5_m4', 'phase_l5_m5', 'phase_l6_m-1', 'phase_l6_m-2', 'phase_l6_m-3', 'phase_l6_m-4', 'phase_l6_m-5', 'phase_l6_m-6', 'phase_l6_m0', 'phase_l6_m1', 'phase_l6_m2', 'phase_l6_m3', 'phase_l6_m4', 'phase_l6_m5', 'phase_l6_m6']\n",
      "All files found ['amp_l2_m-1', 'amp_l2_m-2', 'amp_l2_m0', 'amp_l2_m1', 'amp_l2_m2', 'amp_l3_m-1', 'amp_l3_m-2', 'amp_l3_m-3', 'amp_l3_m0', 'amp_l3_m1', 'amp_l3_m2', 'amp_l3_m3', 'amp_l4_m-1', 'amp_l4_m-2', 'amp_l4_m-3', 'amp_l4_m-4', 'amp_l4_m0', 'amp_l4_m1', 'amp_l4_m2', 'amp_l4_m3', 'amp_l4_m4', 'amp_l5_m-1', 'amp_l5_m-2', 'amp_l5_m-3', 'amp_l5_m-4', 'amp_l5_m-5', 'amp_l5_m0', 'amp_l5_m1', 'amp_l5_m2', 'amp_l5_m3', 'amp_l5_m4', 'amp_l5_m5', 'amp_l6_m-1', 'amp_l6_m-2', 'amp_l6_m-3', 'amp_l6_m-4', 'amp_l6_m-5', 'amp_l6_m-6', 'amp_l6_m0', 'amp_l6_m1', 'amp_l6_m2', 'amp_l6_m3', 'amp_l6_m4', 'amp_l6_m5', 'amp_l6_m6', 'phase_l2_m-1', 'phase_l2_m-2', 'phase_l2_m0', 'phase_l2_m1', 'phase_l2_m2', 'phase_l3_m-1', 'phase_l3_m-2', 'phase_l3_m-3', 'phase_l3_m0', 'phase_l3_m1', 'phase_l3_m2', 'phase_l3_m3', 'phase_l4_m-1', 'phase_l4_m-2', 'phase_l4_m-3', 'phase_l4_m-4', 'phase_l4_m0', 'phase_l4_m1', 'phase_l4_m2', 'phase_l4_m3', 'phase_l4_m4', 'phase_l5_m-1', 'phase_l5_m-2', 'phase_l5_m-3', 'phase_l5_m-4', 'phase_l5_m-5', 'phase_l5_m0', 'phase_l5_m1', 'phase_l5_m2', 'phase_l5_m3', 'phase_l5_m4', 'phase_l5_m5', 'phase_l6_m-1', 'phase_l6_m-2', 'phase_l6_m-3', 'phase_l6_m-4', 'phase_l6_m-5', 'phase_l6_m-6', 'phase_l6_m0', 'phase_l6_m1', 'phase_l6_m2', 'phase_l6_m3', 'phase_l6_m4', 'phase_l6_m5', 'phase_l6_m6']\n",
      "ell_max not provided.\n",
      "modes array not provided. Setting ell_max from file...\n",
      "Chosen ell max 6 Available ell_max 6\n",
      "Constructing the modes list\n",
      "NRTimes not present. Will compute dt auto from mode time axis\n",
      "Reading in modes...\n",
      "Loading 2 -2\n",
      "Creating modes data\n",
      "Default dt is 0.6200000000000045\n",
      "Loading 2 -1\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/vaishakprasad/Documents/Codes/waveformtools/waveformtools/dataIO.py:709: FutureWarning: Unlike other reduction functions (e.g. `skew`, `kurtosis`), the default behavior of `mode` typically preserves the axis it acts along. In SciPy 1.11.0, this behavior will change: the default value of `keepdims` will become False, the `axis` over which the statistic is taken will be eliminated, and the value None will no longer be accepted. Set `keepdims` to True or False to avoid this warning.\n",
      "  dt_auto = mode(np.diff(time_axis))[0][0]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading 2 0\n",
      "Loading 2 1\n",
      "Loading 2 2\n",
      "Loading 3 -3\n",
      "Loading 3 -2\n",
      "Loading 3 -1\n",
      "Loading 3 0\n",
      "Loading 3 1\n",
      "Loading 3 2\n",
      "Loading 3 3\n",
      "Loading 4 -4\n",
      "Loading 4 -3\n",
      "Loading 4 -2\n",
      "Loading 4 -1\n",
      "Loading 4 0\n",
      "Loading 4 1\n",
      "Loading 4 2\n",
      "Loading 4 3\n",
      "Loading 4 4\n",
      "Loading 5 -5\n",
      "Loading 5 -4\n",
      "Loading 5 -3\n",
      "Loading 5 -2\n",
      "Loading 5 -1\n",
      "Loading 5 0\n",
      "Loading 5 1\n",
      "Loading 5 2\n",
      "Loading 5 3\n",
      "Loading 5 4\n",
      "Loading 5 5\n",
      "Loading 6 -6\n",
      "Loading 6 -5\n",
      "Loading 6 -4\n",
      "Loading 6 -3\n",
      "Loading 6 -2\n",
      "Loading 6 -1\n",
      "Loading 6 0\n",
      "Loading 6 1\n",
      "Loading 6 2\n",
      "Loading 6 3\n",
      "Loading 6 4\n",
      "Loading 6 5\n",
      "Loading 6 6\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Name of the simulation\n",
    "label = 'gt strain'\n",
    "\n",
    "# Create a modes array\n",
    "wf_strain_gt = modes_array(label=label, data_dir=strain_data_dir, file_name=strain_file_name)\n",
    "\n",
    "wf_strain_gt.load_modes(ftype='GT', var_type='Strain', centre=True)\n",
    "\n",
    "\n",
    "# Uncomment the below line to save \n",
    "# the modes array as waveformtools hdf file\n",
    "# wf_psi.save_modes(out_file_name='RIT_{label}_mp_psi4_resam_inf_new.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "aa14df98",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'label': 'gt strain',\n",
       " 'data_dir': '/home/vaishakprasad/.cache/MAYA/data',\n",
       " 'file_name': 'GT0001.h5',\n",
       " 'key_format': None,\n",
       " 'ell_max': 6,\n",
       " 'modes_list': [[2, [-2, -1, 0, 1, 2]],\n",
       "  [3, [-3, -2, -1, 0, 1, 2, 3]],\n",
       "  [4, [-4, -3, -2, -1, 0, 1, 2, 3, 4]],\n",
       "  [5, [-5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5]],\n",
       "  [6, [-6, -5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5, 6]]],\n",
       " 'r_ext': inf,\n",
       " 'frequency_axis': None,\n",
       " 'out_file_name': None,\n",
       " 'maxtime': None,\n",
       " 'date': '2023-05-14',\n",
       " 'time': '18:34:33',\n",
       " 'key_ex': None,\n",
       " 'spin_weight': -2,\n",
       " 'actions': 'empty->load_modes->center'}"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wf_strain_gt.get_metadata()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "3a25ddb4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 28024 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 18720 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_modes(wf_strain_gt, nmodes=3, xlim=[-100, 100])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea13277f",
   "metadata": {},
   "source": [
    "<a id='spec'></a>\n",
    "\n",
    "## Sec 4: Loding of SpEC waveform data \\[Back to [top](#toc)\\]\n",
    "\n",
    "\n",
    "\n",
    "Here the $\\Psi_4$ and strain $rh/M$ data from a simulation from the SXS catalogue is loaded info wavefoms modes_array.\n",
    "\n",
    "$$\\label{spec}$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c953c7ac",
   "metadata": {},
   "source": [
    "<a id='sxs_psi4'></a>\n",
    "\n",
    "## Sec 4.1: Load SXS Psi4 data $\\Psi_4$ [Back to [top](#toc)]\n",
    "$$\\label{sxs_psi4}$$\n",
    "\n",
    "Here the $\\Psi_4$ data from a simulation from the sxs catalogue is loaded into a `modes_array`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "9b2909ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set path variables\n",
    "\n",
    "psi4_data_dir = f'{home}/NR/'\n",
    "psi4_file_name = 'rMPsi4_Asymptotic_GeometricUnits_CoM.h5'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "843dcf30",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Passing /home/vaishakprasad/NR//rMPsi4_Asymptotic_GeometricUnits_CoM.h5\n",
      "Loading SpEC data.\n",
      "ell_max not provided.\n",
      "modes array not provided. Setting ell_max from file...\n",
      "Chosen ell max 8 Available ell_max 8\n",
      "Using extrap order 4\n",
      "Constructing the modes list\n",
      "Creating modes data\n",
      "Min dt 0.05718379830432241 and Max dt 0.5000603551985208\n",
      "Resampling at the finest timestep 0.05718379830432241\n",
      "Chosen resampling fineness: finest\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/vaishakprasad/Documents/Codes/waveformtools/waveformtools/dataIO.py:1251: FutureWarning: Unlike other reduction functions (e.g. `skew`, `kurtosis`), the default behavior of `mode` typically preserves the axis it acts along. In SciPy 1.11.0, this behavior will change: the default value of `keepdims` will become False, the `axis` over which the statistic is taken will be eliminated, and the value None will no longer be accepted. Set `keepdims` to True or False to avoid this warning.\n",
      "  dt_auto = mode(np.diff(time_axis))[0][0]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<waveformtools.waveforms.modes_array at 0x7fd8d343d360>"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Name of the simulation\n",
    "label = 'SXS:BBH:0001 Psi4'\n",
    "\n",
    "# Create a modes array\n",
    "wf_psi4_sxs = modes_array(label=label, data_dir=psi4_data_dir, file_name=psi4_file_name)\n",
    "\n",
    "wf_psi4_sxs.load_modes(ftype='SpEC', var_type='Psi4', centre=True)\n",
    "\n",
    "\n",
    "# Uncomment the below line to save \n",
    "# the modes array as waveformtools hdf file\n",
    "# wf_psi.save_modes(out_file_name='RIT_{label}_mp_psi4_resam_inf_new.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "084b79c7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 25232 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 25232 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 14204 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_modes(wf_psi4_sxs, nmodes=3, xlim=[-1800, 100])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1d9fd27",
   "metadata": {},
   "source": [
    "<a id='sxs_strain'></a>\n",
    "\n",
    "## Sec 4.2: Load SXS Strain data $r h /M$ [Back to [top](#toc)]\n",
    "$$\\label{sxs_strain}$$\n",
    "\n",
    "Here the strain $r h/M$ data from a simulation from the sxs catalogue is loaded into a `modes_array`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26e190f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set path variables\n",
    "\n",
    "strain_data_dir = f'{home}/.cache/sxs/SXS:BBH:0001v3/Lev5/'\n",
    "strain_file_name = 'rhOverM_Asymptotic_GeometricUnits_CoM.h5'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "9644407c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Passing /home/vaishakprasad/.cache/sxs/SXS:BBH:0001v3/Lev5//rhOverM_Asymptotic_GeometricUnits_CoM.h5\n",
      "Loading SpEC data.\n",
      "ell_max not provided.\n",
      "modes array not provided. Setting ell_max from file...\n",
      "Chosen ell max 8 Available ell_max 8\n",
      "Using extrap order 4\n",
      "Constructing the modes list\n",
      "Creating modes data\n",
      "Min dt 0.05718379830432241 and Max dt 0.5000603551985208\n",
      "Resampling at the finest timestep 0.05718379830432241\n",
      "Chosen resampling fineness: finest\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<waveformtools.waveforms.modes_array at 0x7fd8d36f5b10>"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Name of the simulation\n",
    "label = 'SXS:BBH:0001 strain'\n",
    "\n",
    "# Create a modes array\n",
    "wf_strain_sxs = modes_array(label=label, data_dir=strain_data_dir, file_name=strain_file_name)\n",
    "\n",
    "wf_strain_sxs.load_modes(ftype='SpEC', var_type='Strain', centre=True)\n",
    "\n",
    "\n",
    "# Uncomment the below line to save \n",
    "# the modes array as waveformtools hdf file\n",
    "# wf_psi.save_modes(out_file_name='RIT_{label}_mp_psi4_resam_inf_new.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "adfc79cd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'label': 'SXS:BBH:0001 strain',\n",
       " 'data_dir': '/home/vaishakprasad/.cache/sxs/SXS:BBH:0001v3/Lev5/',\n",
       " 'file_name': 'rhOverM_Asymptotic_GeometricUnits_CoM.h5',\n",
       " 'key_format': None,\n",
       " 'ell_max': 8,\n",
       " 'modes_list': [[2, [-2, -1, 0, 1, 2]],\n",
       "  [3, [-3, -2, -1, 0, 1, 2, 3]],\n",
       "  [4, [-4, -3, -2, -1, 0, 1, 2, 3, 4]],\n",
       "  [5, [-5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5]],\n",
       "  [6, [-6, -5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5, 6]],\n",
       "  [7, [-7, -6, -5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5, 6, 7]],\n",
       "  [8, [-8, -7, -6, -5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5, 6, 7, 8]]],\n",
       " 'r_ext': inf,\n",
       " 'frequency_axis': None,\n",
       " 'out_file_name': None,\n",
       " 'maxtime': None,\n",
       " 'date': '2023-05-14',\n",
       " 'time': '18:37:38',\n",
       " 'key_ex': None,\n",
       " 'spin_weight': -2,\n",
       " 'actions': 'empty',\n",
       " 'extrap_order': 4}"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wf_strain_sxs.get_metadata()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "99c98da0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 22188 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 14004 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_modes(wf_strain_sxs, nmodes=3, xlim=[-2000, 100])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5041f58c",
   "metadata": {},
   "source": [
    "<a id='spectre'></a>\n",
    "\n",
    "## Sec 5: Loding of SpECTRE waveform data \\[Back to [top](#toc)\\]\n",
    "\n",
    "\n",
    "\n",
    "Here the $\\Psi_4$ and strain $rh/M$ data from a simulation from the SpECTRE CCE catalogue is loaded info wavefoms modes_array.\n",
    "\n",
    "$$\\label{spectre}$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1b6ffd8",
   "metadata": {},
   "source": [
    "<a id='spectre_psi4'></a>\n",
    "\n",
    "## Sec 5.1: Load SpECTRE Psi4 data $\\Psi_4$ [Back to [top](#toc)]\n",
    "$$\\label{spectre_psi4}$$\n",
    "\n",
    "Here the $\\Psi_4$ data from a simulation from the SpECTRE catalogue is loaded into a `modes_array`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "d52270c1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/vaishakprasad/Documents/Codes/waveformtools/waveformtools\n"
     ]
    }
   ],
   "source": [
    "from importlib import reload\n",
    "reload(waveformtools)\n",
    "reload(waveformtools.waveforms)\n",
    "reload(waveformtools.dataIO)\n",
    "from waveformtools.waveforms import modes_array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "b2e4f44b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set path variables\n",
    "\n",
    "psi4_data_dir = f'{home}/NR/'\n",
    "psi4_file_name = 'rMPsi4_Extrapolated_N7_CoM.h5'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17200c43",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Passing /home/vaishakprasad/NR//rMPsi4_Extrapolated_N7_CoM.h5\n",
      "Loading SpECTRE data.\n",
      "Constructing the modes list\n",
      "Creating modes data\n",
      "Min dt 0.011307906874208129 and Max dt 1.0751766069618682\n",
      "Resampling at the default timestep 0.1000003335693691\n",
      "resam_type auto\n"
     ]
    }
   ],
   "source": [
    "# Name of the simulation\n",
    "label = 'SpECTRE Psi4 CCE extrap'\n",
    "\n",
    "# Create a modes array\n",
    "wf_psi4_sp = modes_array(label=label, data_dir=psi4_data_dir, file_name=psi4_file_name)\n",
    "\n",
    "wf_psi4_sp.load_modes(ftype='SpECTRE', var_type='Psi4', centre=True, resam_type='auto')\n",
    "\n",
    "\n",
    "# Uncomment the below line to save \n",
    "# the modes array as waveformtools hdf file\n",
    "# wf_psi.save_modes(out_file_name='RIT_{label}_mp_psi4_resam_inf_new.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "263183bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_modes(wf_psi4_sp, nmodes=3, xlim=[-2000, 100])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22844957",
   "metadata": {},
   "source": [
    "<a id='spectre_strain'></a>\n",
    "\n",
    "## Sec 5.2: Load SpECTRE Strain data $r h /M$ [Back to [top](#toc)]\n",
    "$$\\label{spectre_strain}$$\n",
    "\n",
    "Here the strain $r h/M$ data from a simulation from the SpECTRE CCE catalogue is loaded into a `modes_array`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1d9ec91",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set path variables\n",
    "\n",
    "strain_data_dir = f'{home}/NR/'\n",
    "strain_file_name = 'rhOverM_Extrapolated_N5_CoM_Mem.h5'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5badc01",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Name of the simulation\n",
    "label = 'SpECTRE CCE strain extrap'\n",
    "\n",
    "# Create a modes array\n",
    "wf_strain_sp = modes_array(label=label, data_dir=strain_data_dir, file_name=strain_file_name)\n",
    "\n",
    "wf_strain_sp.load_modes(ftype='SpECTRE', var_type='Strain', centre=True)\n",
    "\n",
    "\n",
    "# Uncomment the below line to save \n",
    "# the modes array as waveformtools hdf file\n",
    "# wf_psi.save_modes(out_file_name='RIT_{label}_mp_psi4_resam_inf_new.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd341ea5",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_modes(wf_strain_sp, nmodes=3, xlim=[-2000, 100])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8e0f535",
   "metadata": {},
   "source": [
    "<a id='transforms'></a>\n",
    "\n",
    "# Sec 6: Transformations to harmonic and angular basis [Back to [top](#toc)]\n",
    "$$\\label{transforms}$$\n",
    "\n",
    "Here the transformation of the Weyl scalar $\\Psi_4$ data from Spin-Weighted Spherical Harmonic (SWSH, $\\ell, m$) basis to coordinate ($\\theta, \\phi$) basis is demonstrated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afba0d12",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a staggered grid on a sphere.\n",
    "from waveformtools import grids\n",
    "# This creates a uniform angular grid stagerred about the north pole.\n",
    "# By default, 4 ghost zones (2 on either sides)\n",
    "# is alloted in each direction.\n",
    "grid_info = grids.spherical_grid(ntheta=125, nphi=244)\n",
    "\n",
    "theta, phi = grid_info.meshgrid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95359222",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# Transform to angular coordinates.\n",
    "waveform_sp = wf2.to_spherical_array(grid_info=grid_info)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73173f5b",
   "metadata": {},
   "source": [
    "<a id='trans1'></a>\n",
    "## Carry out the transformation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c17b586c",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.sin(theta) * np.cos(phi)\n",
    "y = np.sin(theta) * np.sin(phi)\n",
    "z = np.cos(theta)\n",
    "\n",
    "\n",
    "%matplotlib widget\n",
    "import matplotlib\n",
    "from mpl_toolkits.mplot3d import axes3d, Axes3D\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax = Axes3D(fig)\n",
    "\n",
    "# Plot the imaginary part of Psi4 \n",
    "# at timestep\n",
    "tstep = 100 \n",
    "\n",
    "U = waveform_sp.data[:, :, tstep].imag\n",
    "\n",
    "# Assign color dimension\n",
    "color_dimension = U # change to desired fourth dimension\n",
    "\n",
    "minn, maxx = color_dimension.min(), color_dimension.max()\n",
    "norm = matplotlib.colors.Normalize(minn, maxx)\n",
    "\n",
    "m = plt.cm.ScalarMappable(norm=norm, cmap='jet')\n",
    "m.set_array([])\n",
    "\n",
    "fcolors = m.to_rgba(color_dimension)\n",
    "\n",
    "# Plot the surface.\n",
    "surf = ax.plot_surface(x, y, z,\n",
    "                       linewidth=0, antialiased=False, facecolors=fcolors)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "141a6753",
   "metadata": {},
   "source": [
    "<a id='trans2'></a>\n",
    "## Reverse transformation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33467dc6",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "wf2_2 = waveform_sp.to_modes_array()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f53e5017",
   "metadata": {},
   "outputs": [],
   "source": [
    "ell = 2\n",
    "emm = -2\n",
    "delta =  wf2.mode(ell, emm) - wf2_2.mode(ell, emm)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax.plot(wf2.time_axis, wf2.mode(2, 2), label='original', alpha=0.9, c='brown')\n",
    "ax.plot(wf2.time_axis, wf2_2.mode(2, 2), label='recon', linestyle='--', c='blue')\n",
    "ax.set_title('Reconstructed vs original waveforms')\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.set_yscale('log')\n",
    "ax.set_title('Error in transformation')\n",
    "ax.scatter(wf2.time_axis, abs(delta), s=1)\n",
    "plt.show()"
   ]
  }
 ],
 "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.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
