{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a1b2c3d4",
   "metadata": {},
   "source": [
    "# Optimizing Initial Conditions for EOB Waveforms\n",
    "\n",
    "This tutorial demonstrates how to optimize the initial conditions of Effective One Body (EOB) waveforms to best match Numerical Relativity (NR) simulations.\n",
    "\n",
    "The optimizer searches for the best initial energy (E₀) and angular momentum (pₚₕ₀) that minimize the mismatch between EOB and NR waveforms.\n",
    "\n",
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5f6g7h8",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "%config InlineBackend.figure_format = 'retina'\n",
    "\n",
    "import os\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from PyART.analysis.opt_ic import Optimizer\n",
    "from PyART.catalogs.sxs import Waveform_SXS\n",
    "from PyART.catalogs.rit import Waveform_RIT\n",
    "from PyART.analysis.match import Matcher"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "i9j0k1l2",
   "metadata": {},
   "source": [
    "## Load NR Waveform\n",
    "\n",
    "First, we load a numerical relativity waveform that we want to match with EOB:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "m3n4o5p6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load an SXS waveform\n",
    "catalog = 'sxs'\n",
    "sim_id = 180\n",
    "\n",
    "if catalog == 'sxs':\n",
    "    ebbh = Waveform_SXS(\n",
    "        path='./', \n",
    "        download=True, \n",
    "        ID=sim_id, \n",
    "        order=\"Extrapolated_N3.dir\", \n",
    "        ellmax=7,\n",
    "        nu_rescale=True\n",
    "    )\n",
    "    # Remove junk radiation from the beginning\n",
    "    ebbh.cut(200)\n",
    "else:\n",
    "    # For RIT or other catalogs\n",
    "    ebbh = Waveform_RIT(\n",
    "        path='./local_data/rit/', \n",
    "        download=True, \n",
    "        ID=sim_id, \n",
    "        nu_rescale=True\n",
    "    )\n",
    "\n",
    "# Display metadata\n",
    "print('Waveform Metadata:')\n",
    "print('=' * 50)\n",
    "for k, v in ebbh.metadata.items():\n",
    "    print(f'{k:15s} : {v}')\n",
    "print('=' * 50)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "q7r8s9t0",
   "metadata": {},
   "source": [
    "## Configure Optimizer Settings\n",
    "\n",
    "The optimizer requires several settings:\n",
    "1. **Mismatch settings**: Parameters for computing mismatches\n",
    "2. **Minimizer**: Algorithm and parameters for optimization\n",
    "3. **Bounds**: Search ranges for initial conditions\n",
    "4. **Iteration settings**: How to adaptively adjust bounds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "u1v2w3x4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mismatch computation settings\n",
    "mm_settings = {\n",
    "    'cut_second_waveform': True,\n",
    "    'initial_frequency_mm': 10,\n",
    "    'M': 100,  # Total mass in solar masses\n",
    "    'final_frequency_mm': 1024,\n",
    "    'taper_alpha': 0.50,\n",
    "    'taper_start': 0.10,\n",
    "}\n",
    "\n",
    "# Add catalog-specific alignment settings\n",
    "if catalog == 'rit':\n",
    "    mm_settings['pre_align_shift'] = 100.\n",
    "elif catalog == 'sxs':\n",
    "    mm_settings['pre_align_shift'] = 0.\n",
    "\n",
    "# Optimizer settings\n",
    "kind_ic = 'E0pph0'  # Optimize E0 (energy) and pph0 (angular momentum)\n",
    "\n",
    "# Bounds for the optimization\n",
    "bounds = {'E0byM': [None, None], 'pph0': [None, None]}\n",
    "\n",
    "# Adaptive bounds iteration\n",
    "bounds_iter = {\n",
    "    'eps_initial': {'E0byM': 5e-3, 'pph0': 1e-2},\n",
    "    'eps_factors': {'E0byM': 4, 'pph0': 2},\n",
    "    'bad_mm': 1e-2,\n",
    "    'max_iter': 3\n",
    "}\n",
    "\n",
    "# Minimizer configuration (using dual annealing)\n",
    "minimizer = {\n",
    "    'kind': 'dual_annealing',\n",
    "    'opt_maxfun': 100,  # Maximum function evaluations\n",
    "    'opt_max_iter': 1,  # Maximum optimization iterations\n",
    "    'opt_seed': 190521\n",
    "}\n",
    "\n",
    "print(\"Optimizer configured with:\")\n",
    "print(f\"  Kind: {kind_ic}\")\n",
    "print(f\"  Minimizer: {minimizer['kind']}\")\n",
    "print(f\"  Target mass: {mm_settings['M']} Msun\")\n",
    "print(f\"  Good mismatch threshold: 5e-3\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "y5z6a7b8",
   "metadata": {},
   "source": [
    "## Run the Optimizer\n",
    "\n",
    "Now we create the optimizer and let it find the best initial conditions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9d0e1f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create and run optimizer\n",
    "opt = Optimizer(\n",
    "    ebbh,\n",
    "    kind_ic=kind_ic,\n",
    "    mm_settings=mm_settings,\n",
    "    use_nqc=False,\n",
    "    minimizer=minimizer,\n",
    "    opt_good_mm=5e-3,\n",
    "    opt_bounds=bounds,\n",
    "    bounds_iter=bounds_iter,\n",
    "    debug=False,\n",
    "    json_file=None,\n",
    "    overwrite=False\n",
    ")\n",
    "\n",
    "print(\"\\nOptimization Results:\")\n",
    "print('=' * 50)\n",
    "print(f\"Optimized E0/M    : {opt.opt_Waveform.metadata.get('E0byM', 'N/A')}\")\n",
    "print(f\"Optimized pph0    : {opt.opt_Waveform.metadata.get('pph0', 'N/A')}\")\n",
    "print(f\"Final mismatch    : {opt.opt_mismatch:.5e}\")\n",
    "print('=' * 50)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "g3h4i5j6",
   "metadata": {},
   "source": [
    "## Visualize: Mismatch vs Total Mass\n",
    "\n",
    "Let's see how the optimized waveform performs across different total masses:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "k7l8m9n0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute mismatch for a range of masses\n",
    "masses = np.linspace(20, 200, num=19)\n",
    "mm = np.zeros_like(masses)\n",
    "\n",
    "for i, M in enumerate(masses):\n",
    "    mm_settings['M'] = M\n",
    "    matcher = Matcher(ebbh, opt.opt_Waveform, settings=mm_settings)\n",
    "    mm[i] = matcher.mismatch\n",
    "    \n",
    "    if i % 5 == 0:\n",
    "        print(f'M = {M:6.1f} Msun, mismatch = {mm[i]:.3e}')\n",
    "\n",
    "# Plot\n",
    "plt.figure(figsize=(10, 6))\n",
    "plt.plot(masses, mm, linewidth=2, marker='o', markersize=6)\n",
    "plt.yscale('log')\n",
    "plt.xlabel(r'Total Mass $M$ [$M_\\odot$]', fontsize=16)\n",
    "plt.ylabel(r'Mismatch $\\bar{\\mathcal{F}}$', fontsize=16)\n",
    "plt.title('Optimized EOB vs NR Mismatch', fontsize=18)\n",
    "plt.ylim(1e-4, 1e-1)\n",
    "plt.grid(True, alpha=0.3, which='both')\n",
    "plt.axhline(y=5e-3, color='r', linestyle='--', alpha=0.5, label='Good mismatch threshold')\n",
    "plt.axhline(y=1e-2, color='orange', linestyle='--', alpha=0.5, label='Acceptable threshold')\n",
    "plt.legend(fontsize=12)\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "print(f'\\nMinimum mismatch: {mm.min():.3e} at M = {masses[mm.argmin()]:.1f} Msun')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "o1p2q3r4",
   "metadata": {},
   "source": [
    "## Compare Waveforms\n",
    "\n",
    "Let's visually compare the NR and optimized EOB waveforms:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "s5t6u7v8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get merger times\n",
    "nr_mrg, _, _, _ = ebbh.find_max()\n",
    "eob_mrg, _, _, _ = opt.opt_Waveform.find_max()\n",
    "\n",
    "plt.figure(figsize=(14, 5))\n",
    "\n",
    "# Real part\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.plot(ebbh.u - nr_mrg, ebbh.hlm[(2,2)]['real'], \n",
    "         label='NR', linewidth=2, alpha=0.8)\n",
    "plt.plot(opt.opt_Waveform.u - eob_mrg, opt.opt_Waveform.hlm[(2,2)]['real'], \n",
    "         label='Optimized EOB', linewidth=2, alpha=0.8)\n",
    "plt.xlabel('Time (M)', fontsize=14)\n",
    "plt.ylabel(r'Re[$h_{22}$]', fontsize=14)\n",
    "plt.title('Waveform Comparison', fontsize=16)\n",
    "plt.legend(fontsize=12)\n",
    "plt.grid(True, alpha=0.3)\n",
    "plt.xlim([-500, 100])\n",
    "\n",
    "# Amplitude\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.plot(ebbh.u - nr_mrg, ebbh.hlm[(2,2)]['A'], \n",
    "         label='NR', linewidth=2, alpha=0.8)\n",
    "plt.plot(opt.opt_Waveform.u - eob_mrg, opt.opt_Waveform.hlm[(2,2)]['A'], \n",
    "         label='Optimized EOB', linewidth=2, alpha=0.8)\n",
    "plt.xlabel('Time (M)', fontsize=14)\n",
    "plt.ylabel(r'$|h_{22}|$', fontsize=14)\n",
    "plt.title('Amplitude Comparison', fontsize=16)\n",
    "plt.legend(fontsize=12)\n",
    "plt.grid(True, alpha=0.3)\n",
    "plt.xlim([-500, 100])\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "w9x0y1z2",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "This tutorial demonstrated:\n",
    "- Loading NR waveforms from various catalogs (SXS, RIT)\n",
    "- Configuring the optimizer with appropriate settings\n",
    "- Running the optimization to find best-fit initial conditions\n",
    "- Evaluating the optimized waveform across different masses\n",
    "- Visualizing the agreement between NR and optimized EOB\n",
    "\n",
    "### Key Points\n",
    "\n",
    "- The optimizer uses dual annealing by default for robust global optimization\n",
    "- Initial conditions (E₀, pₚₕ₀) significantly affect EOB waveform quality\n",
    "- Mismatches typically vary with total mass\n",
    "- Good mismatches (< 5×10⁻³) indicate excellent EOB-NR agreement\n",
    "\n",
    "### Next Steps\n",
    "\n",
    "- Try different catalogs (RIT, CoRe, ICCUB)\n",
    "- Experiment with different initial condition parameterizations (e0f0 vs E0pph0)\n",
    "- Use NQC (Non-Quasi-Circular) corrections for eccentric binaries\n",
    "- Optimize for multiple modes simultaneously"
   ]
  }
 ],
 "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.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}