diff --git a/LinearStabilityAnalysis/LinearStabilityAnalysis.ipynb b/LinearStabilityAnalysis/LinearStabilityAnalysis.ipynb
index d6ed084..108e744 100644
--- a/LinearStabilityAnalysis/LinearStabilityAnalysis.ipynb
+++ b/LinearStabilityAnalysis/LinearStabilityAnalysis.ipynb
@@ -70,7 +70,9 @@
     "\n",
     "When running this notebook I reccomend increasing the grid-spacing (`dx`) to something larger ($\\sim$`2e3` should still give ok results). The value of `1e2` was used here for producing publication quality figures and is overkill even for that.\n",
     "\n",
-    "If you want to play around with different vorticity profiles this can be done by redefining `zeta` in the below code to return the absolute vorticity profile you fancy."
+    "If you want to play around with different vorticity profiles this can be done by redefining `zeta` in the below code to return the absolute vorticity profile you fancy.\n",
+    "\n",
+    "The analysis can be performed in either hydrostatic mode (as is done in the paper) or non-hydrostatic mode, in which a term in $\\hat{\\omega}^2 d^2_{xx} \\hat{\\psi} / m^2$ is added on the left hand side of the above equation."
    ]
   },
   {
@@ -79,6 +81,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "hydrostatic = True\n",
+    "\n",
     "Lx = 400e3  # Domain width in m\n",
     "dx = 1e2  # Grid spacing in m \n",
     "nx = int(Lx / dx)  # Number of grid points\n",
@@ -145,7 +149,13 @@
     "    m = 2 * np.pi / lambda_val\n",
     "    m2 = np.square(m)\n",
     "    eigen_operator = -N2 / m2 * d2_dx2 + f_zeta\n",
-    "    eigen_val = eigs(eigen_operator, k=1, which='SR', return_eigenvectors=True)\n",
+    "    \n",
+    "    if hydrostatic:\n",
+    "        eigen_val = eigs(eigen_operator, k=1, which='SR', return_eigenvectors=True)\n",
+    "    else:\n",
+    "        generalised_M =  sps.eye(nx) - d2_dx2 / m2  # Identiy - d2/dx2\n",
+    "        eigen_val = eigs(eigen_operator, k=1, M=generalised_M, which='SR', return_eigenvectors=True)\n",
+    "        \n",
     "    return eigen_val\n",
     "\n",
     "# Create empty arrays in which to store the eigenvalues and the eigenfunctions\n",