Updating the Two-Timescale System Notebook

notebooks/L96-two-scale-description.ipynb

@@ -14,27 +14,36 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "In 1996, Lorenz proposed a two-time scale dynamical system, referred to as the Lorenz-96 model (L96), whose behavior mimics the non-linear dynamics of the extratropical atmosphere with simplified representation of multiscale interactions and nonlinear advection [Lorenz (1996)](https://www.ecmwf.int/en/elibrary/10829-predictability-problem-partly-solved). The L96 model consists of a set of equations coupling two scales of variables evolving over slow and fast timescales and arranged around a latitude circle as shown in Fig. 1.\n",
-    "[Lorenz (1996)](https://www.ecmwf.int/en/elibrary/10829-predictability-problem-partly-solved) describes the \"two time-scale\" model in a set of equations (1 and 2) given by:\n",
+    "In 1996, Lorenz proposed a two-time scale dynamical system, referred to as the Lorenz-96 model (L96), whose behavior mimics the non-linear dynamics of the extratropical atmosphere with simplified representation of multiscale interactions and nonlinear advection [Lorenz (1996)](https://www.ecmwf.int/en/elibrary/10829-predictability-problem-partly-solved). The L96 model consists of two sets of equations coupling two sets of variables ($X_k$ and $Y_j,k$), which evolve over two (slow and fast) timescales and are arranged around a latitude circle as shown in Fig. 1. The equations comprising L96 are:\n",
     "\\begin{align}\n",
     "\\frac{d}{dt} X_k\n",
-    "&= - X_{k-1} \\left( X_{k-2} - X_{k+1} \\right) - X_k + F - \\left( \\frac{hc}{b} \\right) \\sum_{j=0}^{J-1} Y_{j,k}\n",
+    "&= - X_{k-1} \\left( X_{k-2} - X_{k+1} \\right) - X_k + F - \\left( \\frac{hc}{b} \\right) \\sum_{j=0}^{J-1} Y_{j,k},\n",
     "\\\\\n",
     "\\frac{d}{dt} Y_{j,k}\n",
     "&= - cbY_{j+1,k} \\left( Y_{j+2,k} - Y_{j-1,k} \\right) - c Y_{j,k} + \\frac{hc}{b} X_k\n",
     "\\end{align}\n",
     "\n",
