add theta to fsle example

Author: AntSimi

examples/07_cube_manipulation/pet_fsle_med.py

@@ -14,7 +14,7 @@
 
 from matplotlib import pyplot as plt
 from numba import njit
-from numpy import arange, empty, isnan, log2, ma, meshgrid, zeros
+from numpy import arange, arctan2, empty, isnan, log2, ma, meshgrid, ones, pi, zeros
 
 from py_eddy_tracker import start_logger
 from py_eddy_tracker.data import get_path
@@ -39,7 +39,7 @@
 # Methods to compute FSLE
 # -----------------------
 @njit(cache=True, fastmath=True)
-def check_p(x, y, g, m, dt, dist_init=0.02, dist_max=0.6):
+def check_p(x, y, flse, theta, m_set, m, dt, dist_init=0.02, dist_max=0.6):
     """
     Check if distance between eastern or northern particle to center particle is bigger than `dist_max`
     """
@@ -60,11 +60,17 @@ def check_p(x, y, g, m, dt, dist_init=0.02, dist_max=0.6):
         de = dxe ** 2 + dye ** 2
 
         if dn >= delta or de >= delta:
-            s1 = dxe ** 2 + dxn ** 2 + dye ** 2 + dyn ** 2
+            s1 = dn + de
+            at1 = 2 * (dxe * dxn + dye * dyn)
+            at2 = de - dn
             s2 = ((dxn + dye) ** 2 + (dxe - dyn) ** 2) * (
                 (dxn - dye) ** 2 + (dxe + dyn) ** 2
             )
-            g[i] = 1 / (2 * dt) * log2(1 / (2 * dist_init ** 2) * (s1 + s2 ** 0.5))
+            flse[i] = 1 / (2 * dt) * log2(1 / (2 * dist_init ** 2) * (s1 + s2 ** 0.5))
+            theta[i] = arctan2(at1, at2 + s2) * 180 / pi
+            # To know where value are set
+            m_set[i] = False
+            # To stop partcile advection
             m[i0], m[i_n], m[i_e] = True, True, True
 
 
@@ -110,10 +116,9 @@ def build_triplet(x, y, step=0.02):
 # Particles
 # ---------
 x0_, y0_ = -5, 30
-lon_p, lat_p = arange(x0_, x0_ + 43, step_grid_out), arange(
-    y0_, y0_ + 16, step_grid_out
-)
-x0, y0 = meshgrid(lon_p, lat_p)
+lon_p = arange(x0_, x0_ + 43, step_grid_out)
+lat_p = arange(y0_, y0_ + 16, step_grid_out)
+y0, x0 = meshgrid(lat_p, lon_p)
 grid_shape = x0.shape
 x0, y0 = x0.reshape(-1), y0.reshape(-1)
 # Identify all particle not on land
@@ -126,6 +131,8 @@ def build_triplet(x, y, step=0.02):
 
 # Array to compute fsle
 fsle = zeros(x0.shape[0], dtype="f4")
+theta = zeros(x0.shape[0], dtype="f4")
+mask = ones(x0.shape[0], dtype="f4")
 x, y = build_triplet(x0, y0, dist_init)
 used = zeros(x.shape[0], dtype="bool")
 
@@ -137,20 +144,23 @@ def build_triplet(x, y, step=0.02):
 for i in range(time_step_by_days * nb_days):
     t, xt, yt = p.__next__()
     dt = t / 86400.0 - t0
-    check_p(xt, yt, fsle, used, dt, dist_max=dist_max, dist_init=dist_init)
+    check_p(xt, yt, fsle, theta, mask, used, dt, dist_max=dist_max, dist_init=dist_init)
 
 # Get index with original_position
 i = ((x0 - x0_) / step_grid_out).astype("i4")
 j = ((y0 - y0_) / step_grid_out).astype("i4")
 fsle_ = empty(grid_shape, dtype="f4")
