diff --git a/src/parcels/_core/kernel.py b/src/parcels/_core/kernel.py index f68c8bf098..09a1fad8d4 100644 --- a/src/parcels/_core/kernel.py +++ b/src/parcels/_core/kernel.py @@ -6,7 +6,6 @@ import numpy as np -from parcels._core.basegrid import GridType from parcels._core.statuscodes import ( StatusCode, _raise_field_interpolation_error, @@ -19,7 +18,6 @@ from parcels._core.warnings import FieldEvalWarning, KernelWarning from parcels._python import assert_same_function_signature from parcels.kernels import ( - AdvectionAnalytical, AdvectionRK4, AdvectionRK45, ) @@ -126,12 +124,12 @@ def check_fieldsets_in_kernels(self, kernel): # TODO v4: this can go into anoth This function is to be called from the derived class when setting up the 'kernel'. """ if self.fieldset is not None: - if kernel is AdvectionAnalytical: - if self._fieldset.U.interp_method != "cgrid_velocity": - raise NotImplementedError("Analytical Advection only works with C-grids") - if self._fieldset.U.grid._gtype not in [GridType.CurvilinearZGrid, GridType.RectilinearZGrid]: - raise NotImplementedError("Analytical Advection only works with Z-grids in the vertical") - elif kernel is AdvectionRK45: + # if kernel is AdvectionAnalytical: + # if self._fieldset.U.interp_method != "cgrid_velocity": + # raise NotImplementedError("Analytical Advection only works with C-grids") + # if self._fieldset.U.grid._gtype not in [GridType.CurvilinearZGrid, GridType.RectilinearZGrid]: + # raise NotImplementedError("Analytical Advection only works with Z-grids in the vertical") + if kernel is AdvectionRK45: if "next_dt" not in [v.name for v in self.pclass.variables]: raise ValueError('ParticleClass requires a "next_dt" for AdvectionRK45 Kernel.') if not hasattr(self.fieldset, "RK45_tol"): diff --git a/src/parcels/_datasets/structured/generated.py b/src/parcels/_datasets/structured/generated.py index 7eba61f559..abb5f6d603 100644 --- a/src/parcels/_datasets/structured/generated.py +++ b/src/parcels/_datasets/structured/generated.py @@ -39,41 +39,56 @@ def simple_UV_dataset(dims=(360, 2, 30, 4), maxdepth=1, mesh="spherical"): ) -def radial_rotation_dataset(xdim=200, ydim=200): # Define 2D flat, square fieldset for testing purposes. +def radial_rotation_dataset(xdim=200, ydim=200, grid_type="A"): # Define 2D flat, square fieldset for testing purposes. lon = np.linspace(0, 60, xdim, dtype=np.float32) lat = np.linspace(0, 60, ydim, dtype=np.float32) x0 = 30.0 # Define the origin to be the centre of the Field. y0 = 30.0 + dx, dy = lon[-1] / xdim, lat[-1] / ydim # Define the grid spacing in x and y directions. + U = np.zeros((2, 1, ydim, xdim), dtype=np.float32) V = np.zeros((2, 1, ydim, xdim), dtype=np.float32) + R = np.zeros((2, 1, ydim, xdim), dtype=np.float32) omega = 2 * np.pi / 86400.0 # Define the rotational period as 1 day. + def calc_r_theta(ln, lt, x0, y0): + r = np.sqrt((ln - x0) ** 2 + (lt - y0) ** 2) + theta = np.arctan2((lt - y0), (ln - x0)) + return r, theta + for i in range(lon.size): for j in range(lat.size): - r = np.sqrt((lon[i] - x0) ** 2 + (lat[j] - y0) ** 2) - assert r >= 0.0 - assert r <= np.sqrt(x0**2 + y0**2) - - theta = np.arctan2((lat[j] - y0), (lon[i] - x0)) - assert abs(theta) <= np.pi + r, theta = calc_r_theta(lon[i], lat[j], x0, y0) + R[:, :, j, i] = r + if grid_type == "A": + r, theta = calc_r_theta(lon[i], lat[j], x0, y0) + U[:, :, j, i] = r * np.sin(theta) * omega + V[:, :, j, i] = -r * np.cos(theta) * omega + elif grid_type == "C": + r, theta = calc_r_theta(lon[i] - dx / 2, lat[j], x0, y0) + U[:, :, j, i] = r * np.sin(theta) * omega - U[:, :, j, i] = r * np.sin(theta) * omega - V[:, :, j, i] = -r * np.cos(theta) * omega + r, theta = calc_r_theta(lon[i], lat[j] - dy / 2, x0, y0) + V[:, :, j, i] = -r * np.cos(theta) * omega return xr.Dataset( - {"U": (["time", "depth", "YG", "XG"], U), "V": (["time", "depth", "YG", "XG"], V)}, + { + "U": (["time", "depth", "YG", "XC"], U), + "V": (["time", "depth", "YC", "XG"], V), + "R": (["time", "depth", "YC", "XC"], R), + }, coords={ "time": (["time"], [np.timedelta64(0, "s"), np.timedelta64(10, "D")], {"axis": "T"}), "depth": (["depth"], np.array([0.0]), {"axis": "Z"}), - "YC": (["YC"], np.arange(ydim) + 0.5, {"axis": "Y"}), - "YG": (["YG"], np.arange(ydim), {"axis": "Y", "c_grid_axis_shift": -0.5}), - "XC": (["XC"], np.arange(xdim) + 0.5, {"axis": "X"}), - "XG": (["XG"], np.arange(xdim), {"axis": "X", "c_grid_axis_shift": -0.5}), - "lat": (["YG"], lat, {"axis": "Y", "c_grid_axis_shift": 0.5}), - "lon": (["XG"], lon, {"axis": "X", "c_grid_axis_shift": -0.5}), + "YC": (["YC"], np.arange(ydim) - 0.5, {"axis": "Y", "c_grid_axis_shift": +0.5}), + "YG": (["YG"], np.arange(ydim), {"axis": "Y"}), + "XC": (["XC"], np.arange(xdim) - 0.5, {"axis": "X", "c_grid_axis_shift": +0.5}), + "XG": (["XG"], np.arange(xdim), {"axis": "X"}), + "lat": (["YG"], lat, {"axis": "Y"}), + "lon": (["XG"], lon, {"axis": "X"}), }, ).pipe( sgrid._attach_sgrid_metadata, @@ -84,7 +99,7 @@ def radial_rotation_dataset(xdim=200, ydim=200): # Define 2D flat, square field node_coordinates=("lon", "lat"), face_dimensions=( sgrid.FaceNodePadding("XC", "XG", sgrid.Padding.LOW), - sgrid.FaceNodePadding("YC", "YG", sgrid.Padding.HIGH), + sgrid.FaceNodePadding("YC", "YG", sgrid.Padding.LOW), ), vertical_dimensions=(sgrid.FaceNodePadding("ZC", "depth", sgrid.Padding.BOTH),), ), diff --git a/src/parcels/interpolators/_xinterpolators.py b/src/parcels/interpolators/_xinterpolators.py index f80f2241a8..01e2e66bd4 100644 --- a/src/parcels/interpolators/_xinterpolators.py +++ b/src/parcels/interpolators/_xinterpolators.py @@ -190,6 +190,105 @@ def interp( return u, v, w +def _get_cgrid_velocities( + vectorfield: VectorField, grid_positions: dict[ptyping.XgridAxis, dict[str, int | float | np.ndarray]] +): + # Helper function to get the edge velocities for a given C-grid vector field and position. + xi, xsi = grid_positions["X"]["index"], grid_positions["X"]["bcoord"] + yi, eta = grid_positions["Y"]["index"], grid_positions["Y"]["bcoord"] + zi, _ = grid_positions["Z"]["index"], grid_positions["Z"]["bcoord"] + ti, tau = grid_positions["T"]["index"], grid_positions["T"]["bcoord"] + + U = vectorfield.U.data + V = vectorfield.V.data + grid = vectorfield.grid + offsets = _get_offsets_dictionary(grid) + tdim, zdim, ydim, xdim = U.shape[0], U.shape[1], U.shape[2], U.shape[3] + lenT = 2 if np.any(tau > 0) else 1 + + 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.is_spherical(): + px = ((px + 180.0) % 360.0) - 180.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), grid.deg2m + ) + 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), grid.deg2m + ) + 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), grid.deg2m + ) + 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), grid.deg2m + ) + + npart = len(xsi) + t_levels = (ti,) if lenT == 1 else (ti, np.clip(ti + 1, 0, tdim - 1)) + + def _compute_corner_data(data, y_levels, x_levels, z_levels=None) -> np.ndarray: + """Gather the two bracketing face values and reduce over time if needed. + + Exactly one of the Z, Y and X axes contributes the two corners. The + other two contribute a single level each. + """ + levels = { + "T": t_levels, + "Z": z_levels if z_levels is not None else (zi,), + "Y": y_levels, + "X": x_levels, + } + axis_dim = grid.get_axis_dim_mapping(data.dims) + corner_data = _gather_corners(data, axis_dim, levels, npart).reshape(lenT, 2, npart) + + 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, :] + return corner_data + + # Compute U velocity: the two corners are the X faces + yi_o = np.clip(yi + offsets["Y"], 0, ydim - 1) + xi_1 = np.clip(xi + 1, 0, xdim - 1) + corner_data = _compute_corner_data(U, y_levels=(yi_o,), x_levels=(xi, xi_1)) + + U0 = corner_data[0, :] * c4 + U1 = corner_data[1, :] * c2 + + # Compute V velocity: the two corners are the Y faces + yi_1 = np.clip(yi + 1, 0, ydim - 1) + xi_o = np.clip(xi + offsets["X"], 0, xdim - 1) + corner_data = _compute_corner_data(V, y_levels=(yi, yi_1), x_levels=(xi_o,)) + + V0 = corner_data[0, :] * c1 + V1 = corner_data[1, :] * c3 + + if vectorfield.W: + W = vectorfield.W.data + + # Compute W velocity: the two corners are the Z faces + yi_o = np.clip(yi + offsets["Y"], 0, ydim - 1) + xi_o = np.clip(xi + offsets["X"], 0, xdim - 1) + zi_0 = np.clip(zi + offsets["Z"], 0, zdim - 1) + zi_1 = np.clip(zi + offsets["Z"] + 1, 0, zdim - 1) + corner_data = _compute_corner_data(W, y_levels=(yi_o,), x_levels=(xi_o,), z_levels=(zi_0, zi_1)) + W0 = corner_data[0, :] + W1 = corner_data[1, :] + else: + W0 = np.zeros_like(U0) + W1 = np.zeros_like(U1) + + return U0, U1, V0, V1, W0, W1, px, py + + class CGrid_Velocity(VectorInterpolator): # noqa: N801 """ Interpolation kernel for velocity fields on a C-Grid. @@ -208,83 +307,13 @@ def interp( Following Delandmeter and Van Sebille (2019), velocity fields should be interpolated only in the direction of the grid cell faces. """ - xi, xsi = grid_positions["X"]["index"], grid_positions["X"]["bcoord"] - yi, eta = grid_positions["Y"]["index"], grid_positions["Y"]["bcoord"] - zi, zeta = grid_positions["Z"]["index"], grid_positions["Z"]["bcoord"] - ti, tau = grid_positions["T"]["index"], grid_positions["T"]["bcoord"] - - U = vectorfield.U.data - V = vectorfield.V.data + _, xsi = grid_positions["X"]["index"], grid_positions["X"]["bcoord"] + _, eta = grid_positions["Y"]["index"], grid_positions["Y"]["bcoord"] + _, zeta = grid_positions["Z"]["index"], grid_positions["Z"]["bcoord"] grid = vectorfield.grid - offsets = _get_offsets_dictionary(grid) - tdim, zdim, ydim, xdim = U.shape[0], U.shape[1], U.shape[2], U.shape[3] - lenT = 2 if np.any(tau > 0) else 1 - - 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.is_spherical(): - px = ((px + 180.0) % 360.0) - 180.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), grid.deg2m - ) - 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), grid.deg2m - ) - 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), grid.deg2m - ) - 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), grid.deg2m - ) - - npart = len(xsi) - t_levels = (ti,) if lenT == 1 else (ti, np.clip(ti + 1, 0, tdim - 1)) - - def _compute_corner_data(data, y_levels, x_levels, z_levels=None) -> np.ndarray: - """Gather the two bracketing face values and reduce over time if needed. - - Exactly one of the Z, Y and X axes contributes the two corners. The - other two contribute a single level each. - """ - levels = { - "T": t_levels, - "Z": z_levels if z_levels is not None else (zi,), - "Y": y_levels, - "X": x_levels, - } - axis_dim = grid.get_axis_dim_mapping(data.dims) - corner_data = _gather_corners(data, axis_dim, levels, npart).reshape(lenT, 2, npart) - - 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, :] - return corner_data - - # Compute U velocity: the two corners are the X faces - yi_o = np.clip(yi + offsets["Y"], 0, ydim - 1) - xi_1 = np.clip(xi + 1, 0, xdim - 1) - corner_data = _compute_corner_data(U, y_levels=(yi_o,), x_levels=(xi, xi_1)) - - U0 = corner_data[0, :] * c4 - U1 = corner_data[1, :] * c2 + U0, U1, V0, V1, W0, W1, px, py = _get_cgrid_velocities(vectorfield, grid_positions) Uvel = (1 - xsi) * U0 + xsi * U1 - - # Compute V velocity: the two corners are the Y faces - yi_1 = np.clip(yi + 1, 0, ydim - 1) - xi_o = np.clip(xi + offsets["X"], 0, xdim - 1) - corner_data = _compute_corner_data(V, y_levels=(yi, yi_1), x_levels=(xi_o,)) - - V0 = corner_data[0, :] * c1 - V1 = corner_data[1, :] * c3 Vvel = (1 - eta) * V0 + eta * V1 if grid._mesh.is_spherical(): @@ -314,16 +343,7 @@ def _compute_corner_data(data, y_levels, x_levels, z_levels=None) -> np.ndarray: v /= conversion if vectorfield.W: - W = vectorfield.W.data - - # Compute W velocity: the two corners are the Z faces - yi_o = np.clip(yi + offsets["Y"], 0, ydim - 1) - xi_o = np.clip(xi + offsets["X"], 0, xdim - 1) - zi_0 = np.clip(zi + offsets["Z"], 0, zdim - 1) - zi_1 = np.clip(zi + offsets["Z"] + 1, 0, zdim - 1) - corner_data = _compute_corner_data(W, y_levels=(yi_o,), x_levels=(xi_o,), z_levels=(zi_0, zi_1)) - - w = corner_data[0, :] * (1 - zeta) + corner_data[1, :] * zeta + w = W0 * (1 - zeta) + W1 * zeta if is_dask_collection(w): w = w.compute() else: diff --git a/src/parcels/kernels/_advection.py b/src/parcels/kernels/_advection.py index 77162aad59..f981902c57 100644 --- a/src/parcels/kernels/_advection.py +++ b/src/parcels/kernels/_advection.py @@ -1,7 +1,5 @@ """Collection of pre-built advection kernels.""" -import math - import numpy as np from parcels._core.statuscodes import StatusCode @@ -162,144 +160,86 @@ def AdvectionAnalytical(particles, fieldset): # pragma: no cover Note that the time-dependent scheme is currently implemented with 'intermediate timesteps' (default 10 per model timestep) and not yet with the full analytical time integration. """ - import numpy as np - import parcels._core.utils.interpolation as i_u + from parcels._core.field import _get_positions + from parcels.interpolators._xinterpolators import CGrid_Velocity, _get_cgrid_velocities tol = 1e-10 - I_s = 10 # number of intermediate time steps + # I_s = 10 # number of intermediate time steps dt = particles.dt - direction = 1.0 if dt > 0 else -1.0 + direction = 1.0 if np.any(dt > 0) else -1.0 withW = True if "W" in [f.name for f in fieldset.fields.values()] else False - withTime = True if len(fieldset.U.grid.time) > 1 else False - tau, zeta, eta, xsi, ti, zi, yi, xi = fieldset.U._search_indices( - particles.z, particles.y, particles.x, particles=particles + + vectorfield = fieldset.UVW if withW else fieldset.UV + if not isinstance(vectorfield.interp_method, CGrid_Velocity): + raise NotImplementedError( + "Analytical advection is only implemented for C-grid velocity fields, " + f"but the fieldset has interp_method={vectorfield.interp_method}" + ) + # withTime = True if len(vectorfield.grid.time) > 1 else False + igrid = vectorfield.igrid + grid = vectorfield.grid + + _, grid_positions = _get_positions( + fieldset.U, particles.t, particles.z, particles.y, particles.x, particles, particles.ei[:, igrid] ) - ds_t = dt - if withTime: - time_i = np.linspace(0, fieldset.U.grid.time[ti + 1] - fieldset.U.grid.time[ti], I_s) - ds_t = min(ds_t, time_i[np.where(particles.time - fieldset.U.grid.time[ti] < time_i)[0][0]]) + _, xsi = grid_positions["X"]["index"], grid_positions["X"]["bcoord"] + _, eta = grid_positions["Y"]["index"], grid_positions["Y"]["bcoord"] + zi, zeta = grid_positions["Z"]["index"], grid_positions["Z"]["bcoord"] + U0, U1, V0, V1, W0, W1, px, py = _get_cgrid_velocities(vectorfield, grid_positions) - if withW: - if abs(xsi - 1) < tol: - if fieldset.U.data[0, zi + 1, yi + 1, xi + 1] > 0: - xi += 1 - xsi = 0 - if abs(eta - 1) < tol: - if fieldset.V.data[0, zi + 1, yi + 1, xi + 1] > 0: - yi += 1 - eta = 0 - if abs(zeta - 1) < tol: - if fieldset.W.data[0, zi + 1, yi + 1, xi + 1] > 0: - zi += 1 - zeta = 0 - else: - if abs(xsi - 1) < tol: - if fieldset.U.data[0, yi + 1, xi + 1] > 0: - xi += 1 - xsi = 0 - if abs(eta - 1) < tol: - if fieldset.V.data[0, yi + 1, xi + 1] > 0: - yi += 1 - eta = 0 - - particles.ei[:] = fieldset.U.ravel_index(zi, yi, xi) - - grid = fieldset.U.grid - if grid._gtype < 2: - px = np.array([grid.x[xi], grid.x[xi + 1], grid.x[xi + 1], grid.x[xi]]) - py = np.array([grid.y[yi], grid.y[yi], grid.y[yi + 1], grid.y[yi + 1]]) - else: - px = np.array([grid.x[yi, xi], grid.x[yi, xi + 1], grid.x[yi + 1, xi + 1], grid.x[yi + 1, xi]]) - py = np.array([grid.y[yi, xi], grid.y[yi, xi + 1], grid.y[yi + 1, xi + 1], grid.y[yi + 1, xi]]) - if grid.mesh == "spherical": - px[0] = px[0] + 360 if px[0] < particles.x - 225 else px[0] - px[0] = px[0] - 360 if px[0] > particles.y + 225 else 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:]) if withW: pz = np.array([grid.depth[zi], grid.depth[zi + 1]]) dz = pz[1] - pz[0] else: dz = 1.0 - c1 = i_u._geodetic_distance(py[0], py[1], px[0], px[1], grid.mesh, np.dot(i_u.phi2D_lin(0.0, xsi), py), grid.deg2m) - c2 = i_u._geodetic_distance(py[1], py[2], px[1], px[2], grid.mesh, np.dot(i_u.phi2D_lin(eta, 1.0), py), grid.deg2m) - c3 = i_u._geodetic_distance(py[2], py[3], px[2], px[3], grid.mesh, np.dot(i_u.phi2D_lin(1.0, xsi), py), grid.deg2m) - c4 = i_u._geodetic_distance(py[3], py[0], px[3], px[0], grid.mesh, np.dot(i_u.phi2D_lin(eta, 0.0), py), grid.deg2m) rad = np.pi / 180.0 - deg2m = grid.deg2m - meshJac = (deg2m * deg2m * math.cos(rad * particles.y)) if grid.mesh == "spherical" else 1 + deg2m = 1852 * 60.0 + meshJac = (deg2m * deg2m * np.cos(rad * particles.y)) if grid._mesh.is_spherical() else 1 dxdy = i_u._compute_jacobian_determinant(py, px, eta, xsi) * meshJac + U0 *= direction * dz + U1 *= direction * dz + V0 *= direction * dz + V1 *= direction * dz if withW: - U0 = direction * fieldset.U.data[ti, zi + 1, yi + 1, xi] * c4 * dz - U1 = direction * fieldset.U.data[ti, zi + 1, yi + 1, xi + 1] * c2 * dz - V0 = direction * fieldset.V.data[ti, zi + 1, yi, xi + 1] * c1 * dz - V1 = direction * fieldset.V.data[ti, zi + 1, yi + 1, xi + 1] * c3 * dz - if withTime: - U0 = U0 * (1 - tau) + tau * direction * fieldset.U.data[ti + 1, zi + 1, yi + 1, xi] * c4 * dz - U1 = U1 * (1 - tau) + tau * direction * fieldset.U.data[ti + 1, zi + 1, yi + 1, xi + 1] * c2 * dz - V0 = V0 * (1 - tau) + tau * direction * fieldset.V.data[ti + 1, zi + 1, yi, xi + 1] * c1 * dz - V1 = V1 * (1 - tau) + tau * direction * fieldset.V.data[ti + 1, zi + 1, yi + 1, xi + 1] * c3 * dz - else: - U0 = direction * fieldset.U.data[ti, yi + 1, xi] * c4 * dz - U1 = direction * fieldset.U.data[ti, yi + 1, xi + 1] * c2 * dz - V0 = direction * fieldset.V.data[ti, yi, xi + 1] * c1 * dz - V1 = direction * fieldset.V.data[ti, yi + 1, xi + 1] * c3 * dz - if withTime: - U0 = U0 * (1 - tau) + tau * direction * fieldset.U.data[ti + 1, yi + 1, xi] * c4 * dz - U1 = U1 * (1 - tau) + tau * direction * fieldset.U.data[ti + 1, yi + 1, xi + 1] * c2 * dz - V0 = V0 * (1 - tau) + tau * direction * fieldset.V.data[ti + 1, yi, xi + 1] * c1 * dz - V1 = V1 * (1 - tau) + tau * direction * fieldset.V.data[ti + 1, yi + 1, xi + 1] * c3 * dz + W0 *= direction * dxdy + W1 *= direction * dxdy def compute_ds(F0, F1, r, direction, tol): # noqa: N803 - up = F0 * (1 - r) + F1 * r - r_target = 1.0 if direction * up >= 0.0 else 0.0 - B = F0 - F1 - delta = -F0 - B = 0 if abs(B) < tol else B - - if abs(B) > tol: - F_r1 = r_target + delta / B - F_r0 = r + delta / B - else: - F_r0, F_r1 = None, None - - if abs(B) < tol and abs(delta) < tol: - ds = float("inf") - elif B == 0: - ds = -(r_target - r) / delta - elif F_r1 * F_r0 < tol: - ds = float("inf") - else: - ds = -1.0 / B * math.log(F_r1 / F_r0) - - if abs(ds) < tol: - ds = float("inf") - return ds, B, delta + with np.errstate(divide="ignore", invalid="ignore"): + up = F0 * (1 - r) + F1 * r + r_target = np.where(direction * up >= 0.0, 1.0, 0.0) + B = F0 - F1 + delta = -F0 + B = np.where(np.abs(B) < tol, np.zeros_like(B), B) + + F_r1 = np.where(np.abs(B) > tol, r_target + delta / B, np.nan) + F_r0 = np.where(np.abs(B) > tol, r + delta / B, np.nan) + + d_s = -1.0 / B * np.log(F_r1 / F_r0) + d_s = np.where(F_r1 * F_r0 < tol, np.inf, d_s) + d_s = np.where(B == 0, -delta * direction / up, d_s) + d_s = np.where((np.abs(B) < tol) & (np.abs(delta) < tol), np.inf, d_s) + + d_s = np.where(d_s < tol, np.inf, d_s) + return d_s, B, delta ds_x, B_x, delta_x = compute_ds(U0, U1, xsi, direction, tol) ds_y, B_y, delta_y = compute_ds(V0, V1, eta, direction, tol) if withW: - W0 = direction * fieldset.W.data[ti, zi, yi + 1, xi + 1] * dxdy - W1 = direction * fieldset.W.data[ti, zi + 1, yi + 1, xi + 1] * dxdy - if withTime: - W0 = W0 * (1 - tau) + tau * direction * fieldset.W.data[ti + 1, zi, yi + 1, xi + 1] * dxdy - W1 = W1 * (1 - tau) + tau * direction * fieldset.W.data[ti + 1, zi + 1, yi + 1, xi + 1] * dxdy ds_z, B_z, delta_z = compute_ds(W0, W1, zeta, direction, tol) else: - ds_z = float("inf") + ds_z = np.inf # take the minimum travel time - s_min = min(abs(ds_x), abs(ds_y), abs(ds_z), abs(ds_t / (dxdy * dz))) + s_min = np.minimum(np.minimum(np.abs(ds_x), np.abs(ds_y)), np.minimum(np.abs(ds_z), np.abs(dt / (dxdy * dz)))) # calculate end position in time s_min def compute_rs(r, B, delta, s_min): # noqa: N803 - if abs(B) < tol: - return -delta * s_min + r - else: - return (r + delta / B) * math.exp(-B * s_min) - delta / B + with np.errstate(divide="ignore", invalid="ignore"): + return np.where(abs(B) < tol, -delta * s_min + r, (r + delta / B) * np.exp(-B * s_min) - delta / B) rs_x = compute_rs(xsi, B_x, delta_x, s_min) rs_y = compute_rs(eta, B_y, delta_y, s_min) @@ -323,7 +263,7 @@ def compute_rs(r, B, delta, s_min): # noqa: N803 rs_z = compute_rs(zeta, B_z, delta_z, s_min) particles.dz += (1.0 - rs_z) * pz[0] + rs_z * pz[1] - particles.z - if particles.dt > 0: - particles.dt = max(direction * s_min * (dxdy * dz), 1e-7).astype("timedelta64[s]") + if direction > 0: + particles.dt = np.maximum(direction * s_min * (dxdy * dz), 1e-7) else: - particles.dt = min(direction * s_min * (dxdy * dz), -1e-7).astype("timedelta64[s]") + particles.dt = np.minimum(direction * s_min * (dxdy * dz), -1e-7) diff --git a/tests/test_advection.py b/tests/test_advection.py index ee009c1c92..ab6428b6b3 100644 --- a/tests/test_advection.py +++ b/tests/test_advection.py @@ -27,7 +27,9 @@ stommel_gyre_dataset, ) from parcels._datasets.structured.generic import datasets_sgrid +from parcels.interpolators import CGrid_Velocity from parcels.kernels import ( + AdvectionAnalytical, AdvectionDiffusionEM, AdvectionDiffusionM1, AdvectionEE, @@ -505,3 +507,52 @@ def test_mitgcm(): 1952691.93845841, ] np.testing.assert_allclose(pset.y, lat_v3, atol=1) + + +def test_analytical_throwserror_on_Agrid(): + ds = simple_UV_dataset(mesh="flat") + fieldset = FieldSet.from_sgrid_conventions(ds, mesh="flat") + pset = ParticleSet(fieldset, x=1, y=1) + + with pytest.raises(NotImplementedError): + pset.execute(AdvectionAnalytical, runtime=1, dt=1) + + +@pytest.mark.parametrize("u", [0.2, -0.3, 0]) +@pytest.mark.parametrize("v", [-0.3, 0, 1]) +@pytest.mark.parametrize("w", [None, 0.05, 0, -0.05]) +@pytest.mark.parametrize("direction", [1, -1]) +def test_uniform_analytical(u, v, w, direction, tmp_parquet): + ds = simple_UV_dataset(mesh="flat") + ds["U"].data[:] = u + ds["V"].data[:] = v + if w is not None: + ds["W"] = xr.full_like(ds["U"], w) + + fieldset = FieldSet.from_sgrid_conventions(ds, mesh="flat") + fieldset.UV.interp_method = CGrid_Velocity() + if w is not None: + fieldset.UVW.interp_method = CGrid_Velocity() + + x0, y0, z0 = 6.1, 6.2, 0.5 + pset = ParticleSet(fieldset, pclass=Particle, x=x0, y=y0, z=z0) + + outfile = ParticleFile(tmp_parquet, outputdt=np.timedelta64(1, "s")) + runtime = 4 + pset.execute( + AdvectionAnalytical, + runtime=np.timedelta64(runtime, "s"), + dt=np.timedelta64(direction, "s"), + output_file=outfile, + ) + assert np.abs(pset.x - x0 - runtime * u * direction) < 1e-6 + assert np.abs(pset.y - y0 - runtime * v * direction) < 1e-6 + if w is not None: + assert np.abs(pset.z - z0 - runtime * w * direction) < 1e-4 + + df = pd.read_parquet(tmp_parquet) + times = (direction * (df["t"] - df["t"][0])).values.astype("timedelta64[s]") + timeref = np.arange(0, 5).astype("timedelta64[s]") + assert np.allclose(times, timeref, atol=np.timedelta64(1, "ms")) + lons = df["x"].values + assert np.allclose(lons, x0 + direction * u * np.arange(0, 5))