-    "The model includes $K$ slow (or low-frequency) variables, denoted as $X_k$, $k=1,\\ldots,K$ and each slow variable $X_k$ is coupled with $J$ fast (or high-frequency) variables $Y_{j,k}$ , $j=1,\\ldots,J$. \n",
-    "The model equations are coupled via the mean term $\\sum_{j=0}^{J-1} Y_{j,k}$, and the coupling depends on three key parameters: $b$, $c$ and $h$. The parameter $b$ determines the magnitude of the nonlinear interactions among the fast variables. The parameter $c$ controls how rapidly the fast variables fluctuate compared to the slow variables. Finally, the parameter $h$ governs the strength of the coupling between the slow and fast variables. \n",
+    "where $X_k$, $k=1,\\ldots,K$, denotes $K$ slow (or low-frequency) variables and $Y_{j,k}$ , $j=1,\\ldots,J$ denotes $J*K$ fast (or high-frequency) variables. Each slow $X_k$ is coupled with $J$ fast $Y_{j,k}$ variables. The slow equations are coupled to the fast equations via a coupling term, $\\sum_{j=0}^{J-1} Y_{j,k}$, which sums over the $J$ fast variables corresponding to a particular $k$. On the other hand, each fast equation is forced by the a coupling term, \\frac{hc}{b} X_k, that depends on the slow variable corresponding that particular $k$. The coupling and the fast time-scale equations depends on three key parameters: $b$, $c$ and $h$. Here $b$ determines the magnitude of the nonlinear interactions among the fast variables, $c$ controls how rapidly the fast variables fluctuate compared to the slow variables and, and $h$ governs the strength of the coupling between the slow and fast variables. Also, the slow time-scale equation is forced by $F$, e.g. [Wilks (2005)](http://rmets.onlinelibrary.wiley.com/doi/abs/10.1256/qj.04.03).\n",
     "\n",
     "The chaotic dynamical system L96 is very useful for testing different numerical methods in atmospheric modeling thanks to its transparency, low computational cost and simplicity compared to Global Climate Models (GCM). The interaction between variables of different scales makes the L96 model of particular interest when evaluating new parameterization methodologies. As such, it was used in assessing different techniques that were later incorporated into GCMs ([Crommelin (2008)](https://journals.ametsoc.org/view/journals/atsc/65/8/2008jas2566.1.xml), [Dorrestijn (2013)](https://royalsocietypublishing.org/doi/10.1098/rsta.2012.0374)).\n",
     "\n",
     "The L96 model has been extensively used as a test bed in various studies including data assimilation approaches ([Law (2016)](https://www.sciencedirect.com/science/article/pii/S0167278915002766?via%3Dihub), [Hatfield (2018)](https://journals.ametsoc.org/view/journals/mwre/146/1/mwr-d-17-0132.1.xml)), stochastic parameterization schemes ([Kwasniok (2012)](https://royalsocietypublishing.org/doi/10.1098/rsta.2011.0384), [Arnold (2013)](https://royalsocietypublishing.org/doi/full/10.1098/rsta.2011.0479), [Chorin (2015)](https://www.pnas.org/doi/10.1073/pnas.1512080112))  and Machine Learning-based parameterization methodologies ([Schneider (2017)](https://agupubs.onlinelibrary.wiley.com/doi/10.1002/2017GL076101), [Dueben (2018)](https://gmd.copernicus.org/articles/11/3999/2018/), [Watson (2019)](https://agupubs.onlinelibrary.wiley.com/doi/full/10.1029/2018MS001597), [Gagne (2020)](https://agupubs.onlinelibrary.wiley.com/doi/full/10.1029/2019MS001896)).\n",
-    "Finally, note that the $F$ term in equation (1) is missing in the original Lorenz 96 paper but the author provides the parameter. All other authors use the term as well, e.g. [Wilks (2005)](http://rmets.onlinelibrary.wiley.com/doi/abs/10.1256/qj.04.03).)\n",
     "\n",
+    "<center>\n",
+    "  <img\n",
+    "    src=\"https://www.researchgate.net/publication/319201436/figure/fig1/AS:869115023589376@1584224577926/Visualisation-of-a-two-scale-Lorenz-96-system-with-J-8-and-K-6-Global-scale-values.png\"\n",
+    "    width=400\n",
+    "  />\n",
+    "</center>\n",
     "\n",
