diff --git a/parcels/_index_search.py b/parcels/_index_search.py index dde987261..6c7f8a26c 100644 --- a/parcels/_index_search.py +++ b/parcels/_index_search.py @@ -83,14 +83,16 @@ def _search_indices_curvilinear_2d( cc = a[1] * b[0] - a[0] * b[1] + x * b[1] - y * a[1] det2 = bb * bb - 4 * aa * cc - det = np.where(det2 > 0, np.sqrt(det2), eta) - eta = np.where(abs(aa) < 1e-12, -cc / bb, np.where(det2 > 0, (-bb + det) / (2 * aa), eta)) - - xsi = np.where( - abs(a[1] + a[3] * eta) < 1e-12, - ((y - py[0]) / (py[1] - py[0]) + (y - py[3]) / (py[2] - py[3])) * 0.5, - (x - a[0] - a[2] * eta) / (a[1] + a[3] * eta), - ) + with np.errstate(divide="ignore", invalid="ignore"): + det = np.where(det2 > 0, np.sqrt(det2), eta) + + eta = np.where(abs(aa) < 1e-12, -cc / bb, np.where(det2 > 0, (-bb + det) / (2 * aa), eta)) + + xsi = np.where( + abs(a[1] + a[3] * eta) < 1e-12, + ((y - py[0]) / (py[1] - py[0]) + (y - py[3]) / (py[2] - py[3])) * 0.5, + (x - a[0] - a[2] * eta) / (a[1] + a[3] * eta), + ) xi = np.where(xsi < -tol, xi - 1, np.where(xsi > 1 + tol, xi + 1, xi)) yi = np.where(eta < -tol, yi - 1, np.where(eta > 1 + tol, yi + 1, yi)) diff --git a/parcels/application_kernels/interpolation.py b/parcels/application_kernels/interpolation.py index 20008c457..467d4280d 100644 --- a/parcels/application_kernels/interpolation.py +++ b/parcels/application_kernels/interpolation.py @@ -4,21 +4,26 @@ from typing import TYPE_CHECKING -import dask.array as dask import numpy as np import xarray as xr +from dask import is_dask_collection + +import parcels.tools.interpolation_utils as i_u if TYPE_CHECKING: - from parcels.field import Field + from parcels.field import Field, VectorField from parcels.uxgrid import _UXGRID_AXES from parcels.xgrid import _XGRID_AXES __all__ = [ + "CGrid_Tracer", + "CGrid_Velocity", "UXPiecewiseConstantFace", "UXPiecewiseLinearNode", "XLinear", "XNearest", "ZeroInterpolator", + "ZeroInterpolator_Vector", ] @@ -36,6 +41,21 @@ def ZeroInterpolator( return 0.0 +def ZeroInterpolator_Vector( + vectorfield: VectorField, + ti: int, + position: dict[str, tuple[int, float | np.ndarray]], + tau: np.float32 | np.float64, + t: np.float32 | np.float64, + z: np.float32 | np.float64, + y: np.float32 | np.float64, + x: np.float32 | np.float64, + applyConversion: bool, +) -> np.float32 | np.float64: + """Template function used for the signature check of the interpolation methods for velocity fields.""" + return 0.0 + + def XLinear( field: Field, ti: int, @@ -53,6 +73,7 @@ def XLinear( axis_dim = field.grid.get_axis_dim_mapping(field.data.dims) data = field.data + tdim, zdim, ydim, xdim = data.shape[0], data.shape[1], data.shape[2], data.shape[3] lenT = 2 if np.any(tau > 0) else 1 lenZ = 2 if np.any(zeta > 0) else 1 @@ -61,22 +82,22 @@ def XLinear( if lenT == 1: ti = np.repeat(ti, lenZ * 4) else: - ti_1 = np.clip(ti + 1, 0, data.shape[0] - 1) + ti_1 = np.clip(ti + 1, 0, tdim - 1) ti = np.concatenate([np.repeat(ti, lenZ * 4), np.repeat(ti_1, lenZ * 4)]) # Depth coordinates: 4 points at zi, 4 at zi+1, repeated for both time levels if lenZ == 1: zi = np.repeat(zi, lenT * 4) else: - zi_1 = np.clip(zi + 1, 0, data.shape[1] - 1) + zi_1 = np.clip(zi + 1, 0, zdim - 1) zi = np.tile(np.array([zi, zi, zi, zi, zi_1, zi_1, zi_1, zi_1]).flatten(), lenT) # Y coordinates: [yi, yi, yi+1, yi+1] for each spatial point, repeated for time/depth - yi_1 = np.clip(yi + 1, 0, data.shape[2] - 1) + yi_1 = np.clip(yi + 1, 0, ydim - 1) yi = np.tile(np.repeat(np.column_stack([yi, yi_1]), 2), (lenT) * (lenZ)) # X coordinates: [xi, xi+1, xi, xi+1] for each spatial point, repeated for time/depth - xi_1 = np.clip(xi + 1, 0, data.shape[3] - 1) + xi_1 = np.clip(xi + 1, 0, xdim - 1) xi = np.tile(np.column_stack([xi, xi_1, xi, xi_1]).flatten(), (lenT) * (lenZ)) # Create DataArrays for indexing @@ -109,7 +130,266 @@ def XLinear( + (1 - xsi) * eta * corner_data[:, 2] + xsi * eta * corner_data[:, 3] ) - return value.compute() if isinstance(value, dask.Array) else value + return value.compute() if is_dask_collection(value) else value + + +def CGrid_Velocity( + vectorfield: VectorField, + ti: int, + position: dict[_XGRID_AXES, tuple[int, float | np.ndarray]], + tau: np.float32 | np.float64, + t: np.float32 | np.float64, + z: np.float32 | np.float64, + y: np.float32 | np.float64, + x: np.float32 | np.float64, + applyConversion: bool, +): + """ + Interpolation kernel for velocity fields on a C-Grid. + Following Delandmeter and Van Sebille (2019), velocity fields should be interpolated + only in the direction of the grid cell faces. + """ + xi, xsi = position["X"] + yi, eta = position["Y"] + zi, zeta = position["Z"] + + U = vectorfield.U.data + V = vectorfield.V.data + grid = vectorfield.grid + tdim, zdim, ydim, xdim = U.shape[0], U.shape[1], U.shape[2], U.shape[3] + + if grid.lon.ndim == 1: + px = np.array([grid.lon[xi], grid.lon[xi + 1], grid.lon[xi + 1], grid.lon[xi]]) + py = np.array([grid.lat[yi], grid.lat[yi], grid.lat[yi + 1], grid.lat[yi + 1]]) + else: + px = np.array([grid.lon[yi, xi], grid.lon[yi, xi + 1], grid.lon[yi + 1, xi + 1], grid.lon[yi + 1, xi]]) + py = np.array([grid.lat[yi, xi], grid.lat[yi, xi + 1], grid.lat[yi + 1, xi + 1], grid.lat[yi + 1, xi]]) + + if grid._mesh == "spherical": + px[0] = np.where(px[0] < x - 225, px[0] + 360, px[0]) + px[0] = np.where(px[0] > x + 225, px[0] - 360, px[0]) + px[1:] = np.where(px[1:] - px[0] > 180, px[1:] - 360, px[1:]) + px[1:] = np.where(-px[1:] + px[0] > 180, px[1:] + 360, px[1:]) + c1 = i_u._geodetic_distance( + py[0], py[1], px[0], px[1], grid._mesh, np.einsum("ij,ji->i", i_u.phi2D_lin(0.0, xsi), py) + ) + c2 = i_u._geodetic_distance( + py[1], py[2], px[1], px[2], grid._mesh, np.einsum("ij,ji->i", i_u.phi2D_lin(eta, 1.0), py) + ) + c3 = i_u._geodetic_distance( + py[2], py[3], px[2], px[3], grid._mesh, np.einsum("ij,ji->i", i_u.phi2D_lin(1.0, xsi), py) + ) + c4 = i_u._geodetic_distance( + py[3], py[0], px[3], px[0], grid._mesh, np.einsum("ij,ji->i", i_u.phi2D_lin(eta, 0.0), py) + ) + + lenT = 2 if np.any(tau > 0) else 1 + + # Create arrays of corner points for xarray.isel + # TODO C grid may not need all xi and yi cornerpoints, so could speed up here? + + # Time coordinates: 4 points at ti, then 4 points at ti+1 + if lenT == 1: + ti_full = np.repeat(ti, 4) + else: + ti_1 = np.clip(ti + 1, 0, tdim - 1) + ti_full = np.concatenate([np.repeat(ti, 4), np.repeat(ti_1, 4)]) + + # Depth coordinates: 4 points at zi, repeated for both time levels + zi_full = np.repeat(zi, lenT * 4) + + # Y coordinates: [yi, yi, yi+1, yi+1] for each spatial point, repeated for time/depth + yi_1 = np.clip(yi + 1, 0, ydim - 1) + yi_full = np.tile(np.repeat(np.column_stack([yi, yi_1]), 2), (lenT)) + # # TODO check why in some cases minus needed here!!! + # yi_minus_1 = np.clip(yi - 1, 0, ydim - 1) + # yi = np.tile(np.repeat(np.column_stack([yi_minus_1, yi]), 2), (lenT)) + + # X coordinates: [xi, xi+1, xi, xi+1] for each spatial point, repeated for time/depth + xi_1 = np.clip(xi + 1, 0, xdim - 1) + xi_full = np.tile(np.column_stack([xi, xi_1, xi, xi_1]).flatten(), (lenT)) + + for data in [U, V]: + axis_dim = grid.get_axis_dim_mapping(data.dims) + + # Create DataArrays for indexing + selection_dict = { + axis_dim["X"]: xr.DataArray(xi_full, dims=("points")), + axis_dim["Y"]: xr.DataArray(yi_full, dims=("points")), + } + if "Z" in axis_dim: + selection_dict[axis_dim["Z"]] = xr.DataArray(zi_full, dims=("points")) + if "time" in data.dims: + selection_dict["time"] = xr.DataArray(ti_full, dims=("points")) + + corner_data = data.isel(selection_dict).data.reshape(lenT, len(xsi), 4) + + if lenT == 2: + tau_full = tau[:, np.newaxis] + corner_data = corner_data[0, :, :] * (1 - tau_full) + corner_data[1, :, :] * tau_full + else: + corner_data = corner_data[0, :, :] + # # See code below for v3 version + # # if self.gridindexingtype == "nemo": + # # U0 = self.U.data[ti, zi, yi + 1, xi] * c4 + # # U1 = self.U.data[ti, zi, yi + 1, xi + 1] * c2 + # # V0 = self.V.data[ti, zi, yi, xi + 1] * c1 + # # V1 = self.V.data[ti, zi, yi + 1, xi + 1] * c3 + # # elif self.gridindexingtype in ["mitgcm", "croco"]: + # # U0 = self.U.data[ti, zi, yi, xi] * c4 + # # U1 = self.U.data[ti, zi, yi, xi + 1] * c2 + # # V0 = self.V.data[ti, zi, yi, xi] * c1 + # # V1 = self.V.data[ti, zi, yi + 1, xi] * c3 + # # TODO Nick can you help use xgcm to fix this implementation? + + # # CROCO and MITgcm grid indexing, + # if data is U: + # U0 = corner_data[:, 0] * c4 + # U1 = corner_data[:, 1] * c2 + # elif data is V: + # V0 = corner_data[:, 0] * c1 + # V1 = corner_data[:, 2] * c3 + # # NEMO grid indexing + if data is U: + U0 = corner_data[:, 2] * c4 + U1 = corner_data[:, 3] * c2 + elif data is V: + V0 = corner_data[:, 1] * c1 + V1 = corner_data[:, 3] * c3 + + U = (1 - xsi) * U0 + xsi * U1 + V = (1 - eta) * V0 + eta * V1 + + deg2m = 1852 * 60.0 + if applyConversion: + meshJac = (deg2m * deg2m * np.cos(np.deg2rad(y))) if grid._mesh == "spherical" else 1 + else: + meshJac = deg2m if grid._mesh == "spherical" else 1 + + jac = i_u._compute_jacobian_determinant(py, px, eta, xsi) * meshJac + + u = ( + (-(1 - eta) * U - (1 - xsi) * V) * px[0] + + ((1 - eta) * U - xsi * V) * px[1] + + (eta * U + xsi * V) * px[2] + + (-eta * U + (1 - xsi) * V) * px[3] + ) / jac + v = ( + (-(1 - eta) * U - (1 - xsi) * V) * py[0] + + ((1 - eta) * U - xsi * V) * py[1] + + (eta * U + xsi * V) * py[2] + + (-eta * U + (1 - xsi) * V) * py[3] + ) / jac + if is_dask_collection(u): + u = u.compute() + v = v.compute() + + # check whether the grid conversion has been applied correctly + xx = (1 - xsi) * (1 - eta) * px[0] + xsi * (1 - eta) * px[1] + xsi * eta * px[2] + (1 - xsi) * eta * px[3] + u = np.where(np.abs((xx - x) / x) > 1e-4, np.nan, u) + + if vectorfield.W: + data = vectorfield.W.data + # Time coordinates: 2 points at ti, then 2 points at ti+1 + if lenT == 1: + ti_full = np.repeat(ti, 2) + else: + ti_1 = np.clip(ti + 1, 0, tdim - 1) + ti_full = np.concatenate([np.repeat(ti, 2), np.repeat(ti_1, 2)]) + + # Depth coordinates: 1 points at zi, repeated for both time levels + zi_1 = np.clip(zi + 1, 0, zdim - 1) + zi_full = np.tile(np.array([zi, zi_1]).flatten(), lenT) + + # Y coordinates: yi+1 for each spatial point, repeated for time/depth + yi_1 = np.clip(yi + 1, 0, ydim - 1) + yi_full = np.tile(yi_1, (lenT) * 2) + + # X coordinates: xi+1 for each spatial point, repeated for time/depth + xi_1 = np.clip(xi + 1, 0, xdim - 1) + xi_full = np.tile(xi_1, (lenT) * 2) + + axis_dim = grid.get_axis_dim_mapping(data.dims) + + # Create DataArrays for indexing + selection_dict = { + axis_dim["X"]: xr.DataArray(xi_full, dims=("points")), + axis_dim["Y"]: xr.DataArray(yi_full, dims=("points")), + axis_dim["Z"]: xr.DataArray(zi_full, dims=("points")), + } + if "time" in data.dims: + selection_dict["time"] = xr.DataArray(ti_full, dims=("points")) + + corner_data = data.isel(selection_dict).data.reshape(lenT, 2, len(xsi)) + + if lenT == 2: + tau_full = tau[np.newaxis, :] + corner_data = corner_data[0, :, :] * (1 - tau_full) + corner_data[1, :, :] * tau_full + else: + corner_data = corner_data[0, :, :] + + w = corner_data[0, :] * (1 - zeta) + corner_data[1, :] * zeta + if is_dask_collection(w): + w = w.compute() + else: + w = np.zeros_like(u) + + return (u, v, w) + + +def CGrid_Tracer( + field: Field, + ti: int, + position: dict[_XGRID_AXES, tuple[int, float | np.ndarray]], + tau: np.float32 | np.float64, + t: np.float32 | np.float64, + z: np.float32 | np.float64, + y: np.float32 | np.float64, + x: np.float32 | np.float64, +): + """Interpolation kernel for tracer fields on a C-Grid. + + Following Delandmeter and Van Sebille (2019), tracer fields should be interpolated + constant over the grid cell + """ + xi, _ = position["X"] + yi, _ = position["Y"] + zi, _ = position["Z"] + + axis_dim = field.grid.get_axis_dim_mapping(field.data.dims) + data = field.data + + lenT = 2 if np.any(tau > 0) else 1 + + if lenT == 2: + ti_1 = np.clip(ti + 1, 0, data.shape[0] - 1) + ti = np.concatenate([np.repeat(ti), np.repeat(ti_1)]) + zi_1 = np.clip(zi + 1, 0, data.shape[1] - 1) + zi = np.concatenate([np.repeat(zi), np.repeat(zi_1)]) + yi_1 = np.clip(yi + 1, 0, data.shape[2] - 1) + yi = np.concatenate([np.repeat(yi), np.repeat(yi_1)]) + xi_1 = np.clip(xi + 1, 0, data.shape[3] - 1) + xi = np.concatenate([np.repeat(xi), np.repeat(xi_1)]) + + # Create DataArrays for indexing + selection_dict = { + axis_dim["X"]: xr.DataArray(xi, dims=("points")), + axis_dim["Y"]: xr.DataArray(yi, dims=("points")), + } + if "Z" in axis_dim: + selection_dict[axis_dim["Z"]] = xr.DataArray(zi, dims=("points")) + if "time" in field.data.dims: + selection_dict["time"] = xr.DataArray(ti, dims=("points")) + + value = data.isel(selection_dict).data.reshape(lenT, len(xi)) + + if lenT == 2: + tau = tau[:, np.newaxis] + value = value[0, :] * (1 - tau) + value[1, :] * tau + else: + value = value[0, :] + + return value.compute() if is_dask_collection(value) else value def XNearest( @@ -172,7 +452,7 @@ def XNearest( else: value = corner_data[0, :] - return value.compute() if isinstance(value, dask.Array) else value + return value.compute() if is_dask_collection(value) else value def UXPiecewiseConstantFace( diff --git a/parcels/field.py b/parcels/field.py index 360e84393..b248414a4 100644 --- a/parcels/field.py +++ b/parcels/field.py @@ -12,7 +12,12 @@ from parcels._core.utils.time import TimeInterval from parcels._reprs import default_repr from parcels._typing import VectorType -from parcels.application_kernels.interpolation import UXPiecewiseLinearNode, XLinear, ZeroInterpolator +from parcels.application_kernels.interpolation import ( + UXPiecewiseLinearNode, + XLinear, + ZeroInterpolator, + ZeroInterpolator_Vector, +) from parcels.particle import KernelParticle from parcels.tools.converters import ( UnitConverter, @@ -282,8 +287,8 @@ def __init__( if vector_interp_method is None: self._vector_interp_method = None else: - _assert_same_function_signature(vector_interp_method, ref=ZeroInterpolator) - self._interp_method = vector_interp_method + _assert_same_function_signature(vector_interp_method, ref=ZeroInterpolator_Vector) + self._vector_interp_method = vector_interp_method def __repr__(self): return f"""<{type(self).__name__}> @@ -298,7 +303,7 @@ def vector_interp_method(self): @vector_interp_method.setter def vector_interp_method(self, method: Callable): - _assert_same_function_signature(method, ref=ZeroInterpolator) + _assert_same_function_signature(method, ref=ZeroInterpolator_Vector) self._vector_interp_method = method def eval(self, time: datetime, z, y, x, particle=None, applyConversion=True): @@ -324,17 +329,19 @@ def eval(self, time: datetime, z, y, x, particle=None, applyConversion=True): w = self.W._interp_method(self.W, ti, position, tau, time, z, y, x) else: w = 0.0 + + if applyConversion: + u = self.U.units.to_target(u, z, y, x) + v = self.V.units.to_target(v, z, y, x) + else: - (u, v, w) = self._vector_interp_method(self, ti, position, time, z, y, x) + (u, v, w) = self._vector_interp_method(self, ti, position, tau, time, z, y, x, applyConversion) for vel in (u, v, w): _update_particle_states_interp_value(particle, vel) - if applyConversion: - u = self.U.units.to_target(u, z, y, x) - v = self.V.units.to_target(v, z, y, x) - if "3D" in self.vector_type: - w = self.W.units.to_target(w, z, y, x) if self.W else 0.0 + if applyConversion and ("3D" in self.vector_type): + w = self.W.units.to_target(w, z, y, x) if self.W else 0.0 if "3D" in self.vector_type: return (u, v, w) diff --git a/parcels/tools/interpolation_utils.py b/parcels/tools/interpolation_utils.py index 71d316e7c..e4893a813 100644 --- a/parcels/tools/interpolation_utils.py +++ b/parcels/tools/interpolation_utils.py @@ -1,4 +1,3 @@ -import math from collections.abc import Callable from typing import Literal @@ -24,10 +23,10 @@ def phi1D_quad(xsi: float) -> list[float]: def phi2D_lin(eta: float, xsi: float) -> list[float]: - phi = [(1-xsi) * (1-eta), - xsi * (1-eta), - xsi * eta , - (1-xsi) * eta ] + phi = np.column_stack([(1-xsi) * (1-eta), + xsi * (1-eta), + xsi * eta , + (1-xsi) * eta ]) return phi @@ -179,18 +178,19 @@ def _geodetic_distance(lat1: float, lat2: float, lon1: float, lon2: float, mesh: if mesh == "spherical": rad = np.pi / 180.0 deg2m = 1852 * 60.0 - return np.sqrt(((lon2 - lon1) * deg2m * math.cos(rad * lat)) ** 2 + ((lat2 - lat1) * deg2m) ** 2) + return np.sqrt(((lon2 - lon1) * deg2m * np.cos(rad * lat)) ** 2 + ((lat2 - lat1) * deg2m) ** 2) else: return np.sqrt((lon2 - lon1) ** 2 + (lat2 - lat1) ** 2) def _compute_jacobian_determinant(py: np.ndarray, px: np.ndarray, eta: float, xsi: float) -> float: - dphidxsi = [eta - 1, 1 - eta, eta, -eta] - dphideta = [xsi - 1, -xsi, xsi, 1 - xsi] + dphidxsi = np.column_stack([eta - 1, 1 - eta, eta, -eta]) + dphideta = np.column_stack([xsi - 1, -xsi, xsi, 1 - xsi]) - dxdxsi = np.dot(px, dphidxsi) - dxdeta = np.dot(px, dphideta) - dydxsi = np.dot(py, dphidxsi) - dydeta = np.dot(py, dphideta) - jac = dxdxsi * dydeta - dxdeta * dydxsi - return jac + dxdxsi_diag = np.einsum("ij,ji->i", dphidxsi, px) + dxdeta_diag = np.einsum("ij,ji->i", dphideta, px) + dydxsi_diag = np.einsum("ij,ji->i", dphidxsi, py) + dydeta_diag = np.einsum("ij,ji->i", dphideta, py) + + jac_diag = dxdxsi_diag * dydeta_diag - dxdeta_diag * dydxsi_diag + return jac_diag diff --git a/parcels/xgrid.py b/parcels/xgrid.py index 6979d6f11..ac6dd89e2 100644 --- a/parcels/xgrid.py +++ b/parcels/xgrid.py @@ -102,6 +102,12 @@ def __init__(self, grid: xgcm.Grid, mesh="flat"): self._spatialhash = None ds = grid._ds + # Set the coordinates for the dataset (needed to be done explicitly for curvilinear grids) + if "lon" in ds: + ds.set_coords("lon") + if "lat" in ds: + ds.set_coords("lat") + if len(set(grid.axes) & {"X", "Y", "Z"}) > 0: # Only if spatial grid is >0D (see #2054 for further development) assert_valid_lat_lon(ds["lat"], ds["lon"], grid.axes) diff --git a/tests/v4/test_advection.py b/tests/v4/test_advection.py index 751c149a3..a8bae9686 100644 --- a/tests/v4/test_advection.py +++ b/tests/v4/test_advection.py @@ -13,7 +13,7 @@ ) from parcels.application_kernels.advection import AdvectionEE, AdvectionRK4, AdvectionRK4_3D, AdvectionRK45 from parcels.application_kernels.advectiondiffusion import AdvectionDiffusionEM, AdvectionDiffusionM1 -from parcels.application_kernels.interpolation import XLinear +from parcels.application_kernels.interpolation import CGrid_Velocity, XLinear from parcels.field import Field, VectorField from parcels.fieldset import FieldSet from parcels.particle import Particle, Variable @@ -21,6 +21,7 @@ from parcels.particleset import ParticleSet from parcels.tools.statuscodes import StatusCode from parcels.xgrid import XGrid +from tests.utils import round_and_hash_float_array kernel = { "EE": AdvectionEE, @@ -356,15 +357,16 @@ def truth_moving(x_0, y_0, t): ("RK45", 0.1), ], ) -@pytest.mark.parametrize("grid_type", ["A"]) # TODO also implement C-grid once available +@pytest.mark.parametrize("grid_type", ["A", "C"]) def test_stommelgyre_fieldset(method, rtol, grid_type): npart = 2 ds = stommel_gyre_dataset(grid_type=grid_type) grid = XGrid.from_dataset(ds) - U = Field("U", ds["U"], grid, interp_method=XLinear) - V = Field("V", ds["V"], grid, interp_method=XLinear) + vector_interp_method = None if grid_type == "A" else CGrid_Velocity + U = Field("U", ds["U"], grid) + V = Field("V", ds["V"], grid) P = Field("P", ds["P"], grid, interp_method=XLinear) - UV = VectorField("UV", U, V) + UV = VectorField("UV", U, V, vector_interp_method=vector_interp_method) fieldset = FieldSet([U, V, P, UV]) dt = np.timedelta64(30, "m") @@ -425,3 +427,133 @@ def UpdateP(particle, fieldset, time): # pragma: no cover pset = ParticleSet(fieldset, pclass=SampleParticle, lon=start_lon, lat=start_lat, time=np.timedelta64(0, "s")) pset.execute([kernel[method], UpdateP], dt=dt, runtime=runtime) np.testing.assert_allclose(pset.p, pset.p_start, rtol=rtol) + + +def test_nemo_curvilinear_fieldset(): + data_folder = parcels.download_example_dataset("NemoCurvilinear_data") + files = data_folder.glob("*.nc4") + ds = xr.open_mfdataset(files, combine="nested", data_vars="minimal", coords="minimal", compat="override") + ds = ( + ds.isel(time_counter=0, drop=True) + .isel(time=0, drop=True) + .isel(z_a=0, drop=True) + .rename({"glamf": "lon", "gphif": "lat", "z": "depth"}) + ) + + xgcm_grid = parcels.xgcm.Grid( + ds, + coords={ + "X": {"left": "x"}, + "Y": {"left": "y"}, + }, + periodic=False, + ) + grid = XGrid(xgcm_grid, mesh="spherical") + + U = parcels.Field("U", ds["U"], grid) + V = parcels.Field("V", ds["V"], grid) + U.units = parcels.GeographicPolar() + V.units = parcels.Geographic() + UV = parcels.VectorField("UV", U, V, vector_interp_method=CGrid_Velocity) + fieldset = parcels.FieldSet([U, V, UV]) + + npart = 20 + lonp = 30 * np.ones(npart) + latp = np.linspace(-70, 88, npart) + runtime = np.timedelta64(12, "h") # TODO increase to 160 days + + def periodicBC(particle, fieldSet, time): # pragma: no cover + particle.dlon = np.where(particle.lon > 180, particle.dlon - 360, particle.dlon) + + pset = parcels.ParticleSet(fieldset, lon=lonp, lat=latp) + pset.execute([AdvectionEE, periodicBC], runtime=runtime, dt=np.timedelta64(6, "h")) + np.testing.assert_allclose(pset.lat_nextloop, latp, atol=1e-1) + + +@pytest.mark.parametrize("method", ["RK4", "RK4_3D"]) +def test_nemo_3D_curvilinear_fieldset(method): + download_dir = parcels.download_example_dataset("NemoNorthSeaORCA025-N006_data") + ufiles = download_dir.glob("*U.nc") + dsu = xr.open_mfdataset(ufiles, decode_times=False, drop_variables=["nav_lat", "nav_lon"]) + dsu = dsu.rename({"time_counter": "time", "uo": "U"}) + + vfiles = download_dir.glob("*V.nc") + dsv = xr.open_mfdataset(vfiles, decode_times=False, drop_variables=["nav_lat", "nav_lon"]) + dsv = dsv.rename({"time_counter": "time", "vo": "V"}) + + wfiles = download_dir.glob("*W.nc") + dsw = xr.open_mfdataset(wfiles, decode_times=False, drop_variables=["nav_lat", "nav_lon"]) + dsw = dsw.rename({"time_counter": "time", "depthw": "depth", "wo": "W"}) + + dsu = dsu.assign_coords(depthu=dsw.depth.values) + dsu = dsu.rename({"depthu": "depth"}) + + dsv = dsv.assign_coords(depthv=dsw.depth.values) + dsv = dsv.rename({"depthv": "depth"}) + + coord_file = f"{download_dir}/coordinates.nc" + dscoord = xr.open_dataset(coord_file, decode_times=False).rename({"glamf": "lon", "gphif": "lat"}) + dscoord = dscoord.isel(time=0, drop=True) + + ds = xr.merge([dsu, dsv, dsw, dscoord]) + ds = ds.drop_vars( + [ + "uos", + "vos", + "nav_lev", + "nav_lon", + "nav_lat", + "tauvo", + "tauuo", + "time_steps", + "gphiu", + "gphiv", + "gphit", + "glamu", + "glamv", + "glamt", + "time_centered_bounds", + "time_counter_bounds", + "time_centered", + ] + ) + ds = ds.drop_vars(["e1f", "e1t", "e1u", "e1v", "e2f", "e2t", "e2u", "e2v"]) + ds["time"] = [np.timedelta64(int(t), "s") + np.datetime64("1900-01-01") for t in ds["time"]] + + ds["W"] *= -1 # Invert W velocity + + xgcm_grid = parcels.xgcm.Grid( + ds, + coords={ + "X": {"left": "x"}, + "Y": {"left": "y"}, + "Z": {"left": "depth"}, + "T": {"center": "time"}, + }, + periodic=False, + ) + grid = XGrid(xgcm_grid, mesh="spherical") + + U = parcels.Field("U", ds["U"], grid) + V = parcels.Field("V", ds["V"], grid) + W = parcels.Field("W", ds["W"], grid) + U.units = parcels.GeographicPolar() + V.units = parcels.Geographic() + UV = parcels.VectorField("UV", U, V, vector_interp_method=CGrid_Velocity) + UVW = parcels.VectorField("UVW", U, V, W, vector_interp_method=CGrid_Velocity) + fieldset = parcels.FieldSet([U, V, W, UV, UVW]) + + npart = 10 + lons = np.linspace(1.9, 3.4, npart) + lats = np.linspace(52.5, 51.6, npart) + pset = parcels.ParticleSet(fieldset, lon=lons, lat=lats, depth=np.ones_like(lons)) + + pset.execute(kernel[method], runtime=np.timedelta64(4, "D"), dt=np.timedelta64(6, "h")) + + if method == "RK4": + np.testing.assert_equal(round_and_hash_float_array([p.lon for p in pset], decimals=5), 29977383852960156017546) + elif method == "RK4_3D": + # TODO check why decimals needs to be so low in RK4_3D (compare to v3) + np.testing.assert_equal( + round_and_hash_float_array([p.depth for p in pset], decimals=1), 29747210774230389239432 + ) diff --git a/tests/v4/test_field.py b/tests/v4/test_field.py index cd71337d7..0a6e10cd5 100644 --- a/tests/v4/test_field.py +++ b/tests/v4/test_field.py @@ -130,7 +130,7 @@ def test_vectorfield_invalid_interpolator(): ds = datasets_structured["ds_2d_left"] grid = XGrid.from_dataset(ds) - def invalid_interpolator_wrong_signature(self, ti, position, tau, t, z, y, invalid): + def invalid_interpolator_wrong_signature(self, ti, position, tau, t, z, y, applyConversion, invalid): return 0.0 # Create component fields