{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# QGS model: Simple run example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reinhold and Pierrehumbert 1982 model version"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This model version is 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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import rc\n",
    "rc('font',**{'family':'serif','sans-serif':['Times'],'size':14})"
   ]
  },
  {
   "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, RungeKuttaTglsIntegrator\n",
    "from qgs.functions.tendencies import create_tendencies\n",
    "from qgs.plotting.util import std_plot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and diagnostics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from qgs.diagnostics.streamfunctions import MiddleAtmosphericStreamfunctionDiagnostic\n",
    "from qgs.diagnostics.variables import VariablesDiagnostic\n",
    "from qgs.diagnostics.multi import MultiDiagnostic"
   ]
  },
  {
   "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 a trajectory on the attractor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "integrator.integrate(0., 100000., dt, ic=ic, write_steps=write_steps)\n",
    "reference_time, reference_traj = integrator.get_trajectories()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "varx = 0\n",
    "vary = 1\n",
    "varz = 2\n",
    "\n",
    "fig = plt.figure(figsize=(10, 8))\n",
    "axi = fig.add_subplot(111, projection='3d')\n",
    "\n",
    "axi.scatter(reference_traj[varx], reference_traj[vary], reference_traj[varz], s=0.2);\n",
    "\n",
    "axi.set_xlabel('$'+model_parameters.latex_var_string[varx]+'$')\n",
    "axi.set_ylabel('$'+model_parameters.latex_var_string[vary]+'$')\n",
    "axi.set_zlabel('$'+model_parameters.latex_var_string[varz]+'$');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "varx = 2\n",
    "vary = 1\n",
    "plt.figure(figsize=(10, 8))\n",
    "\n",
    "plt.plot(reference_traj[varx], reference_traj[vary], marker='o', ms=0.07, ls='')\n",
    "\n",
    "plt.xlabel('$'+model_parameters.latex_var_string[varx]+'$')\n",
    "plt.ylabel('$'+model_parameters.latex_var_string[vary]+'$');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "var = 1\n",
    "plt.figure(figsize=(10, 8))\n",
    "\n",
    "plt.plot(model_parameters.dimensional_time*reference_time, reference_traj[var])\n",
    "\n",
    "plt.xlabel('time (days)')\n",
    "plt.ylabel('$'+model_parameters.latex_var_string[var]+'$');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initial condition sensitivity analysis example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Instantiating a tangent linear integrator with the model tendencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tgls_integrator = RungeKuttaTglsIntegrator()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "tgls_integrator.set_func(f, Df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Integrating with slightly perturbed initial conditions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tangent_ic = 0.0001*np.random.randn(number_of_perturbed_trajectories, model_parameters.ndim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time \n",
    "\n",
    "tgls_integrator.integrate(0., 130., dt=dt, write_steps=write_steps, ic=ic, tg_ic=tangent_ic)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Obtaining the perturbed trajectories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "time, traj, delta = tgls_integrator.get_trajectories()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pert_traj = traj + delta"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Trajectories plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "var = 10\n",
    "plt.figure(figsize=(10, 8))\n",
    "\n",
    "plt.plot(model_parameters.dimensional_time*time, traj[var])\n",
    "plt.plot(model_parameters.dimensional_time*time, pert_traj[:,var].T, ls=':')\n",
    "\n",
    "ax = plt.gca()\n",
    "\n",
    "plt.xlabel('time (days)')\n",
    "plt.ylabel('$'+model_parameters.latex_var_string[var]+'$');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Mean and standard deviation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "var = 1\n",
    "plt.figure(figsize=(10, 8))\n",
    "\n",
    "plt.plot(model_parameters.dimensional_time*time, traj[var])\n",
    "\n",
    "ax = plt.gca()\n",
    "std_plot(model_parameters.dimensional_time*time, np.mean(pert_traj[:,var], axis=0), np.sqrt(np.var(pert_traj[:, var], axis=0)), ax=ax, alpha=0.5)\n",
    "\n",
    "plt.xlabel('time (days)')\n",
    "plt.ylabel('$'+model_parameters.latex_var_string[var]+'$');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Showing the resulting fields (animation)\n",
    "\n",
    "This is an advanced feature showing the time evolution of diagnostic of the model. It shows simultaneously a scatter plot of the variable $\\psi_{{\\rm a}, 2}$ and $\\psi_{{\\rm a}, 3}$, and the geopotential height field at 500 hPa over the orographic height. Please read the documentation for more information."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Creating the diagnostics:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* For the 500hPa geopotential height:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psi = MiddleAtmosphericStreamfunctionDiagnostic(model_parameters, geopotential=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* For the nondimensional variables $\\psi_{{\\rm a}, 2}$ and $\\psi_{{\\rm a}, 3}$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "variable_nondim = VariablesDiagnostic([2, 1], model_parameters, False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# setting also the background\n",
    "background = VariablesDiagnostic([2, 1], model_parameters, False)\n",
    "background.set_data(reference_time, reference_traj)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Creating a multi diagnostic with both:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m = MultiDiagnostic(1,2)\n",
    "m.add_diagnostic(variable_nondim,\n",
    "                 diagnostic_kwargs={'show_time': False, 'background': background},\n",
    "                 plot_kwargs={'ms': 0.2})\n",
    "m.add_diagnostic(psi,\n",
    "                 diagnostic_kwargs={'style': 'contour', 'contour_labels': False},\n",
    "                 plot_kwargs={'colors': 'k'})\n",
    "m.set_data(time, traj)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and show an interactive animation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "rc('font',**{'family':'serif','sans-serif':['Times'],'size':12})\n",
    "m.animate(figsize=(20,6))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "or a movie:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "rc('font',**{'family': 'serif','sans-serif': ['Times'],'size': 12})\n",
    "m.movie(figsize=(20,6), anim_kwargs={'interval': 100, 'frames': len(time)-1})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}