-    "<img src=\"https://www.researchgate.net/publication/319201436/figure/fig1/AS:869115023589376@1584224577926/Visualisation-of-a-two-scale-Lorenz-96-system-with-J-8-and-K-6-Global-scale-values.png\" width=400> *Fig. 1: Visualisation of a two-scale Lorenz '96 system with J = 8 and K = 6. Global-scale variables ($X_k$) are updated based on neighbouring variables and on the local-scale variables ($Y_{j,k}$) associated with the corresponding global-scale variable. Local-scale variabless are updated based on neighbouring variables and the associated global-scale variable. The neighbourhood topology of both local and global-scale variables is circular. Image from [Exploiting the chaotic behaviour of atmospheric models with reconfigurable architectures - Scientific Figure on ResearchGate.](https://www.researchgate.net/figure/Visualisation-of-a-two-scale-Lorenz-96-system-with-J-8-and-K-6-Global-scale-values_fig1_319201436)*\n",
-    "\n"
+    "<!-- <center>\n",
+    "<em>\n",
+    "Fig. 1: Visualisation of a two-scale Lorenz '96 system with J = 8 and K = 6. Global-scale variables ($X_k$) are updated based on neighbouring variables and on the local-scale variables ($Y_{j,k}$) associated with the corresponding global-scale variable. Local-scale variabless are updated based on neighbouring variables and the associated global-scale variable. The neighbourhood topology of both local and global-scale variables is circular. Image from <a href=\"https://www.researchgate.net/figure/Visualisation-of-a-two-scale-Lorenz-96-system-with-J-8-and-K-6-Global-scale-values_fig1_319201436\">Exploiting the chaotic behaviour of atmospheric models with reconfigurable architectures - Scientific Figure on ResearchGate.</a>\n",
+    "\n",
+    "</em>\n",
+    "</center> -->\n",
+    "\n",
+    "<span> <center> *Fig. 1: Visualisation of a two-scale Lorenz '96 system with J = 8 and K = 6. Global-scale variables ($X_k$) are updated based on neighbouring variables and on the local-scale variables ($Y_{j,k}$) associated with the corresponding global-scale variable. Local-scale variabless are updated based on neighbouring variables and the associated global-scale variable. The neighbourhood topology of both local and global-scale variables is circular. Image from [Exploiting the chaotic behaviour of atmospheric models with reconfigurable architectures - Scientific Figure on ResearchGate.](https://www.researchgate.net/figure/Visualisation-of-a-two-scale-Lorenz-96-system-with-J-8-and-K-6-Global-scale-values_fig1_319201436)* </center> </span>\n"
    ]
   },
   {
@@ -50,17 +59,17 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import inspect\n",
-    "\n",
     "import matplotlib.pyplot as plt\n",
-    "import numpy as np"
+    "import numpy as np\n",
+    "\n",
+    "from utils import display_source"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "In the `L96_model` module we provide the function `L96_2t_xdot_ydot` (shown next) that returns the tendencies for the X and Y equations, as well as the tendency to the coupling in the X equation, $-hc/b \\sum_{j=0}^{J-1} Y_{j,k}$."
+    "In the `L96_model` module we provide the function `L96_2t_xdot_ydot` (shown next) that returns the tendencies (sum of the RHS) for the X and Y equations, as well as the tendency corresponding to the coupling in the X equation, $-hc/b \\sum_{j=0}^{J-1} Y_{j,k}$."
    ]
   },
   {
@@ -71,14 +80,14 @@
    "source": [
     "from L96_model import L96_2t_xdot_ydot\n",
     "\n",
-    "print(inspect.getsource(L96_2t_xdot_ydot))"
+    "display_source(L96_2t_xdot_ydot)"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "The tendencies above describe the continuous evolution but to obtain a numerical solution we must integrate in time. In the `L96_model` module we provide the function `integrate_L96_2t` (shown next) that uses fourth-order Runge-Kutta integration. Starting from initial conditions `X0,Y0`, the function returns the the trajectories of `X,Y` sampled at an interval `si`. There is a related function, `integrate_L96_2t_with_coupling`, that in addition to the trajectories, also returns the coupling term, $-hc/b \\sum_{j=0}^{J-1} Y_{j,k}$ at each point in the trajectory."
+    "These tendencies describe the continuous evolution of the L96 model at a particular time. However, to obtain a discrete solution we must integrate numerically in time. In the `L96_model` module, we provide the function `integrate_L96_2t` (shown next) that uses fourth-order Runge-Kutta integration. Starting from the initial conditions `X0,Y0`, the function returns the trajectories of `X,Y` sampled at an interval `si`. There is a related function, `integrate_L96_2t_with_coupling`, that in addition to the trajectories, also returns the coupling term, $-hc/b \\sum_{j=0}^{J-1} Y_{j,k}$ at each point in the trajectory."
    ]
   },
   {
@@ -89,14 +98,14 @@
    "source": [
     "from L96_model import integrate_L96_2t\n",
     "\n",
-    "print(inspect.getsource(integrate_L96_2t))"
+    "display_source(integrate_L96_2t)"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "An example of how to use `integrate_L96_2t` is given below. We choose $K=36$ and $J=10$ (i.e. 36 values of $X$ and 10 values of $Y$ per value of $X$). The choices of $h=1$, $c=10$, and $b=10$ are those used in many studies. We make up an initial conditions (`X(s)=s(s-1)(s+1)` for $s=-1\\ldots 1$, $Y=0$). The value of $F=10$ is sufficient to obtain chaotic behavior here. If you increase $F$, you may need to reduce $dt$ for numerical stability. "
+    "An example of how to use `integrate_L96_2t` is shown below. Here we choose $K=36$ and $J=10$ (i.e. 36 values of $X$ and 10 values of $Y$ per value of $X$). Also, in accordance with previous studies, we set $h=1$, $c=10$, and $b=10$. The initial condition for $X$ is setup as $X(s)=s(s-1)(s+1)$, for $s=-1\\ldots 1$, and the $Y$ is initialized with zeros. The value of $F$ is set to $10$, which is sufficient to obtain chaotic behavior. *Note: if you increase $F$, you may need to reduce $dt$ for numerical stability.* "
    ]
   },
   {
@@ -105,59 +114,90 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "J = 10  # Number of local-scale Y variables per single global-scale X variable\n",
     "K = 36  # Number of globa-scale variables X\n",
+    "J = 10  # Number of local-scale Y variables per single global-scale X variable\n",
     "nt = 1000  # Number of time steps\n",
     "si = 0.005  # Sampling time interval\n",
     "dt = 0.005  # Time step\n",
     "F = 10.0  # Focring\n",
     "h = 1.0  # Coupling coefficient\n",
     "b = 10.0  # ratio of amplitudes\n",
-    "c = 10.0  # time-scale ratio\n",
-    "\n",
-    "\n",
+    "c = 10.0  # time-scale ratio"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
     "def s(k, K):\n",
     "    \"\"\"A non-dimension coordinate from -1..+1 corresponding to k=0..K\"\"\"\n",
-    "    return 2 * (0.5 + k) / K - 1\n",
-    "\n",
-    "\n",
+    "    return 2 * (0.5 + k) / K - 1"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
     "k = np.arange(K)  # For coordinate in plots\n",
     "j = np.arange(J * K)  # For coordinate in plots\n",
+    "\n",
     "# Initial conditions\n",
-    "Xinit = s(k, K) * (s(k, K) - 1) * (s(k, K) + 1)\n",
-    "Yinit = 0 * s(j, J * K) * (s(j, J * K) - 1) * (s(j, J * K) + 1)\n",
+    "X_init = s(k, K) * (s(k, K) - 1) * (s(k, K) + 1)\n",
+    "Y_init = 0 * s(j, J * K) * (s(j, J * K) - 1) * (s(j, J * K) + 1)\n",
     "\n",
     "# \"Run\" model\n",
-    "X, Y, t = integrate_L96_2t(Xinit, Yinit, si, nt, F, h, b, c, dt=dt)\n",
-    "\n",
+    "X, Y, t = integrate_L96_2t(X_init, Y_init, si, nt, F, h, b, c, dt=dt)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "After running the model, we plot the results."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
     "plt.figure(figsize=(12, 10))\n",
     "plt.subplot(221)\n",
+    "\n",
     "# Snapshot of X[k]\n",
     "plt.plot(k, X[-1], label=\"$X_k(t=n_t dt)$\")\n",
     "plt.plot(j / J, Y[-1], label=\"$Y_{j,k}(t=n_t dt)$\")\n",
     "plt.plot(k, Xinit, \"k:\", label=\"$X_k(t=0)$\")\n",
     "plt.plot(j / J, Yinit, \"k:\", label=\"$Y_{j,k}(t=0)$\")\n",
     "plt.legend()\n",
-    "plt.xlabel(\"k, k+j/J\")\n",
+    "plt.xlabel(\"k, k + j/J\")\n",
     "plt.title(\"$X_k, Y_{j,k}$\")\n",
     "plt.subplot(222)\n",
+    "\n",
     "# Sample time-series X[0](t), Y[0](t)\n",
     "plt.plot(t, X[:, 0], label=\"$X_0(t)$\")\n",
     "plt.plot(t, Y[:, 0], label=\"$Y_{0,0}(t)$\")\n",
     "plt.legend()\n",
     "plt.xlabel(\"t\")\n",
     "plt.subplot(223)\n",
+    "\n",
     "# Full model history of X\n",
     "plt.contourf(k, t, X)\n",
     "plt.colorbar(orientation=\"horizontal\")\n",
     "plt.xlabel(\"k\")\n",
     "plt.ylabel(\"t\")\n",
     "plt.title(\"$X_k(t)$\")\n",
     "plt.subplot(224)\n",
+    "\n",
     "# Full model history of Y\n",
     "plt.contourf(j / J, t, Y)\n",
     "plt.colorbar(orientation=\"horizontal\")\n",
-    "plt.xlabel(\"k+j/J\")\n",
+    "plt.xlabel(\"k + j/J\")\n",
     "plt.ylabel(\"t\")\n",
     "plt.title(\"$Y_{j,k}(t)$\");"
    ]
