Uxinterpolation tutorial v4

docs/user_guide/examples/tutorial_interpolation.ipynb

@@ -244,11 +244,221 @@
    "metadata": {},
    "source": [
     "## Interpolation on unstructured grids\n",
-    "```{note}\n",
-    "TODO: add example on simple unstructured fieldset\n",
-    "```\n",
-    "- UXPiecewiseConstantFace\n",
-    "- UXPiecewiseLinearNode"
+    "Parcels v4 supports the use of general circulation model output that is defined on unstructured grids. We include basic interpolators to help you get started, including\n",
+    "- `UXPiecewiseConstantFace` - this interpolator implements piecewise constant interpolation and is appropriate for data that is registered to the face centers of the unstructured grid\n",
+    "- `UXPiecewiseLinearNode` - this interpolator implements barycentric interpolation and is appropriate for data that is registered to the corner vertices of the unstructured grid faces\n",
+    "\n",
+    "To get started, we use a very simple generated `UxArray.UxDataset` that is included with Parcels."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from parcels._datasets.unstructured.generated import simple_small_delaunay\n",
+    "\n",
+    "ds = simple_small_delaunay(nx=4, ny=4)\n",
+    "ds"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Next, we create the `Field` and `Fieldset` objects that will be used later in advancing particles. When creating a `Field` or `VectorField` object for unstructured grid data, we attach a `parcels.UxGrid` object and attach an `interp_method` to each object. For data that is defined on face centers, we use the `UXPiecewiseConstantFace` interpolator and for data that is defined on the face vertices, we use the `UXPiecewiseLinearNode` interpolator. In this example, we  will look specifically at interpolating a tracer field that is defined by the same underlying analytical function, but is defined on both faces and vertices as separate fields."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import parcels\n",
+    "\n",
+    "grid = parcels.UxGrid(ds.uxgrid, ds.nz, mesh=\"flat\")\n",
+    "\n",
+    "Tnode = parcels.Field(\n",
+    "    \"T_node\",\n",
+    "    ds[\"T_node\"],\n",
+    "    grid,\n",
+    "    interp_method=parcels.interpolators.UXPiecewiseLinearNode,\n",
+    ")\n",
+    "Tface = parcels.Field(\n",
+    "    \"T_face\",\n",
+    "    ds[\"T_face\"],\n",
+    "    grid,\n",
+    "    interp_method=parcels.interpolators.UXPiecewiseConstantFace,\n",
+    ")\n",
+    "fieldset = parcels.FieldSet([Tnode, Tface])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We can now create particles to sample data using either the node registered or face registered data."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "SampleParticle = parcels.Particle.add_variable(\n",
+    "    parcels.Variable(\"Tracer\", dtype=np.float32, initial=np.nan)\n",
+    ")\n",
+    "\n",
+    "\n",
+    "def SampleTracer_Node(particles, fieldset):\n",
+    "    particles.Tracer = fieldset.T_node[particles]\n",
+    "\n",
+    "\n",
+    "def SampleTracer_Face(particles, fieldset):\n",
+    "    particles.Tracer = fieldset.T_face[particles]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Now, we can create a `ParticleSet` for each interpolation experiment. In one of the `ParticleSet` objects we use the `SampleTracer_Node` kernel to showcase interpolating with the `UXPiecewiseLinearNode` interpolator for node-registered data. In the other, we use the `SampleTracer_Face` to showcase interpolating with the `UxPiecewiseConstantFace` interpolator for face-registered data."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pset = {}\n",
+    "\n",
+    "xv, yv = np.meshgrid(np.linspace(0.2, 0.8, 8), np.linspace(0.2, 0.9, 8))\n",
+    "pset[\"node\"] = parcels.ParticleSet(\n",
+    "    fieldset, pclass=SampleParticle, lon=xv.flatten(), lat=yv.flatten()\n",
+    ")\n",
+    "pset[\"node\"].execute(\n",
+    "    SampleTracer_Node, runtime=np.timedelta64(1, \"D\"), dt=np.timedelta64(1, \"D\")\n",
+    ")\n",
+    "\n",
+    "pset[\"face\"] = parcels.ParticleSet(\n",
+    "    fieldset, pclass=SampleParticle, lon=xv.flatten(), lat=yv.flatten()\n",
+    ")\n",
+    "pset[\"face\"].execute(\n",
+    "    SampleTracer_Face, runtime=np.timedelta64(1, \"D\"), dt=np.timedelta64(1, \"D\")\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Now, we can plot the node-registered (`T_node`) and face-registered (`T_face`) fields with the values interpolated onto the particles for each field. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "import matplotlib.tri as mtri\n",
+    "import numpy as np\n",
+    "\n",
+    "fig, ax = plt.subplots(1, 2, figsize=(10, 5), constrained_layout=True)\n",
+    "\n",
+    "x = ds.uxgrid.node_lon.values\n",
+    "y = ds.uxgrid.node_lat.values\n",
+    "tri = ds.uxgrid.face_node_connectivity.values\n",
+    "triang = mtri.Triangulation(x, y, triangles=tri)\n",
+    "\n",
+    "\n",
+    "# Shade the triangle faces using smooth shading between the node registered data\n",
+    "T_node = np.squeeze(ds[\"T_node\"][0, 0, :].values)\n",
+    "tpc = ax[0].tripcolor(\n",
+    "    triang,\n",
+    "    T_node,\n",
+    "    shading=\"gouraud\",  # smooth shading\n",
+    "    cmap=\"viridis\",\n",
+    "    vmin=-1.0,\n",
+    "    vmax=1.0,\n",
+    ")\n",
+    "# Plot the element edges\n",
+    "ax[0].triplot(triang, color=\"black\", linewidth=0.5)\n",
+    "\n",
+    "# Plot the node registered data\n",
+    "sc1 = ax[0].scatter(\n",
+    "    x, y, c=T_node, cmap=\"viridis\", s=150, edgecolors=\"black\", vmin=-1.0, vmax=1.0\n",
+    ")\n",
+    "\n",
+    "ax[0].scatter(\n",
+    "    pset[\"node\"].lon,\n",
+    "    pset[\"node\"].lat,\n",
+    "    c=pset[\"node\"].Tracer,\n",
+    "    cmap=\"viridis\",\n",
+    "    edgecolors=\"k\",\n",
+    "    s=50,\n",
+    "    vmin=-1.0,\n",
+    "    vmax=1.0,\n",
+    ")\n",
+    "\n",
+    "cbar = fig.colorbar(sc1, ax=ax[0])\n",
+    "cbar.set_label(\"T\")\n",
+    "\n",
+    "ax[0].set_aspect(\"equal\", \"box\")\n",
+    "ax[0].set_xlabel(\"x\")\n",
+    "ax[0].set_ylabel(\"y\")\n",
+    "ax[0].set_title(\"Node Registered Data\")\n",
+    "\n",
+    "# # Plot the face registered data\n",
+    "T_face = np.squeeze(ds[\"T_face\"][0, 0, :].values)\n",
+    "assert triang.triangles.shape[0] == T_face.shape[0]\n",
+    "tpc2 = ax[1].tripcolor(\n",
+    "    triang,\n",
+    "    facecolors=T_face,\n",
+    "    shading=\"flat\",\n",
+    "    cmap=\"viridis\",\n",
+    "    edgecolors=\"k\",\n",
+    "    linewidth=0.5,\n",
+    "    vmin=-1.0,\n",
+    "    vmax=1.0,\n",
+    ")\n",
+    "# Plot the element edges\n",
+    "ax[1].triplot(triang, color=\"black\", linewidth=0.5)\n",
+    "\n",
+    "# Plot the face registered data\n",
+    "xf = ds.uxgrid.face_lon.values\n",
+    "yf = ds.uxgrid.face_lat.values\n",
+    "\n",
+    "ax[1].scatter(\n",
+    "    pset[\"face\"].lon,\n",
+    "    pset[\"face\"].lat,\n",
+    "    c=pset[\"face\"].Tracer,\n",
+    "    cmap=\"viridis\",\n",
+    "    edgecolors=\"k\",\n",
+    "    s=50,\n",
+    "    vmin=-1.0,\n",
+    "    vmax=1.0,\n",
+    ")\n",
+    "\n",
+    "cbar = fig.colorbar(tpc, ax=ax[1])\n",
+    "cbar.set_label(\"T\")\n",
+    "\n",
+    "ax[1].set_aspect(\"equal\", \"box\")\n",
+    "ax[1].set_xlabel(\"x\")\n",
+    "ax[1].set_ylabel(\"y\")\n",
+    "ax[1].set_title(\"Face Registered Data\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "In both plots the black lines show the edges that define the boundaries between the faces in the unstructured grid. For the node registered field, the background coloring is done using smooth shading based on the value of `T_node` at the corner nodes. For the face registered fields, each face is colored according to the value of `T_face` at the face center. The color of the particles is the value of the function that the particles take on via the corresponding interpolation method."
    ]
   },
   {
@@ -272,7 +482,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "test-notebooks",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
@@ -286,7 +496,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.0"
+   "version": "3.12.11"
   }
  },
  "nbformat": 4,

