{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8b2eeb7d",
   "metadata": {},
   "source": [
    "## Computing Remnant Properties from NRHybSur3dq8 Waveforms\n",
    "\n",
    "This notebook generates waveforms using the `NRHybSur3dq8` surrogate model, computes remnant properties with `gw_remnant`, and compares against the `NRSur3dq8Remnant` fit.\n",
    "\n",
    "Requires `pip install gw_remnant[surrogates]`.\n",
    "\n",
    "Contact: Tousif Islam [tousifislam24@gmail.com]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58f3c540",
   "metadata": {},
   "source": [
    "### 1. Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "fc543290",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "lal.MSUN_SI != Msun\n"
     ]
    }
   ],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\", \"Wswiglal-redir-stdio\")\n",
    "\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": "ae680010",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded NRHybSur3dq8 model\n",
      "Loaded NRSur3dq8Remnant fit.\n"
     ]
    }
   ],
   "source": [
    "sur = gwsurrogate.LoadSurrogate('NRHybSur3dq8')\n",
    "fit = surfinBH.LoadFits('NRSur3dq8Remnant')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d7943289",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compare_remnants(calc, fit, q, chi1, chi2):\n",
    "    \"\"\"Print a side-by-side comparison of gw_remnant vs NRSur3dq8Remnant.\"\"\"\n",
    "    mf, chif, vf, mf_err, chif_err, vf_err = fit.all(q, chi1, chi2)\n",
    "\n",
    "    print(f\"{'Property':<25} {'gw_remnant':>15} {'NRSur3dq8Remnant':>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}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0dc491c",
   "metadata": {},
   "source": [
    "### 2. Non-spinning binary (q = 4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "766841e4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NRHybSur3dq8 time grid: [-29350.61, 135.49] M\n",
      "Output time grid: [-5000.00, 49.90] M\n"
     ]
    }
   ],
   "source": [
    "q = 4.0\n",
    "chi1 = [0, 0, 0]\n",
    "chi2 = [0, 0, 0]\n",
    "\n",
    "times, h = waveform_generator.generate_nrhybsur3dq8(sur, q=q, chi1=chi1, chi2=chi2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "645480ef",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================================\n",
      "Remnant Properties Summary\n",
      "==================================================\n",
      "Mass ratio                    : 4.000\n",
      "Initial mass                  : 1.00000000 M\n",
      "Total energy radiated         : 0.01679011 M\n",
      "Peak luminosity               : 0.00036349\n",
      "Remnant mass                  : 0.97813243 M\n",
      "Remnant spin (dimensionless)  : 0.47088134\n",
      "Remnant spin vector (x,y,z)  : (0.00000000, 0.00000000, 0.47088134)\n",
      "Remnant kick velocity         : 0.00049463 c\n",
      "Remnant kick velocity         : 148.29 km/s\n",
      "Remnant kick vector (x,y,z)  : (0.00008521, 0.00048724, 0.00000000) c\n",
      "Remnant displacement (x,y,z) : (0.00189779, 0.01756277, 0.00000000) M\n",
      "==================================================\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/9k/xlxjfz4d2cz1063spty94p0r0000gq/T/ipykernel_23610/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": "6b71debd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Property                       gw_remnant     NRSur3dq8Remnant\n",
      "--------------------------------------------------------------\n",
      "Remnant mass [M]                 0.978132             0.977807\n",
      "Remnant spin (z)                 0.470881             0.471630\n",
      "Remnant |chi|                    0.470881             0.471630\n",
      "Kick velocity [c]                0.000495             0.000510\n"
     ]
    }
   ],
   "source": [
    "compare_remnants(calc_ns, fit, q, chi1, chi2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65875f70",
   "metadata": {},
   "source": [
    "### 3. Aligned-spin binary (q = 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ceb90c6f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NRHybSur3dq8 time grid: [-27121.80, 135.30] M\n",
      "Output time grid: [-5000.00, 49.90] M\n"
     ]
    }
   ],
   "source": [
    "q = 3.0\n",
    "chi1 = [0, 0, 0.5]\n",
    "chi2 = [0, 0, 0.3]\n",
    "\n",
    "times, h = waveform_generator.generate_nrhybsur3dq8(sur, q=q, chi1=chi1, chi2=chi2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9c2f2ded",
   "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.03452028 M\n",
      "Peak luminosity               : 0.00072088\n",
      "Remnant mass                  : 0.95937439 M\n",
      "Remnant spin (dimensionless)  : 0.76174956\n",
      "Remnant spin vector (x,y,z)  : (0.00000000, 0.00000000, 0.76174956)\n",
      "Remnant kick velocity         : 0.00022711 c\n",
      "Remnant kick velocity         : 68.09 km/s\n",
      "Remnant kick vector (x,y,z)  : (-0.00006610, 0.00021728, 0.00000000) c\n",
      "Remnant displacement (x,y,z) : (-0.00279605, 0.01463853, 0.00000000) M\n",
      "==================================================\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/9k/xlxjfz4d2cz1063spty94p0r0000gq/T/ipykernel_23610/1722731926.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_al = GWRemnantCalculator(time=times, h_dict=h, q=q, chi1=chi1, chi2=chi2)\n"
     ]
    }
   ],
   "source": [
    "calc_al = GWRemnantCalculator(time=times, h_dict=h, q=q, chi1=chi1, chi2=chi2)\n",
    "calc_al.print_remnants()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "895c39ee",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Property                       gw_remnant     NRSur3dq8Remnant\n",
      "--------------------------------------------------------------\n",
      "Remnant mass [M]                 0.959374             0.959083\n",
      "Remnant spin (z)                 0.761750             0.763491\n",
      "Remnant |chi|                    0.761750             0.763491\n",
      "Kick velocity [c]                0.000227             0.000223\n"
     ]
    }
   ],
   "source": [
    "compare_remnants(calc_al, fit, q, chi1, chi2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d0b9dbd",
   "metadata": {},
   "source": [
    "### 4. Anti-aligned-spin binary (q = 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d90f8daf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NRHybSur3dq8 time grid: [-22204.19, 135.11] M\n",
      "Output time grid: [-5000.00, 49.90] M\n"
     ]
    }
   ],
   "source": [
    "q = 2.0\n",
    "chi1 = [0, 0, 0.5]\n",
    "chi2 = [0, 0, -0.5]\n",
    "\n",
    "times, h = waveform_generator.generate_nrhybsur3dq8(sur, q=q, chi1=chi1, chi2=chi2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "4a840b4e",
   "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.03967582 M\n",
      "Peak luminosity               : 0.00092379\n",
      "Remnant mass                  : 0.95352851 M\n",
      "Remnant spin (dimensionless)  : 0.74837908\n",
      "Remnant spin vector (x,y,z)  : (0.00000000, 0.00000000, 0.74837908)\n",
      "Remnant kick velocity         : 0.00030504 c\n",
      "Remnant kick velocity         : 91.45 km/s\n",
      "Remnant kick vector (x,y,z)  : (0.00029094, -0.00009168, 0.00000000) c\n",
      "Remnant displacement (x,y,z) : (0.01533150, -0.00358134, 0.00000000) M\n",
      "==================================================\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/9k/xlxjfz4d2cz1063spty94p0r0000gq/T/ipykernel_23610/2141681472.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_aa = GWRemnantCalculator(time=times, h_dict=h, q=q, chi1=chi1, chi2=chi2)\n"
     ]
    }
   ],
   "source": [
    "calc_aa = GWRemnantCalculator(time=times, h_dict=h, q=q, chi1=chi1, chi2=chi2)\n",
    "calc_aa.print_remnants()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c686a9a7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Property                       gw_remnant     NRSur3dq8Remnant\n",
      "--------------------------------------------------------------\n",
      "Remnant mass [M]                 0.953529             0.953202\n",
      "Remnant spin (z)                 0.748379             0.750797\n",
      "Remnant |chi|                    0.748379             0.750797\n",
      "Kick velocity [c]                0.000305             0.000335\n"
     ]
    }
   ],
   "source": [
    "compare_remnants(calc_aa, fit, q, chi1, chi2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e06693f-47f4-4c7b-af91-6cb5c628a55b",
   "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
}