{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example of modelling: TOI-270"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "plt.rcParams[\"figure.figsize\"] = (10,5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To define a star:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ravest.model import Star\n",
    "\n",
    "star = Star(name=\"TOI-270\", mass=0.386)\n",
    "print(star)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's define some planets. These parameter values for TOI-270 b, c, and d were obtained from [Van Eylen et al. 2021](https://doi.org/10.1093/mnras/stab2143). The full list of parameterisations supported can be seen at `ravest.param.ALLOWED_PARAMETERISATONS`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ravest.model import Planet\n",
    "from ravest.param import ALLOWED_PARAMETERISATIONS, Parameterisation\n",
    "\n",
    "print(\"Allowed parameterisations:\", ALLOWED_PARAMETERISATIONS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parameterisation = Parameterisation(\"P K e w Tc\")\n",
    "\n",
    "planet_b = Planet(letter=\"b\", parameterisation=parameterisation, params={\"P\": 3.3601538, \"K\": 1.27, \"e\": 0.034, \"w\": 0.0,   \"Tc\": 2458387.09505})\n",
    "planet_c = Planet(\"c\", parameterisation, {\"P\": 5.6605731, \"K\": 4.16, \"e\": 0.027, \"w\": 0.2, \"Tc\": 2458389.50285})\n",
    "planet_d = Planet(\"d\", parameterisation, {\"P\": 11.379573, \"K\": 2.56, \"e\": 0.032, \"w\":-0.1, \"Tc\": 2458389.68186})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see what the radial velocity would be due to the effect of one of the individual planets:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "times = np.linspace(2458774, 2458824, 1000)\n",
    "\n",
    "plt.figure()\n",
    "plt.title(planet_d)\n",
    "plt.xlabel(\"Time [days]\")\n",
    "plt.ylabel(\"Radial velocity [m/s]\")\n",
    "plt.plot(times, planet_d.radial_velocity(times))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's add the planets to the star, and use that to model the combined radial velocity that we would see, due to the effect of all of the planets. Notice that we also define a `Trend` object. This can store a constant velocity offset $\\gamma$ ($\\rm{ms}^{-1}$), a linear trend $\\dot{\\gamma}$ ($\\rm{ms}^{-1}/\\rm{day}$), and a quadratic trend $\\ddot{\\gamma}$ ($\\rm{ms}^{-1}/\\rm{day}^2)$, which are often used to account for e.g. instrumental contributions or possible undetected long-period companions. `Trend` also requires a reference/zero-point time $t_0$ for the linear $\\dot\\gamma(t-t_0)$ and quadratic $\\ddot\\gamma(t-t_0)^2$ trend rates to be calculated from. I recommend to use the mean or median of the input times as $t_0$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ravest.model import Trend\n",
    "\n",
    "star.add_planet(planet_b)\n",
    "star.add_planet(planet_c)\n",
    "star.add_planet(planet_d)\n",
    "\n",
    "star.add_trend(Trend(params={\"g\":0, \"gd\":0, \"gdd\":0}, t0=np.mean(times)))\n",
    "\n",
    "plt.figure()\n",
    "plt.title(star)\n",
    "plt.xlabel(\"Time [days]\")\n",
    "plt.ylabel(\"Radial velocity [m/s]\")\n",
    "plt.plot(times, star.radial_velocity(times));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we had some observed data, we can plot the data on top of the system RV model, and look at the residuals and generate phase plots for each planet. As a quick demonstration, I've previously generated some fake data for the TOI-270 system, which we can load in and compare our model to."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fpath = Path.cwd() / \"example_data\" / \"TOI-270.csv\"\n",
    "df = pd.read_csv(fpath)\n",
    "ti = df[\"ti\"].to_numpy()\n",
    "rv = df[\"rv\"].to_numpy()\n",
    "err = df[\"err\"].to_numpy()\n",
    "tel = df[\"tel\"].to_numpy()\n",
    "star.phase_plot(ti, rv, err, tel)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you had some real data and you wanted to fit a planetary model to it, check out the [tutorial notebook on fitting](https://ravest.readthedocs.io/en/latest/Examples/example_fitting.html)."
   ]
  }
 ],
 "metadata": {
  "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}