@@ -168,7 +208,7 @@
    "source": [
     "## A class for numerical integration of the two time-scale L96\n",
     "\n",
-    "For convenience, we provide a class `L96` in `L96_model.py` to reduce the amount of code to use `L96_2t_xdot_ydot`."
+    "For convenience, we also provide a class named `L96` in `L96_model.py` which helps to reduce the amount of code when using the `L96_2t_xdot_ydot` function."
    ]
   },
   {
@@ -241,7 +281,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Re-running the above cell gives the same answer each time because the same initial conditions are used for each invocation or `.run()`. To allow sequences of run, using the end state of one as the starting point for the next, add `store=True` which updates the initial condition stored in the object (`M`) with the final state of the run."
+    "Re-running the above cell gives the same answer each time because the same initial conditions are used for each invocation or `.run()`. To allow sequences of run, using the end state of one as the starting point for the next, add the parameter `store=True` which updates the initial condition stored in the object (`M`) with the final state of the run."
    ]
   },
   {
@@ -257,13 +297,19 @@
     "print(\"Mean absolute difference =\", np.abs(X2 - X).mean() + np.abs(Y2 - Y).mean())"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The entire model creation and execution process demonstrated above can also be implemented in a one-liner version as shown below."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "# The one line version of the above\n",
     "X3, Y3, t = (\n",
     "    L96(36, 10, F=10, dt=0.005)\n",
     "    .set_state(s(M.k, M.K) * (s(M.k, M.K) - 1) * (s(M.k, M.K) + 1), 0 * M.j)\n",
@@ -290,7 +336,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.12"
+   "version": "3.9.16"
   }
  },
  "nbformat": 4,

