{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bbbaafbb-fd7d-4b73-a970-93506ba35d71",
   "metadata": {},
   "source": [
    "# Title page\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf6b8673-a2b5-4b50-ad1b-e90671f57470",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from binarycpython.utils.grid import Population\n",
    "from binarycpython.utils.functions import (\n",
    "    get_help_all,\n",
    "    get_help,\n",
    "    create_hdf5,\n",
    "    output_lines,\n",
    ")\n",
    "from binarycpython.utils.custom_logging_functions import temp_dir\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79ab50b7-591f-4883-af09-116d1835a751",
   "metadata": {},
   "outputs": [],
   "source": [
    "#########################################################\n",
    "# This file serves as an example for running a population.\n",
    "# The use of help(<function>) is a good way to inspect what parameters are there to use\n",
    "#########################################################\n",
    "\n",
    "# Create population object\n",
    "example_pop = Population()\n",
    "\n",
    "# If you want verbosity, set this before other things\n",
    "example_pop.set(verbose=1)\n",
    "\n",
    "# Setting values can be done via .set(<parameter_name>=<value>)\n",
    "# Values that are known to be binary_c_parameters are loaded into bse_options.\n",
    "# Those that are present in the default grid_options are set in grid_options\n",
    "# All other values that you set are put in a custom_options dict\n",
    "example_pop.set(\n",
    "    # binary_c physics options\n",
    "    M_1=10,  # bse_options\n",
    "    separation=0,  # bse_options\n",
    "    orbital_period=45000000080,  # bse_options\n",
    "    max_evolution_time=15000,  # bse_options\n",
    "    eccentricity=0.02,  # bse_options\n",
    "    # Set companion to low mass\n",
    "    M_2=0.08,  # Since in the example we run a single system, we should set the companion mass here. If we donm't do this, the code will complain.\n",
    "    # grid_options\n",
    "    amt_cores=2,  # grid_options\n",
    "    verbose=1,  # verbosity. Not fully configured correctly yet but having it value of 1 prints alot of stuff\n",
    "    # Custom options # TODO: need to be set in grid_options probably\n",
    "    data_dir=os.path.join(\n",
    "        temp_dir(), \"example_python_population_result\"\n",
    "    ),  # custom_options\n",
    "    base_filename=\"example_pop.dat\",  # custom_options\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c986215-93b1-4e30-ad79-f7c397e9ff7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_function(self, output):\n",
    "    # EXAMPLE PARSE_FUNCTION\n",
    "\n",
    "    # extract info from the population instance\n",
    "\n",
    "    # Get some information from the\n",
    "    data_dir = self.custom_options[\"data_dir\"]\n",
    "    base_filename = self.custom_options[\"base_filename\"]\n",
    "\n",
    "    # Check directory, make if necessary\n",
    "    os.makedirs(data_dir, exist_ok=True)\n",
    "\n",
    "    seperator = \" \"\n",
    "\n",
    "    # Create filename\n",
    "    outfilename = os.path.join(data_dir, base_filename)\n",
    "\n",
    "    parameters = [\"time\", \"mass\", \"zams_mass\", \"probability\", \"radius\", \"stellar_type\"]\n",
    "\n",
    "    # Go over the output.\n",
    "    for el in output_lines(output):\n",
    "        headerline = el.split()[0]\n",
    "\n",
    "        # CHeck the header and act accordingly\n",
    "        if headerline == \"MY_STELLAR_DATA\":\n",
    "            values = el.split()[1:]\n",
    "            print(values)\n",
    "\n",
    "            if not len(parameters) == len(values):\n",
    "                print(\"Amount of column names isnt equal to amount of columns\")\n",
    "                raise ValueError\n",
    "\n",
    "            if not os.path.exists(outfilename):\n",
    "                with open(outfilename, \"w\") as f:\n",
    "                    f.write(seperator.join(parameters) + \"\\n\")\n",
    "\n",
    "            with open(outfilename, \"a\") as f:\n",
    "                f.write(seperator.join(values) + \"\\n\")\n",
    "\n",
    "\n",
    "\n",
    "# Creating a parsing function\n",
    "example_pop.set(\n",
    "    parse_function=parse_function,  # Setting the parse function thats used in the evolve_population\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fa5188e-312f-4a05-aeda-0bd4c00629c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Custom logging\n",
    "# Log the moment when the star turns into neutron\n",
    "example_pop.set(\n",
    "    C_logging_code=\"\"\"\n",
    "if(stardata->star[0].stellar_type >= 13)    \n",
    "{\n",
    "    if (stardata->model.time < stardata->model.max_evolution_time)\n",
    "    {\n",
    "        Printf(\"MY_STELLAR_DATA %30.12e %g %g %g %g %d\\\\n\",\n",
    "            // \n",
    "            stardata->model.time, // 1\n",
    "            stardata->star[0].mass, // 2\n",
    "            stardata->common.zero_age.mass[0], // 4\n",
    "            stardata->model.probability, // 5\n",
    "            stardata->star[0].radius, // 6\n",
    "            stardata->star[0].stellar_type // 7\n",
    "      );\n",
    "    };\n",
    "    /* Kill the simulation to save time */\n",
    "    stardata->model.max_evolution_time = stardata->model.time - stardata->model.dtm;\n",
    "};\n",
    "\"\"\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd197154-a8ce-4865-8929-008d3483101a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add grid variables\n",
    "resolution = {\"M_1\": 20, \"q\": 20, \"per\": 40}\n",
    "\n",
    "# Mass\n",
    "example_pop.add_grid_variable(\n",
    "    name=\"lnm1\",\n",
    "    longname=\"Primary mass\",\n",
    "    valuerange=[2, 150],\n",
    "    resolution=\"{}\".format(resolution[\"M_1\"]),\n",
    "    spacingfunc=\"const(math.log(2), math.log(150), {})\".format(resolution[\"M_1\"]),\n",
    "    precode=\"M_1=math.exp(lnm1)\",\n",
    "    probdist=\"three_part_powerlaw(M_1, 0.1, 0.5, 1.0, 150, -1.3, -2.3, -2.3)*M_1\",\n",
    "    dphasevol=\"dlnm1\",\n",
    "    parameter_name=\"M_1\",\n",
    "    condition=\"\",  # Impose a condition on this grid variable. Mostly for a check for yourself\n",
    ")\n",
    "\n",
    "# # Mass ratio\n",
    "# test_pop.add_grid_variable(\n",
    "#     name=\"q\",\n",
    "#     longname=\"Mass ratio\",\n",
    "#     valuerange=[\"0.1/M_1\", 1],\n",
    "#     resolution=\"{}\".format(resolution['q']),\n",
    "#     spacingfunc=\"const(0.1/M_1, 1, {})\".format(resolution['q']),\n",
    "#     probdist=\"flatsections(q, [{'min': 0.1/M_1, 'max': 1.0, 'height': 1}])\",\n",
    "#     dphasevol=\"dq\",\n",
    "#     precode=\"M_2 = q * M_1\",\n",
    "#     parameter_name=\"M_2\",\n",
    "#     condition=\"\",  # Impose a condition on this grid variable. Mostly for a check for yourself\n",
    "# )\n",
    "\n",
    "# #\n",
    "# test_pop.add_grid_variable(\n",
    "#    name=\"log10per\", # in days\n",
    "#    longname=\"log10(Orbital_Period)\",\n",
    "#    valuerange=[0.15, 5.5],\n",
    "#    resolution=\"{}\".format(resolution[\"per\"]),\n",
    "#    spacingfunc=\"const(0.15, 5.5, {})\".format(resolution[\"per\"]),\n",
    "#    precode=\"\"\"orbital_period = 10** log10per\n",
    "# sep = calc_sep_from_period(M_1, M_2, orbital_period)\n",
    "# sep_min = calc_sep_from_period(M_1, M_2, 10**0.15)\n",
    "# sep_max = calc_sep_from_period(M_1, M_2, 10**5.5)\"\"\",\n",
    "#    probdist=\"sana12(M_1, M_2, sep, orbital_period, sep_min, sep_max, math.log10(10**0.15), math.log10(10**5.5), -0.55)\",\n",
    "#    parameter_name=\"orbital_period\",\n",
    "#    dphasevol=\"dlog10per\",\n",
    "# )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c96b151-8c3b-4479-969f-6c2f898497a5",
   "metadata": {},
   "source": [
    "Exporting of all the settings can be done with .export_all_info()\n",
    "on default it exports everything, but can be supressed by turning it off:\n",
    "  population settings (bse_options, grid_options, custom_options), turn off with include_population\n",
    "      settings=False\n",
    "  binary_c_defaults (all the commandline arguments that binary c accepts, and their defaults).\n",
    "      turn off with include_binary_c_defaults=False\n",
    "  include_binary_c_version_info (all the compilation info, and information about the compiled\n",
    "      parameters), turn off with include_binary_c_version_info=False\n",
    "  include_binary_c_help_all (all the help information for all the binary_c parameters),\n",
    "      turn off with include_binary_c_help_all=Fase\n",
    "On default it will write this to the custom_options['data_dir'], but that can be overriden by\n",
    "  setting use_datadir=False and providing an outfile=<>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ba83443c-64b4-4126-90b3-b1ca76c0f73d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'/tmp/binary_c_python/example_python_population_result/example_pop_settings.json'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "example_pop.export_all_info()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef5b51a1-e56d-4bec-a5c7-b9d5bceeceba",
   "metadata": {},
   "source": [
    "Executing a single system\n",
    "This uses the M_1 orbital period etc set with the set function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c0a9aed7-1fc9-4fe3-8bcf-4e987bfc00a9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Creating the code for the shared library for the custom logging\n",
      "['2.773586668293e+01', '1.33526', '10', '1', '1.72498e-05', '13']\n",
      "None\n"
     ]
    }
   ],
   "source": [
    "output = example_pop.evolve_single()\n",
    "print(output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ea376c1-1e92-45af-8cab-9d7fdca564eb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "## Executing a population\n",
    "## This uses the values generated by the grid_variables\n",
    "example_pop.evolve()  # TODO: update this function call"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce937b07-7292-4bb1-b2e7-d69e51bdff6e",
   "metadata": {
    "tags": []
   },
   "source": [
    "Wrapping up the results to an hdf5 file can be done by using the create_hdf5\n",
    "(<directory containing data and settings>) This function takes the settings file\n",
    "(ending in _settings.json) and the data files (ending in .dat) from the data_dir\n",
    "and packing them into an hdf5 file, which is then written into the same data_dir directory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ad444138-f854-4fdb-8e18-49e35ac9c7e2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Creating /tmp/binary_c_python/example_python_population_result/example_pop.hdf5\n",
      "Adding settings to HDF5 file\n",
      "Adding data to HDF5 file\n"
     ]
    }
   ],
   "source": [
    "create_hdf5(data_dir=example_pop.custom_options[\"data_dir\"], name=\"example_pop.hdf5\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34da69bc-fe60-4a9d-b07a-f6cd2f216964",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
