diff --git a/parcels/_datasets/unstructured/generic.py b/parcels/_datasets/unstructured/generic.py index 91704adeb2..ac9542638d 100644 --- a/parcels/_datasets/unstructured/generic.py +++ b/parcels/_datasets/unstructured/generic.py @@ -24,6 +24,10 @@ def _stommel_gyre_delaunay(): lon, lat = np.meshgrid(np.linspace(0, 60.0, Nx, dtype=np.float32), np.linspace(0, 60.0, Nx, 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 = ( @@ -40,9 +44,10 @@ def _stommel_gyre_delaunay(): uxgrid.attrs["Conventions"] = "UGRID-1.0" # Define arrays U (zonal), V (meridional) and P (sea surface height) - U = np.zeros((1, 1, lat.size), dtype=np.float64) - V = np.zeros((1, 1, lat.size), dtype=np.float64) - P = np.zeros((1, 1, lat.size), dtype=np.float64) + U = np.zeros((1, nz1, lat.size), dtype=np.float64) + V = np.zeros((1, nz1, lat.size), dtype=np.float64) + W = np.zeros((1, nz, lat.size), dtype=np.float64) + P = np.zeros((1, nz1, lat.size), dtype=np.float64) for i, (x, y) in enumerate(zip(lon_flat, lat_flat, strict=False)): xi = x / 60.0 @@ -72,7 +77,20 @@ def _stommel_gyre_delaunay(): dims=["time", "nz1", "n_node"], coords=dict( time=(["time"], [TIME[0]]), - nz1=(["nz1"], [0]), + nz1=(["nz1"], zc), + ), + attrs=dict( + description="meridional velocity", units="m/s", location="node", 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" @@ -85,12 +103,12 @@ def _stommel_gyre_delaunay(): dims=["time", "nz1", "n_node"], coords=dict( time=(["time"], [TIME[0]]), - nz1=(["nz1"], [0]), + nz1=(["nz1"], zc), ), attrs=dict(description="pressure", units="N/m^2", location="node", mesh="delaunay", Conventions="UGRID-1.0"), ) - return ux.UxDataset({"U": u, "V": v, "p": p}, uxgrid=uxgrid) + return ux.UxDataset({"U": u, "V": v, "W": w, "p": p}, uxgrid=uxgrid) def _fesom2_square_delaunay_uniform_z_coordinate(): diff --git a/parcels/application_kernels/interpolation.py b/parcels/application_kernels/interpolation.py index a1abf14064..18efd02763 100644 --- a/parcels/application_kernels/interpolation.py +++ b/parcels/application_kernels/interpolation.py @@ -26,9 +26,8 @@ def UXPiecewiseConstantFace( This interpolation method is appropriate for fields that are face registered, such as u,v in FESOM. """ - # TODO joe : handle vertical interpolation - zi, fi = field.unravel_index(ei) - return field.data[ti, zi, fi] + zi, fi = field.grid.unravel_index(ei) + return field.data.values[ti, zi, fi] def UXPiecewiseLinearNode( @@ -43,11 +42,20 @@ def UXPiecewiseLinearNode( x: np.float32 | np.float64, ): """ - Piecewise linear interpolation kernel for node registered data. This - interpolation method is appropriate for fields that are node registered - such as the vertical velocity w in FESOM. + 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 + and piecewise linear interpolation in the vertical direction. """ - # TODO joe : handle vertical interpolation - zi, fi = field.unravel_index(ei) - node_ids = field.data.uxgrid.face_node_connectivity[fi, :] - return np.dot(field.data[ti, zi, node_ids], bcoords) + k, fi = field.grid.unravel_index(ei) + node_ids = field.grid.uxgrid.face_node_connectivity[fi, :] + # 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 = np.dot(field.data.values[ti, k, node_ids], bcoords) + fzkp1 = np.dot(field.data.values[ti, k + 1, node_ids], bcoords) + + # Then, do piecewise linear interpolation in the vertical direction + zk = field.grid.z.values[k] + zkp1 = field.grid.z.values[k + 1] + return (fzk * (zkp1 - z) + fzkp1 * (z - zk)) / (zkp1 - zk) # Linear interpolation in the vertical direction diff --git a/parcels/field.py b/parcels/field.py index 5b66ffcce9..d1377db0d5 100644 --- a/parcels/field.py +++ b/parcels/field.py @@ -316,7 +316,7 @@ def eval(self, time: datetime, z, y, x, particle=None, applyConversion=True): try: tau, ti = _search_time_index(self, time) - bcoords, _ei = self.grid.search(self, z, y, x, ei=_ei) + bcoords, _ei = self.grid.search(z, y, x, ei=_ei) value = self._interp_method(self, ti, _ei, bcoords, tau, time, z, y, x) if np.isnan(value): diff --git a/parcels/uxgrid.py b/parcels/uxgrid.py index abdca45931..eccd51aa58 100644 --- a/parcels/uxgrid.py +++ b/parcels/uxgrid.py @@ -15,8 +15,25 @@ class UxGrid(BaseGrid): for interpolation on unstructured grids. """ - def __init__(self, grid: ux.grid.Grid) -> UxGrid: + def __init__(self, grid: ux.grid.Grid, z: ux.UxDataArray) -> UxGrid: + """ + Initializes the UxGrid with a uxarray grid and vertical coordinate array. + + Parameters + ---------- + grid : ux.grid.Grid + The uxarray grid object containing the unstructured grid data. + z : ux.UxDataArray + A 1D array of vertical coordinates (depths) associated with the layer interface heights (not the mid-layer depths). + While uxarray allows nz to be spatially and temporally varying, the parcels.UxGrid class considers the case where + the vertical coordinate is constant in time and space. This implies flat bottom topography and no moving ALE vertical grid. + """ self.uxgrid = grid + if not isinstance(z, ux.UxDataArray): + raise TypeError("z must be an instance of ux.UxDataArray") + if z.ndim != 1: + raise ValueError("z must be a 1D array of vertical coordinates") + self.z = z def search( self, z: float, y: float, x: float, ei: int | None = None, search2D: bool = False @@ -24,32 +41,46 @@ def search( tol = 1e-10 def try_face(fid): - # TODO : Vertical search is not implemented yet, so we assume z is not used. bcoords, err = self.uxgrid._get_barycentric_coordinates(y, x, fid) if (bcoords >= 0).all() and (bcoords <= 1).all() and err < tol: - return bcoords, self.ravel_index(0, fid) # Z and time indices are 0 for now + return bcoords, fid return None, None + def find_vertical_index() -> int: + if search2D: + return 0 + else: + nz = self.z.shape[0] + if nz == 1: + return 0 + zf = self.z.values + # Return zi such that zf[zi] <= z < zf[zi+1] + zi = np.searchsorted(zf, z, side="right") - 1 # Search assumes that z is positive and increasing with i + if zi < 0 or zi >= nz - 1: + raise FieldOutOfBoundError(z, y, x) + return zi + + zi = find_vertical_index() # Find the vertical cell center nearest to z + if ei is not None: - zi, fi = self.unravel_index(ei) - bcoords, ei_new = try_face(fi) + _, fi = self.unravel_index(ei) + bcoords, fi_new = try_face(fi) if bcoords is not None: - return bcoords, ei_new - + return bcoords, self.ravel_index(zi, fi_new) # Try neighbors of current face for neighbor in self.uxgrid.face_face_connectivity[fi, :]: if neighbor == -1: continue - bcoords, ei_new = try_face(neighbor) + bcoords, fi_new = try_face(neighbor) if bcoords is not None: - return bcoords, ei_new + return bcoords, self.ravel_index(zi, fi_new) # Global fallback using spatial hash fi, bcoords = self.uxgrid.get_spatial_hash().query([[x, y]]) if fi == -1: raise FieldOutOfBoundError(z, y, x) - return bcoords, self.ravel_index(zi, fi) + return bcoords[0], self.ravel_index(zi, fi[0]) def _get_barycentric_coordinates(self, y, x, fi): """Checks if a point is inside a given face id on a UxGrid.""" diff --git a/tests/v4/test_field.py b/tests/v4/test_field.py index fc481c624b..6dbe77f85b 100644 --- a/tests/v4/test_field.py +++ b/tests/v4/test_field.py @@ -3,7 +3,7 @@ import uxarray as ux import xarray as xr -from parcels import Field, xgcm +from parcels import Field, UXPiecewiseConstantFace, UXPiecewiseLinearNode, xgcm 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 @@ -33,7 +33,10 @@ def test_field_init_param_types(): pytest.param(ux.UxDataArray(), XGrid(xgcm.Grid(datasets_structured["ds_2d_left"])), id="uxdata-grid"), pytest.param( xr.DataArray(), - UxGrid(datasets_unstructured["stommel_gyre_delaunay"].uxgrid), + UxGrid( + datasets_unstructured["stommel_gyre_delaunay"].uxgrid, + z=datasets_unstructured["stommel_gyre_delaunay"].coords["nz"], + ), id="xarray-uxgrid", ), ], @@ -110,6 +113,38 @@ def test_vectorfield_init_different_time_intervals(): ... +def test_field_unstructured_z_linear(): + """Tests correctness of piecewise constant and piecewise linear interpolation methods on an unstructured grid with a vertical coordinate. + The example dataset is a FESOM2 square Delaunay grid with uniform z-coordinate. Cell centered and layer registered data are defined to be + 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) + + # Change the pressure values to be linearly dependent on the vertical coordinate + for k, z in enumerate(ds.coords["nz1"]): + 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"]): + ds["W"].values[:, k, :] = z + + grid = UxGrid(ds.uxgrid, z=ds.coords["nz"]) + # 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) + + # Test above first cell center - for piecewise constant, should return the depth of the first cell center + assert np.isclose(P.eval(time=ds.time[0].values, z=10.0, y=30.0, x=30.0, applyConversion=False), 55.555557) + # Test below first cell center, but in the first layer - for piecewise constant, should return the depth of the first cell center + assert np.isclose(P.eval(time=ds.time[0].values, z=65.0, y=30.0, x=30.0, applyConversion=False), 55.555557) + # Test bottom layer - for piecewise constant, should return the depth of the of the bottom layer cell center + assert np.isclose(P.eval(time=ds.time[0].values, z=900.0, y=30.0, x=30.0, applyConversion=False), 944.44445801) + + W = Field(name="W", data=ds.W, grid=grid, interp_method=UXPiecewiseLinearNode) + assert np.isclose(W.eval(time=ds.time[0].values, z=10.0, y=30.0, x=30.0, applyConversion=False), 10.0) + assert np.isclose(W.eval(time=ds.time[0].values, z=65.0, y=30.0, x=30.0, applyConversion=False), 65.0) + assert np.isclose(W.eval(time=ds.time[0].values, z=900.0, y=30.0, x=30.0, applyConversion=False), 900.0) + + def test_field_unstructured_grid_creation(): ... diff --git a/tests/v4/test_uxarray_fieldset.py b/tests/v4/test_uxarray_fieldset.py index bf4313129f..fbb2cc3840 100644 --- a/tests/v4/test_uxarray_fieldset.py +++ b/tests/v4/test_uxarray_fieldset.py @@ -37,13 +37,13 @@ def uv_fesom_channel(ds_fesom_channel) -> VectorField: U=Field( name="U", data=ds_fesom_channel.U, - grid=UxGrid(ds_fesom_channel.uxgrid), + grid=UxGrid(ds_fesom_channel.uxgrid, z=ds_fesom_channel.coords["nz"]), interp_method=UXPiecewiseConstantFace, ), V=Field( name="V", data=ds_fesom_channel.V, - grid=UxGrid(ds_fesom_channel.uxgrid), + grid=UxGrid(ds_fesom_channel.uxgrid, z=ds_fesom_channel.coords["nz"]), interp_method=UXPiecewiseConstantFace, ), ) @@ -57,17 +57,20 @@ def uvw_fesom_channel(ds_fesom_channel) -> VectorField: U=Field( name="U", data=ds_fesom_channel.U, - grid=UxGrid(ds_fesom_channel.uxgrid), + grid=UxGrid(ds_fesom_channel.uxgrid, z=ds_fesom_channel.coords["nz"]), interp_method=UXPiecewiseConstantFace, ), V=Field( name="V", data=ds_fesom_channel.V, - grid=UxGrid(ds_fesom_channel.uxgrid), + grid=UxGrid(ds_fesom_channel.uxgrid, z=ds_fesom_channel.coords["nz"]), interp_method=UXPiecewiseConstantFace, ), W=Field( - name="W", data=ds_fesom_channel.W, grid=UxGrid(ds_fesom_channel.uxgrid), interp_method=UXPiecewiseLinearNode + name="W", + data=ds_fesom_channel.W, + grid=UxGrid(ds_fesom_channel.uxgrid, z=ds_fesom_channel.coords["nz"]), + interp_method=UXPiecewiseLinearNode, ), ) return UVW @@ -114,7 +117,6 @@ def test_fesom_channel(ds_fesom_channel, uvw_fesom_channel): pset.execute(endtime=timedelta(days=1), dt=timedelta(hours=1)) -@pytest.mark.xfail(reason="https://github.com/OceanParcels/Parcels/pull/2026#issuecomment-2945609874") # TODO: Fix def test_fesom2_square_delaunay_uniform_z_coordinate_eval(): """ Test the evaluation of a fieldset with a FESOM2 square Delaunay grid and uniform z-coordinate. @@ -124,14 +126,14 @@ def test_fesom2_square_delaunay_uniform_z_coordinate_eval(): ds = datasets_unstructured["fesom2_square_delaunay_uniform_z_coordinate"] UVW = VectorField( name="UVW", - U=Field(name="U", data=ds.U, grid=UxGrid(ds.uxgrid), interp_method=UXPiecewiseConstantFace), - V=Field(name="V", data=ds.V, grid=UxGrid(ds.uxgrid), interp_method=UXPiecewiseConstantFace), - W=Field(name="W", data=ds.W, grid=UxGrid(ds.uxgrid), interp_method=UXPiecewiseLinearNode), + U=Field(name="U", data=ds.U, grid=UxGrid(ds.uxgrid, z=ds.coords["nz"]), interp_method=UXPiecewiseConstantFace), + V=Field(name="V", data=ds.V, grid=UxGrid(ds.uxgrid, z=ds.coords["nz"]), interp_method=UXPiecewiseConstantFace), + W=Field(name="W", data=ds.W, grid=UxGrid(ds.uxgrid, z=ds.coords["nz"]), interp_method=UXPiecewiseLinearNode), ) - P = Field(name="p", data=ds.p, grid=UxGrid(ds.uxgrid), interp_method=UXPiecewiseConstantFace) + P = Field(name="p", data=ds.p, grid=UxGrid(ds.uxgrid, z=ds.coords["nz"]), interp_method=UXPiecewiseConstantFace) fieldset = FieldSet([UVW, P, UVW.U, UVW.V, UVW.W]) assert fieldset.U.eval(time=ds.time[0].values, z=1.0, y=30.0, x=30.0, applyConversion=False) == 1.0 assert fieldset.V.eval(time=ds.time[0].values, z=1.0, y=30.0, x=30.0, applyConversion=False) == 1.0 - # assert fieldset.W.eval(time=ds.time[0].values, z=1.0, y=30.0, x=30.0, applyConversion=False) == 0.0 - assert fieldset.P.eval(time=ds.time[0].values, z=1.0, y=30.0, x=30.0, applyConversion=False) == 1.0 + assert fieldset.W.eval(time=ds.time[0].values, z=1.0, y=30.0, x=30.0, applyConversion=False) == 0.0 + assert fieldset.p.eval(time=ds.time[0].values, z=1.0, y=30.0, x=30.0, applyConversion=False) == 1.0