notebooks/L96_model.py

@@ -61,12 +61,6 @@ def L96_2t_xdot_ydot(X, Y, F, h, b, c):
     Ysummed = Y.reshape((K, J)).sum(axis=-1)
 
     Xdot = np.roll(X, 1) * (np.roll(X, -1) - np.roll(X, 2)) - X + F - hcb * Ysummed
-    #     for k in range(K):
-    #         Xdot[k] = ( X[(k+1)%K] - X[k-2] ) * X[k-1] - X[k] + F - hcb * Ysummed[k]
-
-    # for j in range(JK):
-    #        k = j//J
-    #        Ydot[j] = -c * b * Y[(j+1)%JK] * ( Y[(j+2)%JK] - Y[j-1] ) - c * Y[j] + hcb * X[k]
     Ydot = (
         -c * b * np.roll(Y, -1) * (np.roll(Y, -2) - np.roll(Y, 1))
         - c * Y

notebooks/utils.py

@@ -0,0 +1,14 @@
+import inspect
+
+from IPython.display import HTML, display
+from pygments import highlight
+from pygments.formatters import HtmlFormatter
+from pygments.lexers import PythonLexer
+
+
+def display_source(non_builtin_object):
+    """Displays the source code of the given object with syntax highlighting."""
+    code = inspect.getsource(non_builtin_object)
+    html = highlight(code, PythonLexer(), HtmlFormatter(style="colorful"))
+    stylesheet = f"<style>{HtmlFormatter().get_style_defs('.highlight')}</style>"
+    display(HTML(f"{stylesheet}{html}"))