{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# QGS model: MAOOAM with dynamic reference temperatures"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Coupled ocean-atmosphere 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 ocean with 8 modes, with dynamic reference temperature evolution in both components. \n",
    "\n",
    "More details can be found in the articles:\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",
    "* TBD\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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import rc\n",
    "rc('font',**{'family':'serif','sans-serif':['Times'],'size':12})"
   ]
  },
  {
   "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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and diagnostics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from qgs.diagnostics.streamfunctions import MiddleAtmosphericStreamfunctionDiagnostic, OceanicLayerStreamfunctionDiagnostic\n",
    "from qgs.diagnostics.temperatures import MiddleAtmosphericTemperatureDiagnostic, OceanicLayerTemperatureDiagnostic\n",
    "from qgs.diagnostics.variables import VariablesDiagnostic, GeopotentialHeightDifferenceDiagnostic\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 = 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}, dynamic_T=True)\n",
    "\n",
    "# Modes definition: with dynamic reference temperature, the user has to use the symbolic modes\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",
    "model_parameters.set_oceanic_basin_fourier_modes(2, 4, mode=\"symbolic\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setting MAOOAM parameters according to the publication linked above\n",
    "model_parameters.set_params({'kd': 0.0290, 'kdp': 0.0290, 'r': 1.e-7,\n",
    "                             'h': 136.5, 'd': 1.1e-7})\n",
    "model_parameters.atemperature_params.set_params({'eps': 0.7, 'hlambda': 15.06})\n",
    "model_parameters.gotemperature_params.set_params({'gamma': 5.6e8})"
   ]
  },
  {
   "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(103., 0)\n",
    "model_parameters.atemperature_params.set_insolation(103., 1)\n",
    "model_parameters.gotemperature_params.set_insolation(310., 0)\n",
    "model_parameters.gotemperature_params.set_insolation(310., 1)"
   ]
  },
  {
   "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": {
    "scrolled": true
   },
   "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",
    "## Might take several minutes, depending on your cpu computational power.\n",
    "ic = np.random.rand(model_parameters.ndim)*0.01\n",
    "ic[29] = 3.  # Setting reasonable initial reference temperature\n",
    "ic[10] = 1.5\n",
    "integrator.integrate(0., 2000000.1, 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., 3000000., dt, ic=ic, write_steps=write_steps)\n",
    "reference_time, reference_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 = 22\n",
    "vary = 31\n",
    "varz = 0\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]+'$', labelpad=12.)\n",
    "axi.set_ylabel('$'+model_parameters.latex_var_string[vary]+'$', labelpad=12.)\n",
    "axi.set_zlabel('$'+model_parameters.latex_var_string[varz]+'$', labelpad=12.);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "varx = 22\n",
    "vary = 31\n",
    "plt.figure(figsize=(10, 8))\n",
    "\n",
    "plt.plot(reference_traj[varx], reference_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 = 10\n",
    "plt.figure(figsize=(10, 8))\n",
    "\n",
    "plt.plot(model_parameters.dimensional_time*reference_time, 2* model_parameters.temperature_scaling * reference_traj[var], label='$'+model_parameters.latex_var_string[var]+'$ [K]')\n",
    "plt.plot(model_parameters.dimensional_time*reference_time, model_parameters.temperature_scaling * reference_traj[29], label='$'+model_parameters.latex_var_string[29]+'$ [K]')\n",
    "\n",
    "\n",
    "plt.xlabel('time (days)')\n",
    "#plt.ylabel();\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "var = 22\n",
    "plt.figure(figsize=(10, 8))\n",
    "\n",
    "plt.plot(model_parameters.dimensional_time*reference_time, reference_traj[var])\n",
    "\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 variables $\\psi_{{\\rm a}, 1}$, $\\psi_{{\\rm o}, 2}$ and $\\delta T_{{\\rm o}, 2}$, with the corresponding atmospheric and oceanic streamfunctions and temperature at 500 hPa. 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_a = MiddleAtmosphericStreamfunctionDiagnostic(model_parameters, geopotential=True, delta_x=0.2, delta_y=0.2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* For the 500hPa atmospheric temperature:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "theta_a = MiddleAtmosphericTemperatureDiagnostic(model_parameters, delta_x=0.2, delta_y=0.2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* For the ocean streamfunction:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psi_o = OceanicLayerStreamfunctionDiagnostic(model_parameters, delta_x=0.2, delta_y=0.2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* For the ocean temperature:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "theta_o = OceanicLayerTemperatureDiagnostic(model_parameters, delta_x=0.2, delta_y=0.2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* For the nondimensional variables $\\psi_{{\\rm a}, 1}$, $\\psi_{{\\rm o}, 2}$ and $\\delta T_{{\\rm o}, 2}$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "variable_nondim = VariablesDiagnostic([22, 31, 0], model_parameters, False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* For the geopotential height difference between North and South:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "geopot_dim = GeopotentialHeightDifferenceDiagnostic([[[np.pi/model_parameters.scale_params.n, np.pi/4], [np.pi/model_parameters.scale_params.n, 3*np.pi/4]]],\n",
    "                                                    model_parameters, True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# setting also the background\n",
    "background = VariablesDiagnostic([22, 31, 0], model_parameters, False)\n",
    "background.set_data(reference_time, reference_traj)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Selecting a subset of the data to plot:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stride = 10\n",
    "time = reference_time[10000:10000+5200*stride:stride]\n",
    "traj = reference_traj[:, 10000:10000+5200*stride:stride]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Creating a multi diagnostic with all the diagnostics:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m = MultiDiagnostic(2,3)\n",
    "m.add_diagnostic(geopot_dim, diagnostic_kwargs={'style':'moving-timeserie'})\n",
    "m.add_diagnostic(theta_a, diagnostic_kwargs={'show_time':False})\n",
    "m.add_diagnostic(theta_o, diagnostic_kwargs={'show_time':False})\n",
    "m.add_diagnostic(variable_nondim, diagnostic_kwargs={'show_time':False, 'background': background, 'style':'3Dscatter'}, plot_kwargs={'ms': 0.2})\n",
    "m.add_diagnostic(psi_a, diagnostic_kwargs={'show_time':False})\n",
    "m.add_diagnostic(psi_o, diagnostic_kwargs={'show_time':False})\n",
    "\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=(23,12))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "or a movie (may takes some minutes to compute):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "rc('font',**{'family':'serif','sans-serif':['Times'],'size':12})\n",
    "m.movie(figsize=(23.5,12), anim_kwargs={'interval': 100, 'frames':2000})"
   ]
  },
  {
   "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
}