-used_ = zeros(grid_shape, dtype="bool")
-fsle_[j, i] = fsle
-used_[j, i] = used[::3]
+theta_ = empty(grid_shape, dtype="f4")
+mask_ = ones(grid_shape, dtype="bool")
+fsle_[i, j] = fsle
+theta_[i, j] = theta
+mask_[i, j] = mask
 # Create a grid object
 fsle_custom = RegularGridDataset.with_array(
     coordinates=("lon", "lat"),
     datas=dict(
-        fsle=ma.array(fsle_.T, mask=~used_.T),
+        fsle=ma.array(fsle_, mask=mask_),
+        theta=ma.array(theta_, mask=mask_),
         lon=lon_p,
         lat=lat_p,
     ),
@@ -169,3 +179,15 @@ def build_triplet(x, y, step=0.02):
 m = fsle_custom.display(ax, 1 / fsle_custom.grid("fsle"), **kw)
 ax.grid()
 cb = plt.colorbar(m, cax=fig.add_axes([0.94, 0.05, 0.01, 0.9]))
+# %%
+# Display Theta
+# -------------
+fig = plt.figure(figsize=(13, 5), dpi=150)
+ax = fig.add_axes([0.03, 0.03, 0.90, 0.94])
+ax.set_xlim(-6, 36.5), ax.set_ylim(30, 46)
+ax.set_aspect("equal")
+ax.set_title("Theta from finite size lyapunov exponent", weight="bold")
+kw = dict(cmap="Spectral_r", vmin=-180, vmax=180)
+m = fsle_custom.display(ax, fsle_custom.grid("theta"), **kw)
+ax.grid()
+cb = plt.colorbar(m, cax=fig.add_axes([0.94, 0.05, 0.01, 0.9]))

notebooks/python_module/07_cube_manipulation/pet_fsle_med.ipynb

@@ -26,7 +26,7 @@
       },
       "outputs": [],
       "source": [
-        "from matplotlib import pyplot as plt\nfrom numba import njit\nfrom numpy import arange, empty, isnan, log2, ma, meshgrid, zeros\n\nfrom py_eddy_tracker import start_logger\nfrom py_eddy_tracker.data import get_path\nfrom py_eddy_tracker.dataset.grid import GridCollection, RegularGridDataset\n\nstart_logger().setLevel(\"ERROR\")"
+        "from matplotlib import pyplot as plt\nfrom numba import njit\nfrom numpy import arange, arctan2, empty, isnan, log2, ma, meshgrid, ones, pi, zeros\n\nfrom py_eddy_tracker import start_logger\nfrom py_eddy_tracker.data import get_path\nfrom py_eddy_tracker.dataset.grid import GridCollection, RegularGridDataset\n\nstart_logger().setLevel(\"ERROR\")"
       ]
     },
     {
@@ -62,7 +62,7 @@
       },
       "outputs": [],
       "source": [
-        "@njit(cache=True, fastmath=True)\ndef check_p(x, y, g, m, dt, dist_init=0.02, dist_max=0.6):\n    \"\"\"\n    Check if distance between eastern or northern particle to center particle is bigger than `dist_max`\n    \"\"\"\n    nb_p = x.shape[0] // 3\n    delta = dist_max ** 2\n    for i in range(nb_p):\n        i0 = i * 3\n        i_n = i0 + 1\n        i_e = i0 + 2\n        # If particle already set, we skip\n        if m[i0] or m[i_n] or m[i_e]:\n            continue\n        # Distance with north\n        dxn, dyn = x[i0] - x[i_n], y[i0] - y[i_n]\n        dn = dxn ** 2 + dyn ** 2\n        # Distance with east\n        dxe, dye = x[i0] - x[i_e], y[i0] - y[i_e]\n        de = dxe ** 2 + dye ** 2\n\n        if dn >= delta or de >= delta:\n            s1 = dxe ** 2 + dxn ** 2 + dye ** 2 + dyn ** 2\n            s2 = ((dxn + dye) ** 2 + (dxe - dyn) ** 2) * (\n                (dxn - dye) ** 2 + (dxe + dyn) ** 2\n            )\n            g[i] = 1 / (2 * dt) * log2(1 / (2 * dist_init ** 2) * (s1 + s2 ** 0.5))\n            m[i0], m[i_n], m[i_e] = True, True, True\n\n\n@njit(cache=True)\ndef build_triplet(x, y, step=0.02):\n    \"\"\"\n    Triplet building for each position we add east and north point with defined step\n    \"\"\"\n    nb_x = x.shape[0]\n    x_ = empty(nb_x * 3, dtype=x.dtype)\n    y_ = empty(nb_x * 3, dtype=y.dtype)\n    for i in range(nb_x):\n        i0 = i * 3\n        i_n, i_e = i0 + 1, i0 + 2\n        x__, y__ = x[i], y[i]\n        x_[i0], y_[i0] = x__, y__\n        x_[i_n], y_[i_n] = x__, y__ + step\n        x_[i_e], y_[i_e] = x__ + step, y__\n    return x_, y_"
+        "@njit(cache=True, fastmath=True)\ndef check_p(x, y, flse, theta, m_set, m, dt, dist_init=0.02, dist_max=0.6):\n    \"\"\"\n    Check if distance between eastern or northern particle to center particle is bigger than `dist_max`\n    \"\"\"\n    nb_p = x.shape[0] // 3\n    delta = dist_max ** 2\n    for i in range(nb_p):\n        i0 = i * 3\n        i_n = i0 + 1\n        i_e = i0 + 2\n        # If particle already set, we skip\n        if m[i0] or m[i_n] or m[i_e]:\n            continue\n        # Distance with north\n        dxn, dyn = x[i0] - x[i_n], y[i0] - y[i_n]\n        dn = dxn ** 2 + dyn ** 2\n        # Distance with east\n        dxe, dye = x[i0] - x[i_e], y[i0] - y[i_e]\n        de = dxe ** 2 + dye ** 2\n\n        if dn >= delta or de >= delta:\n            s1 = dn + de\n            at1 = 2 * (dxe * dxn + dye * dyn)\n            at2 = de - dn\n            s2 = ((dxn + dye) ** 2 + (dxe - dyn) ** 2) * (\n                (dxn - dye) ** 2 + (dxe + dyn) ** 2\n            )\n            flse[i] = 1 / (2 * dt) * log2(1 / (2 * dist_init ** 2) * (s1 + s2 ** 0.5))\n            theta[i] = arctan2(at1, at2 + s2) * 180 / pi\n            # To know where value are set\n            m_set[i] = False\n            # To stop partcile advection\n            m[i0], m[i_n], m[i_e] = True, True, True\n\n\n@njit(cache=True)\ndef build_triplet(x, y, step=0.02):\n    \"\"\"\n    Triplet building for each position we add east and north point with defined step\n    \"\"\"\n    nb_x = x.shape[0]\n    x_ = empty(nb_x * 3, dtype=x.dtype)\n    y_ = empty(nb_x * 3, dtype=y.dtype)\n    for i in range(nb_x):\n        i0 = i * 3\n        i_n, i_e = i0 + 1, i0 + 2\n        x__, y__ = x[i], y[i]\n        x_[i0], y_[i0] = x__, y__\n        x_[i_n], y_[i_n] = x__, y__ + step\n        x_[i_e], y_[i_e] = x__ + step, y__\n    return x_, y_"
       ]
     },
     {
@@ -80,7 +80,7 @@
       },
       "outputs": [],
       "source": [
-        "# Step in degrees for ouput\nstep_grid_out = 1/25.\n# Initial separation in degrees\ndist_init = 1/50.\n# Final separation in degrees\ndist_max = 1/5.\n# Time of start\nt0 = 20268\n# Number of time step by days\ntime_step_by_days = 5\n# Maximal time of advection\n# Here we limit because our data cube cover only 3 month\nnb_days = 85\n# Backward or forward\nbackward=True"
+        "# Step in degrees for ouput\nstep_grid_out = 1 / 25.0\n# Initial separation in degrees\ndist_init = 1 / 50.0\n# Final separation in degrees\ndist_max = 1 / 5.0\n# Time of start\nt0 = 20268\n# Number of time step by days\ntime_step_by_days = 5\n# Maximal time of advection\n# Here we limit because our data cube cover only 3 month\nnb_days = 85\n# Backward or forward\nbackward = True"
       ]
     },
     {
@@ -98,7 +98,7 @@
       },
       "outputs": [],
       "source": [
-        "x0_, y0_ = -5, 30\nlon_p, lat_p = arange(x0_, x0_ + 43, step_grid_out), arange(y0_, y0_ + 16, step_grid_out)\nx0, y0 = meshgrid(lon_p, lat_p)\ngrid_shape = x0.shape\nx0, y0 = x0.reshape(-1), y0.reshape(-1)\n# Identify all particle not on land\nm = ~isnan(c[t0].interp(\"adt\", x0, y0))\nx0, y0 = x0[m], y0[m]"
+        "x0_, y0_ = -5, 30\nlon_p = arange(x0_, x0_ + 43, step_grid_out)\nlat_p = arange(y0_, y0_ + 16, step_grid_out)\ny0, x0 = meshgrid(lat_p, lon_p)\ngrid_shape = x0.shape\nx0, y0 = x0.reshape(-1), y0.reshape(-1)\n# Identify all particle not on land\nm = ~isnan(c[t0].interp(\"adt\", x0, y0))\nx0, y0 = x0[m], y0[m]"
       ]
     },
     {
@@ -116,7 +116,7 @@
       },
       "outputs": [],
       "source": [
-        "# Array to compute fsle\nfsle = zeros(x0.shape[0], dtype=\"f4\")\nx, y = build_triplet(x0, y0, dist_init)\nused = zeros(x.shape[0], dtype=\"bool\")\n\n# advection generator\nkw = dict(t_init=t0, nb_step=1, backward=backward, mask_particule=used)\np = c.advect(x, y, \"u\", \"v\", time_step=86400 / time_step_by_days, **kw)\n\n# We check at each step of advection if particle distance is over `dist_max`\nfor i in range(time_step_by_days * nb_days):\n    t, xt, yt = p.__next__()\n    dt = t / 86400.0 - t0\n    check_p(xt, yt, fsle, used, dt, dist_max=dist_max, dist_init=dist_init)\n\n# Get index with original_position\ni = ((x0 - x0_) / step_grid_out).astype(\"i4\")\nj = ((y0 - y0_) / step_grid_out).astype(\"i4\")\nfsle_ = empty(grid_shape, dtype=\"f4\")\nused_ = zeros(grid_shape, dtype=\"bool\")\nfsle_[j, i] = fsle\nused_[j, i] = used[::3]\n# Create a grid object\nfsle_custom = RegularGridDataset.with_array(\n    coordinates=(\"lon\", \"lat\"),\n    datas=dict(\n        fsle=ma.array(fsle_.T, mask=~used_.T),\n        lon=lon_p,\n        lat=lat_p,\n    ),\n    centered=True,\n)"
+        "# Array to compute fsle\nfsle = zeros(x0.shape[0], dtype=\"f4\")\ntheta = zeros(x0.shape[0], dtype=\"f4\")\nmask = ones(x0.shape[0], dtype=\"f4\")\nx, y = build_triplet(x0, y0, dist_init)\nused = zeros(x.shape[0], dtype=\"bool\")\n\n# advection generator\nkw = dict(t_init=t0, nb_step=1, backward=backward, mask_particule=used)\np = c.advect(x, y, \"u\", \"v\", time_step=86400 / time_step_by_days, **kw)\n\n# We check at each step of advection if particle distance is over `dist_max`\nfor i in range(time_step_by_days * nb_days):\n    t, xt, yt = p.__next__()\n    dt = t / 86400.0 - t0\n    check_p(xt, yt, fsle, theta, mask, used, dt, dist_max=dist_max, dist_init=dist_init)\n\n# Get index with original_position\ni = ((x0 - x0_) / step_grid_out).astype(\"i4\")\nj = ((y0 - y0_) / step_grid_out).astype(\"i4\")\nfsle_ = empty(grid_shape, dtype=\"f4\")\ntheta_ = empty(grid_shape, dtype=\"f4\")\nmask_ = ones(grid_shape, dtype=\"bool\")\nfsle_[i, j] = fsle\ntheta_[i, j] = theta\nmask_[i, j] = mask\n# Create a grid object\nfsle_custom = RegularGridDataset.with_array(\n    coordinates=(\"lon\", \"lat\"),\n    datas=dict(\n        fsle=ma.array(fsle_, mask=mask_),\n        theta=ma.array(theta_, mask=mask_),\n        lon=lon_p,\n        lat=lat_p,\n    ),\n    centered=True,\n)"
       ]
     },
     {
@@ -136,6 +136,24 @@
       "source": [
         "fig = plt.figure(figsize=(13, 5), dpi=150)\nax = fig.add_axes([0.03, 0.03, 0.90, 0.94])\nax.set_xlim(-6, 36.5), ax.set_ylim(30, 46)\nax.set_aspect(\"equal\")\nax.set_title(\"Finite size lyapunov exponent\", weight=\"bold\")\nkw = dict(cmap=\"viridis_r\", vmin=-15, vmax=0)\nm = fsle_custom.display(ax, 1 / fsle_custom.grid(\"fsle\"), **kw)\nax.grid()\ncb = plt.colorbar(m, cax=fig.add_axes([0.94, 0.05, 0.01, 0.9]))"
       ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Display Theta\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "fig = plt.figure(figsize=(13, 5), dpi=150)\nax = fig.add_axes([0.03, 0.03, 0.90, 0.94])\nax.set_xlim(-6, 36.5), ax.set_ylim(30, 46)\nax.set_aspect(\"equal\")\nax.set_title(\"Theta from finite size lyapunov exponent\", weight=\"bold\")\nkw = dict(cmap=\"Spectral_r\", vmin=-180, vmax=180)\nm = fsle_custom.display(ax, fsle_custom.grid(\"theta\"), **kw)\nax.grid()\ncb = plt.colorbar(m, cax=fig.add_axes([0.94, 0.05, 0.01, 0.9]))"
+      ]
     }
   ],
   "metadata": {