{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1f2c324c",
   "metadata": {},
   "source": [
    "## Computing Remnant Properties from NRSur7dq4 Waveforms\n",
    "\n",
    "This notebook generates waveforms using the `NRSur7dq4` surrogate model, computes remnant properties with `gw_remnant`, and compares against the `NRSur7dq4Remnant` fit.\n",
    "\n",
    "Requires `pip install gw_remnant[surrogates]`.\n",
    "\n",
    "Contact: Tousif Islam [tousifislam24@gmail.com]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac19b950",
   "metadata": {},
   "source": [
    "### 1. Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4609d499",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "lal.MSUN_SI != Msun\n"
     ]
    }
   ],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\", \"Wswiglal-redir-stdio\")\n",
    "import numpy as np\n",
    "import gwsurrogate\n",
    "import surfinBH\n",
    "import gwtools\n",
    "from gw_remnant.gw_remnant_calculator import GWRemnantCalculator\n",
    "from gw_remnant.gw_utils import waveform_generator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "21cbcb10",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded NRSur7dq4 model\n",
      "Loaded NRSur7dq4Remnant fit.\n"
     ]
    }
   ],
   "source": [
    "sur = gwsurrogate.LoadSurrogate('NRSur7dq4')\n",
    "fit = surfinBH.LoadFits('NRSur7dq4Remnant')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6ec02b5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compare_remnants(calc, fit, q, chi1, chi2, omega0):\n",
    "    \"\"\"Print a side-by-side comparison of gw_remnant vs NRSur7dq4Remnant.\"\"\"\n",
    "    mf, chif, vf, mf_err, chif_err, vf_err = fit.all(q, chi1, chi2, omega0=omega0)\n",
    "\n",
    "    print(f\"{'Property':<25} {'gw_remnant':>15} {'NRSur7dq4Remnant':>20}\")\n",
    "    print(\"-\" * 62)\n",
    "    print(f\"{'Remnant mass [M]':<25} {calc.remnant_mass:>15.6f} {mf:>20.6f}\")\n",
    "    print(f\"{'Remnant spin (z)':<25} {calc.remnant_spin:>15.6f} {chif[2]:>20.6f}\")\n",
    "    print(f\"{'Remnant |chi|':<25} {np.linalg.norm(calc.remnant_spin_vector):>15.6f} {np.linalg.norm(chif):>20.6f}\")\n",
    "    print(f\"{'Kick velocity [c]':<25} {calc.remnant_kick:>15.6f} {np.linalg.norm(vf):>20.6f}\")\n",
    "    print(f\"{'Spin vector':<25} {np.array2string(calc.remnant_spin_vector, precision=4):>15}\")\n",
    "    print(f\"{'':25} {np.array2string(chif, precision=4):>15}  (NRSur7dq4Remnant)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9fd98e16",
   "metadata": {},
   "source": [
    "### 2. Non-spinning binary (q = 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3d0e41ee",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NRSur7dq4 time grid: [-4299.55, 100.45] M\n",
      "Output time grid: [-4000.00, 49.90] M\n"
     ]
    }
   ],
   "source": [
    "q = 2.0\n",
    "chi1 = [0, 0, 0]\n",
    "chi2 = [0, 0, 0]\n",
    "\n",
    "times, h = waveform_generator.generate_nrsur7dq4(sur, q=q, chi1=chi1, chi2=chi2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "02de73b4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================================\n",
      "Remnant Properties Summary\n",
      "==================================================\n",
      "Mass ratio                    : 2.000\n",
      "Initial mass                  : 1.00000000 M\n",
      "Total energy radiated         : 0.03164319 M\n",
      "Peak luminosity               : 0.00076585\n",
      "Remnant mass                  : 0.96137688 M\n",
      "Remnant spin (dimensionless)  : 0.61992975\n",
      "Remnant spin vector (x,y,z)  : (0.00001354, 0.00003122, 0.61992975)\n",
      "Remnant kick velocity         : 0.00049130 c\n",
      "Remnant kick velocity         : 147.29 km/s\n",
      "Remnant kick vector (x,y,z)  : (-0.00048506, 0.00007767, -0.00000779) c\n",
      "Remnant displacement (x,y,z) : (-0.01430924, 0.00316052, -0.00027950) M\n",
      "==================================================\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/9k/xlxjfz4d2cz1063spty94p0r0000gq/T/ipykernel_23607/3385016684.py:1: UserWarning: Tips: If you are using NR waveforms, ensure that they do not contain junk radiation, as that is known to corrupt the remnant property estimation.\n",
      "  calc_ns = GWRemnantCalculator(time=times, h_dict=h, q=q, chi1=chi1, chi2=chi2)\n"
     ]
    }
   ],
   "source": [
    "calc_ns = GWRemnantCalculator(time=times, h_dict=h, q=q, chi1=chi1, chi2=chi2)\n",
    "calc_ns.print_remnants()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ee3ff42b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded NRSur7dq4 model\n",
      "Property                       gw_remnant     NRSur7dq4Remnant\n",
      "--------------------------------------------------------------\n",
      "Remnant mass [M]                 0.961377             0.961193\n",
      "Remnant spin (z)                 0.619930             0.623422\n",
      "Remnant |chi|                    0.619930             0.623422\n",
      "Kick velocity [c]                0.000491             0.000488\n",
      "Spin vector               [1.3544e-05 3.1220e-05 6.1993e-01]\n",
      "                          [ 5.2284e-05 -1.0361e-04  6.2342e-01]  (NRSur7dq4Remnant)\n"
     ]
    }
   ],
   "source": [
    "omega0 = abs(np.gradient(0.5 * gwtools.phase(h[(2, 2)])) / np.gradient(times))[0]\n",
    "compare_remnants(calc_ns, fit, q, chi1, chi2, omega0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b53554c4",
   "metadata": {},
   "source": [
    "### 3. Non-precessing (aligned-spin) binary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "7b21c17e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NRSur7dq4 time grid: [-4299.47, 100.53] M\n",
      "Output time grid: [-4000.00, 49.90] M\n"
     ]
    }
   ],
   "source": [
    "q = 2.0\n",
    "chi1 = [0, 0, 0.6]\n",
    "chi2 = [0, 0, -0.3]\n",
    "\n",
    "times, h = waveform_generator.generate_nrsur7dq4(sur, q=q, chi1=chi1, chi2=chi2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "468ad8c1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================================\n",
      "Remnant Properties Summary\n",
      "==================================================\n",
      "Mass ratio                    : 2.000\n",
      "Initial mass                  : 1.00000000 M\n",
      "Total energy radiated         : 0.04418731 M\n",
      "Peak luminosity               : 0.00101209\n",
      "Remnant mass                  : 0.94855348 M\n",
      "Remnant spin (dimensionless)  : 0.78789787\n",
      "Remnant spin vector (x,y,z)  : (-0.00001578, -0.00002466, 0.78789787)\n",
      "Remnant kick velocity         : 0.00026158 c\n",
      "Remnant kick velocity         : 78.42 km/s\n",
      "Remnant kick vector (x,y,z)  : (-0.00024390, 0.00009340, 0.00001461) c\n",
      "Remnant displacement (x,y,z) : (-0.00487550, 0.00188224, 0.00061485) M\n",
      "==================================================\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/9k/xlxjfz4d2cz1063spty94p0r0000gq/T/ipykernel_23607/3374374379.py:1: UserWarning: Tips: If you are using NR waveforms, ensure that they do not contain junk radiation, as that is known to corrupt the remnant property estimation.\n",
      "  calc_np = GWRemnantCalculator(time=times, h_dict=h, q=q, chi1=chi1, chi2=chi2)\n"
     ]
    }
   ],
   "source": [
    "calc_np = GWRemnantCalculator(time=times, h_dict=h, q=q, chi1=chi1, chi2=chi2)\n",
    "calc_np.print_remnants()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c0803345",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Property                       gw_remnant     NRSur7dq4Remnant\n",
      "--------------------------------------------------------------\n",
      "Remnant mass [M]                 0.948553             0.948758\n",
      "Remnant spin (z)                 0.787898             0.792294\n",
      "Remnant |chi|                    0.787898             0.792294\n",
      "Kick velocity [c]                0.000262             0.000287\n",
      "Spin vector               [-1.5779e-05 -2.4665e-05  7.8790e-01]\n",
      "                          [-5.8305e-05 -2.7152e-04  7.9229e-01]  (NRSur7dq4Remnant)\n"
     ]
    }
   ],
   "source": [
    "omega0 = abs(np.gradient(0.5 * gwtools.phase(h[(2, 2)])) / np.gradient(times))[0]\n",
    "compare_remnants(calc_np, fit, q, chi1, chi2, omega0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1bca93a",
   "metadata": {},
   "source": [
    "### 4. Precessing binary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "dffb8d8f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NRSur7dq4 time grid: [-4304.26, 95.74] M\n",
      "Output time grid: [-4000.00, 49.90] M\n"
     ]
    }
   ],
   "source": [
    "q = 3.0\n",
    "chi1 = [0.5, 0.0, 0.3]\n",
    "chi2 = [0.0, 0.4, -0.2]\n",
    "\n",
    "times, h = waveform_generator.generate_nrsur7dq4(sur, q=q, chi1=chi1, chi2=chi2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "96d2860f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================================\n",
      "Remnant Properties Summary\n",
      "==================================================\n",
      "Mass ratio                    : 3.000\n",
      "Initial mass                  : 1.00000000 M\n",
      "Total energy radiated         : 0.02909093 M\n",
      "Peak luminosity               : 0.00066007\n",
      "Remnant mass                  : 0.96467585 M\n",
      "Remnant spin (dimensionless)  : 0.68281051\n",
      "Remnant spin vector (x,y,z)  : (0.21062174, 0.04002432, 0.68281051)\n",
      "Remnant kick velocity         : 0.00077546 c\n",
      "Remnant kick velocity         : 232.48 km/s\n",
      "Remnant kick vector (x,y,z)  : (0.00064341, 0.00014168, 0.00040902) c\n",
      "Remnant displacement (x,y,z) : (0.02673383, -0.00266597, 0.00121992) M\n",
      "==================================================\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/9k/xlxjfz4d2cz1063spty94p0r0000gq/T/ipykernel_23607/4242708690.py:1: UserWarning: Tips: If you are using NR waveforms, ensure that they do not contain junk radiation, as that is known to corrupt the remnant property estimation.\n",
      "  calc_p = GWRemnantCalculator(time=times, h_dict=h, q=q, chi1=chi1, chi2=chi2)\n"
     ]
    }
   ],
   "source": [
    "calc_p = GWRemnantCalculator(time=times, h_dict=h, q=q, chi1=chi1, chi2=chi2)\n",
    "calc_p.print_remnants()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "6e8c751d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Property                       gw_remnant     NRSur7dq4Remnant\n",
      "--------------------------------------------------------------\n",
      "Remnant mass [M]                 0.964676             0.964583\n",
      "Remnant spin (z)                 0.682811             0.689538\n",
      "Remnant |chi|                    0.715677             0.718883\n",
      "Kick velocity [c]                0.000775             0.001510\n",
      "Spin vector               [0.2106 0.04   0.6828]\n",
      "                          [0.2    0.0365 0.6895]  (NRSur7dq4Remnant)\n"
     ]
    }
   ],
   "source": [
    "omega0 = abs(np.gradient(0.5 * gwtools.phase(h[(2, 2)])) / np.gradient(times))[0]\n",
    "compare_remnants(calc_p, fit, q, chi1, chi2, omega0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99fbb317-11e8-45a5-b684-83356768ee1e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "kitp-py310",
   "language": "python",
   "name": "kitp-py310"
  },
  "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.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}