Merge pull request AntSimi#128 from AntSimi/particles_drift

Particles drift

Author: AntSimi

examples/07_cube_manipulation/README.rst

@@ -1,6 +1,2 @@
 Time grid computation
 =====================
-
-.. warning::
-
-    Time grid is under development, API could move quickly!

examples/07_cube_manipulation/pet_particles_drift.py

@@ -0,0 +1,46 @@
+"""
+Build path of particle drifting
+===============================
+
+"""
+
+from matplotlib import pyplot as plt
+from numpy import arange, meshgrid
+
+from py_eddy_tracker import start_logger
+from py_eddy_tracker.data import get_demo_path
+from py_eddy_tracker.dataset.grid import GridCollection
+
+start_logger().setLevel("ERROR")
+
+# %%
+# Load data cube
+c = GridCollection.from_netcdf_cube(
+    get_demo_path("dt_med_allsat_phy_l4_2005T2.nc"),
+    "longitude",
+    "latitude",
+    "time",
+    heigth="adt",
+)
+
+# %%
+# Advection properties
+nb_days, step_by_day = 10, 6
+nb_time = step_by_day * nb_days
+kw_p = dict(nb_step=1, time_step=86400 / step_by_day)
+t0 = 20210
+
+# %%
+# Get paths
+x0, y0 = meshgrid(arange(32, 35, 0.5), arange(32.5, 34.5, 0.5))
+x0, y0 = x0.reshape(-1), y0.reshape(-1)
+t, x, y = c.path(x0, y0, "u", "v", t_init=t0, **kw_p, nb_time=nb_time)
+
+# %%
+# Plot paths
+ax = plt.figure(figsize=(9, 6)).add_subplot(111, aspect="equal")
+ax.plot(x0, y0, "k.", ms=20)
+ax.plot(x, y, lw=3)
+ax.set_title("10 days particle paths")
+ax.set_xlim(31, 35), ax.set_ylim(32, 34.5)
+ax.grid()

notebooks/python_module/07_cube_manipulation/pet_particles_drift.ipynb

@@ -0,0 +1,126 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "%matplotlib inline"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "\n# Build path of particle drifting\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "from matplotlib import pyplot as plt\nfrom numpy import arange, meshgrid\n\nfrom py_eddy_tracker import start_logger\nfrom py_eddy_tracker.data import get_demo_path\nfrom py_eddy_tracker.dataset.grid import GridCollection\n\nstart_logger().setLevel(\"ERROR\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "Load data cube\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "c = GridCollection.from_netcdf_cube(\n    get_demo_path(\"dt_med_allsat_phy_l4_2005T2.nc\"),\n    \"longitude\",\n    \"latitude\",\n    \"time\",\n    heigth=\"adt\",\n)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "Advection properties\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "nb_days, step_by_day = 10, 6\nnb_time = step_by_day * nb_days\nkw_p = dict(nb_step=1, time_step=86400 / step_by_day)\nt0 = 20210"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "Get paths\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "x0, y0 = meshgrid(arange(32, 35, 0.5), arange(32.5, 34.5, 0.5))\nx0, y0 = x0.reshape(-1), y0.reshape(-1)\nt, x, y = c.path(x0, y0, \"u\", \"v\", t_init=t0, **kw_p, nb_time=nb_time)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "Plot paths\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "ax = plt.figure(figsize=(9, 6)).add_subplot(111, aspect=\"equal\")\nax.plot(x0, y0, \"k.\", ms=20)\nax.plot(x, y, lw=3)\nax.set_title(\"10 days particle paths\")\nax.set_xlim(31, 35), ax.set_ylim(32, 34.5)\nax.grid()"
+      ]
+    }
+  ],
+  "metadata": {
+    "kernelspec": {
+      "display_name": "Python 3",
+      "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.7.7"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

src/py_eddy_tracker/dataset/grid.py

@@ -2517,6 +2517,28 @@ def get_previous_time_step(self, t_init):
             logger.debug(f"i={i}, t={t}, dataset={dataset}")
             yield t, dataset
 
+    def path(self, x0, y0, *args, nb_time=2, **kwargs):
+        """
+        At each call it will update position in place with u & v field
+
+        :param array x0: Longitude of obs to move
+        :param array y0: Latitude of obs to move
+        :param int nb_time: Number of iteration for particle
+        :param dict kwargs: look at :py:meth:`GridCollection.advect`
+
+        :return: t,x,y
+
+        .. minigallery:: py_eddy_tracker.GridCollection.path
+        """
+        particles = self.advect(x0.copy(), y0.copy(), *args, **kwargs)
+        t = empty(nb_time + 1, dtype="f8")
+        x = empty((nb_time + 1, x0.size), dtype=x0.dtype)
+        y = empty(x.shape, dtype=y0.dtype)
+        t[0], x[0], y[0] = kwargs.get("t_init"), x0, y0
+        for i in range(nb_time):
+            t[i + 1], x[i + 1], y[i + 1] = particles.__next__()
+        return t, x, y
+
 
 @njit(cache=True)
 def advect_t(x_g, y_g, u_g0, v_g0, m_g0, u_g1, v_g1, m_g1, x, y, m, weigths, half_w=0):