diff --git a/docs/user_guide/examples/tutorial_interpolation.ipynb b/docs/user_guide/examples/tutorial_interpolation.ipynb index b64ed72c7..e1233ee25 100644 --- a/docs/user_guide/examples/tutorial_interpolation.ipynb +++ b/docs/user_guide/examples/tutorial_interpolation.ipynb @@ -260,8 +260,10 @@ "source": [ "## Interpolators on unstructured grids\n", "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", + "- `UxConstantFaceConstantZC` - this interpolator implements piecewise constant interpolation in both the lateral and vertical directions. It is appropriate for data that is registered to the face centers of the unstructured grid and centered on vertical layers.\n", + "- `UxLinearNodeConstantZC` - this interpolator implements barycentric interpolation in the lateral direction and piecewise constant in the vertical direction. It is appropriate for data that is registered to the corner vertices of the unstructured grid faces and centered on vertical layers.\n", + "- `UxLinearNodeLinearZF` - this interpolator implements barycentric interpolation in the lateral direction and piecewise linear interpolation in the vertical direction. It is appropriate for data that is registered to the corner vertices of the unstructured grid faces and on vertical layer interfaces.\n", + "- `UxConstantFaceLinearZF` - this interpolator implements piecewise constant interpolation in the lateral direction and piecewise linear interpolation in the vertical direction. It is appropriate for data that is registered to the face centers of the unstructured grid and centered on vertical layers\n", "\n", "To get started, we use a very simple generated `UxArray.UxDataset` that is included with Parcels." ] @@ -282,7 +284,7 @@ "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 {py:obj}`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." + "Next, we create the`Fieldset` object that will be used later in advancing particles. When using the `FieldSet.from_uxdataset()` method, Parcels takes care of automatically assigning the appropriate interpolator based on the coordinates of each data variable." ] }, { @@ -295,21 +297,7 @@ "\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])" + "fieldset = parcels.FieldSet.from_uxdataset(ds, mesh=\"flat\")" ] }, { @@ -342,7 +330,7 @@ "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." + "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 `UxLinearNodeLinearZF` interpolator for node-registered data. In the other, we use the `SampleTracer_Face` to showcase interpolating with the `UxConstantFaceConstantZC` interpolator for face-registered data." ] }, { @@ -496,7 +484,7 @@ ], "metadata": { "kernelspec": { - "display_name": "test-notebooks", + "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, @@ -510,7 +498,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.13.9" + "version": "3.12.11" } }, "nbformat": 4, diff --git a/docs/user_guide/examples/tutorial_nestedgrids.ipynb b/docs/user_guide/examples/tutorial_nestedgrids.ipynb index 7cc465100..8920ae767 100644 --- a/docs/user_guide/examples/tutorial_nestedgrids.ipynb +++ b/docs/user_guide/examples/tutorial_nestedgrids.ipynb @@ -312,7 +312,7 @@ " \"face_polygon\": (\n", " (\n", " \"time\",\n", - " \"nz\",\n", + " \"zf\",\n", " \"n_face\",\n", " ),\n", " face_poly[np.newaxis, np.newaxis, :],\n", @@ -325,7 +325,7 @@ " },\n", " coords={\n", " \"time\": np.array([np.timedelta64(0, \"ns\")]),\n", - " \"nz\": np.array([0]),\n", + " \"zf\": np.array([0]),\n", " \"n_node\": np.arange(n_node),\n", " \"n_face\": np.arange(n_face),\n", " },\n", @@ -338,8 +338,8 @@ " \"GridID\",\n", " uxda,\n", " # TODO note that here we need to use mesh=\"flat\" otherwise the hashing doesn't work. See https://github.com/Parcels-code/Parcels/pull/2439#discussion_r2627664010\n", - " parcels.UxGrid(uxda.uxgrid, z=uxda[\"nz\"], mesh=\"flat\"),\n", - " interp_method=parcels.interpolators.UxPiecewiseConstantFace,\n", + " parcels.UxGrid(uxda.uxgrid, z=uxda[\"zf\"], mesh=\"flat\"),\n", + " interp_method=parcels.interpolators.UxConstantFaceConstantZC,\n", ")\n", "fieldset = parcels.FieldSet([GridID])" ] @@ -559,7 +559,7 @@ ], "metadata": { "kernelspec": { - "display_name": "docs", + "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, diff --git a/docs/user_guide/examples_v3/tutorial_stommel_uxarray.ipynb b/docs/user_guide/examples_v3/tutorial_stommel_uxarray.ipynb index 4b51c34c6..4d41f8b75 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.UxConstantFaceConstantZC`. 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 UxConstantFaceConstantZC\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=UxConstantFaceConstantZC,\n", ")\n", "V = Field(\n", " name=\"V\",\n", " data=ds.V,\n", " grid=grid,\n", - " interp_method=UxPiecewiseConstantFace,\n", + " interp_method=UxConstantFaceConstantZC,\n", ")\n", "P = Field(\n", " name=\"P\",\n", " data=ds.p,\n", " grid=grid,\n", - " interp_method=UxPiecewiseConstantFace,\n", + " interp_method=UxConstantFaceConstantZC,\n", ")" ] }, diff --git a/src/parcels/_core/field.py b/src/parcels/_core/field.py index 988636f13..405578013 100644 --- a/src/parcels/_core/field.py +++ b/src/parcels/_core/field.py @@ -60,9 +60,8 @@ class Field: Notes ----- The xarray.DataArray or uxarray.UxDataArray object contains the field data and metadata. - - * dims: (time, [nz1 | nz], [face_lat | node_lat | edge_lat], [face_lon | node_lon | edge_lon]) - * attrs: (location, mesh, mesh) + * dims: (time, [zc | zf], [face_lat | node_lat | edge_lat], [face_lon | node_lon | edge_lon]) + * attrs: (location, mesh) When using a xarray.DataArray object: @@ -177,10 +176,10 @@ def zdim(self): if type(self.data) is xr.DataArray: return self.grid.zdim else: - if "nz1" in self.data.dims: - return self.data.sizes["nz1"] - elif "nz" in self.data.dims: - return self.data.sizes["nz"] + if "zc" in self.data.dims: + return self.data.sizes["zc"] + elif "zf" in self.data.dims: + return self.data.sizes["zf"] else: return 0 @@ -403,9 +402,9 @@ def _assert_valid_uxdataarray(data: ux.UxDataArray): uxarray.UxDataArray object. """ # Validate dimensions - if not ("nz1" in data.dims or "nz" in data.dims): + if not ("zc" in data.dims or "zf" in data.dims): raise ValueError( - "Field is missing a 'nz1' or 'nz' dimension in the field's metadata. " + "Field is missing a 'zc' or 'zf' dimension in the field's metadata. " "This attribute is required for xarray.DataArray objects." ) @@ -424,23 +423,6 @@ def _assert_valid_uxdataarray(data: ux.UxDataArray): "This attribute is required for xarray.DataArray objects." ) - _assert_valid_uxgrid(data.uxgrid) - - -def _assert_valid_uxgrid(grid): - """Verifies that all the required attributes are present in the uxarray.UxDataArray.UxGrid object.""" - if "Conventions" not in grid.attrs.keys(): - raise ValueError( - "Field is missing a 'Conventions' attribute in the field's metadata. " - "This attribute is required for uxarray.UxDataArray objects." - ) - if grid.attrs["Conventions"] != "UGRID-1.0": - raise ValueError( - "Field has a 'Conventions' attribute that is not 'UGRID-1.0'. " - "This attribute is required for uxarray.UxDataArray objects." - "See https://ugrid-conventions.github.io/ugrid-conventions/ for more information." - ) - def _assert_compatible_combination(data: xr.DataArray | ux.UxDataArray, grid: ux.Grid | XGrid): if isinstance(data, ux.UxDataArray): diff --git a/src/parcels/_core/fieldset.py b/src/parcels/_core/fieldset.py index 1477479d8..830e2d29b 100644 --- a/src/parcels/_core/fieldset.py +++ b/src/parcels/_core/fieldset.py @@ -19,7 +19,15 @@ 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 ( + UxConstantFaceConstantZC, + UxConstantFaceLinearZF, + UxLinearNodeConstantZC, + UxLinearNodeLinearZF, + XConstantField, + XLinear, + ZeroInterpolator, +) if TYPE_CHECKING: from parcels._core.basegrid import BaseGrid @@ -239,35 +247,39 @@ def from_copernicusmarine(ds: xr.Dataset): ) return FieldSet.from_sgrid_conventions(ds, mesh="spherical") - def from_fesom2(ds: ux.UxDataset): - """Create a FieldSet from a FESOM2 uxarray.UxDataset. + def from_uxdataset(ds: ux.UxDataset, mesh: str = "spherical"): + """Create a FieldSet from a Parcels compliant uxarray.UxDataset. + The main requirements for a uxDataset are naming conventions for vertical grid dimensions & coordinates + + zf - Name for coordinate and dimension for vertical positions at layer interfaces + zc - Name for coordinate and dimension for vertical positions at layer centers Parameters ---------- ds : uxarray.UxDataset - uxarray.UxDataset as obtained from the uxarray package. + uxarray.UxDataset as obtained from the uxarray package but with appropriate named vertical dimensions Returns ------- FieldSet FieldSet object containing the fields from the dataset that can be used for a Parcels simulation. """ - ds = ds.copy() ds_dims = list(ds.dims) - if not all(dim in ds_dims for dim in ["time", "nz", "nz1"]): + if not all(dim in ds_dims for dim in ["time", "zf", "zc"]): raise ValueError( - f"Dataset missing one of the required dimensions 'time', 'nz', or 'nz1'. Found dimensions {ds_dims}" + f"Dataset missing one of the required dimensions 'time', 'zf', or 'zc' for uxDataset. Found dimensions {ds_dims}" ) - grid = UxGrid(ds.uxgrid, z=ds.coords["nz"], mesh="spherical") - ds = _discover_fesom2_U_and_V(ds) + + grid = UxGrid(ds.uxgrid, z=ds.coords["zf"], mesh=mesh) + ds = _discover_ux_U_and_V(ds) fields = {} if "U" in ds.data_vars and "V" in ds.data_vars: fields["U"] = Field("U", ds["U"], grid, _select_uxinterpolator(ds["U"])) - fields["V"] = Field("V", ds["V"], grid, _select_uxinterpolator(ds["U"])) + fields["V"] = Field("V", ds["V"], grid, _select_uxinterpolator(ds["V"])) if "W" in ds.data_vars: - fields["W"] = Field("W", ds["W"], grid, _select_uxinterpolator(ds["U"])) + fields["W"] = Field("W", ds["W"], grid, _select_uxinterpolator(ds["W"])) fields["UVW"] = VectorField("UVW", fields["U"], fields["V"], fields["W"]) else: fields["UV"] = VectorField("UV", fields["U"], fields["V"]) @@ -277,6 +289,61 @@ def from_fesom2(ds: ux.UxDataset): return FieldSet(list(fields.values())) + def from_fesom2(ds: ux.UxDataset, mesh: str = "spherical"): + """Create a FieldSet from a FESOM2 uxarray.UxDataset. + + Parameters + ---------- + ds : uxarray.UxDataset + uxarray.UxDataset as obtained from the uxarray package. + + Returns + ------- + FieldSet + FieldSet object containing the fields from the dataset that can be used for a Parcels simulation. + """ + ds = ds.copy() + ds_dims = list(ds.dims) + if not all(dim in ds_dims for dim in ["time", "nz", "nz1"]): + raise ValueError( + f"Dataset missing one of the required dimensions 'time', 'nz', or 'nz1' for FESOM data. Found dimensions {ds_dims}" + ) + ds = ds.rename( + { + "nz": "zf", # Vertical Interface + "nz1": "zc", # Vertical Center + } + ).set_index(zf="zf", zc="zc") + + return FieldSet.from_uxdataset(ds, mesh=mesh) + + def from_icon(ds: ux.UxDataset, mesh: str = "spherical"): + """Create a FieldSet from a ICON uxarray.UxDataset. + + Parameters + ---------- + ds : uxarray.UxDataset + uxarray.UxDataset as obtained from the uxarray package. + + Returns + ------- + FieldSet + FieldSet object containing the fields from the dataset that can be used for a Parcels simulation. + """ + ds = ds.copy() + ds_dims = list(ds.dims) + if not all(dim in ds_dims for dim in ["time", "depth", "depth_2"]): + raise ValueError( + f"Dataset missing one of the required dimensions 'time', 'depth', or 'depth_2' for ICON data. Found dimensions {ds_dims}" + ) + ds = ds.rename( + { + "depth_2": "zf", # Vertical Interface + "depth": "zc", # Vertical Center + } + ).set_index(zf="zf", zc="zc") + return FieldSet.from_uxdataset(ds, mesh=mesh) + def from_sgrid_conventions( ds: xr.Dataset, mesh: Mesh ): # TODO: Update mesh to be discovered from the dataset metadata @@ -471,13 +538,13 @@ def _discover_copernicusmarine_U_and_V(ds: xr.Dataset) -> xr.Dataset: return ds -def _discover_fesom2_U_and_V(ds: ux.UxDataset) -> ux.UxDataset: +def _discover_ux_U_and_V(ds: ux.UxDataset) -> ux.UxDataset: # Common variable names for U and V found in UxDatasets - common_fesom_UV = [("unod", "vnod"), ("u", "v")] - common_fesom_W = ["w"] + common_ux_UV = [("unod", "vnod"), ("u", "v")] + common_ux_W = ["w"] if "W" not in ds: - for common_W in common_fesom_W: + for common_W in common_ux_W: if common_W in ds: ds = _ds_rename_using_standard_names(ds, {common_W: "W"}) break @@ -489,7 +556,7 @@ def _discover_fesom2_U_and_V(ds: ux.UxDataset) -> ux.UxDataset: "Dataset has only one of the two variables 'U' and 'V'. Please rename the appropriate variable in your dataset to have both 'U' and 'V' for Parcels simulation." ) - for common_U, common_V in common_fesom_UV: + for common_U, common_V in common_ux_UV: if common_U in ds: if common_V not in ds: raise ValueError( @@ -526,16 +593,20 @@ def _ds_rename_using_standard_names(ds: xr.Dataset | ux.UxDataset, name_dict: di 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, - # (nz,n_node): node/corner laterally, layer interfaces vertically — barycentric lateral & linear vertical - "nz,n_node": UxPiecewiseLinearNode, + # (zc,n_face): face-center laterally, layer centers vertically — piecewise constant + "zc,n_face": UxConstantFaceConstantZC, + # (zc,n_node): node/corner laterally, layer centers vertically — barycentric lateral & piecewise constant vertical + "zc,n_node": UxLinearNodeConstantZC, + # (zf,n_node): node/corner laterally, layer interfaces vertically — barycentric lateral & linear vertical + "zf,n_node": UxLinearNodeLinearZF, + # (zf,n_face): face-center laterally, layer interfaces vertically — piecewise constant lateral & linear vertical + "zf,n_face": UxConstantFaceLinearZF, } # Extract only spatial dimensions, neglecting time da_spatial_dims = tuple(d for d in da.dims if d not in ("time",)) if len(da_spatial_dims) != 2: raise ValueError( - "Fields on unstructured grids must have two spatial dimensions, one vertical (nz or nz1) and one lateral (n_face, n_edge, or n_node)" + "Fields on unstructured grids must have two spatial dimensions, one vertical (zf or zc) and one lateral (n_face, n_edge, or n_node)" ) # Construct key (string) for mapping to interpolator @@ -543,7 +614,7 @@ def _select_uxinterpolator(da: ux.UxDataArray): vdim = None ldim = None for d in da_spatial_dims: - if d in ("nz", "nz1"): + if d in ("zf", "zc"): vdim = d if d in ("n_face", "n_node"): ldim = d @@ -553,4 +624,8 @@ def _select_uxinterpolator(da: ux.UxDataArray): if key in supported_uxinterp_mapping.keys(): return supported_uxinterp_mapping[key] - return None + logger.warning( + f"Interpolator not found for UxDataArray {da.name} with spatial dimensions {da_spatial_dims}. Setting default interpolator to ZeroInterpolator" + ) + + return ZeroInterpolator diff --git a/src/parcels/_core/index_search.py b/src/parcels/_core/index_search.py index c2dd18351..a81a17da2 100644 --- a/src/parcels/_core/index_search.py +++ b/src/parcels/_core/index_search.py @@ -255,7 +255,7 @@ def uxgrid_point_in_cell(grid, y: np.ndarray, x: np.ndarray, yi: np.ndarray, xi: points = ptilde - pdotnhat[:, None] * nhat + face_vertices[:, 0, :] else: - nids = grid.uxgrid.face_node_connectivity[xi].values + nids = grid.uxgrid.face_node_connectivity[xi, :].values face_vertices = np.stack( ( grid.uxgrid.node_lon[nids.ravel()].values.reshape(nids.shape), @@ -283,7 +283,7 @@ def _triangle_area(A, B, C): if A.shape[-1] == 2: # 2D case: cross product reduces to scalar z-component cross = d1[..., 0] * d2[..., 1] - d1[..., 1] * d2[..., 0] - area = 0.5 * np.abs(cross) + area = 0.5 * cross elif A.shape[-1] == 3: # 3D case: full vector cross product cross = np.cross(d1, d2) @@ -307,8 +307,8 @@ def _barycentric_coordinates(nodes, points, min_area=1e-8): nodes : numpy.ndarray Polygon verties per query of shape (M, 3, 2/3) where M is the number of query points. The second dimension corresponds to the number of vertices - The last dimension can be either 2 or 3, where 3 corresponds to the (z, y, x) coordinates of each vertex and 2 corresponds to the - (lat, lon) coordinates of each vertex. + The last dimension can be either 2 or 3, where 3 corresponds to the (x, y, z) coordinates of each vertex and 2 corresponds to the + (lon, lat) coordinates of each vertex. points : numpy.ndarray Spherical coordinates of the point (M,2/3) where M is the number of query points. @@ -321,24 +321,21 @@ def _barycentric_coordinates(nodes, points, min_area=1e-8): """ M, K = nodes.shape[: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) + vi0 = nodes[:, 0, :] # (M,K,2) + vi1 = nodes[:, 1, :] # (M,K,2) + vi2 = nodes[:, 2, :] # (M,K,2) - # a0 = area(v_{i-1}, v_i, v_{i+1}) - a0 = _triangle_area(vim1, vi, vi1) # (M,K) + a = _triangle_area(vi0, vi1, vi2) # (M,K) - # a1 = area(P, v_{i-1}, v_i); a2 = area(P, v_i, v_{i+1}) P = points - a1 = _triangle_area(P, vi1, vim1) - a2 = _triangle_area(P, vim1, vi) - a3 = _triangle_area(P, vi, vi1) + a0 = _triangle_area(P, vi1, vi2) + a1 = _triangle_area(P, vi2, vi0) + a2 = _triangle_area(P, vi0, vi1) - # wi = a0 / (a1c * a2c) # (M,K) bcoords = np.zeros((M, K)) - bcoords[:, 0] = a1 / a0 - bcoords[:, 1] = a2 / a0 - bcoords[:, 2] = a3 / a0 + bcoords[:, 0] = a0 / a + bcoords[:, 1] = a1 / a + bcoords[:, 2] = a2 / a return bcoords diff --git a/src/parcels/_core/utils/unstructured.py b/src/parcels/_core/utils/unstructured.py index 26a339a6d..8e971b05d 100644 --- a/src/parcels/_core/utils/unstructured.py +++ b/src/parcels/_core/utils/unstructured.py @@ -1,6 +1,6 @@ DIM_TO_VERTICAL_LOCATION_MAP = { - "nz1": "center", - "nz": "face", + "zc": "center", + "zf": "face", } diff --git a/src/parcels/_datasets/unstructured/generated.py b/src/parcels/_datasets/unstructured/generated.py index 5382b31dd..69b425ea2 100644 --- a/src/parcels/_datasets/unstructured/generated.py +++ b/src/parcels/_datasets/unstructured/generated.py @@ -20,8 +20,6 @@ def simple_small_delaunay(nx=10, ny=10): 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) @@ -36,12 +34,12 @@ def simple_small_delaunay(nx=10, ny=10): 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) + U = np.zeros((1, zc.size, uxgrid.n_face), dtype=np.float64) + V = np.zeros((1, zc.size, uxgrid.n_face), dtype=np.float64) + W = np.zeros((1, zf.size, uxgrid.n_node), dtype=np.float64) + P = np.zeros((1, zc.size, 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) + Tface = np.zeros((1, zc.size, 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) @@ -50,7 +48,7 @@ def simple_small_delaunay(nx=10, ny=10): 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) + Tnode = np.zeros((1, zc.size, 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) @@ -58,10 +56,10 @@ def simple_small_delaunay(nx=10, ny=10): data=U, name="U", uxgrid=uxgrid, - dims=["time", "nz1", "n_face"], + dims=["time", "zc", "n_face"], coords=dict( time=(["time"], [TIME[0]]), - nz1=(["nz1"], zc), + zc=(["zc"], zc), ), attrs=dict( description="zonal velocity", units="m/s", location="face", mesh="delaunay", Conventions="UGRID-1.0" @@ -71,10 +69,10 @@ def simple_small_delaunay(nx=10, ny=10): data=V, name="V", uxgrid=uxgrid, - dims=["time", "nz1", "n_face"], + dims=["time", "zc", "n_face"], coords=dict( time=(["time"], [TIME[0]]), - nz1=(["nz1"], zc), + zc=(["zc"], zc), ), attrs=dict( description="meridional velocity", units="m/s", location="face", mesh="delaunay", Conventions="UGRID-1.0" @@ -84,10 +82,10 @@ def simple_small_delaunay(nx=10, ny=10): data=W, name="W", uxgrid=uxgrid, - dims=["time", "nz", "n_node"], + dims=["time", "zf", "n_node"], coords=dict( time=(["time"], [TIME[0]]), - nz=(["nz"], zf), + zf=(["zf"], zf), ), attrs=dict( description="meridional velocity", units="m/s", location="node", mesh="delaunay", Conventions="UGRID-1.0" @@ -97,10 +95,10 @@ def simple_small_delaunay(nx=10, ny=10): data=P, name="p", uxgrid=uxgrid, - dims=["time", "nz1", "n_face"], + dims=["time", "zc", "n_face"], coords=dict( time=(["time"], [TIME[0]]), - nz1=(["nz1"], zc), + zc=(["zc"], zc), ), attrs=dict(description="pressure", units="N/m^2", location="face", mesh="delaunay", Conventions="UGRID-1.0"), ) @@ -109,10 +107,10 @@ def simple_small_delaunay(nx=10, ny=10): data=Tface, name="T_face", uxgrid=uxgrid, - dims=["time", "nz1", "n_face"], + dims=["time", "zc", "n_face"], coords=dict( time=(["time"], [TIME[0]]), - nz1=(["nz1"], zc), + zc=(["zc"], zc), ), attrs=dict( description="Tracer field sampled on face centers", @@ -126,10 +124,10 @@ def simple_small_delaunay(nx=10, ny=10): data=Tnode, name="T_node", uxgrid=uxgrid, - dims=["time", "nz", "n_node"], + dims=["time", "zc", "n_node"], coords=dict( time=(["time"], [TIME[0]]), - nz1=(["nz"], zf), + zc=(["zc"], zc), ), attrs=dict( description="Tracer field sampled on face vertices", diff --git a/src/parcels/_datasets/unstructured/generic.py b/src/parcels/_datasets/unstructured/generic.py index 1d5b456dc..05088c9cb 100644 --- a/src/parcels/_datasets/unstructured/generic.py +++ b/src/parcels/_datasets/unstructured/generic.py @@ -21,10 +21,10 @@ def _stommel_gyre_delaunay(): The velocity field provides a slow moving interior circulation and a western boundary current. All fields are placed on the vertices of the grid and at the element vertical faces. """ - lon, lat = np.meshgrid(np.linspace(0, 60.0, Nx, dtype=np.float32), np.linspace(0, 60.0, Nx, dtype=np.float32)) + lon, lat = np.meshgrid(np.linspace(0, 60.0, Nx, dtype=np.float64), np.linspace(0, 60.0, Nx, dtype=np.float64)) lon_flat = lon.ravel() lat_flat = lat.ravel() - zf = np.linspace(0.0, 1000.0, 2, endpoint=True, dtype=np.float32) # Vertical element faces + zf = np.linspace(0.0, 1000.0, 2, endpoint=True, dtype=np.float64) # Vertical element faces zc = 0.5 * (zf[:-1] + zf[1:]) # Vertical element centers nz = zf.size nz1 = zc.size @@ -120,10 +120,10 @@ def _fesom2_square_delaunay_uniform_z_coordinate(): The pressure field is constant. All fields are placed on location consistent with FESOM2 variable placement conventions """ - lon, lat = np.meshgrid(np.linspace(0, 60.0, Nx, dtype=np.float32), np.linspace(0, 60.0, Nx, dtype=np.float32)) + lon, lat = np.meshgrid(np.linspace(0, 60.0, Nx, dtype=np.float64), np.linspace(0, 60.0, Nx, dtype=np.float64)) lon_flat = lon.ravel() lat_flat = lat.ravel() - zf = np.linspace(0.0, 1000.0, 10, endpoint=True, dtype=np.float32) # Vertical element faces + zf = np.linspace(0.0, 1000.0, 10, endpoint=True, dtype=np.float64) # Vertical element faces zc = 0.5 * (zf[:-1] + zf[1:]) # Vertical element centers nz = zf.size nz1 = zc.size @@ -218,12 +218,12 @@ def _fesom2_square_delaunay_antimeridian(): All fields are placed on location consistent with FESOM2 variable placement conventions """ lon, lat = np.meshgrid( - np.linspace(-210.0, -150.0, Nx, dtype=np.float32), np.linspace(-40.0, 40.0, Nx, dtype=np.float32) + np.linspace(-210.0, -150.0, Nx, dtype=np.float64), np.linspace(-40.0, 40.0, Nx, dtype=np.float64) ) # wrap longitude from [-180,180] lon_flat = lon.ravel() lat_flat = lat.ravel() - zf = np.linspace(0.0, 1000.0, 10, endpoint=True, dtype=np.float32) # Vertical element faces + zf = np.linspace(0.0, 1000.0, 10, endpoint=True, dtype=np.float64) # Vertical element faces zc = 0.5 * (zf[:-1] + zf[1:]) # Vertical element centers nz = zf.size nz1 = zc.size @@ -311,8 +311,109 @@ def _fesom2_square_delaunay_antimeridian(): return ux.UxDataset({"U": u, "V": v, "W": w, "p": p}, uxgrid=uxgrid) +def _icon_square_delaunay_uniform_z_coordinate(): + """ + Delaunay grid with uniform z-coordinate, mimicking an ICON dataset. + This dataset consists of a square domain with closed boundaries, where the grid is generated using Delaunay triangulation. + The bottom topography is flat and uniform, and the vertical grid spacing is constant with 10 layers spanning [0,1000.0] + The lateral velocity field components are non-zero constant, and the vertical velocity component is zero. + The pressure field is constant. + All fields are face registered and at vertical layer centers, except for the vertical velocity component, which is + at vertical layer interfaces. + """ + lon, lat = np.meshgrid(np.linspace(0, 60.0, Nx, dtype=np.float64), np.linspace(0, 60.0, Nx, dtype=np.float64)) + lon_flat = lon.ravel() + lat_flat = lat.ravel() + zf = np.linspace(0.0, 1000.0, 10, endpoint=True, dtype=np.float64) # Vertical element faces + zc = 0.5 * (zf[:-1] + zf[1:]) # Vertical element centers + + # mask any point on one of the boundaries + mask = ( + np.isclose(lon_flat, 0.0) | np.isclose(lon_flat, 60.0) | np.isclose(lat_flat, 0.0) | np.isclose(lat_flat, 60.0) + ) + + boundary_points = np.flatnonzero(mask) + + uxgrid = ux.Grid.from_points( + (lon_flat, lat_flat), + method="regional_delaunay", + boundary_points=boundary_points, + ) + + # Define arrays U (zonal), V (meridional) and P (sea surface height) + U = np.ones( + (T, zc.size, uxgrid.n_face), dtype=np.float64 + ) # Lateral velocity is on the element centers and face centers + V = np.ones( + (T, zc.size, uxgrid.n_face), dtype=np.float64 + ) # Lateral velocity is on the element centers and face centers + W = np.zeros( + (T, zf.size, uxgrid.n_face), dtype=np.float64 + ) # Vertical velocity is on the element faces and face vertices + P = np.ones((T, zc.size, uxgrid.n_node), dtype=np.float64) # Pressure is on the element centers and face vertices + + U[0, :, :] = zc[:, None] * uxgrid.face_lon.values[None, :] + V[0, :, :] = zc[:, None] * uxgrid.face_lat.values[None, :] + W[0, :, :] = zf[:, None] * uxgrid.face_lon.values[None, :] * uxgrid.face_lat.values[None, :] + P[0, :, :] = zc[:, None] * (uxgrid.node_lon.values[None, :] + uxgrid.node_lat.values[None, :]) + + u = ux.UxDataArray( + data=U, + name="U", + uxgrid=uxgrid, + dims=["time", "depth", "n_face"], + coords=dict( + time=(["time"], TIME), + depth=(["depth"], 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", "depth", "n_face"], + coords=dict( + time=(["time"], TIME), + depth=(["depth"], 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", "depth_2", "n_face"], + coords=dict( + time=(["time"], TIME), + depth_2=(["depth_2"], zf), + ), + attrs=dict( + description="vertical velocity", units="m/s", location="node", mesh="delaunay", Conventions="UGRID-1.0" + ), + ) + p = ux.UxDataArray( + data=P, + name="p", + uxgrid=uxgrid, + dims=["time", "depth", "n_node"], + coords=dict( + time=(["time"], TIME), + depth=(["depth"], zc), + ), + attrs=dict(description="pressure", units="N/m^2", location="node", mesh="delaunay", Conventions="UGRID-1.0"), + ) + + return ux.UxDataset({"U": u, "V": v, "W": w, "p": p}, uxgrid=uxgrid) + + datasets = { "stommel_gyre_delaunay": _stommel_gyre_delaunay(), "fesom2_square_delaunay_uniform_z_coordinate": _fesom2_square_delaunay_uniform_z_coordinate(), "fesom2_square_delaunay_antimeridian": _fesom2_square_delaunay_antimeridian(), + "icon_square_delaunay_uniform_z_coordinate": _icon_square_delaunay_uniform_z_coordinate(), } diff --git a/src/parcels/interpolators.py b/src/parcels/interpolators.py index 8e23d4952..a5fd8266c 100644 --- a/src/parcels/interpolators.py +++ b/src/parcels/interpolators.py @@ -18,8 +18,8 @@ __all__ = [ "CGrid_Tracer", "CGrid_Velocity", - "UxPiecewiseConstantFace", - "UxPiecewiseLinearNode", + "UxConstantFaceConstantZC", + "UxLinearNodeLinearZF", "XConstantField", "XFreeslip", "XLinear", @@ -641,30 +641,69 @@ def XLinearInvdistLandTracer( return values.compute() if is_dask_collection(values) else values -def UxPiecewiseConstantFace( +def UxConstantFaceConstantZC( particle_positions: dict[str, float | np.ndarray], grid_positions: dict[_UXGRID_AXES, dict[str, int | float | np.ndarray]], field: Field, ): - """ - Piecewise constant interpolation kernel for face registered data. - This interpolation method is appropriate for fields that are - face registered, such as u,v in FESOM. - """ + """Piecewise constant interpolation kernel for face registered data that is vertically centered (on zc points)""" return field.data.values[ grid_positions["T"]["index"], grid_positions["Z"]["index"], grid_positions["FACE"]["index"] ] -def UxPiecewiseLinearNode( +def UxConstantFaceLinearZF( + particle_positions: dict[str, float | np.ndarray], + grid_positions: dict[_UXGRID_AXES, dict[str, int | float | np.ndarray]], + field: Field, +): + """ + Piecewise constant interpolation (lateral) with linear vertical interpolation kernel for face registered data + that is located at vertical interface levels (on zf points) + """ + ti = grid_positions["T"]["index"] + zi, fi = grid_positions["Z"]["index"], grid_positions["FACE"]["index"] + z = particle_positions["z"] + + # The zi refers to the vertical layer index. The field in this routine are assumed to be defined at the vertical interface levels. + # For interface zi, the interface indices are [zi, zi+1], so we need to use the values at zi and zi+1. + # First, do barycentric interpolation in the lateral direction for each interface level + fzk = field.data.values[ti, zi, fi] + fzkp1 = field.data.values[ti, zi + 1, fi] + + # Then, do piecewise linear interpolation in the vertical direction + zk = field.grid.z.values[zi] + zkp1 = field.grid.z.values[zi + 1] + return (fzk * (zkp1 - z) + fzkp1 * (z - zk)) / (zkp1 - zk) # Linear interpolation in the vertical direction + + +def UxLinearNodeConstantZC( + particle_positions: dict[str, float | np.ndarray], + grid_positions: dict[_UXGRID_AXES, dict[str, int | float | np.ndarray]], + field: Field, +): + """ + Piecewise linear interpolation kernel for node registered data that is vertically centered (zc points). + Effectively, it applies barycentric interpolation in the lateral direction + and piecewise constant interpolation in the vertical direction. + """ + ti = grid_positions["T"]["index"] + zi, fi = grid_positions["Z"]["index"], grid_positions["FACE"]["index"] + bcoords = grid_positions["FACE"]["bcoord"] + node_ids = field.grid.uxgrid.face_node_connectivity[fi, :].values + return np.sum( + field.data.values[ti[:, None], zi[:, None], node_ids] * bcoords, axis=-1 + ) # Linear interpolation in the vertical direction + + +def UxLinearNodeLinearZF( particle_positions: dict[str, float | np.ndarray], grid_positions: dict[_UXGRID_AXES, dict[str, int | float | np.ndarray]], field: Field, ): """ - Piecewise linear interpolation kernel for node registered data located at vertical interface levels. - This interpolation method is appropriate for fields that are node registered such as the vertical - velocity W in FESOM2. Effectively, it applies barycentric interpolation in the lateral direction + Piecewise linear interpolation kernel for node registered data located at vertical interface levels (zf points). + Effectively, it applies barycentric interpolation in the lateral direction and piecewise linear interpolation in the vertical direction. """ ti = grid_positions["T"]["index"] diff --git a/tests/test_field.py b/tests/test_field.py index 7795c953d..afc3d26b0 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 UxConstantFaceConstantZC, UxLinearNodeConstantZC, UxLinearNodeLinearZF, XLinear def test_field_init_param_types(): @@ -186,18 +186,24 @@ def test_field_unstructured_z_linear(): linear functions of the vertical coordinate. This allows for testing of exactness of the interpolation methods. """ ds = datasets_unstructured["fesom2_square_delaunay_uniform_z_coordinate"].copy(deep=True) + ds = ds.rename( + { + "nz": "zf", # Vertical Interface + "nz1": "zc", # Vertical Center + } + ).set_index(zf="zf", zc="zc") # Change the pressure values to be linearly dependent on the vertical coordinate - for k, z in enumerate(ds.coords["nz1"]): + for k, z in enumerate(ds.coords["zc"]): ds["p"].values[:, k, :] = z # Change the vertical velocity values to be linearly dependent on the vertical coordinate - for k, z in enumerate(ds.coords["nz"]): + for k, z in enumerate(ds.coords["zf"]): ds["W"].values[:, k, :] = z - grid = UxGrid(ds.uxgrid, z=ds.coords["nz"], mesh="flat") + grid = UxGrid(ds.uxgrid, z=ds.coords["zf"], 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=UxLinearNodeConstantZC) # Test above first cell center - for piecewise constant, should return the depth of the first cell center assert np.isclose( @@ -215,7 +221,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=UxLinearNodeLinearZF) assert np.isclose( W.eval(time=[0], z=[10.0], y=[30.0], x=[30.0], applyConversion=False), 10.0, @@ -232,10 +238,16 @@ def test_field_unstructured_z_linear(): def test_field_constant_in_time(): """Tests field evaluation for a field with no time interval (i.e., constant in time).""" - ds = datasets_unstructured["stommel_gyre_delaunay"] - grid = UxGrid(ds.uxgrid, z=ds.coords["nz"], mesh="flat") + ds = datasets_unstructured["stommel_gyre_delaunay"].copy(deep=True) + ds = ds.rename( + { + "nz": "zf", # Vertical Interface + "nz1": "zc", # Vertical Center + } + ).set_index(zf="zf", zc="zc") + grid = UxGrid(ds.uxgrid, z=ds.coords["zf"], 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=UxConstantFaceConstantZC) # 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_fieldset.py b/tests/test_fieldset.py index efb3d0092..c45e07c35 100644 --- a/tests/test_fieldset.py +++ b/tests/test_fieldset.py @@ -322,6 +322,14 @@ def test_fieldset_from_copernicusmarine_with_W(caplog): assert "renamed it to 'W'" in caplog.text +def test_fieldset_from_icon(): + ds = datasets_unstructured["icon_square_delaunay_uniform_z_coordinate"] + fieldset = FieldSet.from_icon(ds) + assert "U" in fieldset.fields + assert "V" in fieldset.fields + assert "UVW" in fieldset.fields + + def test_fieldset_from_fesom2(): ds = datasets_unstructured["stommel_gyre_delaunay"] fieldset = FieldSet.from_fesom2(ds) diff --git a/tests/test_particleset_execute.py b/tests/test_particleset_execute.py index c998677bb..71d4cc121 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 UxConstantFaceConstantZC, UxLinearNodeLinearZF, XLinear from parcels.kernels import AdvectionEE, AdvectionRK2, AdvectionRK4, AdvectionRK4_3D, AdvectionRK45 from tests.common_kernels import DoNothing from tests.utils import DEFAULT_PARTICLES @@ -478,25 +478,31 @@ def SetLat2(p): def test_uxstommelgyre_pset_execute(): - ds = datasets_unstructured["stommel_gyre_delaunay"] - grid = UxGrid(grid=ds.uxgrid, z=ds.coords["nz"], mesh="spherical") + ds = datasets_unstructured["stommel_gyre_delaunay"].copy(deep=True) + ds = ds.rename( + { + "nz": "zf", # Vertical Interface + "nz1": "zc", # Vertical Center + } + ).set_index(zf="zf", zc="zc") + grid = UxGrid(grid=ds.uxgrid, z=ds.coords["zf"], mesh="spherical") U = Field( name="U", data=ds.U, grid=grid, - interp_method=UxPiecewiseConstantFace, + interp_method=UxConstantFaceConstantZC, ) V = Field( name="V", data=ds.V, grid=grid, - interp_method=UxPiecewiseConstantFace, + interp_method=UxConstantFaceConstantZC, ) P = Field( name="P", data=ds.p, grid=grid, - interp_method=UxPiecewiseConstantFace, + interp_method=UxConstantFaceConstantZC, ) UV = VectorField(name="UV", U=U, V=V) fieldset = FieldSet([UV, UV.U, UV.V, P]) @@ -518,31 +524,37 @@ def test_uxstommelgyre_pset_execute(): def test_uxstommelgyre_multiparticle_pset_execute(): - ds = datasets_unstructured["stommel_gyre_delaunay"] - grid = UxGrid(grid=ds.uxgrid, z=ds.coords["nz"], mesh="spherical") + ds = datasets_unstructured["stommel_gyre_delaunay"].copy(deep=True) + ds = ds.rename( + { + "nz": "zf", # Vertical Interface + "nz1": "zc", # Vertical Center + } + ).set_index(zf="zf", zc="zc") + grid = UxGrid(grid=ds.uxgrid, z=ds.coords["zf"], mesh="spherical") U = Field( name="U", data=ds.U, grid=grid, - interp_method=UxPiecewiseConstantFace, + interp_method=UxConstantFaceConstantZC, ) V = Field( name="V", data=ds.V, grid=grid, - interp_method=UxPiecewiseConstantFace, + interp_method=UxConstantFaceConstantZC, ) W = Field( name="W", data=ds.W, grid=grid, - interp_method=UxPiecewiseLinearNode, + interp_method=UxLinearNodeLinearZF, ) P = Field( name="P", data=ds.p, grid=grid, - interp_method=UxPiecewiseConstantFace, + interp_method=UxConstantFaceConstantZC, ) UVW = VectorField(name="UVW", U=U, V=V, W=W) fieldset = FieldSet([UVW, UVW.U, UVW.V, UVW.W, P]) @@ -569,19 +581,19 @@ def test_uxstommelgyre_pset_execute_output(): name="U", data=ds.U, grid=grid, - interp_method=UxPiecewiseConstantFace, + interp_method=UxConstantFaceConstantZC, ) V = Field( name="V", data=ds.V, grid=grid, - interp_method=UxPiecewiseConstantFace, + interp_method=UxConstantFaceConstantZC, ) P = Field( name="P", data=ds.p, grid=grid, - interp_method=UxPiecewiseConstantFace, + interp_method=UxConstantFaceConstantZC, ) 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 05aa3c5b2..bad99b0a0 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, + UxConstantFaceConstantZC, + UxLinearNodeLinearZF, ) @@ -28,6 +28,12 @@ def ds_fesom_channel() -> ux.UxDataset: f"{fesom_path}/w.fesom_channel.nc", ] ds = ux.open_mfdataset(grid_path, data_path).rename_vars({"u": "U", "v": "V", "w": "W"}) + ds = ds.rename( + { + "nz": "zf", # Vertical Interface + "nz1": "zc", # Vertical Center + } + ).set_index(zf="zf", zc="zc") return ds @@ -38,14 +44,14 @@ def uv_fesom_channel(ds_fesom_channel) -> VectorField: U=Field( name="U", data=ds_fesom_channel.U, - grid=UxGrid(ds_fesom_channel.uxgrid, z=ds_fesom_channel.coords["nz"], mesh="flat"), - interp_method=UxPiecewiseConstantFace, + grid=UxGrid(ds_fesom_channel.uxgrid, z=ds_fesom_channel.coords["zf"], mesh="flat"), + interp_method=UxConstantFaceConstantZC, ), 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, + grid=UxGrid(ds_fesom_channel.uxgrid, z=ds_fesom_channel.coords["zf"], mesh="flat"), + interp_method=UxConstantFaceConstantZC, ), ) return UV @@ -58,20 +64,20 @@ def uvw_fesom_channel(ds_fesom_channel) -> VectorField: U=Field( name="U", data=ds_fesom_channel.U, - grid=UxGrid(ds_fesom_channel.uxgrid, z=ds_fesom_channel.coords["nz"], mesh="flat"), - interp_method=UxPiecewiseConstantFace, + grid=UxGrid(ds_fesom_channel.uxgrid, z=ds_fesom_channel.coords["zf"], mesh="flat"), + interp_method=UxConstantFaceConstantZC, ), 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, + grid=UxGrid(ds_fesom_channel.uxgrid, z=ds_fesom_channel.coords["zf"], mesh="flat"), + interp_method=UxConstantFaceConstantZC, ), 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, + grid=UxGrid(ds_fesom_channel.uxgrid, z=ds_fesom_channel.coords["zf"], mesh="flat"), + interp_method=UxLinearNodeLinearZF, ), ) return UVW @@ -101,8 +107,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 = UxConstantFaceConstantZC + fieldset.V.interp_method = UxConstantFaceConstantZC def test_fesom2_square_delaunay_uniform_z_coordinate_eval(): @@ -112,14 +118,20 @@ def test_fesom2_square_delaunay_uniform_z_coordinate_eval(): Since the underlying data is constant, we can check that the values are as expected. """ ds = datasets_unstructured["fesom2_square_delaunay_uniform_z_coordinate"] - grid = UxGrid(ds.uxgrid, z=ds.coords["nz"], mesh="flat") + ds = ds.rename( + { + "nz": "zf", # Vertical Interface + "nz1": "zc", # Vertical Center + } + ).set_index(zf="zf", zc="zc") + grid = UxGrid(ds.uxgrid, z=ds.coords["zf"], 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=UxConstantFaceConstantZC), + V=Field(name="V", data=ds.V, grid=grid, interp_method=UxConstantFaceConstantZC), + W=Field(name="W", data=ds.W, grid=grid, interp_method=UxLinearNodeLinearZF), ) - P = Field(name="p", data=ds.p, grid=grid, interp_method=UxPiecewiseLinearNode) + P = Field(name="p", data=ds.p, grid=grid, interp_method=UxLinearNodeLinearZF) fieldset = FieldSet([UVW, P, UVW.U, UVW.V, UVW.W]) assert np.isclose( @@ -154,12 +166,18 @@ def test_fesom2_square_delaunay_antimeridian_eval(): Ensures that the fieldset can be created and evaluated correctly. Since the underlying data is constant, we can check that the values are as expected. """ - ds = datasets_unstructured["fesom2_square_delaunay_antimeridian"] + ds = datasets_unstructured["fesom2_square_delaunay_antimeridian"].copy(deep=True) + ds = ds.rename( + { + "nz": "zf", # Vertical Interface + "nz1": "zc", # Vertical Center + } + ).set_index(zf="zf", zc="zc") P = Field( name="p", data=ds.p, - grid=UxGrid(ds.uxgrid, z=ds.coords["nz"], mesh="spherical"), - interp_method=UxPiecewiseLinearNode, + grid=UxGrid(ds.uxgrid, z=ds.coords["zf"], mesh="spherical"), + interp_method=UxLinearNodeLinearZF, ) fieldset = FieldSet([P]) @@ -167,3 +185,46 @@ def test_fesom2_square_delaunay_antimeridian_eval(): assert np.isclose(fieldset.p.eval(time=[0], z=[1.0], y=[30.0], x=[-180.0], applyConversion=False), 1.0) assert np.isclose(fieldset.p.eval(time=[0], z=[1.0], y=[30.0], x=[180.0], applyConversion=False), 1.0) assert np.isclose(fieldset.p.eval(time=[0], z=[1.0], y=[30.0], x=[170.0], applyConversion=False), 1.0) + + +def test_icon_evals(): + ds = datasets_unstructured["icon_square_delaunay_uniform_z_coordinate"].copy(deep=True) + fieldset = FieldSet.from_icon(ds, mesh="flat") + + # Query points, are chosen to be just a fraction off from the center of a cell for testing + # This generic dataset has an effective lateral grid-spacing of 3 degrees and vertical grid + # spacing of 100m - shifting by 1/10 of a degree laterally and 10m vertically should keep us + # within the cell and make for easy exactness checking of constant and linear interpolation + xc = ds.uxgrid.face_lon.values + yc = ds.uxgrid.face_lat.values + zc = 0.0 * xc + ds.depth.values[1] # Make zc the same length as xc + + tq = 0.0 * xc + xq = xc + 0.1 + yq = yc + 0.1 + zq = zc + 10.0 + + # The exact function for U is U=z*x . The U variable is center registered both laterally and + # vertically. In this case, piecewise constant interpolation is expected in both directions. + # The expected value for interpolation is then just computed using the cell center locations + assert np.allclose(fieldset.U.eval(time=tq, z=zq, y=yq, x=xq, applyConversion=False), zc * xc) + + # The exact function for V is V=z*y . The V variable is center registered both laterally and + # vertically. In this case, piecewise constant interpolation is expected in both directions + # The expected value for interpolation is then just computed using the cell center locations + assert np.allclose(fieldset.V.eval(time=tq, z=zq, y=yq, x=xq, applyConversion=False), zc * yc) + + # The exact function for W is W=z*x*y . The W variable is center registered laterally and + # interface registered vertically. In this case, piecewise constant interpolation is expected + # laterally, while piecewise linear is expected vertically. + # The expected value for interpolation is then just computed using the cell center locations + # for the latitude and longitude, and the query point for the vertical interpolation + assert np.allclose(fieldset.W.eval(time=tq, z=zq, y=yq, x=xq, applyConversion=False), zq * yc * xc) + + # The exact function for P is P=z*(x+y) . The P variable is node registered laterally and + # center registered vertically. In this case, barycentric interpolation is expected + # laterally and piecewise constant is expected vertically + # Since barycentric interpolation is exact for functions f=a*x+b*y laterally, the expected + # value for interpolation is then just computed using query point locations + # for the latitude and longitude, and the layer centers vertically. + assert np.allclose(fieldset.p.eval(time=tq, z=zq, y=yq, x=xq, applyConversion=False), zc * (xq + yq)) diff --git a/tests/test_uxgrid.py b/tests/test_uxgrid.py index 7b7312c24..2241e68cf 100644 --- a/tests/test_uxgrid.py +++ b/tests/test_uxgrid.py @@ -3,10 +3,13 @@ from parcels import UxGrid from parcels._datasets.unstructured.generic import datasets as uxdatasets +GENERIC_Z_COORDS = ["nz", "zf", "depth_2"] + @pytest.mark.parametrize("uxds", [pytest.param(uxds, id=key) for key, uxds in uxdatasets.items()]) def test_uxgrid_init_on_generic_datasets(uxds): - UxGrid(uxds.uxgrid, z=uxds.coords["nz"], mesh="flat") + vertical_coord = next((z_coord for z_coord in uxds.coords if z_coord in GENERIC_Z_COORDS), None) + UxGrid(uxds.uxgrid, z=uxds.coords[vertical_coord], mesh="flat") @pytest.mark.parametrize("uxds", [uxdatasets["stommel_gyre_delaunay"]]) @@ -26,5 +29,5 @@ def test_uxgrid_mesh(uxds, mesh): def test_xgrid_get_axis_dim(uxds): grid = UxGrid(uxds.uxgrid, z=uxds.coords["nz"], mesh="flat") - assert grid.get_axis_dim("FACE") == 721 + assert grid.get_axis_dim("FACE") == 718 assert grid.get_axis_dim("Z") == 2 diff --git a/tests/utils/test_unstructured.py b/tests/utils/test_unstructured.py index e8c296fec..143220335 100644 --- a/tests/utils/test_unstructured.py +++ b/tests/utils/test_unstructured.py @@ -8,14 +8,14 @@ def test_get_vertical_location_from_dims(): # Test with nz1 dimension - assert get_vertical_location_from_dims(("nz1", "time")) == "center" + assert get_vertical_location_from_dims(("zc", "time")) == "center" # Test with nz dimension - assert get_vertical_location_from_dims(("nz", "time")) == "face" + assert get_vertical_location_from_dims(("zf", "time")) == "face" # Test with both dimensions with pytest.raises(ValueError): - get_vertical_location_from_dims(("nz1", "nz", "time")) + get_vertical_location_from_dims(("zc", "zf", "time")) # Test with no vertical dimension with pytest.raises(ValueError): @@ -24,10 +24,10 @@ def test_get_vertical_location_from_dims(): def test_get_vertical_dim_name_from_location(): # Test with center location - assert get_vertical_dim_name_from_location("center") == "nz1" + assert get_vertical_dim_name_from_location("center") == "zc" # Test with face location - assert get_vertical_dim_name_from_location("face") == "nz" + assert get_vertical_dim_name_from_location("face") == "zf" with pytest.raises(KeyError): get_vertical_dim_name_from_location("invalid_location")