{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a1",
   "metadata": {},
   "source": [
    "## Computing Remnant Properties from pySEOBNR Waveforms\n",
    "\n",
    "This notebook generates waveforms using the `pySEOBNR` effective-one-body models (SEOBNRv5HM, SEOBNRv5PHM, SEOBNRv5EHM), computes remnant properties with `gw_remnant`, and compares against NR surrogate remnant fits.\n",
    "\n",
    "Requires `pip install gw_remnant[eob]`.\n",
    "\n",
    "Contact: Tousif Islam [tousifislam24@gmail.com]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1",
   "metadata": {},
   "source": [
    "### 1. Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c1",
   "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 lal\n",
    "import surfinBH\n",
    "import gwtools\n",
    "from pyseobnr import GenerateWaveform\n",
    "from gw_remnant.gw_remnant_calculator import GWRemnantCalculator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded NRSur3dq8Remnant fit.\n",
      "Loaded NRSur7dq4Remnant fit.\n"
     ]
    }
   ],
   "source": [
    "fit_3dq8 = surfinBH.LoadFits('NRSur3dq8Remnant')\n",
    "fit_7dq4 = surfinBH.LoadFits('NRSur7dq4Remnant')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "M_TOTAL = 70.0   # total mass in Msun (arbitrary, cancels in geometric units)\n",
    "DISTANCE = 100.0  # distance in Mpc (arbitrary, cancels in geometric units)\n",
    "\n",
    "\n",
    "def generate_seobnr(approximant, q, chi1, chi2, e_ref=0.0):\n",
    "    \"\"\"Generate a pySEOBNR waveform and return in geometric units (M=1).\"\"\"\n",
    "    m1 = M_TOTAL * q / (1 + q)\n",
    "    m2 = M_TOTAL / (1 + q)\n",
    "\n",
    "    params = {\n",
    "        \"mass1\": m1, \"mass2\": m2,\n",
    "        \"spin1x\": chi1[0], \"spin1y\": chi1[1], \"spin1z\": chi1[2],\n",
    "        \"spin2x\": chi2[0], \"spin2y\": chi2[1], \"spin2z\": chi2[2],\n",
    "        \"approximant\": approximant,\n",
    "        \"f22_start\": 20.0,\n",
    "        \"distance\": DISTANCE,\n",
    "        \"deltaT\": 1.0 / 4096.0,\n",
    "    }\n",
    "    if e_ref > 0:\n",
    "        params[\"eccentricity\"] = e_ref\n",
    "\n",
    "    gen = GenerateWaveform(params)\n",
    "    t_si, hlm_si = gen.generate_td_modes()\n",
    "\n",
    "    # Convert SI -> geometric units (M=1)\n",
    "    Mtot_s = M_TOTAL * lal.MTSUN_SI\n",
    "    D_m = DISTANCE * 1e6 * lal.PC_SI\n",
    "    Mtot_m = M_TOTAL * lal.MRSUN_SI\n",
    "\n",
    "    t = t_si / Mtot_s\n",
    "    h_dict = {(l, m): -h * D_m / Mtot_m for (l, m), h in hlm_si.items()}\n",
    "\n",
    "    print(f\"{approximant} time grid: [{t[0]:.1f}, {t[-1]:.1f}] M\")\n",
    "    return t, h_dict\n",
    "\n",
    "\n",
    "def compare_remnants(calc, fit, q, chi1, chi2, omega0=None, fit_name=\"fit\"):\n",
    "    \"\"\"Print a side-by-side comparison of gw_remnant vs a surfinBH remnant fit.\"\"\"\n",
    "    if omega0 is not None:\n",
    "        mf, chif, vf, *_ = fit.all(q, chi1, chi2, omega0=omega0)\n",
    "    else:\n",
    "        mf, chif, vf, *_ = fit.all(q, chi1, chi2)\n",
    "\n",
    "    print(f\"{'Property':<25} {'gw_remnant':>15} {fit_name:>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)}\")\n",
    "    print(f\"{'':25} {np.array2string(chif, precision=4)}  ({fit_name})\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2",
   "metadata": {},
   "source": [
    "---\n",
    "### 2. Non-spinning (SEOBNRv5HM, q = 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SEOBNRv5HM time grid: [-2329.5, 347.8] M\n"
     ]
    }
   ],
   "source": [
    "q = 3.0\n",
    "chi1 = [0, 0, 0]\n",
    "chi2 = [0, 0, 0]\n",
    "\n",
    "times, h = generate_seobnr('SEOBNRv5HM', q, chi1, chi2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c2",
   "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.02140288 M\n",
      "Peak luminosity               : 0.00050797\n",
      "Remnant mass                  : 0.97180206 M\n",
      "Remnant spin (dimensionless)  : 0.53837510\n",
      "Remnant spin vector (x,y,z)  : (-0.00000000, 0.00000000, 0.53837510)\n",
      "Remnant kick velocity         : 0.00050084 c\n",
      "Remnant kick velocity         : 150.15 km/s\n",
      "Remnant kick vector (x,y,z)  : (0.00046378, 0.00018906, 0.00000000) c\n",
      "Remnant displacement (x,y,z) : (0.16253014, 0.05912293, 0.00000000) M\n",
      "==================================================\n",
      "Property                       gw_remnant     NRSur3dq8Remnant\n",
      "--------------------------------------------------------------\n",
      "Remnant mass [M]                 0.971802             0.971233\n",
      "Remnant spin (z)                 0.538375             0.540600\n",
      "Remnant |chi|                    0.538375             0.540600\n",
      "Kick velocity [c]                0.000501             0.000513\n",
      "Spin vector               [-7.4205e-20  6.1728e-21  5.3838e-01]\n",
      "                          [0.     0.     0.5406]  (NRSur3dq8Remnant)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/9k/xlxjfz4d2cz1063spty94p0r0000gq/T/ipykernel_23625/3118498222.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()\n",
    "compare_remnants(calc_ns, fit_3dq8, q, chi1, chi2, fit_name='NRSur3dq8Remnant')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3",
   "metadata": {},
   "source": [
    "---\n",
    "### 3. Aligned-spin (SEOBNRv5HM, q = 2, chi1z = 0.5, chi2z = -0.3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SEOBNRv5HM time grid: [-2257.0, 372.2] M\n"
     ]
    }
   ],
   "source": [
    "q = 2.0\n",
    "chi1 = [0, 0, 0.5]\n",
    "chi2 = [0, 0, -0.3]\n",
    "\n",
    "times, h = generate_seobnr('SEOBNRv5HM', q, chi1, chi2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c3",
   "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.03913652 M\n",
      "Peak luminosity               : 0.00092268\n",
      "Remnant mass                  : 0.95265535 M\n",
      "Remnant spin (dimensionless)  : 0.75606931\n",
      "Remnant spin vector (x,y,z)  : (-0.00000000, -0.00000000, 0.75606931)\n",
      "Remnant kick velocity         : 0.00019383 c\n",
      "Remnant kick velocity         : 58.11 km/s\n",
      "Remnant kick vector (x,y,z)  : (0.00019118, 0.00003196, -0.00000000) c\n",
      "Remnant displacement (x,y,z) : (0.07562675, 0.00941414, -0.00000000) M\n",
      "==================================================\n",
      "Property                       gw_remnant     NRSur3dq8Remnant\n",
      "--------------------------------------------------------------\n",
      "Remnant mass [M]                 0.952655             0.951674\n",
      "Remnant spin (z)                 0.756069             0.762036\n",
      "Remnant |chi|                    0.756069             0.762036\n",
      "Kick velocity [c]                0.000194             0.000277\n",
      "Spin vector               [-1.8888e-20 -5.4327e-21  7.5607e-01]\n",
      "                          [0.    0.    0.762]  (NRSur3dq8Remnant)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/9k/xlxjfz4d2cz1063spty94p0r0000gq/T/ipykernel_23625/604263903.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()\n",
    "compare_remnants(calc_al, fit_3dq8, q, chi1, chi2, fit_name='NRSur3dq8Remnant')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4",
   "metadata": {},
   "source": [
    "---\n",
    "### 4. Precessing (SEOBNRv5PHM, q = 2)\n",
    "\n",
    "**Note:** The PN expressions assume the initial orbital angular momentum is aligned with the z-axis. For precessing binaries this assumption breaks down, so the remnant spin is slightly off even though the remnant mass and kick velocity remain accurate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SEOBNRv5PHM time grid: [-2163.8, 365.5] M\n"
     ]
    }
   ],
   "source": [
    "q = 2.0\n",
    "chi1 = [0.4, 0.0, 0.3]\n",
    "chi2 = [0.0, 0.3, -0.2]\n",
    "\n",
    "times, h = generate_seobnr('SEOBNRv5PHM', q, chi1, chi2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c4",
   "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.03532244 M\n",
      "Peak luminosity               : 0.00085427\n",
      "Remnant mass                  : 0.95652174 M\n",
      "Remnant spin (dimensionless)  : 0.70962809\n",
      "Remnant spin vector (x,y,z)  : (0.14346747, 0.00813686, 0.70962809)\n",
      "Remnant kick velocity         : 0.00031231 c\n",
      "Remnant kick velocity         : 93.63 km/s\n",
      "Remnant kick vector (x,y,z)  : (-0.00024042, -0.00007480, -0.00018478) c\n",
      "Remnant displacement (x,y,z) : (-0.08045489, -0.02178497, -0.06562347) M\n",
      "==================================================\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/9k/xlxjfz4d2cz1063spty94p0r0000gq/T/ipykernel_23625/469808673.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_pr = GWRemnantCalculator(time=times, h_dict=h, q=q, chi1=chi1, chi2=chi2)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded NRSur7dq4 model\n",
      "Property                       gw_remnant     NRSur7dq4Remnant\n",
      "--------------------------------------------------------------\n",
      "Remnant mass [M]                 0.956522             0.954669\n",
      "Remnant spin (z)                 0.709628             0.709370\n",
      "Remnant |chi|                    0.724031             0.723153\n",
      "Kick velocity [c]                0.000312             0.002488\n",
      "Spin vector               [0.1435 0.0081 0.7096]\n",
      "                          [0.1382 0.0254 0.7094]  (NRSur7dq4Remnant)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded NRSur7dq4 model\n",
      "Property                       gw_remnant     NRSur7dq4Remnant\n",
      "--------------------------------------------------------------\n",
      "Remnant mass [M]                 0.956522             0.954669\n",
      "Remnant spin (z)                 0.709628             0.709370\n",
      "Remnant |chi|                    0.724031             0.723153\n",
      "Kick velocity [c]                0.000312             0.002488\n",
      "Spin vector               [0.1435 0.0081 0.7096]\n",
      "                          [0.1382 0.0254 0.7094]  (NRSur7dq4Remnant)\n"
     ]
    }
   ],
   "source": [
    "calc_pr = GWRemnantCalculator(time=times, h_dict=h, q=q, chi1=chi1, chi2=chi2)\n",
    "calc_pr.print_remnants()\n",
    "\n",
    "omega0 = abs(np.gradient(0.5 * gwtools.phase(h[(2, 2)])) / np.gradient(times))[0]\n",
    "compare_remnants(calc_pr, fit_7dq4, q, chi1, chi2, omega0=omega0,\n",
    "                 fit_name='NRSur7dq4Remnant')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5",
   "metadata": {},
   "source": [
    "---\n",
    "### 5. Eccentric (SEOBNRv5EHM, q = 2, e = 0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SEOBNRv5EHM time grid: [-1958.1, 352.4] M\n"
     ]
    }
   ],
   "source": [
    "q = 2.0\n",
    "chi1 = [0, 0, 0]\n",
    "chi2 = [0, 0, 0]\n",
    "e_ref = 0.1\n",
    "\n",
    "times, h = generate_seobnr('SEOBNRv5EHM', q, chi1, chi2, e_ref=e_ref)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c5",
   "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.03039380 M\n",
      "Peak luminosity               : 0.00075077\n",
      "Remnant mass                  : 0.96167657 M\n",
      "Remnant spin (dimensionless)  : 0.62364773\n",
      "Remnant spin vector (x,y,z)  : (-0.00000000, -0.00000000, 0.62364773)\n",
      "Remnant kick velocity         : 0.00047126 c\n",
      "Remnant kick velocity         : 141.28 km/s\n",
      "Remnant kick vector (x,y,z)  : (-0.00036052, 0.00030349, -0.00000000) c\n",
      "Remnant displacement (x,y,z) : (-0.10889326, 0.10845557, -0.00000000) M\n",
      "==================================================\n",
      "Property                       gw_remnant     NRSur3dq8Remnant\n",
      "--------------------------------------------------------------\n",
      "Remnant mass [M]                 0.961677             0.961195\n",
      "Remnant spin (z)                 0.623648             0.623424\n",
      "Remnant |chi|                    0.623648             0.623424\n",
      "Kick velocity [c]                0.000471             0.000445\n",
      "Spin vector               [-1.0200e-20 -4.5642e-22  6.2365e-01]\n",
      "                          [0.     0.     0.6234]  (NRSur3dq8Remnant)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/9k/xlxjfz4d2cz1063spty94p0r0000gq/T/ipykernel_23625/495517983.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_ecc = GWRemnantCalculator(time=times, h_dict=h, q=q, chi1=chi1, chi2=chi2,\n"
     ]
    }
   ],
   "source": [
    "calc_ecc = GWRemnantCalculator(time=times, h_dict=h, q=q, chi1=chi1, chi2=chi2,\n",
    "                               e_ref=e_ref)\n",
    "calc_ecc.print_remnants()\n",
    "compare_remnants(calc_ecc, fit_3dq8, q, chi1, chi2, fit_name='NRSur3dq8Remnant')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61c54ca0-ab1f-44aa-99aa-11459d20ad52",
   "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
}