diff --git a/docs/user_guide/examples/tutorial_interpolation.ipynb b/docs/user_guide/examples/tutorial_interpolation.ipynb index 6180c70ce..5bf5706fe 100644 --- a/docs/user_guide/examples/tutorial_interpolation.ipynb +++ b/docs/user_guide/examples/tutorial_interpolation.ipynb @@ -131,7 +131,10 @@ " fieldset, pclass=SampleParticle, lon=xv.flatten(), lat=yv.flatten()\n", " )\n", " pset[p_interp.__name__].execute(\n", - " SampleP, runtime=np.timedelta64(1, \"D\"), dt=np.timedelta64(1, \"D\")\n", + " SampleP,\n", + " runtime=np.timedelta64(1, \"D\"),\n", + " dt=np.timedelta64(1, \"D\"),\n", + " verbose_progress=False,\n", " )" ] }, @@ -239,26 +242,245 @@ "```" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Interpolation at boundaries\n", + "In some cases, we need to implement specific boundary conditions, for example to prevent particles from \n", + "getting \"stuck\" near land. [This guide](../examples_v3/documentation_unstuck_Agrid.ipynb) describes \n", + "how to implement this in parcels using `parcels.interpolators.XFreeslip` and `parcels.interpolators.XPartialslip`." + ] + }, { "cell_type": "markdown", "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": [ - "## Interpolation at boundaries\n", - "In some cases, we need to implement specific boundary conditions, for example to prevent particles from \n", - "getting \"stuck\" near land. [This guide](../examples_v3/documentation_unstuck_Agrid.ipynb) describes \n", - "how to implement this in parcels using `parcels.interpolators.XFreeslip` and `parcels.interpolators.XPartialslip`." + "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 numpy as np\n", + "\n", + "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", + "\n", + "pset[\"node\"] = parcels.ParticleSet(\n", + " fieldset, pclass=SampleParticle, lon=xv.flatten(), lat=yv.flatten()\n", + ")\n", + "pset[\"node\"].execute(\n", + " SampleTracer_Node,\n", + " runtime=np.timedelta64(1, \"D\"),\n", + " dt=np.timedelta64(1, \"D\"),\n", + " verbose_progress=False,\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,\n", + " runtime=np.timedelta64(1, \"D\"),\n", + " dt=np.timedelta64(1, \"D\"),\n", + " verbose_progress=False,\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 +494,7 @@ ], "metadata": { "kernelspec": { - "display_name": "test-notebooks", + "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, @@ -286,7 +508,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.0" + "version": "3.12.11" } }, "nbformat": 4, diff --git a/docs/user_guide/examples_v3/tutorial_stommel_uxarray.ipynb b/docs/user_guide/examples_v3/tutorial_stommel_uxarray.ipynb index bb0b3fd6a..4b51c34c6 100644 --- a/docs/user_guide/examples_v3/tutorial_stommel_uxarray.ipynb +++ b/docs/user_guide/examples_v3/tutorial_stommel_uxarray.ipynb @@ -112,7 +112,7 @@ "\n", "In Parcels, grid searching is conducted with respect to the faces. In other words, when a grid index `ei` is provided to an interpolation method, this refers the face index `fi` at vertical layer `zi` (when unraveled). Within the interpolation method, the `field.grid.uxgrid.face_node_connectivity` attribute can be used to obtain the node indices that surround the face. Using these connectivity tables is necessary for properly indexing node registered data.\n", "\n", - "For the example Stommel gyre dataset in this tutorial, the `u` and `v` velocity components are face registered (similar to FESOM). Parcels includes a nearest neighbor interpolator for face registered unstructured grid data through `Parcels.application_kernels.interpolation.UXPiecewiseConstantFace`. Below, we create the `Field`s `U` and `V` and associate them with the `UxGrid` we created previously and this interpolation method." + "For the example Stommel gyre dataset in this tutorial, the `u` and `v` velocity components are face registered (similar to FESOM). Parcels includes a nearest neighbor interpolator for face registered unstructured grid data through `Parcels.application_kernels.interpolation.UxPiecewiseConstantFace`. Below, we create the `Field`s `U` and `V` and associate them with the `UxGrid` we created previously and this interpolation method." ] }, { @@ -121,26 +121,26 @@ "metadata": {}, "outputs": [], "source": [ - "from parcels.application_kernels.interpolation import UXPiecewiseConstantFace\n", + "from parcels.application_kernels.interpolation import UxPiecewiseConstantFace\n", "from parcels.field import Field\n", "\n", "U = Field(\n", " name=\"U\",\n", " data=ds.U,\n", " grid=grid,\n", - " interp_method=UXPiecewiseConstantFace,\n", + " interp_method=UxPiecewiseConstantFace,\n", ")\n", "V = Field(\n", " name=\"V\",\n", " data=ds.V,\n", " grid=grid,\n", - " interp_method=UXPiecewiseConstantFace,\n", + " interp_method=UxPiecewiseConstantFace,\n", ")\n", "P = Field(\n", " name=\"P\",\n", " data=ds.p,\n", " grid=grid,\n", - " interp_method=UXPiecewiseConstantFace,\n", + " interp_method=UxPiecewiseConstantFace,\n", ")" ] }, diff --git a/src/parcels/_core/fieldset.py b/src/parcels/_core/fieldset.py index 4346fb233..b3d958990 100644 --- a/src/parcels/_core/fieldset.py +++ b/src/parcels/_core/fieldset.py @@ -18,7 +18,7 @@ from parcels._core.xgrid import _DEFAULT_XGCM_KWARGS, XGrid from parcels._logger import logger from parcels._typing import Mesh -from parcels.interpolators import UXPiecewiseConstantFace, UXPiecewiseLinearNode, XConstantField, XLinear +from parcels.interpolators import UxPiecewiseConstantFace, UxPiecewiseLinearNode, XConstantField, XLinear if TYPE_CHECKING: from parcels._core.basegrid import BaseGrid @@ -459,9 +459,9 @@ def _select_uxinterpolator(da: ux.UxDataArray): """Selects the appropriate uxarray interpolator for a given uxarray UxDataArray""" supported_uxinterp_mapping = { # (nz1,n_face): face-center laterally, layer centers vertically — piecewise constant - "nz1,n_face": UXPiecewiseConstantFace, + "nz1,n_face": UxPiecewiseConstantFace, # (nz,n_node): node/corner laterally, layer interfaces vertically — barycentric lateral & linear vertical - "nz,n_node": UXPiecewiseLinearNode, + "nz,n_node": UxPiecewiseLinearNode, } # Extract only spatial dimensions, neglecting time da_spatial_dims = tuple(d for d in da.dims if d not in ("time",)) diff --git a/src/parcels/_core/index_search.py b/src/parcels/_core/index_search.py index 173d14226..a225621fc 100644 --- a/src/parcels/_core/index_search.py +++ b/src/parcels/_core/index_search.py @@ -211,9 +211,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 ------- @@ -226,10 +226,10 @@ def uxgrid_point_in_cell(grid, y: np.ndarray, x: np.ndarray, yi: np.ndarray, xi: lon_rad = np.deg2rad(x) lat_rad = np.deg2rad(y) x_cart, y_cart, z_cart = _latlon_rad_to_xyz(lat_rad, lon_rad) - points = np.column_stack((x_cart.flatten(), y_cart.flatten(), z_cart.flatten())) + points = np.column_stack((x_cart.flatten(), y_cart.flatten(), z_cart.flatten())) # (M,3) # 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), @@ -238,8 +238,22 @@ def uxgrid_point_in_cell(grid, y: np.ndarray, x: np.ndarray, yi: np.ndarray, xi: ), axis=-1, ) + + # Get projection points onto element plane + # for the projection, all points are computed relative to v0 + r1 = np.squeeze(face_vertices[:, 1, :] - face_vertices[:, 0, :]) # (M,3) + r2 = np.squeeze(face_vertices[:, 2, :] - face_vertices[:, 0, :]) # (M,3) + nhat = np.cross(r1, r2) + norm = np.linalg.norm(nhat, axis=-1) + nhat = nhat / norm[:, None] + # Calculate the component of the points in the direction of nhat + ptilde = points - np.squeeze(face_vertices[:, 0, :]) + pdotnhat = np.sum(ptilde * nhat, axis=-1) + # Reconstruct points with normal component removed. + points = ptilde - pdotnhat[:, None] * nhat + np.squeeze(face_vertices[:, 0, :]) + 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), @@ -251,10 +265,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 @@ -280,8 +295,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 ---------- @@ -302,29 +319,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 diff --git a/src/parcels/_core/uxgrid.py b/src/parcels/_core/uxgrid.py index a3fc9e240..5e2f628a8 100644 --- a/src/parcels/_core/uxgrid.py +++ b/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") diff --git a/src/parcels/_datasets/unstructured/generated.py b/src/parcels/_datasets/unstructured/generated.py new file mode 100644 index 000000000..5382b31dd --- /dev/null +++ b/src/parcels/_datasets/unstructured/generated.py @@ -0,0 +1,143 @@ +import math + +import numpy as np +import uxarray as ux +import xarray as xr + +T = 2 +vmax = 1.0 +delta = 0.1 +TIME = xr.date_range("2000", "2001", T) + + +def simple_small_delaunay(nx=10, ny=10): + """ + Data on a small Delaunay grid. The naming convention of the dataset and grid is consistent with what is + provided by UXArray when reading in FESOM2 datasets. + """ + lon, lat = np.meshgrid(np.linspace(0, 1.0, nx, dtype=np.float32), np.linspace(0, 1.0, ny, dtype=np.float32)) + lon_flat = lon.ravel() + lat_flat = lat.ravel() + zf = np.linspace(0.0, 1000.0, 2, endpoint=True, dtype=np.float32) # Vertical element faces + zc = 0.5 * (zf[:-1] + zf[1:]) # Vertical element centers + nz = zf.size + nz1 = zc.size + + # mask any point on one of the boundaries + mask = np.isclose(lon_flat, 0.0) | np.isclose(lon_flat, 1.0) | np.isclose(lat_flat, 0.0) | np.isclose(lat_flat, 1.0) + + boundary_points = np.flatnonzero(mask) + + uxgrid = ux.Grid.from_points( + (lon_flat, lat_flat), + method="regional_delaunay", + boundary_points=boundary_points, + ) + uxgrid.attrs["Conventions"] = "UGRID-1.0" + + # Define arrays U (zonal), V (meridional), W (vertical), and P (sea surface height) + U = np.zeros((1, nz1, uxgrid.n_face), dtype=np.float64) + V = np.zeros((1, nz1, uxgrid.n_face), dtype=np.float64) + W = np.zeros((1, nz, uxgrid.n_node), dtype=np.float64) + P = np.zeros((1, nz1, uxgrid.n_face), dtype=np.float64) + # Define Tface, a ficticious tracer field on the face centroids + Tface = np.zeros((1, nz1, uxgrid.n_face), dtype=np.float64) + + for i, (x, y) in enumerate(zip(uxgrid.face_lon, uxgrid.face_lat, strict=False)): + P[0, :, i] = -vmax * delta * (1 - x) * (math.exp(-x / delta) - 1) * np.sin(math.pi * y) + U[0, :, i] = -vmax * (1 - math.exp(-x / delta) - x) * np.cos(math.pi * y) + V[0, :, i] = vmax * ((2.0 - x) * math.exp(-x / delta) - 1) * np.sin(math.pi * y) + Tface[0, :, i] = np.sin(math.pi * y) * np.cos(math.pi * x) + + # Define Tnode, the same ficticious tracer field as above but on the face corner vertices + Tnode = np.zeros((1, nz, uxgrid.n_node), dtype=np.float64) + for i, (x, y) in enumerate(zip(uxgrid.node_lon, uxgrid.node_lat, strict=False)): + Tnode[0, :, i] = np.sin(math.pi * y) * np.cos(math.pi * x) + + u = ux.UxDataArray( + data=U, + name="U", + uxgrid=uxgrid, + dims=["time", "nz1", "n_face"], + coords=dict( + time=(["time"], [TIME[0]]), + nz1=(["nz1"], zc), + ), + attrs=dict( + description="zonal velocity", units="m/s", location="face", mesh="delaunay", Conventions="UGRID-1.0" + ), + ) + v = ux.UxDataArray( + data=V, + name="V", + uxgrid=uxgrid, + dims=["time", "nz1", "n_face"], + coords=dict( + time=(["time"], [TIME[0]]), + nz1=(["nz1"], zc), + ), + attrs=dict( + description="meridional velocity", units="m/s", location="face", mesh="delaunay", Conventions="UGRID-1.0" + ), + ) + w = ux.UxDataArray( + data=W, + name="W", + uxgrid=uxgrid, + dims=["time", "nz", "n_node"], + coords=dict( + time=(["time"], [TIME[0]]), + nz=(["nz"], zf), + ), + attrs=dict( + description="meridional velocity", units="m/s", location="node", mesh="delaunay", Conventions="UGRID-1.0" + ), + ) + p = ux.UxDataArray( + data=P, + name="p", + uxgrid=uxgrid, + dims=["time", "nz1", "n_face"], + coords=dict( + time=(["time"], [TIME[0]]), + nz1=(["nz1"], zc), + ), + attrs=dict(description="pressure", units="N/m^2", location="face", mesh="delaunay", Conventions="UGRID-1.0"), + ) + + tface = ux.UxDataArray( + data=Tface, + name="T_face", + uxgrid=uxgrid, + dims=["time", "nz1", "n_face"], + coords=dict( + time=(["time"], [TIME[0]]), + nz1=(["nz1"], zc), + ), + attrs=dict( + description="Tracer field sampled on face centers", + units="None", + location="face", + mesh="delaunay", + Conventions="UGRID-1.0", + ), + ) + tnode = ux.UxDataArray( + data=Tnode, + name="T_node", + uxgrid=uxgrid, + dims=["time", "nz", "n_node"], + coords=dict( + time=(["time"], [TIME[0]]), + nz1=(["nz"], zf), + ), + attrs=dict( + description="Tracer field sampled on face vertices", + units="None", + location="node", + mesh="delaunay", + Conventions="UGRID-1.0", + ), + ) + + return ux.UxDataset({"U": u, "V": v, "W": w, "p": p, "T_face": tface, "T_node": tnode}, uxgrid=uxgrid) diff --git a/src/parcels/interpolators.py b/src/parcels/interpolators.py index abd707eea..8e23d4952 100644 --- a/src/parcels/interpolators.py +++ b/src/parcels/interpolators.py @@ -18,8 +18,8 @@ __all__ = [ "CGrid_Tracer", "CGrid_Velocity", - "UXPiecewiseConstantFace", - "UXPiecewiseLinearNode", + "UxPiecewiseConstantFace", + "UxPiecewiseLinearNode", "XConstantField", "XFreeslip", "XLinear", @@ -641,7 +641,7 @@ def XLinearInvdistLandTracer( return values.compute() if is_dask_collection(values) else values -def UXPiecewiseConstantFace( +def UxPiecewiseConstantFace( particle_positions: dict[str, float | np.ndarray], grid_positions: dict[_UXGRID_AXES, dict[str, int | float | np.ndarray]], field: Field, @@ -656,7 +656,7 @@ def UXPiecewiseConstantFace( ] -def UXPiecewiseLinearNode( +def UxPiecewiseLinearNode( particle_positions: dict[str, float | np.ndarray], grid_positions: dict[_UXGRID_AXES, dict[str, int | float | np.ndarray]], field: Field, diff --git a/tests/test_field.py b/tests/test_field.py index 1cd2ae37e..7795c953d 100644 --- a/tests/test_field.py +++ b/tests/test_field.py @@ -9,7 +9,7 @@ from parcels._datasets.structured.generic import T as T_structured from parcels._datasets.structured.generic import datasets as datasets_structured from parcels._datasets.unstructured.generic import datasets as datasets_unstructured -from parcels.interpolators import UXPiecewiseConstantFace, UXPiecewiseLinearNode, XLinear +from parcels.interpolators import UxPiecewiseConstantFace, UxPiecewiseLinearNode, XLinear def test_field_init_param_types(): @@ -197,7 +197,7 @@ def test_field_unstructured_z_linear(): grid = UxGrid(ds.uxgrid, z=ds.coords["nz"], mesh="flat") # Note that the vertical coordinate is required to be the position of the layer interfaces ("nz"), not the mid-layers ("nz1") - P = Field(name="p", data=ds.p, grid=grid, interp_method=UXPiecewiseConstantFace) + P = Field(name="p", data=ds.p, grid=grid, interp_method=UxPiecewiseConstantFace) # Test above first cell center - for piecewise constant, should return the depth of the first cell center assert np.isclose( @@ -215,7 +215,7 @@ def test_field_unstructured_z_linear(): 944.44445801, ) - W = Field(name="W", data=ds.W, grid=grid, interp_method=UXPiecewiseLinearNode) + W = Field(name="W", data=ds.W, grid=grid, interp_method=UxPiecewiseLinearNode) assert np.isclose( W.eval(time=[0], z=[10.0], y=[30.0], x=[30.0], applyConversion=False), 10.0, @@ -235,7 +235,7 @@ def test_field_constant_in_time(): ds = datasets_unstructured["stommel_gyre_delaunay"] grid = UxGrid(ds.uxgrid, z=ds.coords["nz"], mesh="flat") # Note that the vertical coordinate is required to be the position of the layer interfaces ("nz"), not the mid-layers ("nz1") - P = Field(name="p", data=ds.p, grid=grid, interp_method=UXPiecewiseConstantFace) + P = Field(name="p", data=ds.p, grid=grid, interp_method=UxPiecewiseConstantFace) # Assert that the field can be evaluated at any time, and returns the same value time = np.datetime64("2000-01-01T00:00:00") diff --git a/tests/test_particleset_execute.py b/tests/test_particleset_execute.py index 3e4397aed..60be3da19 100644 --- a/tests/test_particleset_execute.py +++ b/tests/test_particleset_execute.py @@ -23,7 +23,7 @@ from parcels._datasets.structured.generated import simple_UV_dataset from parcels._datasets.structured.generic import datasets as datasets_structured from parcels._datasets.unstructured.generic import datasets as datasets_unstructured -from parcels.interpolators import UXPiecewiseConstantFace, UXPiecewiseLinearNode, XLinear +from parcels.interpolators import UxPiecewiseConstantFace, UxPiecewiseLinearNode, XLinear from parcels.kernels import AdvectionEE, AdvectionRK2, AdvectionRK4, AdvectionRK4_3D, AdvectionRK45 from tests.common_kernels import DoNothing from tests.utils import DEFAULT_PARTICLES @@ -490,19 +490,19 @@ def test_uxstommelgyre_pset_execute(): name="U", data=ds.U, grid=grid, - interp_method=UXPiecewiseConstantFace, + interp_method=UxPiecewiseConstantFace, ) V = Field( name="V", data=ds.V, grid=grid, - interp_method=UXPiecewiseConstantFace, + interp_method=UxPiecewiseConstantFace, ) P = Field( name="P", data=ds.p, grid=grid, - interp_method=UXPiecewiseConstantFace, + interp_method=UxPiecewiseConstantFace, ) UV = VectorField(name="UV", U=U, V=V) fieldset = FieldSet([UV, UV.U, UV.V, P]) @@ -519,8 +519,8 @@ def test_uxstommelgyre_pset_execute(): runtime=np.timedelta64(10, "m"), dt=np.timedelta64(60, "s"), ) - np.testing.assert_allclose(pset[0].lon, 29.997648, atol=1e-3) - np.testing.assert_allclose(pset[0].lat, 4.998691, atol=1e-3) + np.testing.assert_allclose(pset[0].lon, 29.997387, atol=1e-3) + np.testing.assert_allclose(pset[0].lat, 4.998546, atol=1e-3) def test_uxstommelgyre_multiparticle_pset_execute(): @@ -530,25 +530,25 @@ def test_uxstommelgyre_multiparticle_pset_execute(): name="U", data=ds.U, grid=grid, - interp_method=UXPiecewiseConstantFace, + interp_method=UxPiecewiseConstantFace, ) V = Field( name="V", data=ds.V, grid=grid, - interp_method=UXPiecewiseConstantFace, + interp_method=UxPiecewiseConstantFace, ) W = Field( name="W", data=ds.W, grid=grid, - interp_method=UXPiecewiseLinearNode, + interp_method=UxPiecewiseLinearNode, ) P = Field( name="P", data=ds.p, grid=grid, - interp_method=UXPiecewiseConstantFace, + interp_method=UxPiecewiseConstantFace, ) UVW = VectorField(name="UVW", U=U, V=V, W=W) fieldset = FieldSet([UVW, UVW.U, UVW.V, UVW.W, P]) @@ -575,19 +575,19 @@ def test_uxstommelgyre_pset_execute_output(): name="U", data=ds.U, grid=grid, - interp_method=UXPiecewiseConstantFace, + interp_method=UxPiecewiseConstantFace, ) V = Field( name="V", data=ds.V, grid=grid, - interp_method=UXPiecewiseConstantFace, + interp_method=UxPiecewiseConstantFace, ) P = Field( name="P", data=ds.p, grid=grid, - interp_method=UXPiecewiseConstantFace, + interp_method=UxPiecewiseConstantFace, ) UV = VectorField(name="UV", U=U, V=V) fieldset = FieldSet([UV, UV.U, UV.V, P]) diff --git a/tests/test_uxarray_fieldset.py b/tests/test_uxarray_fieldset.py index 570af9d34..05aa3c5b2 100644 --- a/tests/test_uxarray_fieldset.py +++ b/tests/test_uxarray_fieldset.py @@ -13,8 +13,8 @@ ) from parcels._datasets.unstructured.generic import datasets as datasets_unstructured from parcels.interpolators import ( - UXPiecewiseConstantFace, - UXPiecewiseLinearNode, + UxPiecewiseConstantFace, + UxPiecewiseLinearNode, ) @@ -39,13 +39,13 @@ def uv_fesom_channel(ds_fesom_channel) -> VectorField: name="U", data=ds_fesom_channel.U, grid=UxGrid(ds_fesom_channel.uxgrid, z=ds_fesom_channel.coords["nz"], mesh="flat"), - interp_method=UXPiecewiseConstantFace, + interp_method=UxPiecewiseConstantFace, ), V=Field( name="V", data=ds_fesom_channel.V, grid=UxGrid(ds_fesom_channel.uxgrid, z=ds_fesom_channel.coords["nz"], mesh="flat"), - interp_method=UXPiecewiseConstantFace, + interp_method=UxPiecewiseConstantFace, ), ) return UV @@ -59,19 +59,19 @@ def uvw_fesom_channel(ds_fesom_channel) -> VectorField: name="U", data=ds_fesom_channel.U, grid=UxGrid(ds_fesom_channel.uxgrid, z=ds_fesom_channel.coords["nz"], mesh="flat"), - interp_method=UXPiecewiseConstantFace, + interp_method=UxPiecewiseConstantFace, ), V=Field( name="V", data=ds_fesom_channel.V, grid=UxGrid(ds_fesom_channel.uxgrid, z=ds_fesom_channel.coords["nz"], mesh="flat"), - interp_method=UXPiecewiseConstantFace, + interp_method=UxPiecewiseConstantFace, ), W=Field( name="W", data=ds_fesom_channel.W, grid=UxGrid(ds_fesom_channel.uxgrid, z=ds_fesom_channel.coords["nz"], mesh="flat"), - interp_method=UXPiecewiseLinearNode, + interp_method=UxPiecewiseLinearNode, ), ) return UVW @@ -101,8 +101,8 @@ def test_set_interp_methods(ds_fesom_channel, uv_fesom_channel): assert (fieldset.V.data == ds_fesom_channel.V).all() # Set the interpolation method for each field - fieldset.U.interp_method = UXPiecewiseConstantFace - fieldset.V.interp_method = UXPiecewiseConstantFace + fieldset.U.interp_method = UxPiecewiseConstantFace + fieldset.V.interp_method = UxPiecewiseConstantFace def test_fesom2_square_delaunay_uniform_z_coordinate_eval(): @@ -115,17 +115,37 @@ def test_fesom2_square_delaunay_uniform_z_coordinate_eval(): grid = UxGrid(ds.uxgrid, z=ds.coords["nz"], mesh="flat") UVW = VectorField( name="UVW", - U=Field(name="U", data=ds.U, grid=grid, interp_method=UXPiecewiseConstantFace), - V=Field(name="V", data=ds.V, grid=grid, interp_method=UXPiecewiseConstantFace), - W=Field(name="W", data=ds.W, grid=grid, interp_method=UXPiecewiseLinearNode), + U=Field(name="U", data=ds.U, grid=grid, interp_method=UxPiecewiseConstantFace), + V=Field(name="V", data=ds.V, grid=grid, interp_method=UxPiecewiseConstantFace), + W=Field(name="W", data=ds.W, grid=grid, interp_method=UxPiecewiseLinearNode), ) - P = Field(name="p", data=ds.p, grid=grid, interp_method=UXPiecewiseLinearNode) + P = Field(name="p", data=ds.p, grid=grid, interp_method=UxPiecewiseLinearNode) fieldset = FieldSet([UVW, P, UVW.U, UVW.V, UVW.W]) - assert fieldset.U.eval(time=[0], z=[1.0], y=[30.0], x=[30.0], applyConversion=False) == 1.0 - assert fieldset.V.eval(time=[0], z=[1.0], y=[30.0], x=[30.0], applyConversion=False) == 1.0 - assert fieldset.W.eval(time=[0], z=[1.0], y=[30.0], x=[30.0], applyConversion=False) == 0.0 - assert fieldset.p.eval(time=[0], z=[1.0], y=[30.0], x=[30.0], applyConversion=False) == 1.0 + assert np.isclose( + fieldset.U.eval(time=[0.0], z=[1.0], y=[30.0], x=[30.0], applyConversion=False), + 1.0, + rtol=1e-3, + atol=1e-6, + ) + assert np.isclose( + fieldset.V.eval(time=[0.0], z=[1.0], y=[30.0], x=[30.0], applyConversion=False), + 1.0, + rtol=1e-3, + atol=1e-6, + ) + assert np.isclose( + fieldset.W.eval(time=[0.0], z=[1.0], y=[30.0], x=[30.0], applyConversion=False), + 0.0, + rtol=1e-3, + atol=1e-6, + ) + assert np.isclose( + fieldset.p.eval(time=[0.0], z=[1.0], y=[30.0], x=[30.0], applyConversion=False), + 1.0, + rtol=1e-3, + atol=1e-6, + ) def test_fesom2_square_delaunay_antimeridian_eval(): @@ -139,7 +159,7 @@ def test_fesom2_square_delaunay_antimeridian_eval(): name="p", data=ds.p, grid=UxGrid(ds.uxgrid, z=ds.coords["nz"], mesh="spherical"), - interp_method=UXPiecewiseLinearNode, + interp_method=UxPiecewiseLinearNode, ) fieldset = FieldSet([P])