{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# QGS model: MAOSOAM"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Coupled ocean-atmosphere channel model version"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This model version is a 2-layer channel QG atmosphere truncated at wavenumber 2 coupled, both by friction and heat exchange, to a shallow water **channel** ocean also truncated at wavenumber 2. \n",
    "\n",
    "More details can be found in the articles:\n",
    "*  Vannitsem, S., Solé‐Pomies, R. and De Cruz, L. (2019). *Routes to long‐term atmospheric predictability in reduced‐order coupled ocean–atmosphere systems ‐ Impact of the ocean basin boundary conditions.* Quarterly Journal of the Royal Meteorological Society, **145**: 2791– 2805. [doi.org/10.1002/qj.3594](https://doi.org/10.1002/qj.3594)\n",
    "* Vannitsem, S., Demaeyer, J., De Cruz, L., & Ghil, M. (2015). *Low-frequency variability and heat transport in a low-order nonlinear coupled ocean–atmosphere model*. Physica D: Nonlinear Phenomena, **309**, 71-85. [doi:10.1016/j.physd.2015.07.006](https://doi.org/10.1016/j.physd.2015.07.006)\n",
    "* De Cruz, L., Demaeyer, J. and Vannitsem, S. (2016). *The Modular Arbitrary-Order Ocean-Atmosphere Model: MAOOAM v1.0*, Geosci. Model Dev., **9**, 2793-2808. [doi:10.5194/gmd-9-2793-2016](https://doi.org/10.5194/gmd-9-2793-2016)\n",
    "\n",
    "or in the documentation and on [readthedocs](https://qgs.readthedocs.io/en/latest/files/model/maooam_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.basis.fourier import contiguous_channel_basis\n",
    "from qgs.integrators.integrator import RungeKuttaIntegrator\n",
    "from qgs.functions.tendencies import create_tendencies"
   ]
  },
  {
   "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 = 100\n",
    "\n",
    "number_of_trajectories = 1"
   ]
  },
  {
   "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({'n': 1.5})\n",
    "\n",
    "# Mode truncation at the wavenumber 2 in both x and y spatial\n",
    "# coordinates for the atmosphere\n",
    "model_parameters.set_atmospheric_channel_fourier_modes(2, 2, mode=\"symbolic\")\n",
    "# Mode truncation at the wavenumber 2 in the x and at the \n",
    "# wavenumber 4 in the y spatial coordinates for the ocean\n",
    "ocean_basis = contiguous_channel_basis(2, 2, 1.5)\n",
    "model_parameters.set_oceanic_modes(ocean_basis)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setting MAOOAM parameters according to the publication linked above\n",
    "model_parameters.set_params({'phi0_npi': 0.3056, 'kd': 0.026778245344758034, 'kdp': 0.026778245344758034, 'r': 1.e-8,\n",
    "                             'h': 1000.0, 'd': 1.6e-8, 'f0': 1.195e-4, 'sigma': 0.14916, 'n':1.7})\n",
    "model_parameters.atemperature_params.set_params({'eps': 0.76, 'T0': 270.,\n",
    "                                                 'hlambda': 16.064})\n",
    "model_parameters.gotemperature_params.set_params({'gamma': 4e9, 'T0': 285.})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Setting the short-wave radiation component as in the publication above: $C_{\\text{a},1}$ and $C_{\\text{o},1}$ \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_parameters.atemperature_params.set_insolation(350/3., 0)\n",
    "model_parameters.gotemperature_params.set_insolation(350, 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Printing the model's parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_parameters.print_params()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Creating the tendencies function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "## Might take several minutes, depending on the number of cpus you have.\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 from an initial condition on the attractors obtained after a long transient integration time "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ic = np.array([ 2.34980646e-02, -5.91652353e-03,  3.20923307e-03, -1.08916714e-03,\n",
    "       -1.13188144e-03, -5.14454554e-03,  1.50294902e-02, -2.20518843e-04,\n",
    "        4.55325496e-03, -1.18748859e-03,  2.27043688e-02,  4.29437410e-04,\n",
    "        3.74041445e-03, -1.78681895e-03, -1.71853500e-03,  3.68921542e-04,\n",
    "       -6.42748591e-04, -2.81188015e-03, -2.14109639e-03, -1.41736652e-03,\n",
    "        3.24489725e-09,  3.97502699e-05, -7.47489713e-05,  9.89194512e-06,\n",
    "        5.52902699e-06,  6.43875197e-05, -6.95990073e-05,  1.21618381e-04,\n",
    "        7.08494425e-05, -1.11255308e-04,  4.13406579e-02, -7.90716982e-03,\n",
    "        1.33752621e-02,  1.66742520e-02,  6.29900201e-03,  1.76761107e-02,\n",
    "       -5.40207169e-02,  1.29814807e-02, -4.35142923e-02, -7.62511906e-03])"
   ]
  },
  {
   "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., 1000000., dt, ic=ic, write_steps=write_steps)\n",
    "time, traj = integrator.get_trajectories()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plotting the result in 3D and 2D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "varx = 20\n",
    "vary = 30\n",
    "varz = 0\n",
    "\n",
    "fig = plt.figure(figsize=(10, 8))\n",
    "axi = fig.gca(projection='3d')\n",
    "\n",
    "axi.scatter(traj[varx], traj[vary], 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 = 30\n",
    "vary = 10\n",
    "plt.figure(figsize=(10, 8))\n",
    "\n",
    "plt.plot(traj[varx], traj[vary], marker='o', ms=0.1, 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 = 30\n",
    "plt.figure(figsize=(10, 8))\n",
    "\n",
    "plt.plot(model_parameters.dimensional_time*time, traj[var])\n",
    "\n",
    "plt.xlabel('time (days)')\n",
    "plt.ylabel('$'+model_parameters.latex_var_string[var]+'$');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "var = 10\n",
    "plt.figure(figsize=(10, 8))\n",
    "\n",
    "plt.plot(model_parameters.dimensional_time*time, traj[var])\n",
    "\n",
    "plt.xlabel('time (days)')\n",
    "plt.ylabel('$'+model_parameters.latex_var_string[var]+'$');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "var = 20\n",
    "plt.figure(figsize=(10, 8))\n",
    "\n",
    "plt.plot(model_parameters.dimensional_time*time, traj[var])\n",
    "\n",
    "plt.xlabel('time (days)')\n",
    "plt.ylabel('$'+model_parameters.latex_var_string[var]+'$');"
   ]
  }
 ],
 "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
}