src/parcels/_core/index_search.py

@@ -209,9 +209,9 @@ def uxgrid_point_in_cell(grid, y: np.ndarray, x: np.ndarray, yi: np.ndarray, xi:
     x : np.ndarray
         Array of longitudes of the points to check.
     yi : np.ndarray
-        Array of face indices corresponding to the points.
-    xi : np.ndarray
         Not used, but included for compatibility with other search functions.
+    xi : np.ndarray
+        Array of face indices corresponding to the points.
 
     Returns
     -------
@@ -227,7 +227,7 @@ def uxgrid_point_in_cell(grid, y: np.ndarray, x: np.ndarray, yi: np.ndarray, xi:
         points = np.column_stack((x_cart.flatten(), y_cart.flatten(), z_cart.flatten()))
 
         # Get the vertex indices for each face
-        nids = grid.uxgrid.face_node_connectivity[yi].values
+        nids = grid.uxgrid.face_node_connectivity[xi].values
         face_vertices = np.stack(
             (
                 grid.uxgrid.node_x[nids.ravel()].values.reshape(nids.shape),
@@ -237,7 +237,7 @@ def uxgrid_point_in_cell(grid, y: np.ndarray, x: np.ndarray, yi: np.ndarray, xi:
             axis=-1,
         )
     else:
-        nids = grid.uxgrid.face_node_connectivity[yi].values
+        nids = grid.uxgrid.face_node_connectivity[xi].values
         face_vertices = np.stack(
             (
                 grid.uxgrid.node_lon[nids.ravel()].values.reshape(nids.shape),
@@ -249,11 +249,11 @@ def uxgrid_point_in_cell(grid, y: np.ndarray, x: np.ndarray, yi: np.ndarray, xi:
 
     M = len(points)
 
-    is_in_cell = np.zeros(M, dtype=np.int32)
-
     coords = _barycentric_coordinates(face_vertices, points)
-    is_in_cell = np.where(np.all((coords >= -1e-6) & (coords <= 1 + 1e-6), axis=1), 1, 0)
 
+    is_in_cell = np.zeros(M, dtype=np.int32)
+    is_in_cell = np.where(np.all((coords >= -1e-6), axis=1), 1, 0)
+    is_in_cell &= np.isclose(np.sum(coords, axis=1), 1.0, rtol=1e-3, atol=1e-6)
     return is_in_cell, coords
 
 
@@ -278,8 +278,10 @@ def _triangle_area(A, B, C):
 def _barycentric_coordinates(nodes, points, min_area=1e-8):
     """
     Compute the barycentric coordinates of a point P inside a convex polygon using area-based weights.
-    So that this method generalizes to n-sided polygons, we use the Waschpress points as the generalized
-    barycentric coordinates, which is only valid for convex polygons.
+    This method currently only works for triangular cells (K=3)
+
+    TODO: So that this method generalizes to n-sided polygons, we can use the Waschpress points as the generalized
+    barycentric coordinates, which is valid for convex polygons.
 
     Parameters
     ----------
@@ -300,29 +302,24 @@ def _barycentric_coordinates(nodes, points, min_area=1e-8):
     """
     M, K = nodes.shape[:2]
 
-    # roll(-1) to get vi+1, roll(+1) to get vi-1
-    vi = nodes  # (M,K,2)
-    vi1 = np.roll(nodes, shift=-1, axis=1)  # (M,K,2)
-    vim1 = np.roll(nodes, shift=+1, axis=1)  # (M,K,2)
+    vi = np.squeeze(nodes[:, 0, :])  # (M,K,2)
+    vi1 = np.squeeze(nodes[:, 1, :])  # (M,K,2)
+    vim1 = np.squeeze(nodes[:, 2, :])  # (M,K,2)
 
     # a0 = area(v_{i-1}, v_i, v_{i+1})
     a0 = _triangle_area(vim1, vi, vi1)  # (M,K)
 
     # a1 = area(P, v_{i-1}, v_i); a2 = area(P, v_i, v_{i+1})
-    P = points[:, None, :]  # (M,1,2) -> (M,K,2)
-    a1 = _triangle_area(P, vim1, vi)
-    a2 = _triangle_area(P, vi, vi1)
-
-    # clamp tiny denominators for stability
-    a1c = np.maximum(a1, min_area)
-    a2c = np.maximum(a2, min_area)
-
-    wi = a0 / (a1c * a2c)  # (M,K)
-
-    sum_wi = wi.sum(axis=1, keepdims=True)  # (M,1)
-    # Avoid 0/0: if sum_wi==0 (degenerate), keep zeros
-    with np.errstate(invalid="ignore", divide="ignore"):
-        bcoords = wi / sum_wi
+    P = points
+    a1 = _triangle_area(P, vi1, vim1)
+    a2 = _triangle_area(P, vim1, vi)
+    a3 = _triangle_area(P, vi, vi1)
+
+    # wi = a0 / (a1c * a2c)  # (M,K)
+    bcoords = np.zeros((M, K))
+    bcoords[:, 0] = a1 / a0
+    bcoords[:, 1] = a2 / a0
+    bcoords[:, 2] = a3 / a0
 
     return bcoords
 

src/parcels/_core/uxgrid.py

@@ -33,6 +33,8 @@ def __init__(self, grid: ux.grid.Grid, z: ux.UxDataArray, mesh) -> UxGrid:
         mesh : str
             The type of mesh used for the grid. Either "flat" or "spherical".
         """
+        if grid.n_max_face_nodes > 3:
+            raise ValueError("Provided ux.grid.Grid must contain only triangular cells (n_max_face_nodes=3)")
         self.uxgrid = grid
         if not isinstance(z, ux.UxDataArray):
             raise TypeError("z must be an instance of ux.UxDataArray")