{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# QGS model: Simple run example with comparison of the Covariant Lyapunov vectors computation method (see last section)"
   ]
  },
  {
   "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)"
   ]
  },
  {
   "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": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import sys, os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "sys.path.extend([os.path.abspath('../../')])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "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": {
    "scrolled": false
   },
   "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": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "np.random.seed(210217)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Importing the model's modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "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 the Lyapunovs Estimators"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from qgs.toolbox.lyapunov import LyapunovsEstimator, CovariantLyapunovsEstimator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from qgs.toolbox.lyapunov import _compute_backward_lyap_traj_jit, _compute_forward_lyap_traj_jit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Systems definition"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "General parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "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": {
    "scrolled": false
   },
   "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": {
    "scrolled": false
   },
   "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": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "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": {
    "scrolled": false
   },
   "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": {
    "scrolled": false
   },
   "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": {
    "scrolled": false
   },
   "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": {
    "scrolled": false
   },
   "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": {
    "scrolled": false
   },
   "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": {
    "scrolled": false
   },
   "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": [
    "## Comparing Covariant Lyapunov vectors computation\n",
    "\n",
    "Here we compare the two methods used to compute the CLVs: the method 0 (Ginelli et al. algorithm) and the method 1 (subspaces intersection method). These methods are described in:\n",
    "\n",
    "* **Method 0:**\n",
    "* **Method 1:**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Covariant Lyapunovs Estimator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "clvint = CovariantLyapunovsEstimator()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Computing the CLVs with the Ginelli et al. algorithm (method 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "clvint.set_func(f, Df)\n",
    "clvint.compute_clvs(0., 10000., 40000., 50000., 0.1, 0.1, ic, write_steps=1)\n",
    "ctl0, ctraj0, cexp0, cvec0 = clvint.get_clvs()\n",
    "clvint.terminate()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plotting the spectrum for reference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(15, 4))\n",
    "\n",
    "mean_exp = np.mean(cexp0, axis=-1)\n",
    "\n",
    "x_pos = np.arange(1.,model_parameters.ndim+1,1)\n",
    "\n",
    "plt.bar(x_pos, mean_exp)\n",
    "\n",
    "plt.vlines(x_pos, -0.55, np.minimum(0.,mean_exp)-0.035, linestyles='dashdot', colors='tab:gray')\n",
    "\n",
    "plt.xticks(x_pos, map(str,range(1, model_parameters.ndim+1,1)))\n",
    "yt=[-0.5,-0.4,-0.3,-0.2,-0.1,0.,0.1]\n",
    "plt.yticks(yt, map(str,yt))\n",
    "\n",
    "plt.xlim(x_pos[0]-1., x_pos[-1]+1.)\n",
    "plt.ylim(np.min(mean_exp)-0.1, np.max(mean_exp)+0.1)\n",
    "\n",
    "plt.ylabel(\"Lyapunov exponent\");\n",
    "plt.xlabel(\"Index of the Lyapunov exponent\");\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Computing the CLVs with the subspaces intersection method along the same trajectory (method 1 done manually using hidden routine)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Computing the BLVs and FLVs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "pretime = ctl0[:100001]\n",
    "time = ctl0[100000:200001]\n",
    "posttime = ctl0[200000:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "backtraj = ctraj0[..., :200001][np.newaxis, ...]\n",
    "forwtraj = ctraj0[..., 100000:][np.newaxis, ...]\n",
    "cvec0 = cvec0[..., 100000:200001]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "ftraj, fexp, fvec = _compute_forward_lyap_traj_jit(f, Df, time, posttime, forwtraj, 0.1, model_parameters.ndim, 1, False, 1, clvint.b, clvint.c, clvint.a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "btraj, bexp, bvec = _compute_backward_lyap_traj_jit(f, Df, pretime, time, backtraj, 0.1, model_parameters.ndim, 1, False, 1, clvint.b, clvint.c, clvint.a)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Computing the subspaces intersections"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "ctraj1 = forwtraj[..., :100001]\n",
    "n_records = ctraj1.shape[-1]\n",
    "cvec1 = np.zeros((model_parameters.ndim, model_parameters.ndim, n_records))\n",
    "i_traj = 0\n",
    "for ti in range(n_records):\n",
    "    for j in range(model_parameters.ndim):\n",
    "        u, z, w = np.linalg.svd(bvec[i_traj, :, :j+1, ti].T @ fvec[i_traj, :, :model_parameters.ndim-j, ti])\n",
    "        basis = bvec[i_traj, :, :j+1, ti] @ u\n",
    "        cvec1[:, j, ti] = basis[:, 0]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Showing the first CLVs obtained by both method at a given time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Obtained by method 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "cvec0[:,0,0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Obtained by method 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "cvec1[:,0,0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plotting component by component the difference between each vector obtained by the two different methods"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each component is plotted with a different color"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "vars = slice(0, model_parameters.ndim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(20, int(model_parameters.ndim*8/2)), constrained_layout=False)\n",
    "grid = fig.add_gridspec(int(model_parameters.ndim/2), 2)\n",
    "axs = grid.subplots()\n",
    "\n",
    "for vec, ax in enumerate(axs.flatten()):\n",
    "    ax.plot(model_parameters.dimensional_time*time, (np.abs(cvec0[vars,vec,:])-np.abs(cvec1[vars,vec,:])).T);\n",
    "    ax.set_xlabel('time (days)')\n",
    "    ax.set_title('CLV '+str(vec+1))\n",
    "    "
   ]
  }
 ],
 "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
}