{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# QGS model: Example of DiffEqPy usage"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preamble"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook shows how to use the [DiffEqPy](https://github.com/SciML/diffeqpy) package to integrate the qgs model and compare with the qgs Runge-Kutta integrator."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**You must finish the install of DiffEqPy manually before running this notebook!**\n",
    "\n",
    "The DiffEqPy package first installation step is done by Anaconda in the qgs environment but then you must install [Julia](https://julialang.org/downloads/) and follow the final manual installation instruction found in the link above.\n",
    "\n",
    "These can be summed up as opening a terminal and doing:\n",
    "```\n",
    "conda activate qgs\n",
    "python\n",
    "```\n",
    "and then inside the Python command line interface do:\n",
    "\n",
    "```\n",
    ">>> import diffeqpy\n",
    ">>> diffeqpy.install()\n",
    "```\n",
    "which will then finalize the installation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reinhold and Pierrehumbert 1982 model version"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook use the model version with a simple 2-layer channel QG atmosphere truncated at wavenumber 2 on a beta-plane with a simple orography (a montain and a valley). \n",
    "\n",
    "More detail can be found in the articles:\n",
    "* Reinhold, B. B., & Pierrehumbert, R. T. (1982). *Dynamics of weather regimes: Quasi-stationary waves and blocking*. Monthly Weather Review, **110** (9), 1105-1145. [doi:10.1175/1520-0493(1982)110%3C1105:DOWRQS%3E2.0.CO;2](https://doi.org/10.1175/1520-0493(1982)110%3C1105:DOWRQS%3E2.0.CO;2)\n",
    "* Cehelsky, P., & Tung, K. K. (1987). *Theories of multiple equilibria and weather regimes—A critical reexamination. Part II: Baroclinic two-layer models*. Journal of the atmospheric sciences, **44** (21), 3282-3303. [doi:10.1175/1520-0469(1987)044%3C3282%3ATOMEAW%3E2.0.CO%3B2](https://doi.org/10.1175/1520-0469(1987)044%3C3282%3ATOMEAW%3E2.0.CO%3B2)\n",
    "\n",
    "or in the documentation and on [readthedocs](https://qgs.readthedocs.io/en/latest/files/model/oro_model.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Modules import"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, setting the path and loading of some modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys, os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sys.path.extend([os.path.abspath('../')])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import Axes3D"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Initializing the random number generator (for reproducibility). -- Disable if needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(210217)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Importing the model's modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from qgs.params.params import QgParams\n",
    "from qgs.integrators.integrator import RungeKuttaIntegrator\n",
    "from qgs.functions.tendencies import create_tendencies\n",
    "from qgs.plotting.util import std_plot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Importing Julia DifferentialEquations.jl package"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from julia.api import Julia\n",
    "jl = Julia(compiled_modules=False)\n",
    "from diffeqpy import de"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Importing also Numba"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from numba import jit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ignoring Numba performance warnings. -- Comment to disable if you want to see them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Systems definition"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "General parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Time parameters\n",
    "dt = 0.1\n",
    "# Saving the model state n steps\n",
    "write_steps = 5\n",
    "\n",
    "number_of_trajectories = 1\n",
    "number_of_perturbed_trajectories = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Setting some model parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model parameters instantiation with some non-default specs\n",
    "model_parameters = QgParams({'phi0_npi': np.deg2rad(50.)/np.pi, 'hd':0.3})\n",
    "# Mode truncation at the wavenumber 2 in both x and y spatial coordinate\n",
    "model_parameters.set_atmospheric_channel_fourier_modes(2, 2)\n",
    "\n",
    "# Changing (increasing) the orography depth and the meridional temperature gradient\n",
    "model_parameters.ground_params.set_orography(0.4, 1)\n",
    "model_parameters.atemperature_params.set_thetas(0.2, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Printing the model's parameters\n",
    "model_parameters.print_params()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Creating the tendencies function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "f, Df = create_tendencies(model_parameters)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Time integration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Defining an integrator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "integrator = RungeKuttaIntegrator()\n",
    "integrator.set_func(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Start on a random initial condition and integrate over a transient time to obtain an initial condition on the attractors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "ic = np.random.rand(model_parameters.ndim)*0.1\n",
    "integrator.integrate(0., 200000., dt, ic=ic, write_steps=0)\n",
    "time, ic = integrator.get_trajectories()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now integrate to obtain with the qgs RK4 integrator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "integrator.integrate(0., 200., dt, ic=ic, write_steps=write_steps)\n",
    "time, traj = integrator.get_trajectories()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And also with the DifferentialEquations ODE Solver"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# defining a function with the correct header\n",
    "\n",
    "@jit\n",
    "def f_jl(x,p,t):\n",
    "    u = f(t,x)\n",
    "    return u"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# Defining the problem and integrating\n",
    "prob = de.ODEProblem(f_jl, ic, (0., 200.))\n",
    "sol = de.solve(prob, saveat=write_steps * dt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plotting the result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "var = 10\n",
    "jtraj = np.array(sol.u).T\n",
    "plt.figure(figsize=(10, 8))\n",
    "\n",
    "plt.plot(model_parameters.dimensional_time*time, traj[var], label=\"Qgs RK4 integrator\")\n",
    "plt.plot(model_parameters.dimensional_time*time, jtraj[var], label=\"DiffEqPy default integrator\")\n",
    "\n",
    "plt.legend()\n",
    "plt.xlabel('time (days)')\n",
    "plt.ylabel('$'+model_parameters.latex_var_string[var]+'$');"
   ]
  }
 ],
 "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}