{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evolve individual binaries 🐞"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It can be extremely useful to evolve a single binary from specific initial conditions or to start at a specific evolutionary state.\n",
    "This is useful for debugging binaries and checking if your custom steps or flow are correctly working.\n",
    "\n",
    "This tutorial will cover:\n",
    "\n",
    "- How to initialise a zero-age main sequence binary (ZAMS).\n",
    "- How to re-evolve a binary in an existing population.\n",
    "- How to evolve a binary from specific initial states.\n",
    "\n",
    "\n",
    "## Evolve a ZAMS binary\n",
    "\n",
    "To evolve a binary from ZAMS, we will:\n",
    "\n",
    "\n",
    "1. Load the standard simulation properties.\n",
    "2. Load the steps. \n",
    "3. Initialise the binary.\n",
    "4. Evolve it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "from posydon.config import PATH_TO_POSYDON\n",
    "\n",
    "path_to_params = os.path.join(PATH_TO_POSYDON, \"posydon/popsyn/population_params_default.ini\")\n",
    "shutil.copyfile(path_to_params, './population_params.ini')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load the function to load the simulation properties from the ini file\n",
    "from posydon.popsyn.io import simprop_kwargs_from_ini\n",
    "from posydon.binary_evol.simulationproperties import SimulationProperties\n",
    "\n",
    "# Load the simulation properties from the default ini file. \n",
    "sim_kwargs = simprop_kwargs_from_ini('population_params.ini')\n",
    "# manually add the metallicity to each step that requires it\n",
    "metallicity = {'metallicity':1}\n",
    "\n",
    "sim_kwargs['step_HMS_HMS'][1].update(metallicity)\n",
    "sim_kwargs['step_CO_HeMS'][1].update(metallicity)\n",
    "sim_kwargs['step_CO_HMS_RLO'][1].update(metallicity)\n",
    "sim_kwargs['step_CO_HeMS_RLO'][1].update(metallicity)\n",
    "sim_kwargs['step_detached'][1].update(metallicity)\n",
    "sim_kwargs['step_disrupted'][1].update(metallicity)\n",
    "sim_kwargs['step_merged'][1].update(metallicity)\n",
    "sim_kwargs['step_initially_single'][1].update(metallicity)\n",
    "\n",
    "sim_prop = SimulationProperties(**sim_kwargs)\n",
    "# Load the steps and required data\n",
    "sim_prop.load_steps(verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load the binary and single star classes\n",
    "from posydon.binary_evol.singlestar import SingleStar\n",
    "from posydon.binary_evol.binarystar import BinaryStar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "STAR1 = SingleStar(**{'mass': 30.782576, \n",
    "                      'state': 'H-rich_Core_H_burning'})\n",
    "STAR2 = SingleStar(**{'mass':20.273864,\n",
    "                      'state': 'H-rich_Core_H_burning'})\n",
    "\n",
    "BINARY = BinaryStar(STAR1, STAR2,  \n",
    "                    **{'time': 0.0, 'state': 'detached', 'event': 'ZAMS', 'orbital_period':3513.150157, 'eccentricity': 0.0},\n",
    "                    properties = sim_prop)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note: depending on the evolution of the binary, you might get some warnings about the Roche lobe calculation.\n",
    "BINARY.evolve()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You're now able to access the binary and evolutionary information. Using the function, `BinaryStar.to_df()` we can inspect the evolution of the systm.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# we will add the step names to the dataframe\n",
    "col = ['time', 'step_names', 'state', 'event', 'orbital_period', 'eccentricity', 'S1_state', 'S2_state', 'S1_mass', 'S2_mass']\n",
    "\n",
    "BINARY.to_df(extra_columns={'step_names':'string'})[col]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Re-evolving the Binary\n",
    "\n",
    "The above binary might have been disrupted in the first or second SN due to the strength of the supernova kick.\n",
    "You can restore the binary completely or to a specific state and re-evolve it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# restore to the original state\n",
    "BINARY.restore()\n",
    "BINARY.to_df()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BINARY.evolve()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Even thought we reset the binary, the natal kick properties stored in the natal kick array are not reset. This is because we want to be able to re-evolve the binary to the same final state. In case you want to reset the natal kick array values, then set the list to `None` values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# kick magnitude km/s, azimuthal angle rad, polar angle rad and mean anomaly rad (TODO check documentation)\n",
    "print(BINARY.star_1.natal_kick_array)\n",
    "print(BINARY.star_2.natal_kick_array)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This will be exactly the same evolution as the previous one\n",
    "BINARY.to_df(extra_columns={'step_names':'string'})[col]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Changing the kick\n",
    "\n",
    "We will restore the `BINARY` to its initial state and increase the first kick velocity to disrupt the binary after the first supernova."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BINARY.restore()\n",
    "BINARY.to_df()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# [velocity, azimuthal angle, polar angle, phase]\n",
    "BINARY.star_1.natal_kick_array = [1000., 0.858129274334538, 1.9157148786534735, 1.8675467897282945]\n",
    "BINARY.star_2.natal_kick_array = [None, None, None, None] # default\n",
    "print(BINARY.star_1.natal_kick_array)\n",
    "print(BINARY.star_2.natal_kick_array)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BINARY.evolve()\n",
    "BINARY.to_df(extra_columns={'step_names':'string'})[col]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Evolve from specific step\n",
    "\n",
    "You can provide the row number (starting at 0), to reset the binary to that specific evolutionary phase.\n",
    "This can be useful for debugging a specific evolutionary state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "BINARY.restore(3)\n",
    "BINARY.to_df()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading a binary from an existing population"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assume we run the population somewhere else and you have the h5 file associated with it, and you want to load the population in memory and re-evolve a binary. \n",
    "We will have to load the `population_params.ini` file associated with the Population file to re-evolve the binary in the same way.\n",
    "\n",
    "We will use the `BBH_contact.h5` file in the example dataset and use the standard `population_params.ini` SimulationProperties from earlier in this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from posydon.popsyn.synthetic_population import Population\n",
    "\n",
    "\n",
    "data_path = f'PATH/TO/TUTORIAL/POPULATIONS/'\n",
    "\n",
    "pop = Population(f'{data_path}/BBH_contact.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 143,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Because the original population contains two extre columns. These have to be defined manually here.\n",
    "# Otherwise, the reseting of the binaries fails.\n",
    "BINARY = BinaryStar.from_df(pop.history[0],\n",
    "                            extra_columns={'step_names':'string','step_times':'float'})\n",
    "\n",
    "# you have to set the simulation properties for the binary\n",
    "BINARY.properties = sim_prop\n",
    "\n",
    "# Set the natal kick arrays from the oneline to get the exact same evolution as the original binary\n",
    "BINARY.star_1.natal_kick_array = pop.oneline[0][['S1_natal_kick_array_0', 'S1_natal_kick_array_1', 'S1_natal_kick_array_2', 'S1_natal_kick_array_3']].values.flatten()\n",
    "BINARY.star_2.natal_kick_array = pop.oneline[0][['S2_natal_kick_array_0', 'S2_natal_kick_array_1', 'S2_natal_kick_array_2', 'S2_natal_kick_array_3']].values.flatten()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BINARY.to_df(extra_columns={'step_names':'string'})[col]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 146,
   "metadata": {},
   "outputs": [],
   "source": [
    "BINARY.restore()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BINARY.evolve()\n",
    "BINARY.to_df(extra_columns={'step_names':'string'})[col]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evolving a Binary Star starting from an arbitrary state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you want to evolve a custom binary star, you can do so by crafting yourself the `BinaryStar` object and providing the `SimulationProperties` object. For example, let's evolve a neutron star with a low mass helium star in Roche lobe overflow.\n",
    "And start a binary in the detached step, which requires a bit of additional input data.\n",
    "\n",
    "We will use the standard SimulationProperties, loaded earlier in this notebook.\n",
    "\n",
    "\n",
    "<div class=\"alert alert-warning\"><b>Caution!</b> \n",
    "\n",
    "While you can evolve a binary from an arbitrary state, you will need to provide data on the internal structure of the star, if you're starting in the detached step.\n",
    "Otherwise, the matching to a single star model will not work!\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# we initialise the binary to be tight and Roche lobe overflowing in a circular orbit.\n",
    "binary = BinaryStar(star_1 = SingleStar(**{'state' : 'NS',\n",
    "                                           'mass' : 1.1,\n",
    "                                           'spin' : 0.,}),\n",
    "                    star_2 = SingleStar(**{'state' : 'H-rich_Core_H_burning',\n",
    "                                           'mass' : 2.5,\n",
    "                                           'natal_kick_array' : [10., 0., 0., 0.]}),\n",
    "                    **{'time' : 0.,\n",
    "                       'state' : 'RLO2',\n",
    "                       'event' : 'oRLO2',\n",
    "                       'orbital_period' : 1.,\n",
    "                       'eccentricity' : 0.},\n",
    "                    properties = sim_prop,\n",
    "                    )\n",
    "binary.evolve()\n",
    "binary.to_df(extra_columns={'step_names':'string'})[col]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# we perform the same evolution, but with a stripped helium star to go to the CO_HeMS_RLO grid\n",
    "binary = BinaryStar(star_1 = SingleStar(**{'state' : 'NS',\n",
    "                                           'mass' : 1.1,\n",
    "                                           'spin' : 0.,}),\n",
    "                    star_2 = SingleStar(**{'state' : 'stripped_He_Core_He_burning',\n",
    "                                           'mass' : 2.5,\n",
    "                                           'natal_kick_array' : [10., 0., 0., 0.]}),\n",
    "                    **{'time' : 0.,\n",
    "                       'state' : 'RLO2',\n",
    "                       'event' : 'oRLO2',\n",
    "                       'orbital_period' : 1.,\n",
    "                       'eccentricity' : 0.},\n",
    "                    properties = sim_prop,\n",
    "                    )\n",
    "binary.evolve()\n",
    "binary.to_df(extra_columns={'step_names':'string'})[col]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will change the binary to a much wider system and have it start in the detached state.\n",
    "This requires additional starting information to allow the star to be matched to the single star grids."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 163,
   "metadata": {},
   "outputs": [],
   "source": [
    "from posydon.utils.constants import Zsun\n",
    "from posydon.utils.common_functions import orbital_separation_from_period\n",
    "\n",
    "import numpy as np\n",
    "Z = 1.0\n",
    "\n",
    "# Setup the central abundances\n",
    "zams_table = {2.: 2.915e-01,\n",
    "                      1.: 2.703e-01,\n",
    "                      0.45: 2.586e-01,\n",
    "                      0.2: 2.533e-01,\n",
    "                      0.1: 2.511e-01,\n",
    "                      0.01: 2.492e-01,\n",
    "                      0.001: 2.49e-01,\n",
    "                      0.0001: 2.49e-01}\n",
    "\n",
    "Y = zams_table[Z]\n",
    "Z = Z*Zsun\n",
    "X = 1 - Y - Z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "STAR1 = SingleStar(**{'mass':1.2,\n",
    "                      'state': 'NS'})\n",
    "\n",
    "STAR2 = SingleStar(**{'mass': 17.782576, \n",
    "                      'state': 'H-rich_Core_H_burning',\n",
    "                      # add the metallicity and central abundances\n",
    "                      'metallicity':Z, \n",
    "                      'center_h1':X,\n",
    "                      'center_he4':Y, \n",
    "                      # add a numerical value for the radius\n",
    "                      'log_R': np.nan,\n",
    "                      # add the helium core mass\n",
    "                      'he_core_mass': 0.0,\n",
    "                      })\n",
    "\n",
    "binary = BinaryStar(STAR1, STAR2,\n",
    "                    **{'time' : 0.,\n",
    "                       'state' : 'detached',\n",
    "                       'event' : None,\n",
    "                       'orbital_period' : 5000.,\n",
    "                        # calculate the separation; current bug that the separation is not automatically calculated if orbital period is given\n",
    "                       'separation': orbital_separation_from_period(5000., 17.782576, 1.2),\n",
    "                       'eccentricity' : 0.},\n",
    "                    properties = sim_prop,\n",
    "                    )\n",
    "binary.evolve()\n",
    "binary.to_df(extra_columns={'step_names':'string'})[col]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "posydon_env",
   "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}