diff --git a/parcels/_core/utils/time.py b/parcels/_core/utils/time.py index 64623eba15..4473e87e49 100644 --- a/parcels/_core/utils/time.py +++ b/parcels/_core/utils/time.py @@ -48,6 +48,10 @@ def __init__(self, left: T, right: T) -> None: def __contains__(self, item: T) -> bool: return self.left <= item <= self.right + def is_all_time_in_interval(self, time): + item = np.atleast_1d(time) + return (self.left <= item).all() and (item <= self.right).all() + def __repr__(self) -> str: return f"TimeInterval(left={self.left!r}, right={self.right!r})" diff --git a/parcels/_index_search.py b/parcels/_index_search.py index 7c1f11ef2d..f95c215b0d 100644 --- a/parcels/_index_search.py +++ b/parcels/_index_search.py @@ -39,36 +39,14 @@ def _search_time_index(field: Field, time: datetime): if the sampled value is outside the time value range. """ if field.time_interval is None: - return 0, 0 + return np.zeros(shape=time.shape, dtype=np.float32), np.zeros(shape=time.shape, dtype=np.int32) - if time not in field.time_interval: + if not field.time_interval.is_all_time_in_interval(time): _raise_time_extrapolation_error(time, field=None) - time_index = field.data.time <= time - - if time_index.all(): - # If given time > last known field time, use - # the last field frame without interpolation - ti = len(field.data.time) - 1 - - elif np.logical_not(time_index).all(): - # If given time < any time in the field, use - # the first field frame without interpolation - ti = 0 - else: - ti = int(time_index.argmin() - 1) if time_index.any() else 0 - if len(field.data.time) == 1: - tau = 0 - elif ti == len(field.data.time) - 1: - tau = 1 - else: - tau = ( - (time - field.data.time[ti]).dt.total_seconds().values - / (field.data.time[ti + 1] - field.data.time[ti]).dt.total_seconds().values - if field.data.time[ti] != field.data.time[ti + 1] - else 0 - ) - return tau, ti + ti = np.searchsorted(field.data.time.data, time, side="right") - 1 + tau = (time - field.data.time.data[ti]) / (field.data.time.data[ti + 1] - field.data.time.data[ti]) + return np.atleast_1d(tau), np.atleast_1d(ti) def search_indices_vertical_z(depth, gridindexingtype: GridIndexingType, z: float): @@ -273,13 +251,13 @@ def _search_indices_rectilinear( def _search_indices_curvilinear_2d( grid: XGrid, y: float, x: float, yi_guess: int | None = None, xi_guess: int | None = None -): +): # TODO fix typing instructions to make clear that y, x etc need to be ndarrays yi, xi = yi_guess, xi_guess if yi is None or xi is None: faces = grid.get_spatial_hash().query(np.column_stack((y, x))) yi, xi = faces[0] - xsi = eta = -1.0 + xsi = eta = -1.0 * np.ones(len(x), dtype=float) invA = np.array( [ [1, 0, 0, 0], @@ -303,7 +281,7 @@ def _search_indices_curvilinear_2d( # if y < field.lonlat_minmax[2] or y > field.lonlat_minmax[3]: # _raise_field_out_of_bound_error(z, y, x) - while xsi < -tol or xsi > 1 + tol or eta < -tol or eta > 1 + tol: + while np.any(xsi < -tol) or np.any(xsi > 1 + tol) or np.any(eta < -tol) or np.any(eta > 1 + tol): 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]]) @@ -313,40 +291,29 @@ def _search_indices_curvilinear_2d( aa = a[3] * b[2] - a[2] * b[3] bb = a[3] * b[0] - a[0] * b[3] + a[1] * b[2] - a[2] * b[1] + x * b[3] - y * a[3] cc = a[1] * b[0] - a[0] * b[1] + x * b[1] - y * a[1] - if abs(aa) < 1e-12: # Rectilinear cell, or quasi - eta = -cc / bb - else: - det2 = bb * bb - 4 * aa * cc - if det2 > 0: # so, if det is nan we keep the xsi, eta from previous iter - det = np.sqrt(det2) - eta = (-bb + det) / (2 * aa) - if abs(a[1] + a[3] * eta) < 1e-12: # this happens when recti cell rotated of 90deg - xsi = ((y - py[0]) / (py[1] - py[0]) + (y - py[3]) / (py[2] - py[3])) * 0.5 - else: - xsi = (x - a[0] - a[2] * eta) / (a[1] + a[3] * eta) - if xsi < 0 and eta < 0 and xi == 0 and yi == 0: - _raise_field_out_of_bound_error(0, y, x) - if xsi > 1 and eta > 1 and xi == grid.xdim - 1 and yi == grid.ydim - 1: - _raise_field_out_of_bound_error(0, y, x) - if xsi < -tol: - xi -= 1 - elif xsi > 1 + tol: - xi += 1 - if eta < -tol: - yi -= 1 - elif eta > 1 + tol: - yi += 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), + ) + + 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)) + (yi, xi) = _reconnect_bnd_indices(yi, xi, grid.ydim, grid.xdim, grid.mesh) it += 1 if it > maxIterSearch: print(f"Correct cell not found after {maxIterSearch} iterations") _raise_field_out_of_bound_error(0, y, x) - xsi = max(0.0, xsi) - eta = max(0.0, eta) - xsi = min(1.0, xsi) - eta = min(1.0, eta) + xsi = np.where(xsi < 0.0, 0.0, np.where(xsi > 1.0, 1.0, xsi)) + eta = np.where(eta < 0.0, 0.0, np.where(eta > 1.0, 1.0, eta)) - if not ((0 <= xsi <= 1) and (0 <= eta <= 1)): + if np.any((xsi < 0) | (xsi > 1) | (eta < 0) | (eta > 1)): _raise_field_sampling_error(y, x) return (yi, eta, xi, xsi) @@ -442,20 +409,12 @@ def _search_indices_curvilinear(field, time, z, y, x, ti, particle=None, search2 def _reconnect_bnd_indices(yi: int, xi: int, ydim: int, xdim: int, sphere_mesh: bool): - if xi < 0: - if sphere_mesh: - xi = xdim - 2 - else: - xi = 0 - if xi > xdim - 2: - if sphere_mesh: - xi = 0 - else: - xi = xdim - 2 - if yi < 0: - yi = 0 - if yi > ydim - 2: - yi = ydim - 2 - if sphere_mesh: - xi = xdim - xi + xi = np.where(xi < 0, (xdim - 2) if sphere_mesh else 0, xi) + xi = np.where(xi > xdim - 2, 0 if sphere_mesh else (xdim - 2), xi) + + xi = np.where(yi > ydim - 2, xdim - xi if sphere_mesh else xi, xi) + + yi = np.where(yi < 0, 0, yi) + yi = np.where(yi > ydim - 2, ydim - 2, yi) + return yi, xi diff --git a/parcels/application_kernels/advection.py b/parcels/application_kernels/advection.py index d21111d98e..928da7d53a 100644 --- a/parcels/application_kernels/advection.py +++ b/parcels/application_kernels/advection.py @@ -21,11 +21,11 @@ def AdvectionRK4(particle, fieldset, time): # pragma: no cover dt = particle.dt / np.timedelta64(1, "s") # TODO: improve API for converting dt to seconds (u1, v1) = fieldset.UV[particle] lon1, lat1 = (particle.lon + u1 * 0.5 * dt, particle.lat + v1 * 0.5 * dt) - (u2, v2) = fieldset.UV[time + 0.5 * particle.dt, particle.depth, lat1, lon1, particle] + (u2, v2) = fieldset.UV[particle.time + 0.5 * particle.dt, particle.depth, lat1, lon1, particle] lon2, lat2 = (particle.lon + u2 * 0.5 * dt, particle.lat + v2 * 0.5 * dt) - (u3, v3) = fieldset.UV[time + 0.5 * particle.dt, particle.depth, lat2, lon2, particle] + (u3, v3) = fieldset.UV[particle.time + 0.5 * particle.dt, particle.depth, lat2, lon2, particle] lon3, lat3 = (particle.lon + u3 * dt, particle.lat + v3 * dt) - (u4, v4) = fieldset.UV[time + particle.dt, particle.depth, lat3, lon3, particle] + (u4, v4) = fieldset.UV[particle.time + particle.dt, particle.depth, lat3, lon3, particle] particle.dlon += (u1 + 2 * u2 + 2 * u3 + u4) / 6.0 * dt particle.dlat += (v1 + 2 * v2 + 2 * v3 + v4) / 6.0 * dt @@ -37,15 +37,15 @@ def AdvectionRK4_3D(particle, fieldset, time): # pragma: no cover lon1 = particle.lon + u1 * 0.5 * dt lat1 = particle.lat + v1 * 0.5 * dt dep1 = particle.depth + w1 * 0.5 * dt - (u2, v2, w2) = fieldset.UVW[time + 0.5 * particle.dt, dep1, lat1, lon1, particle] + (u2, v2, w2) = fieldset.UVW[particle.time + 0.5 * particle.dt, dep1, lat1, lon1, particle] lon2 = particle.lon + u2 * 0.5 * dt lat2 = particle.lat + v2 * 0.5 * dt dep2 = particle.depth + w2 * 0.5 * dt - (u3, v3, w3) = fieldset.UVW[time + 0.5 * particle.dt, dep2, lat2, lon2, particle] + (u3, v3, w3) = fieldset.UVW[particle.time + 0.5 * particle.dt, dep2, lat2, lon2, particle] lon3 = particle.lon + u3 * dt lat3 = particle.lat + v3 * dt dep3 = particle.depth + w3 * dt - (u4, v4, w4) = fieldset.UVW[time + particle.dt, dep3, lat3, lon3, particle] + (u4, v4, w4) = fieldset.UVW[particle.time + particle.dt, dep3, lat3, lon3, particle] particle.dlon += (u1 + 2 * u2 + 2 * u3 + u4) / 6 * dt particle.dlat += (v1 + 2 * v2 + 2 * v3 + v4) / 6 * dt particle.ddepth += (w1 + 2 * w2 + 2 * w3 + w4) / 6 * dt @@ -56,35 +56,35 @@ def AdvectionRK4_3D_CROCO(particle, fieldset, time): # pragma: no cover This kernel assumes the vertical velocity is the 'w' field from CROCO output and works on sigma-layers. """ dt = particle.dt / np.timedelta64(1, "s") # TODO: improve API for converting dt to seconds - sig_dep = particle.depth / fieldset.H[time, 0, particle.lat, particle.lon] + sig_dep = particle.depth / fieldset.H[particle.time, 0, particle.lat, particle.lon] - (u1, v1, w1) = fieldset.UVW[time, particle.depth, particle.lat, particle.lon, particle] - w1 *= sig_dep / fieldset.H[time, 0, particle.lat, particle.lon] + (u1, v1, w1) = fieldset.UVW[particle.time, particle.depth, particle.lat, particle.lon, particle] + w1 *= sig_dep / fieldset.H[particle.time, 0, particle.lat, particle.lon] lon1 = particle.lon + u1 * 0.5 * dt lat1 = particle.lat + v1 * 0.5 * dt sig_dep1 = sig_dep + w1 * 0.5 * dt - dep1 = sig_dep1 * fieldset.H[time, 0, lat1, lon1] + dep1 = sig_dep1 * fieldset.H[particle.time, 0, lat1, lon1] - (u2, v2, w2) = fieldset.UVW[time + 0.5 * particle.dt, dep1, lat1, lon1, particle] - w2 *= sig_dep1 / fieldset.H[time, 0, lat1, lon1] + (u2, v2, w2) = fieldset.UVW[particle.time + 0.5 * particle.dt, dep1, lat1, lon1, particle] + w2 *= sig_dep1 / fieldset.H[particle.time, 0, lat1, lon1] lon2 = particle.lon + u2 * 0.5 * dt lat2 = particle.lat + v2 * 0.5 * dt sig_dep2 = sig_dep + w2 * 0.5 * dt - dep2 = sig_dep2 * fieldset.H[time, 0, lat2, lon2] + dep2 = sig_dep2 * fieldset.H[particle.time, 0, lat2, lon2] - (u3, v3, w3) = fieldset.UVW[time + 0.5 * particle.dt, dep2, lat2, lon2, particle] - w3 *= sig_dep2 / fieldset.H[time, 0, lat2, lon2] + (u3, v3, w3) = fieldset.UVW[particle.time + 0.5 * particle.dt, dep2, lat2, lon2, particle] + w3 *= sig_dep2 / fieldset.H[particle.time, 0, lat2, lon2] lon3 = particle.lon + u3 * dt lat3 = particle.lat + v3 * dt sig_dep3 = sig_dep + w3 * dt - dep3 = sig_dep3 * fieldset.H[time, 0, lat3, lon3] + dep3 = sig_dep3 * fieldset.H[particle.time, 0, lat3, lon3] - (u4, v4, w4) = fieldset.UVW[time + particle.dt, dep3, lat3, lon3, particle] - w4 *= sig_dep3 / fieldset.H[time, 0, lat3, lon3] + (u4, v4, w4) = fieldset.UVW[particle.time + particle.dt, dep3, lat3, lon3, particle] + w4 *= sig_dep3 / fieldset.H[particle.time, 0, lat3, lon3] lon4 = particle.lon + u4 * dt lat4 = particle.lat + v4 * dt sig_dep4 = sig_dep + w4 * dt - dep4 = sig_dep4 * fieldset.H[time, 0, lat4, lon4] + dep4 = sig_dep4 * fieldset.H[particle.time, 0, lat4, lon4] particle.dlon += (u1 + 2 * u2 + 2 * u3 + u4) / 6 * dt particle.dlat += (v1 + 2 * v2 + 2 * v3 + v4) / 6 * dt @@ -115,14 +115,7 @@ def AdvectionRK45(particle, fieldset, time): # pragma: no cover Time-step dt is halved if error is larger than fieldset.RK45_tol, and doubled if error is smaller than 1/10th of tolerance. """ - dt = particle.next_dt / np.timedelta64(1, "s") # TODO: improve API for converting dt to seconds - if dt > fieldset.RK45_max_dt: - dt = fieldset.RK45_max_dt - particle.next_dt = fieldset.RK45_max_dt * np.timedelta64(1, "s") - if dt < fieldset.RK45_min_dt: - particle.next_dt = fieldset.RK45_min_dt * np.timedelta64(1, "s") - return StatusCode.Repeat - particle.dt = particle.next_dt + dt = particle.dt / np.timedelta64(1, "s") # TODO: improve API for converting dt to seconds c = [1.0 / 4.0, 3.0 / 8.0, 12.0 / 13.0, 1.0, 1.0 / 2.0] A = [ @@ -137,42 +130,58 @@ def AdvectionRK45(particle, fieldset, time): # pragma: no cover (u1, v1) = fieldset.UV[particle] lon1, lat1 = (particle.lon + u1 * A[0][0] * dt, particle.lat + v1 * A[0][0] * dt) - (u2, v2) = fieldset.UV[time + c[0] * particle.dt, particle.depth, lat1, lon1, particle] + (u2, v2) = fieldset.UV[particle.time + c[0] * particle.dt, particle.depth, lat1, lon1, particle] lon2, lat2 = ( particle.lon + (u1 * A[1][0] + u2 * A[1][1]) * dt, particle.lat + (v1 * A[1][0] + v2 * A[1][1]) * dt, ) - (u3, v3) = fieldset.UV[time + c[1] * particle.dt, particle.depth, lat2, lon2, particle] + (u3, v3) = fieldset.UV[particle.time + c[1] * particle.dt, particle.depth, lat2, lon2, particle] lon3, lat3 = ( particle.lon + (u1 * A[2][0] + u2 * A[2][1] + u3 * A[2][2]) * dt, particle.lat + (v1 * A[2][0] + v2 * A[2][1] + v3 * A[2][2]) * dt, ) - (u4, v4) = fieldset.UV[time + c[2] * particle.dt, particle.depth, lat3, lon3, particle] + (u4, v4) = fieldset.UV[particle.time + c[2] * particle.dt, particle.depth, lat3, lon3, particle] lon4, lat4 = ( particle.lon + (u1 * A[3][0] + u2 * A[3][1] + u3 * A[3][2] + u4 * A[3][3]) * dt, particle.lat + (v1 * A[3][0] + v2 * A[3][1] + v3 * A[3][2] + v4 * A[3][3]) * dt, ) - (u5, v5) = fieldset.UV[time + c[3] * particle.dt, particle.depth, lat4, lon4, particle] + (u5, v5) = fieldset.UV[particle.time + c[3] * particle.dt, particle.depth, lat4, lon4, particle] lon5, lat5 = ( particle.lon + (u1 * A[4][0] + u2 * A[4][1] + u3 * A[4][2] + u4 * A[4][3] + u5 * A[4][4]) * dt, particle.lat + (v1 * A[4][0] + v2 * A[4][1] + v3 * A[4][2] + v4 * A[4][3] + v5 * A[4][4]) * dt, ) - (u6, v6) = fieldset.UV[time + c[4] * particle.dt, particle.depth, lat5, lon5, particle] + (u6, v6) = fieldset.UV[particle.time + c[4] * particle.dt, particle.depth, lat5, lon5, particle] lon_4th = (u1 * b4[0] + u2 * b4[1] + u3 * b4[2] + u4 * b4[3] + u5 * b4[4]) * dt lat_4th = (v1 * b4[0] + v2 * b4[1] + v3 * b4[2] + v4 * b4[3] + v5 * b4[4]) * dt lon_5th = (u1 * b5[0] + u2 * b5[1] + u3 * b5[2] + u4 * b5[3] + u5 * b5[4] + u6 * b5[5]) * dt lat_5th = (v1 * b5[0] + v2 * b5[1] + v3 * b5[2] + v4 * b5[3] + v5 * b5[4] + v6 * b5[5]) * dt - kappa = math.sqrt(math.pow(lon_5th - lon_4th, 2) + math.pow(lat_5th - lat_4th, 2)) - if (kappa <= fieldset.RK45_tol) or (math.fabs(dt) < math.fabs(fieldset.RK45_min_dt)): - particle.dlon += lon_4th - particle.dlat += lat_4th - if (kappa <= fieldset.RK45_tol / 10) and (math.fabs(dt * 2) <= math.fabs(fieldset.RK45_max_dt)): - particle.next_dt *= 2 - else: - particle.next_dt /= 2 - return StatusCode.Repeat + kappa = np.sqrt(np.pow(lon_5th - lon_4th, 2) + np.pow(lat_5th - lat_4th, 2)) + + good_particles = (kappa <= fieldset.RK45_tol) | (np.fabs(dt) <= np.fabs(fieldset.RK45_min_dt)) + particle.dlon += np.where(good_particles, lon_5th, 0) + particle.dlat += np.where(good_particles, lat_5th, 0) + + increase_dt_particles = ( + good_particles & (kappa <= fieldset.RK45_tol / 10) & (np.fabs(dt * 2) <= np.fabs(fieldset.RK45_max_dt)) + ) + particle.dt = np.where(increase_dt_particles, particle.dt * 2, particle.dt) + particle.dt = np.where( + particle.dt > fieldset.RK45_max_dt * np.timedelta64(1, "s"), + fieldset.RK45_max_dt * np.timedelta64(1, "s"), + particle.dt, + ) + particle.state = np.where(good_particles, StatusCode.Success, particle.state) + + repeat_particles = np.invert(good_particles) + particle.dt = np.where(repeat_particles, particle.dt / 2, particle.dt) + particle.dt = np.where( + particle.dt < fieldset.RK45_min_dt * np.timedelta64(1, "s"), + fieldset.RK45_min_dt * np.timedelta64(1, "s"), + particle.dt, + ) + particle.state = np.where(repeat_particles, StatusCode.Repeat, particle.state) def AdvectionAnalytical(particle, fieldset, time): # pragma: no cover diff --git a/parcels/application_kernels/advectiondiffusion.py b/parcels/application_kernels/advectiondiffusion.py index 15362a6c07..afae2b5188 100644 --- a/parcels/application_kernels/advectiondiffusion.py +++ b/parcels/application_kernels/advectiondiffusion.py @@ -3,9 +3,6 @@ See `this tutorial <../examples/tutorial_diffusion.ipynb>`__ for a detailed explanation. """ -import math -import random - import numpy as np __all__ = ["AdvectionDiffusionEM", "AdvectionDiffusionM1", "DiffusionUniformKh"] @@ -28,21 +25,21 @@ def AdvectionDiffusionM1(particle, fieldset, time): # pragma: no cover """ dt = particle.dt / np.timedelta64(1, "s") # TODO: improve API for converting dt to seconds # Wiener increment with zero mean and std of sqrt(dt) - dWx = random.normalvariate(0, math.sqrt(math.fabs(dt))) - dWy = random.normalvariate(0, math.sqrt(math.fabs(dt))) + dWx = np.random.normal(0, np.sqrt(np.fabs(dt))) + dWy = np.random.normal(0, np.sqrt(np.fabs(dt))) - Kxp1 = fieldset.Kh_zonal[time, particle.depth, particle.lat, particle.lon + fieldset.dres] - Kxm1 = fieldset.Kh_zonal[time, particle.depth, particle.lat, particle.lon - fieldset.dres] + Kxp1 = fieldset.Kh_zonal[particle.time, particle.depth, particle.lat, particle.lon + fieldset.dres, particle] + Kxm1 = fieldset.Kh_zonal[particle.time, particle.depth, particle.lat, particle.lon - fieldset.dres, particle] dKdx = (Kxp1 - Kxm1) / (2 * fieldset.dres) - u, v = fieldset.UV[time, particle.depth, particle.lat, particle.lon] - bx = math.sqrt(2 * fieldset.Kh_zonal[time, particle.depth, particle.lat, particle.lon]) + u, v = fieldset.UV[particle.time, particle.depth, particle.lat, particle.lon, particle] + bx = np.sqrt(2 * fieldset.Kh_zonal[particle.time, particle.depth, particle.lat, particle.lon, particle]) - Kyp1 = fieldset.Kh_meridional[time, particle.depth, particle.lat + fieldset.dres, particle.lon] - Kym1 = fieldset.Kh_meridional[time, particle.depth, particle.lat - fieldset.dres, particle.lon] + Kyp1 = fieldset.Kh_meridional[particle.time, particle.depth, particle.lat + fieldset.dres, particle.lon, particle] + Kym1 = fieldset.Kh_meridional[particle.time, particle.depth, particle.lat - fieldset.dres, particle.lon, particle] dKdy = (Kyp1 - Kym1) / (2 * fieldset.dres) - by = math.sqrt(2 * fieldset.Kh_meridional[time, particle.depth, particle.lat, particle.lon]) + by = np.sqrt(2 * fieldset.Kh_meridional[particle.time, particle.depth, particle.lat, particle.lon, particle]) # Particle positions are updated only after evaluating all terms. particle.dlon += u * dt + 0.5 * dKdx * (dWx**2 + dt) + bx * dWx @@ -64,22 +61,22 @@ def AdvectionDiffusionEM(particle, fieldset, time): # pragma: no cover """ dt = particle.dt / np.timedelta64(1, "s") # Wiener increment with zero mean and std of sqrt(dt) - dWx = random.normalvariate(0, math.sqrt(math.fabs(dt))) - dWy = random.normalvariate(0, math.sqrt(math.fabs(dt))) + dWx = np.random.normal(0, np.sqrt(np.fabs(dt))) + dWy = np.random.normal(0, np.sqrt(np.fabs(dt))) - u, v = fieldset.UV[time, particle.depth, particle.lat, particle.lon] + u, v = fieldset.UV[particle.time, particle.depth, particle.lat, particle.lon, particle] - Kxp1 = fieldset.Kh_zonal[time, particle.depth, particle.lat, particle.lon + fieldset.dres] - Kxm1 = fieldset.Kh_zonal[time, particle.depth, particle.lat, particle.lon - fieldset.dres] + Kxp1 = fieldset.Kh_zonal[particle.time, particle.depth, particle.lat, particle.lon + fieldset.dres, particle] + Kxm1 = fieldset.Kh_zonal[particle.time, particle.depth, particle.lat, particle.lon - fieldset.dres, particle] dKdx = (Kxp1 - Kxm1) / (2 * fieldset.dres) ax = u + dKdx - bx = math.sqrt(2 * fieldset.Kh_zonal[time, particle.depth, particle.lat, particle.lon]) + bx = np.sqrt(2 * fieldset.Kh_zonal[particle.time, particle.depth, particle.lat, particle.lon, particle]) - Kyp1 = fieldset.Kh_meridional[time, particle.depth, particle.lat + fieldset.dres, particle.lon] - Kym1 = fieldset.Kh_meridional[time, particle.depth, particle.lat - fieldset.dres, particle.lon] + Kyp1 = fieldset.Kh_meridional[particle.time, particle.depth, particle.lat + fieldset.dres, particle.lon, particle] + Kym1 = fieldset.Kh_meridional[particle.time, particle.depth, particle.lat - fieldset.dres, particle.lon, particle] dKdy = (Kyp1 - Kym1) / (2 * fieldset.dres) ay = v + dKdy - by = math.sqrt(2 * fieldset.Kh_meridional[time, particle.depth, particle.lat, particle.lon]) + by = np.sqrt(2 * fieldset.Kh_meridional[particle.time, particle.depth, particle.lat, particle.lon, particle]) # Particle positions are updated only after evaluating all terms. particle.dlon += ax * dt + bx * dWx @@ -106,11 +103,13 @@ def DiffusionUniformKh(particle, fieldset, time): # pragma: no cover """ dt = particle.dt / np.timedelta64(1, "s") # Wiener increment with zero mean and std of sqrt(dt) - dWx = random.normalvariate(0, math.sqrt(math.fabs(dt))) - dWy = random.normalvariate(0, math.sqrt(math.fabs(dt))) + dWx = np.random.normal(0, np.sqrt(np.fabs(dt))) + dWy = np.random.normal(0, np.sqrt(np.fabs(dt))) + + print(particle) - bx = math.sqrt(2 * fieldset.Kh_zonal[particle]) - by = math.sqrt(2 * fieldset.Kh_meridional[particle]) + bx = np.sqrt(2 * fieldset.Kh_zonal[particle]) + by = np.sqrt(2 * fieldset.Kh_meridional[particle]) particle.dlon += bx * dWx particle.dlat += by * dWy diff --git a/parcels/application_kernels/interpolation.py b/parcels/application_kernels/interpolation.py index a41b76afce..1df817a2a9 100644 --- a/parcels/application_kernels/interpolation.py +++ b/parcels/application_kernels/interpolation.py @@ -4,99 +4,38 @@ from typing import TYPE_CHECKING +import dask.array as dask import numpy as np - -from parcels.field import Field -from parcels.tools.statuscodes import ( - FieldOutOfBoundError, -) +import xarray as xr if TYPE_CHECKING: + from parcels.field import Field from parcels.uxgrid import _UXGRID_AXES from parcels.xgrid import _XGRID_AXES __all__ = [ "UXPiecewiseConstantFace", "UXPiecewiseLinearNode", - "XBiLinear", - "XBiLinearPeriodic", - "XTriLinear", + "XLinear", + "ZeroInterpolator", ] -def XBiLinear( - 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, -): - """Bilinear interpolation on a regular grid.""" - xi, xsi = position["X"] - yi, eta = position["Y"] - zi, _ = position["Z"] - - data = field.data.data[:, zi, yi : yi + 2, xi : xi + 2] - if tau > 0: - data = (1 - tau) * data[ti, :, :] + tau * data[ti + 1, :, :] - else: - data = data[ti, :, :] - - return ( - (1 - xsi) * (1 - eta) * data[0, 0] - + xsi * (1 - eta) * data[0, 1] - + xsi * eta * data[1, 1] - + (1 - xsi) * eta * data[1, 0] - ) - - -def XBiLinearPeriodic( +def ZeroInterpolator( field: Field, ti: int, - position: dict[_XGRID_AXES, tuple[int, float | np.ndarray]], + 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, -): - """Bilinear interpolation on a regular grid with periodic boundary conditions in horizontal directions.""" - xi, xsi = position["X"] - yi, eta = position["Y"] - zi, _ = position["Z"] - - if xi < 0: - xi = 0 - xsi = (x - field.grid.lon[xi]) / (field.grid.lon[xi + 1] - field.grid.lon[xi]) - if yi < 0: - yi = 0 - eta = (y - field.grid.lat[yi]) / (field.grid.lat[yi + 1] - field.grid.lat[yi]) - - data = field.data.data[:, zi, yi : yi + 2, xi : xi + 2] - data = (1 - tau) * data[ti, :, :] + tau * data[ti + 1, :, :] - - xsi = 0 if not np.isfinite(xsi) else xsi - eta = 0 if not np.isfinite(eta) else eta - - if xsi > 0 and eta > 0: - return ( - (1 - xsi) * (1 - eta) * data[0, 0] - + xsi * (1 - eta) * data[0, 1] - + xsi * eta * data[1, 1] - + (1 - xsi) * eta * data[1, 0] - ) - elif xsi > 0 and eta == 0: - return (1 - xsi) * data[0, 0] + xsi * data[0, 1] - elif xsi == 0 and eta > 0: - return (1 - eta) * data[0, 0] + eta * data[1, 0] - else: - return data[0, 0] +) -> np.float32 | np.float64: + """Template function used for the signature check of the lateral interpolation methods.""" + return 0.0 -def XTriLinear( +def XLinear( field: Field, ti: int, position: dict[_XGRID_AXES, tuple[int, float | np.ndarray]], @@ -111,32 +50,65 @@ def XTriLinear( yi, eta = position["Y"] zi, zeta = position["Z"] - if zi < 0 or xi < 0 or yi < 0: - raise FieldOutOfBoundError + axis_dim = field.grid.get_axis_dim_mapping(field.data.dims) + data = field.data - data = field.data.data[:, zi : zi + 2, yi : yi + 2, xi : xi + 2] - data = (1 - tau) * data[ti, :, :, :] + tau * data[ti + 1, :, :, :] - if zeta > 0: - data = (1 - zeta) * data[0, :, :] + zeta * data[1, :, :] + lenT = 2 if np.any(tau > 0) else 1 + lenZ = 2 if np.any(zeta > 0) else 1 + + # Time coordinates: 8 points at ti, then 8 points at ti+1 + if lenT == 1: + ti = np.repeat(ti, lenZ * 4) + else: + ti_1 = np.clip(ti + 1, 0, data.shape[0] - 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: - data = data[0, :, :] - - xsi = 0 if not np.isfinite(xsi) else xsi - eta = 0 if not np.isfinite(eta) else eta - - if xsi > 0 and eta > 0: - return ( - (1 - xsi) * (1 - eta) * data[0, 0] - + xsi * (1 - eta) * data[0, 1] - + xsi * eta * data[1, 1] - + (1 - xsi) * eta * data[1, 0] - ) - elif xsi > 0 and eta == 0: - return (1 - xsi) * data[0, 0] + xsi * data[0, 1] - elif xsi == 0 and eta > 0: - return (1 - eta) * data[0, 0] + eta * data[1, 0] + zi_1 = np.clip(zi + 1, 0, data.shape[1] - 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 = 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 = np.tile(np.column_stack([xi, xi_1, xi, xi_1]).flatten(), (lenT) * (lenZ)) + + # 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 data.dims: + selection_dict["time"] = xr.DataArray(ti, dims=("points")) + + corner_data = data.isel(selection_dict).data.reshape(lenT, lenZ, len(xsi), 4) + + if lenT == 2: + tau = tau[np.newaxis, :, np.newaxis] + corner_data = corner_data[0, :, :, :] * (1 - tau) + corner_data[1, :, :, :] * tau else: - return data[0, 0] + corner_data = corner_data[0, :, :, :] + + if lenZ == 2: + zeta = zeta[:, np.newaxis] + corner_data = corner_data[0, :, :] * (1 - zeta) + corner_data[1, :, :] * zeta + else: + corner_data = corner_data[0, :, :] + + value = ( + (1 - xsi) * (1 - eta) * corner_data[:, 0] + + xsi * (1 - eta) * corner_data[:, 1] + + (1 - xsi) * eta * corner_data[:, 2] + + xsi * eta * corner_data[:, 3] + ) + return value.compute() if isinstance(value, dask.Array) else value def UXPiecewiseConstantFace( diff --git a/parcels/field.py b/parcels/field.py index 1228b0a1f6..7c2170ac97 100644 --- a/parcels/field.py +++ b/parcels/field.py @@ -16,6 +16,7 @@ VectorType, assert_valid_mesh, ) +from parcels.application_kernels.interpolation import UXPiecewiseLinearNode, XLinear, ZeroInterpolator from parcels.particle import KernelParticle from parcels.tools.converters import ( UnitConverter, @@ -23,13 +24,10 @@ ) from parcels.tools.statuscodes import ( AllParcelsErrorCodes, - FieldOutOfBoundError, - FieldOutOfBoundSurfaceError, - FieldSamplingError, - _raise_field_out_of_bound_error, + StatusCode, ) from parcels.uxgrid import UxGrid -from parcels.xgrid import XGrid, _transpose_xfield_data_to_tzyx +from parcels.xgrid import LEFT_OUT_OF_BOUNDS, RIGHT_OUT_OF_BOUNDS, XGrid, _transpose_xfield_data_to_tzyx from ._index_search import _search_time_index @@ -52,23 +50,9 @@ def _deal_with_errors(error, key, vector_type: VectorType): return 0 -def ZeroInterpolator( - field: Field, - 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, -) -> np.float32 | np.float64: - """Template function used for the signature check of the lateral interpolation methods.""" - return 0.0 - - _DEFAULT_INTERPOLATOR_MAPPING = { - XGrid: ZeroInterpolator, # TODO v4: Update these to better defaults - UxGrid: ZeroInterpolator, + XGrid: XLinear, + UxGrid: UXPiecewiseLinearNode, } @@ -258,18 +242,11 @@ def eval(self, time: datetime, z, y, x, particle=None, applyConversion=True): # else: # _ei = particle.ei[self.igrid] - try: - tau, ti = _search_time_index(self, time) - position = self.grid.search(z, y, x, ei=_ei) - value = self._interp_method(self, ti, position, tau, time, z, y, x) - - if np.isnan(value): - # Detect Out-of-bounds sampling and raise exception - _raise_field_out_of_bound_error(z, y, x) + tau, ti = _search_time_index(self, time) + position = self.grid.search(z, y, x, ei=_ei) + _update_particle_states(particle, position) - except (FieldSamplingError, FieldOutOfBoundError, FieldOutOfBoundSurfaceError) as e: - e.add_note(f"Error interpolating field '{self.name}'.") - raise e + value = self._interp_method(self, ti, position, tau, time, z, y, x) if applyConversion: value = self.units.to_target(value, z, y, x) @@ -345,31 +322,28 @@ def eval(self, time: datetime, z, y, x, particle=None, applyConversion=True): # else: # _ei = particle.ei[self.igrid] - try: - tau, ti = _search_time_index(self.U, time) - position = self.grid.search(z, y, x, ei=_ei) - if self._vector_interp_method is None: - u = self.U._interp_method(self.U, ti, position, tau, time, z, y, x) - v = self.V._interp_method(self.V, ti, position, tau, time, z, y, x) - if "3D" in self.vector_type: - w = self.W._interp_method(self.W, ti, position, tau, time, z, y, x) - else: - (u, v, w) = self._vector_interp_method(self, ti, position, time, z, y, x) + tau, ti = _search_time_index(self.U, time) + position = self.grid.search(z, y, x, ei=_ei) + _update_particle_states(particle, position) - 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 self._vector_interp_method is None: + u = self.U._interp_method(self.U, ti, position, tau, time, z, y, x) + v = self.V._interp_method(self.V, ti, position, tau, time, z, y, x) + if "3D" in self.vector_type: + w = self.W._interp_method(self.W, ti, position, tau, time, z, y, x) + else: + (u, v, w) = self._vector_interp_method(self, ti, position, time, z, y, x) + 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: - return (u, v, w) - else: - return (u, v) + w = self.W.units.to_target(w, z, y, x) if self.W else 0.0 - except (FieldSamplingError, FieldOutOfBoundError, FieldOutOfBoundSurfaceError) as e: - e.add_note(f"Error interpolating field '{self.name}'.") - raise e + if "3D" in self.vector_type: + return (u, v, w) + else: + return (u, v) def __getitem__(self, key): try: @@ -381,6 +355,17 @@ def __getitem__(self, key): return _deal_with_errors(error, key, vector_type=self.vector_type) +def _update_particle_states(particle, position): + """Update the particle states based on the position dictionary.""" + if particle and "X" in position: # TODO also support uxgrid search + particle.state = np.where(position["X"][0] < 0, StatusCode.ErrorOutOfBounds, particle.state) + particle.state = np.where(position["Y"][0] < 0, StatusCode.ErrorOutOfBounds, particle.state) + particle.state = np.where(position["Z"][0] == RIGHT_OUT_OF_BOUNDS, StatusCode.ErrorOutOfBounds, particle.state) + particle.state = np.where( + position["Z"][0] == LEFT_OUT_OF_BOUNDS, StatusCode.ErrorThroughSurface, particle.state + ) + + def _assert_valid_uxdataarray(data: ux.UxDataArray): """Verifies that all the required attributes are present in the xarray.DataArray or uxarray.UxDataArray object. diff --git a/parcels/kernel.py b/parcels/kernel.py index 7bd4a59d3d..4cd8f37778 100644 --- a/parcels/kernel.py +++ b/parcels/kernel.py @@ -15,10 +15,10 @@ from parcels.basegrid import GridType from parcels.tools.statuscodes import ( StatusCode, - TimeExtrapolationError, _raise_field_out_of_bound_error, _raise_field_out_of_bound_surface_error, _raise_field_sampling_error, + _raise_time_extrapolation_error, ) from parcels.tools.warnings import KernelWarning @@ -213,105 +213,76 @@ def from_list(cls, fieldset, ptype, pyfunc_list): return cls(fieldset, ptype, pyfunc_list) def execute(self, pset, endtime, dt): - """Execute this Kernel over a ParticleSet for several timesteps.""" - pset._data["state"][:] = StatusCode.Evaluate + """Execute this Kernel over a ParticleSet for several timesteps. - if abs(dt) < np.timedelta64(1000, "ns"): # TODO still needed? - warnings.warn( - "'dt' is too small, causing numerical accuracy limit problems. Please chose a higher 'dt' and rather scale the 'time' axis of the field accordingly. (related issue #762)", - RuntimeWarning, - stacklevel=2, - ) + Parameters + ---------- + pset : + object of (sub-)type ParticleSet + endtime : + endtime of this overall kernel evaluation step + dt : + computational integration timestep from pset.execute + """ + compute_time_direction = 1 if dt > 0 else -1 + + pset._data["state"][:] = StatusCode.Evaluate if not self._positionupdate_kernels_added: self.add_positionupdate_kernels() self._positionupdate_kernels_added = True - for p in pset: - self.evaluate_particle(p, endtime) - if p.state == StatusCode.StopAllExecution: - return StatusCode.StopAllExecution + while (len(pset) > 0) and np.any(np.isin(pset.state, [StatusCode.Evaluate, StatusCode.Repeat])): + time_to_endtime = compute_time_direction * (endtime - pset.time_nextloop) - # Remove all particles that signalled deletion - self.remove_deleted(pset) - - # Identify particles that threw errors - n_error = pset._num_error_particles - - while n_error > 0: - for i in pset._error_particles: - p = pset[i] - if p.state == StatusCode.StopExecution: - return - if p.state == StatusCode.StopAllExecution: - return StatusCode.StopAllExecution - if p.state == StatusCode.Repeat: - p.state = StatusCode.Evaluate - elif p.state == StatusCode.ErrorTimeExtrapolation: - raise TimeExtrapolationError(p.time) - elif p.state == StatusCode.ErrorOutOfBounds: - _raise_field_out_of_bound_error(p.depth, p.lat, p.lon) - elif p.state == StatusCode.ErrorThroughSurface: - _raise_field_out_of_bound_surface_error(p.depth, p.lat, p.lon) - elif p.state == StatusCode.Error: - _raise_field_sampling_error(p.depth, p.lat, p.lon) - elif p.state == StatusCode.Delete: - pass - else: - warnings.warn( - f"Deleting particle {p.trajectory} because of non-recoverable error", - RuntimeWarning, - stacklevel=2, - ) - p.delete() + if all(time_to_endtime <= 0): + return StatusCode.Success - # Remove all particles that signalled deletion - self.remove_deleted(pset) # Generalizable version! + # adapt dt to end exactly on endtime + if compute_time_direction == 1: + pset.dt = np.maximum(np.minimum(pset.dt, time_to_endtime), 0) + else: + pset.dt = np.minimum(np.maximum(pset.dt, -time_to_endtime), 0) - # Re-execute Kernels to retry particles with StatusCode.Repeat - for p in pset: - self.evaluate_particle(p, endtime) + # run kernels for all particles that need to be evaluated + evaluate_particles = (pset.state == StatusCode.Evaluate) & (pset.dt != 0) + for f in self._pyfuncs: + f(pset[evaluate_particles], self._fieldset, None) - n_error = pset._num_error_particles + # check for particles that have to be repeated + repeat_particles = pset.state == StatusCode.Repeat + while np.any(repeat_particles): + f(pset[repeat_particles], self._fieldset, None) + repeat_particles = pset.state == StatusCode.Repeat - def evaluate_particle(self, p, endtime): - """Execute the kernel evaluation of for an individual particle. + # revert to original dt (unless in RK45 mode) + if not hasattr(self.fieldset, "RK45_tol"): + pset._data["dt"][:] = dt - Parameters - ---------- - p : - object of (sub-)type Particle - endtime : - endtime of this overall kernel evaluation step - dt : - computational integration timestep - """ - sign_dt = 1 if p.dt >= 0 else -1 - while p.state in [StatusCode.Evaluate, StatusCode.Repeat]: - if sign_dt * (endtime - p.time_nextloop) <= 0: - return p - - pre_dt = p.dt - try: # Use next_dt from AdvectionRK45 if it is set - if abs(endtime - p.time_nextloop) < abs(p.next_dt) - np.timedelta64(1000, "ns"): - p.next_dt = sign_dt * (endtime - p.time_nextloop) - except KeyError: - if sign_dt * (endtime - p.time_nextloop) <= p.dt: - p.dt = sign_dt * (endtime - p.time_nextloop) - res = None - for f in self._pyfuncs: - res_tmp = f(p, self._fieldset, p.time_nextloop) - if res_tmp is not None: # TODO v4: Remove once all kernels return StatusCode - res = res_tmp - if res in [StatusCode.StopExecution, StatusCode.Repeat]: - break - - if res is None: - if p.state == StatusCode.Success: - if sign_dt * (p.time - endtime) > 0: - p.state = StatusCode.Evaluate - else: - p.state = res + # Reset particle state for particles that signalled success and have not reached endtime yet + particles_to_evaluate = (pset.state == StatusCode.Success) & (time_to_endtime > 0) + pset[particles_to_evaluate].state = StatusCode.Evaluate + + # delete particles that signalled deletion + self.remove_deleted(pset) + + # check and throw errors + if np.any(pset.state == StatusCode.StopAllExecution): + return StatusCode.StopAllExecution - p.dt = pre_dt - return p + errors_to_throw = { + StatusCode.ErrorTimeExtrapolation: _raise_time_extrapolation_error, + StatusCode.ErrorOutOfBounds: _raise_field_out_of_bound_error, + StatusCode.ErrorThroughSurface: _raise_field_out_of_bound_surface_error, + StatusCode.Error: _raise_field_sampling_error, + } + + for error_code, error_func in errors_to_throw.items(): + if np.any(pset.state == error_code): + inds = pset.state == error_code + if error_code == StatusCode.ErrorTimeExtrapolation: + error_func(pset[inds].time) + else: + error_func(pset[inds].depth, pset[inds].lat, pset[inds].lon) + + return pset diff --git a/parcels/particle.py b/parcels/particle.py index 546f53cf62..8a51e3680f 100644 --- a/parcels/particle.py +++ b/parcels/particle.py @@ -138,6 +138,13 @@ def __setattr__(self, name, value): else: self._data[name][self._index] = value + def __getitem__(self, index): + self._index = index + return self + + def __len__(self): + return len(self._index) + def _assert_no_duplicate_variable_names(*, existing_vars: list[Variable], new_vars: list[Variable]): existing_names = {var.name for var in existing_vars} diff --git a/parcels/particleset.py b/parcels/particleset.py index 2647800e11..871d151b35 100644 --- a/parcels/particleset.py +++ b/parcels/particleset.py @@ -176,6 +176,14 @@ def __getitem__(self, index): """Get a single particle by index.""" return KernelParticle(self._data, index=index) + def __setattr__(self, name, value): + if name in ["_data"]: + object.__setattr__(self, name, value) + elif isinstance(self._data, dict) and name in self._data.keys(): + self._data[name][:] = value + else: + object.__setattr__(self, name, value) + @staticmethod def lonlatdepth_dtype_from_field_interp_method(field): # TODO update this when now interp methods are implemented @@ -517,11 +525,11 @@ def execute( raise TypeError("The runtime must be a np.timedelta64 object") else: - if not np.isnat(self._data["time_nextloop"]).any(): + if not np.isnat(self.time_nextloop).any(): if sign_dt > 0: - start_time = self._data["time_nextloop"].min() + start_time = self.time_nextloop.min() else: - start_time = self._data["time_nextloop"].max() + start_time = self.time_nextloop.max() else: if sign_dt > 0: start_time = self.fieldset.time_interval.left @@ -572,9 +580,7 @@ def execute( next_time = end_time # TODO update to min(next_output, end_time) when ParticleFile works else: next_time = end_time # TODO update to max(next_output, end_time) when ParticleFile works - res = self._kernel.execute(self, endtime=next_time, dt=dt) - if res == StatusCode.StopAllExecution: - return StatusCode.StopAllExecution + self._kernel.execute(self, endtime=next_time, dt=dt) # TODO: Handle IO timing based of timedelta or datetime objects if next_output: diff --git a/parcels/tools/converters.py b/parcels/tools/converters.py index 17311d2055..68c3810b6c 100644 --- a/parcels/tools/converters.py +++ b/parcels/tools/converters.py @@ -1,6 +1,6 @@ from __future__ import annotations -from math import cos, pi +from math import pi import numpy as np import numpy.typing as npt @@ -62,10 +62,10 @@ class GeographicPolar(UnitConverter): target_unit = "degree" def to_target(self, value, z, y, x): - return value / 1000.0 / 1.852 / 60.0 / cos(y * pi / 180) + return value / 1000.0 / 1.852 / 60.0 / np.cos(y * pi / 180) def to_source(self, value, z, y, x): - return value * 1000.0 * 1.852 * 60.0 * cos(y * pi / 180) + return value * 1000.0 * 1.852 * 60.0 * np.cos(y * pi / 180) class GeographicSquare(UnitConverter): @@ -90,10 +90,10 @@ class GeographicPolarSquare(UnitConverter): target_unit = "degree2" def to_target(self, value, z, y, x): - return value / pow(1000.0 * 1.852 * 60.0 * cos(y * pi / 180), 2) + return value / pow(1000.0 * 1.852 * 60.0 * np.cos(y * pi / 180), 2) def to_source(self, value, z, y, x): - return value * pow(1000.0 * 1.852 * 60.0 * cos(y * pi / 180), 2) + return value * pow(1000.0 * 1.852 * 60.0 * np.cos(y * pi / 180), 2) unitconverters_map = { diff --git a/parcels/uxgrid.py b/parcels/uxgrid.py index cc622feaf7..7f0048843f 100644 --- a/parcels/uxgrid.py +++ b/parcels/uxgrid.py @@ -5,8 +5,8 @@ import numpy as np import uxarray as ux -from parcels.field import FieldOutOfBoundError from parcels.spatialhash import _barycentric_coordinates +from parcels.tools.statuscodes import FieldOutOfBoundError from parcels.xgrid import _search_1d_array from .basegrid import BaseGrid @@ -95,14 +95,14 @@ def try_face(fid): return bcoords, self.ravel_index(zi, neighbor) # Global fallback as last ditch effort - face_ids = self.uxgrid.get_faces_containing_point([x, y], return_counts=False)[0] + points = np.column_stack((x, y)) + face_ids = self.uxgrid.get_faces_containing_point(points, return_counts=False)[0] fi = face_ids[0] if len(face_ids) > 0 else -1 if fi == -1: raise FieldOutOfBoundError(z, y, x) bcoords = try_face(fi) if bcoords is None: raise FieldOutOfBoundError(z, y, x) - return {"Z": (zi, zeta), "FACE": (fi, bcoords)} def _get_barycentric_coordinates_latlon(self, y, x, fi): @@ -118,9 +118,10 @@ def _get_barycentric_coordinates_latlon(self, y, x, fi): ) ) - coord = np.deg2rad([x, y]) + coord = np.deg2rad(np.column_stack((x, y))) bcoord = np.asarray(_barycentric_coordinates(nodes, coord)) - err = abs(np.dot(bcoord, nodes[:, 0]) - coord[0]) + abs(np.dot(bcoord, nodes[:, 1]) - coord[1]) + proj_coord = np.matmul(np.transpose(nodes), bcoord) + err = np.linalg.norm(proj_coord - coord) return bcoord, err def _get_barycentric_coordinates_cartesian(self, y, x, fi): diff --git a/parcels/xgrid.py b/parcels/xgrid.py index 77a49d7406..bac507d8e1 100644 --- a/parcels/xgrid.py +++ b/parcels/xgrid.py @@ -10,7 +10,6 @@ from parcels._index_search import _search_indices_curvilinear_2d from parcels.basegrid import BaseGrid from parcels.spatialhash import SpatialHash -from parcels.tools.statuscodes import FieldOutOfBoundError, FieldOutOfBoundSurfaceError _XGRID_AXES = Literal["X", "Y", "Z"] _XGRID_AXES_ORDERING: Sequence[_XGRID_AXES] = "ZYX" @@ -23,6 +22,9 @@ _DEFAULT_XGCM_KWARGS = {"periodic": False} +LEFT_OUT_OF_BOUNDS = -2 +RIGHT_OUT_OF_BOUNDS = -1 + def get_cell_count_along_dim(axis: xgcm.Axis) -> int: first_coord = list(axis.coords.items())[0] @@ -277,15 +279,6 @@ def search(self, z, y, x, ei=None): zi, zeta = _search_1d_array(ds.depth.values, z) else: zi, zeta = 0, 0.0 - if zi == -1: - if zeta < 0: - raise FieldOutOfBoundError( - f"Depth {z} is out of bounds for the grid with depth values {ds.depth.values}." - ) - elif zeta > 1: - raise FieldOutOfBoundSurfaceError( - f"Depth {z} is out of the surface for the grid with depth values {ds.depth.values}." - ) if ds.lon.ndim == 1: yi, eta = _search_1d_array(ds.lat.values, y) @@ -319,6 +312,38 @@ def _fpoint_info(self): return axis_position_mapping + def get_axis_dim_mapping(self, dims: list[str]) -> dict[_XGRID_AXES, str]: + """ + Maps xarray dimension names to their corresponding axis (X, Y, Z). + + WARNING: This API is unstable and subject to change in future versions. + + Parameters + ---------- + dims : list[str] + List of xarray dimension names + + Returns + ------- + dict[_XGRID_AXES, str] + Dictionary mapping axes (X, Y, Z) to their corresponding dimension names + + Examples + -------- + >>> grid.get_axis_dim_mapping(['time', 'lat', 'lon']) + {'Y': 'lat', 'X': 'lon'} + + Notes + ----- + Only returns mappings for spatial axes (X, Y, Z) that are present in the grid. + """ + result = {} + for dim in dims: + axis = get_axis_from_dim_name(self.xgcm_grid.axes, dim) + if axis in self.axes: # Only include spatial axes (X, Y, Z) + result[cast(_XGRID_AXES, axis)] = dim + return result + def get_spatial_hash( self, reconstruct=False, @@ -477,7 +502,6 @@ def _search_1d_array( Searches for the particle location in a 1D array and returns barycentric coordinate along dimension. Assumptions: - - particle position x is within the bounds of the array - array is strictly monotonically increasing. Parameters @@ -489,17 +513,30 @@ def _search_1d_array( Returns ------- - int - Index of the element just before the position x in the array. - float + array of int + Index of the element just before the position x in the array. Note that this index is -2 if the index is left out of bounds and -1 if the index is right out of bounds. + array of float Barycentric coordinate. """ # TODO v4: We probably rework this to deal with 0D arrays before this point (as we already know field dimensionality) if len(arr) < 2: - return 0, 0.0 - i = np.argmin(arr <= x) - 1 - bcoord = (x - arr[i]) / (arr[i + 1] - arr[i]) - return i, bcoord + return np.zeros(shape=x.shape, dtype=np.int32), np.zeros_like(x) + index = np.searchsorted(arr, x, side="right") - 1 + # Use broadcasting to avoid repeated array access + arr_index = arr[index] + arr_next = arr[np.clip(index + 1, 1, len(arr) - 1)] # Ensure we don't go out of bounds + bcoord = (x - arr_index) / (arr_next - arr_index) + + # TODO check how we can avoid searchsorted when grid spacing is uniform + # dx = arr[1] - arr[0] + # index = ((x - arr[0]) / dx).astype(int) + # index = np.clip(index, 0, len(arr) - 2) + # bcoord = (x - arr[index]) / dx + + index = np.where(x < arr[0], LEFT_OUT_OF_BOUNDS, index) + index = np.where(x >= arr[-1], RIGHT_OUT_OF_BOUNDS, index) + + return np.atleast_1d(index), np.atleast_1d(bcoord) def _convert_center_pos_to_fpoint( diff --git a/tests/test_kernel_execution.py b/tests/test_kernel_execution.py index e9e72e62ca..6a61f29860 100644 --- a/tests/test_kernel_execution.py +++ b/tests/test_kernel_execution.py @@ -2,13 +2,10 @@ import pytest from parcels import ( - FieldOutOfBoundError, FieldSet, Particle, ParticleSet, - StatusCode, ) -from tests.common_kernels import DeleteParticle, DoNothing, MoveEast, MoveNorth from tests.utils import create_fieldset_unit_mesh @@ -54,132 +51,6 @@ def SampleP(particle, fieldset, time): # pragma: no cover assert np.allclose(lons[0], 0.2) -@pytest.mark.parametrize( - "start, end, substeps, dt", - [ - (0.0, 10.0, 1, 1.0), - (0.0, 10.0, 4, 1.0), - (0.0, 10.0, 1, 3.0), - (2.0, 16.0, 5, 3.0), - (20.0, 10.0, 4, -1.0), - (20.0, -10.0, 7, -2.0), - ], -) -def test_execution_endtime(fieldset_unit_mesh, start, end, substeps, dt): - npart = 10 - pset = ParticleSet( - fieldset_unit_mesh, pclass=Particle, time=start, lon=np.linspace(0, 1, npart), lat=np.linspace(1, 0, npart) - ) - pset.execute(DoNothing, endtime=end, dt=dt) - assert np.allclose(pset.time_nextloop, end) - - -@pytest.mark.parametrize( - "start, end, substeps, dt", - [ - (0.0, 10.0, 1, 1.0), - (0.0, 10.0, 4, 1.0), - (0.0, 10.0, 1, 3.0), - (2.0, 16.0, 5, 3.0), - (20.0, 10.0, 4, -1.0), - (20.0, -10.0, 7, -2.0), - ], -) -def test_execution_runtime(fieldset_unit_mesh, start, end, substeps, dt): - npart = 10 - pset = ParticleSet( - fieldset_unit_mesh, pclass=Particle, time=start, lon=np.linspace(0, 1, npart), lat=np.linspace(1, 0, npart) - ) - t_step = abs(end - start) / substeps - for _ in range(substeps): - pset.execute(DoNothing, runtime=t_step, dt=dt) - assert np.allclose(pset.time_nextloop, end) - - -def test_execution_fail_out_of_bounds(fieldset_unit_mesh): - npart = 10 - - def MoveRight(particle, fieldset, time): # pragma: no cover - tmp1, tmp2 = fieldset.UV[time, particle.depth, particle.lat, particle.lon + 0.1, particle] - particle.dlon += 0.1 - - pset = ParticleSet(fieldset_unit_mesh, pclass=Particle, lon=np.linspace(0, 1, npart), lat=np.linspace(1, 0, npart)) - with pytest.raises(FieldOutOfBoundError): - pset.execute(MoveRight, endtime=10.0, dt=1.0) - assert len(pset) == npart - assert (pset.lon - 1.0 > -1.0e12).all() - - -def test_execution_recover_out_of_bounds(fieldset_unit_mesh): - npart = 2 - - def MoveRight(particle, fieldset, time): # pragma: no cover - tmp1, tmp2 = fieldset.UV[time, particle.depth, particle.lat, particle.lon + 0.1, particle] - particle.dlon += 0.1 - - def MoveLeft(particle, fieldset, time): # pragma: no cover - if particle.state == StatusCode.ErrorOutOfBounds: - particle.dlon -= 1.0 - particle.state = StatusCode.Success - - lon = np.linspace(0.05, 0.95, npart) - lat = np.linspace(1, 0, npart) - pset = ParticleSet(fieldset_unit_mesh, pclass=Particle, lon=lon, lat=lat) - pset.execute([MoveRight, MoveLeft], endtime=11.0, dt=1.0) - assert len(pset) == npart - assert np.allclose(pset.lon, lon, rtol=1e-5) - assert np.allclose(pset.lat, lat, rtol=1e-5) - - -def test_execution_check_all_errors(fieldset_unit_mesh): - def MoveRight(particle, fieldset, time): # pragma: no cover - tmp1, tmp2 = fieldset.UV[time, particle.depth, particle.lat, particle.lon, particle] - - def RecoverAllErrors(particle, fieldset, time): # pragma: no cover - if particle.state > 4: - particle.state = StatusCode.Delete - - pset = ParticleSet(fieldset_unit_mesh, pclass=Particle, lon=10, lat=0) - pset.execute([MoveRight, RecoverAllErrors], endtime=11.0, dt=1.0) - assert len(pset) == 0 - - -def test_execution_check_stopallexecution(fieldset_unit_mesh): - def addoneLon(particle, fieldset, time): # pragma: no cover - particle.dlon += 1 - - if particle.lon + particle.dlon >= 10: - particle.state = StatusCode.StopAllExecution - - pset = ParticleSet(fieldset_unit_mesh, pclass=Particle, lon=[0, 1], lat=[0, 0]) - pset.execute(addoneLon, endtime=20.0, dt=1.0) - assert pset[0].lon == 9 - assert pset[0].time == 9 - assert pset[1].lon == 1 - assert pset[1].time == 0 - - -def test_execution_delete_out_of_bounds(fieldset_unit_mesh): - npart = 10 - - def MoveRight(particle, fieldset, time): # pragma: no cover - tmp1, tmp2 = fieldset.UV[time, particle.depth, particle.lat, particle.lon + 0.1, particle] - particle.dlon += 0.1 - - lon = np.linspace(0.05, 0.95, npart) - lat = np.linspace(1, 0, npart) - pset = ParticleSet(fieldset_unit_mesh, pclass=Particle, lon=lon, lat=lat) - pset.execute([MoveRight, DeleteParticle], endtime=10.0, dt=1.0) - assert len(pset) == 0 - - -def test_kernel_add_no_new_variables(fieldset_unit_mesh): - pset = ParticleSet(fieldset_unit_mesh, pclass=Particle, lon=[0.5], lat=[0.5]) - pset.execute(pset.Kernel(MoveEast) + pset.Kernel(MoveNorth), endtime=2.0, dt=1.0) - assert np.allclose(pset.lon, 0.6, rtol=1e-5) - assert np.allclose(pset.lat, 0.6, rtol=1e-5) - - def test_multi_kernel_duplicate_varnames(fieldset_unit_mesh): # Testing for merging of two Kernels with the same variable declared # Should throw a warning, but go ahead regardless diff --git a/tests/utils.py b/tests/utils.py index fc792d73b6..551f82edd9 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -10,7 +10,7 @@ import parcels from parcels._datasets.structured.generated import simple_UV_dataset -from parcels.application_kernels.interpolation import XBiLinear +from parcels.application_kernels.interpolation import XLinear from parcels.field import Field, VectorField from parcels.fieldset import FieldSet from parcels.xgrid import _FIELD_DATA_ORDERING, XGrid, get_axis_from_dim_name @@ -75,8 +75,8 @@ def create_fieldset_zeros_conversion(mesh_type="spherical", xdim=200, ydim=100) ds["lon"].data = np.linspace(-1e6 * mesh_conversion, 1e6 * mesh_conversion, xdim) ds["lat"].data = np.linspace(-1e6 * mesh_conversion, 1e6 * mesh_conversion, ydim) grid = XGrid.from_dataset(ds) - U = Field("U", ds["U"], grid, mesh_type=mesh_type, interp_method=XBiLinear) - V = Field("V", ds["V"], grid, mesh_type=mesh_type, interp_method=XBiLinear) + U = Field("U", ds["U"], grid, mesh_type=mesh_type, interp_method=XLinear) + V = Field("V", ds["V"], grid, mesh_type=mesh_type, interp_method=XLinear) UV = VectorField("UV", U, V) return FieldSet([U, V, UV]) diff --git a/tests/v4/test_advection.py b/tests/v4/test_advection.py index 5244501475..7b209ab30c 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 XBiLinear, XBiLinearPeriodic, XTriLinear +from parcels.application_kernels.interpolation import XLinear from parcels.field import Field, VectorField from parcels.fieldset import FieldSet from parcels.particle import Particle, Variable @@ -38,8 +38,8 @@ def test_advection_zonal(mesh_type, npart=10): ds = simple_UV_dataset(mesh_type=mesh_type) ds["U"].data[:] = 1.0 grid = XGrid.from_dataset(ds) - U = Field("U", ds["U"], grid, mesh_type=mesh_type, interp_method=XBiLinear) - V = Field("V", ds["V"], grid, mesh_type=mesh_type, interp_method=XBiLinear) + U = Field("U", ds["U"], grid, mesh_type=mesh_type, interp_method=XLinear) + V = Field("V", ds["V"], grid, mesh_type=mesh_type, interp_method=XLinear) UV = VectorField("UV", U, V) fieldset = FieldSet([U, V, UV]) @@ -54,8 +54,7 @@ def test_advection_zonal(mesh_type, npart=10): def periodicBC(particle, fieldset, time): particle.total_dlon += particle.dlon - particle.lon = np.fmod(particle.lon, fieldset.U.grid.lon[-1]) - particle.lat = np.fmod(particle.lat, fieldset.U.grid.lat[-1]) + particle.lon = np.fmod(particle.lon, 2) def test_advection_zonal_periodic(): @@ -64,17 +63,25 @@ def test_advection_zonal_periodic(): ds["lon"].data = np.array([0, 2]) ds["lat"].data = np.array([0, 2]) + # add a halo + halo = ds.isel(XG=0) + halo.lon.values = ds.lon.values[1] + 1 + halo.XG.values = ds.XG.values[1] + 2 + ds = xr.concat([ds, halo], dim="XG") + grid = XGrid.from_dataset(ds) - U = Field("U", ds["U"], grid, interp_method=XBiLinearPeriodic) - V = Field("V", ds["V"], grid, interp_method=XBiLinearPeriodic) + U = Field("U", ds["U"], grid, interp_method=XLinear) + V = Field("V", ds["V"], grid, interp_method=XLinear) UV = VectorField("UV", U, V) fieldset = FieldSet([U, V, UV]) PeriodicParticle = Particle.add_variable(Variable("total_dlon", initial=0)) - pset = ParticleSet(fieldset, pclass=PeriodicParticle, lon=[0.5], lat=[0.5]) + startlon = np.array([0.5, 0.4]) + pset = ParticleSet(fieldset, pclass=PeriodicParticle, lon=startlon, lat=[0.5, 0.5]) pset.execute([AdvectionEE, periodicBC], runtime=np.timedelta64(40, "s"), dt=np.timedelta64(1, "s")) - np.testing.assert_allclose(pset.total_dlon[0], 4, atol=1e-5) - np.testing.assert_allclose(pset.lon_nextloop[0], 0.5, atol=1e-5) + np.testing.assert_allclose(pset.total_dlon, 4, atol=1e-5) + np.testing.assert_allclose(pset.lon_nextloop, startlon, atol=1e-5) + np.testing.assert_allclose(pset.lat_nextloop, 0.5, atol=1e-5) def test_horizontal_advection_in_3D_flow(npart=10): @@ -82,9 +89,9 @@ def test_horizontal_advection_in_3D_flow(npart=10): ds = simple_UV_dataset(mesh_type="flat") ds["U"].data[:] = 1.0 grid = XGrid.from_dataset(ds) - U = Field("U", ds["U"], grid, interp_method=XTriLinear) + U = Field("U", ds["U"], grid, interp_method=XLinear) U.data[:, 0, :, :] = 0.0 # Set U to 0 at the surface - V = Field("V", ds["V"], grid, interp_method=XTriLinear) + V = Field("V", ds["V"], grid, interp_method=XLinear) UV = VectorField("UV", U, V) fieldset = FieldSet([U, V, UV]) @@ -100,28 +107,32 @@ def test_horizontal_advection_in_3D_flow(npart=10): def test_advection_3D_outofbounds(direction, wErrorThroughSurface): ds = simple_UV_dataset(mesh_type="flat") grid = XGrid.from_dataset(ds) - U = Field("U", ds["U"], grid, interp_method=XTriLinear) - U.data[:] = 0.01 # Set U to 0 at the surface - V = Field("V", ds["V"], grid, interp_method=XTriLinear) - W = Field("W", ds["V"], grid, interp_method=XTriLinear) # Use V as W for testing + U = Field("U", ds["U"], grid, interp_method=XLinear) + U.data[:] = 0.01 # Set U to small value (to avoid horizontal out of bounds) + V = Field("V", ds["V"], grid, interp_method=XLinear) + W = Field("W", ds["V"], grid, interp_method=XLinear) # Use V as W for testing W.data[:] = -1.0 if direction == "up" else 1.0 UVW = VectorField("UVW", U, V, W) UV = VectorField("UV", U, V) fieldset = FieldSet([U, V, W, UVW, UV]) def DeleteParticle(particle, fieldset, time): # pragma: no cover - if particle.state == StatusCode.ErrorOutOfBounds or particle.state == StatusCode.ErrorThroughSurface: - particle.state = StatusCode.Delete + particle.state = np.where(particle.state == StatusCode.ErrorOutOfBounds, StatusCode.Delete, particle.state) + particle.state = np.where(particle.state == StatusCode.ErrorThroughSurface, StatusCode.Delete, particle.state) def SubmergeParticle(particle, fieldset, time): # pragma: no cover - if particle.state == StatusCode.ErrorThroughSurface: - dt = particle.dt / np.timedelta64(1, "s") - (u, v) = fieldset.UV[particle] - particle.dlon = u * dt - particle.dlat = v * dt - particle.ddepth = 0.0 - particle.depth = 0 - particle.state = StatusCode.Evaluate + if len(particle.state) == 0: + return + inds = np.argwhere(particle.state == StatusCode.ErrorThroughSurface).flatten() + if len(inds) == 0: + return + dt = particle.dt / np.timedelta64(1, "s") + (u, v) = fieldset.UV[particle[inds]] + particle[inds].dlon = u * dt + particle[inds].dlat = v * dt + particle[inds].ddepth = 0.0 + particle[inds].depth = 0 + particle[inds].state = StatusCode.Evaluate kernels = [AdvectionRK4_3D] if wErrorThroughSurface: @@ -175,11 +186,11 @@ def test_length1dimensions(u, v, w): # TODO: Refactor this test to be more read ds["W"] = (["time", "depth", "YG", "XG"], W) grid = XGrid.from_dataset(ds) - U = Field("U", ds["U"], grid, interp_method=XTriLinear) - V = Field("V", ds["V"], grid, interp_method=XTriLinear) + U = Field("U", ds["U"], grid, interp_method=XLinear) + V = Field("V", ds["V"], grid, interp_method=XLinear) fields = [U, V, VectorField("UV", U, V)] if w: - W = Field("W", ds["W"], grid, interp_method=XTriLinear) + W = Field("W", ds["W"], grid, interp_method=XLinear) fields.append(VectorField("UVW", U, V, W)) fieldset = FieldSet(fields) @@ -198,8 +209,8 @@ def test_length1dimensions(u, v, w): # TODO: Refactor this test to be more read def test_radialrotation(npart=10): ds = radial_rotation_dataset() grid = XGrid.from_dataset(ds) - U = parcels.Field("U", ds["U"], grid, mesh_type="flat", interp_method=XBiLinear) - V = parcels.Field("V", ds["V"], grid, mesh_type="flat", interp_method=XBiLinear) + U = parcels.Field("U", ds["U"], grid, mesh_type="flat", interp_method=XLinear) + V = parcels.Field("V", ds["V"], grid, mesh_type="flat", interp_method=XLinear) UV = parcels.VectorField("UV", U, V) fieldset = parcels.FieldSet([U, V, UV]) @@ -227,17 +238,17 @@ def test_radialrotation(npart=10): ("AdvDiffM1", 1e-2), ("RK4", 1e-5), ("RK4_3D", 1e-5), - ("RK45", 1e-5), + ("RK45", 1e-4), ], ) def test_moving_eddy(method, rtol): ds = moving_eddy_dataset() grid = XGrid.from_dataset(ds) - U = Field("U", ds["U"], grid, interp_method=XBiLinear) - V = Field("V", ds["V"], grid, interp_method=XBiLinear) + U = Field("U", ds["U"], grid, interp_method=XLinear) + V = Field("V", ds["V"], grid, interp_method=XLinear) if method == "RK4_3D": # Using W to test 3D advection (assuming same velocity as V) - W = Field("W", ds["V"], grid, interp_method=XTriLinear) + W = Field("W", ds["V"], grid, interp_method=XLinear) UVW = VectorField("UVW", U, V, W) fieldset = FieldSet([U, V, W, UVW]) else: @@ -246,23 +257,18 @@ def test_moving_eddy(method, rtol): if method in ["AdvDiffEM", "AdvDiffM1"]: # Add zero diffusivity field for diffusion kernels ds["Kh"] = (["time", "depth", "YG", "XG"], np.full(ds["U"].shape, 0)) - fieldset.add_field(Field("Kh", ds["Kh"], grid, interp_method=XBiLinear), "Kh_zonal") - fieldset.add_field(Field("Kh", ds["Kh"], grid, interp_method=XBiLinear), "Kh_meridional") + fieldset.add_field(Field("Kh", ds["Kh"], grid, interp_method=XLinear), "Kh_zonal") + fieldset.add_field(Field("Kh", ds["Kh"], grid, interp_method=XLinear), "Kh_meridional") fieldset.add_constant("dres", 0.1) start_lon, start_lat, start_depth = 12000, 12500, 12500 - dt = np.timedelta64(3, "m") + dt = np.timedelta64(30, "m") if method == "RK45": - # Use RK45Particles to set next_dt - RK45Particles = Particle.add_variable(Variable("next_dt", initial=dt, dtype="timedelta64[s]")) - fieldset.add_constant("RK45_tol", 1e-3) + fieldset.add_constant("RK45_tol", rtol) - pclass = RK45Particles if method == "RK45" else Particle - pset = ParticleSet( - fieldset, pclass=pclass, lon=start_lon, lat=start_lat, depth=start_depth, time=np.timedelta64(0, "s") - ) - pset.execute(kernel[method], dt=dt, endtime=np.timedelta64(6, "h")) + pset = ParticleSet(fieldset, lon=start_lon, lat=start_lat, depth=start_depth, time=np.timedelta64(0, "s")) + pset.execute(kernel[method], dt=dt, endtime=np.timedelta64(1, "h")) def truth_moving(x_0, y_0, t): t /= np.timedelta64(1, "s") @@ -280,30 +286,27 @@ def truth_moving(x_0, y_0, t): @pytest.mark.parametrize( "method, rtol", [ - ("EE", 1e-2), + ("EE", 1e-1), ("RK4", 1e-5), - ("RK45", 1e-3), + ("RK45", 1e-4), ], ) def test_decaying_moving_eddy(method, rtol): ds = decaying_moving_eddy_dataset() grid = XGrid.from_dataset(ds) - U = Field("U", ds["U"], grid, interp_method=XBiLinear) - V = Field("V", ds["V"], grid, interp_method=XBiLinear) + U = Field("U", ds["U"], grid, interp_method=XLinear) + V = Field("V", ds["V"], grid, interp_method=XLinear) UV = VectorField("UV", U, V) fieldset = FieldSet([U, V, UV]) start_lon, start_lat = 10000, 10000 - dt = np.timedelta64(2, "m") + dt = np.timedelta64(60, "m") if method == "RK45": - # Use RK45Particles to set next_dt - RK45Particles = Particle.add_variable(Variable("next_dt", initial=dt, dtype="timedelta64[s]")) - fieldset.add_constant("RK45_tol", 1e-3) + fieldset.add_constant("RK45_tol", rtol) + fieldset.add_constant("RK45_min_dt", 10 * 60) - pclass = RK45Particles if method == "RK45" else Particle - - pset = ParticleSet(fieldset, pclass=pclass, lon=start_lon, lat=start_lat, time=np.timedelta64(0, "s")) + pset = ParticleSet(fieldset, lon=start_lon, lat=start_lat, time=np.timedelta64(0, "s")) pset.execute(kernel[method], dt=dt, endtime=np.timedelta64(1, "D")) def truth_moving(x_0, y_0, t): @@ -325,57 +328,79 @@ def truth_moving(x_0, y_0, t): np.testing.assert_allclose(pset.lat_nextloop, exp_lat, rtol=rtol) -# TODO decrease atol for these tests once the C-grid is implemented @pytest.mark.parametrize( - "method, atol", + "method, rtol", [ - ("RK4", 1), - ("RK45", 1), + ("RK4", 0.1), + ("RK45", 0.1), ], ) @pytest.mark.parametrize("grid_type", ["A"]) # TODO also implement C-grid once available -@pytest.mark.parametrize("flowfield", ["stommel_gyre", "peninsula"]) -def test_gyre_flowfields(method, grid_type, atol, flowfield): +def test_stommelgyre_fieldset(method, rtol, grid_type): npart = 2 - if flowfield == "peninsula": - ds = peninsula_dataset(grid_type=grid_type) - start_lat = np.linspace(3e3, 47e3, npart) - start_lon = 3e3 * np.ones_like(start_lat) - runtime = np.timedelta64(1, "D") - else: - ds = stommel_gyre_dataset(grid_type=grid_type) - start_lon = np.linspace(10e3, 100e3, npart) - start_lat = np.ones_like(start_lon) * 5000e3 - runtime = np.timedelta64(2, "D") + ds = stommel_gyre_dataset(grid_type=grid_type) grid = XGrid.from_dataset(ds) - U = Field("U", ds["U"], grid, interp_method=XBiLinear) - V = Field("V", ds["V"], grid, interp_method=XBiLinear) - P = Field("P", ds["P"], grid, interp_method=XBiLinear) + U = Field("U", ds["U"], grid, interp_method=XLinear) + V = Field("V", ds["V"], grid, interp_method=XLinear) + P = Field("P", ds["P"], grid, interp_method=XLinear) UV = VectorField("UV", U, V) fieldset = FieldSet([U, V, P, UV]) - dt = np.timedelta64(1, "m") # TODO check these settings (and possibly increase) + dt = np.timedelta64(30, "m") + runtime = np.timedelta64(1, "D") + start_lon = np.linspace(10e3, 100e3, npart) + start_lat = np.ones_like(start_lon) * 5000e3 if method == "RK45": - # Use RK45Particles to set next_dt - SampleParticle = Particle.add_variable( - [ - Variable("p", initial=0.0, dtype=np.float32), - Variable("p_start", initial=0.0, dtype=np.float32), - Variable("next_dt", initial=dt, dtype="timedelta64[s]"), - ] - ) - fieldset.add_constant("RK45_tol", 1e-3) - else: - SampleParticle = Particle.add_variable( - [Variable("p", initial=0.0, dtype=np.float32), Variable("p_start", initial=0.0, dtype=np.float32)] - ) + fieldset.add_constant("RK45_tol", rtol) + + SampleParticle = Particle.add_variable( + [Variable("p", initial=0.0, dtype=np.float32), Variable("p_start", initial=0.0, dtype=np.float32)] + ) + + def UpdateP(particle, fieldset, time): # pragma: no cover + particle.p = fieldset.P[particle.time, particle.depth, particle.lat, particle.lon] + particle.p_start = np.where(particle.time == 0, particle.p, particle.p_start) + + 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) + + +@pytest.mark.parametrize( + "method, rtol", + [ + ("RK4", 5e-3), + ("RK45", 1e-4), + ], +) +@pytest.mark.parametrize("grid_type", ["A"]) # TODO also implement C-grid once available +def test_peninsula_fieldset(method, rtol, grid_type): + npart = 2 + ds = peninsula_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) + P = Field("P", ds["P"], grid, interp_method=XLinear) + UV = VectorField("UV", U, V) + fieldset = FieldSet([U, V, P, UV]) + + dt = np.timedelta64(30, "m") + runtime = np.timedelta64(1, "D") + start_lat = np.linspace(3e3, 47e3, npart) + start_lon = 3e3 * np.ones_like(start_lat) + + if method == "RK45": + fieldset.add_constant("RK45_tol", rtol) + + SampleParticle = Particle.add_variable( + [Variable("p", initial=0.0, dtype=np.float32), Variable("p_start", initial=0.0, dtype=np.float32)] + ) def UpdateP(particle, fieldset, time): # pragma: no cover - if time == 0: - particle.p_start = fieldset.P[particle.time, particle.depth, particle.lat, particle.lon] particle.p = fieldset.P[particle.time, particle.depth, particle.lat, particle.lon] + particle.p_start = np.where(particle.time == 0, particle.p, particle.p_start) 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_start, pset.p, atol=atol) + np.testing.assert_allclose(pset.p, pset.p_start, rtol=rtol) diff --git a/tests/v4/test_diffusion.py b/tests/v4/test_diffusion.py index 3519af4d0e..74cfa04292 100644 --- a/tests/v4/test_diffusion.py +++ b/tests/v4/test_diffusion.py @@ -6,7 +6,7 @@ from parcels._datasets.structured.generated import simple_UV_dataset from parcels.application_kernels import AdvectionDiffusionEM, AdvectionDiffusionM1, DiffusionUniformKh -from parcels.application_kernels.interpolation import XBiLinear +from parcels.application_kernels.interpolation import XLinear from parcels.field import Field, VectorField from parcels.fieldset import FieldSet from parcels.particle import Particle, Variable @@ -25,19 +25,19 @@ def test_fieldKh_Brownian(mesh_type): ds["lon"].data = np.array([-1e6, 1e6]) ds["lat"].data = np.array([-1e6, 1e6]) grid = XGrid.from_dataset(ds) - U = Field("U", ds["U"], grid, mesh_type=mesh_type, interp_method=XBiLinear) - V = Field("V", ds["V"], grid, mesh_type=mesh_type, interp_method=XBiLinear) + U = Field("U", ds["U"], grid, mesh_type=mesh_type, interp_method=XLinear) + V = Field("V", ds["V"], grid, mesh_type=mesh_type, interp_method=XLinear) ds["Kh_zonal"] = (["time", "depth", "YG", "XG"], np.full((2, 1, 2, 2), kh_zonal)) ds["Kh_meridional"] = (["time", "depth", "YG", "XG"], np.full((2, 1, 2, 2), kh_meridional)) - Kh_zonal = Field("Kh_zonal", ds["Kh_zonal"], grid=grid, mesh_type=mesh_type, interp_method=XBiLinear) - Kh_meridional = Field("Kh_meridional", ds["Kh_meridional"], grid=grid, mesh_type=mesh_type, interp_method=XBiLinear) + Kh_zonal = Field("Kh_zonal", ds["Kh_zonal"], grid=grid, mesh_type=mesh_type, interp_method=XLinear) + Kh_meridional = Field("Kh_meridional", ds["Kh_meridional"], grid=grid, mesh_type=mesh_type, interp_method=XLinear) UV = VectorField("UV", U, V) fieldset = FieldSet([U, V, UV, Kh_zonal, Kh_meridional]) npart = 100 runtime = np.timedelta64(2, "h") - random.seed(1234) + np.random.seed(1234) pset = ParticleSet(fieldset=fieldset, lon=np.zeros(npart), lat=np.zeros(npart)) pset.execute(pset.Kernel(DiffusionUniformKh), runtime=runtime, dt=np.timedelta64(1, "h")) @@ -62,8 +62,8 @@ def test_fieldKh_SpatiallyVaryingDiffusion(mesh_type, kernel): ds["lon"].data = np.linspace(-1e6, 1e6, xdim) ds["lat"].data = np.linspace(-1e6, 1e6, ydim) grid = XGrid.from_dataset(ds) - U = Field("U", ds["U"], grid, mesh_type=mesh_type, interp_method=XBiLinear) - V = Field("V", ds["V"], grid, mesh_type=mesh_type, interp_method=XBiLinear) + U = Field("U", ds["U"], grid, mesh_type=mesh_type, interp_method=XLinear) + V = Field("V", ds["V"], grid, mesh_type=mesh_type, interp_method=XLinear) Kh = np.zeros((ydim, xdim), dtype=np.float32) for x in range(xdim): @@ -71,22 +71,22 @@ def test_fieldKh_SpatiallyVaryingDiffusion(mesh_type, kernel): ds["Kh_zonal"] = (["time", "depth", "YG", "XG"], np.full((2, 1, ydim, xdim), Kh)) ds["Kh_meridional"] = (["time", "depth", "YG", "XG"], np.full((2, 1, ydim, xdim), Kh)) - Kh_zonal = Field("Kh_zonal", ds["Kh_zonal"], grid=grid, mesh_type=mesh_type, interp_method=XBiLinear) - Kh_meridional = Field("Kh_meridional", ds["Kh_meridional"], grid=grid, mesh_type=mesh_type, interp_method=XBiLinear) + Kh_zonal = Field("Kh_zonal", ds["Kh_zonal"], grid=grid, mesh_type=mesh_type, interp_method=XLinear) + Kh_meridional = Field("Kh_meridional", ds["Kh_meridional"], grid=grid, mesh_type=mesh_type, interp_method=XLinear) UV = VectorField("UV", U, V) fieldset = FieldSet([U, V, UV, Kh_zonal, Kh_meridional]) fieldset.add_constant("dres", float(ds["lon"][1] - ds["lon"][0])) - npart = 100 + npart = 10000 - random.seed(1636) + np.random.seed(1636) pset = ParticleSet(fieldset=fieldset, lon=np.zeros(npart), lat=np.zeros(npart)) pset.execute(pset.Kernel(kernel), runtime=np.timedelta64(4, "h"), dt=np.timedelta64(1, "h")) tol = 2000 * mesh_conversion # effectively 2000 m errors (because of low numbers of particles) assert np.allclose(np.mean(pset.lon), 0, atol=tol) assert np.allclose(np.mean(pset.lat), 0, atol=tol) - assert stats.skew(pset.lon) > stats.skew(pset.lat) + assert abs(stats.skew(pset.lon)) > abs(stats.skew(pset.lat)) @pytest.mark.parametrize("lambd", [1, 5]) @@ -95,16 +95,16 @@ def test_randomexponential(lambd): npart = 1000 # Rate parameter for random.expovariate - fieldset.lambd = lambd + fieldset.add_constant("lambd", lambd) # Set random seed - random.seed(1234) + np.random.seed(1234) pset = ParticleSet(fieldset=fieldset, lon=np.zeros(npart), lat=np.zeros(npart), depth=np.zeros(npart)) def vertical_randomexponential(particle, fieldset, time): # pragma: no cover # Kernel for random exponential variable in depth direction - particle.depth = random.expovariate(fieldset.lambd) + particle.depth = np.random.exponential(scale=1 / fieldset.lambd, size=len(particle)) pset.execute(vertical_randomexponential, runtime=np.timedelta64(1, "s"), dt=np.timedelta64(1, "s")) @@ -123,7 +123,7 @@ def test_randomvonmises(mu, kappa): fieldset.kappa = kappa # Set random seed - random.seed(1234) + np.random.seed(1234) AngleParticle = Particle.add_variable(Variable("angle")) pset = ParticleSet( @@ -131,14 +131,12 @@ def test_randomvonmises(mu, kappa): ) def vonmises(particle, fieldset, time): # pragma: no cover - particle.angle = random.vonmisesvariate(fieldset.mu, fieldset.kappa) + particle.angle = np.array([random.vonmisesvariate(fieldset.mu, fieldset.kappa) for _ in range(len(particle))]) pset.execute(vonmises, runtime=np.timedelta64(1, "s"), dt=np.timedelta64(1, "s")) - angles = np.array([p.angle for p in pset]) - - assert np.allclose(np.mean(angles), mu, atol=0.1) + assert np.allclose(np.mean(pset.angle), mu, atol=0.1) vonmises_mean = stats.vonmises.mean(kappa=kappa, loc=mu) - assert np.allclose(np.mean(angles), vonmises_mean, atol=0.1) + assert np.allclose(np.mean(pset.angle), vonmises_mean, atol=0.1) vonmises_var = stats.vonmises.var(kappa=kappa, loc=mu) - assert np.allclose(np.var(angles), vonmises_var, atol=0.1) + assert np.allclose(np.var(pset.angle), vonmises_var, atol=0.1) diff --git a/tests/v4/test_index_search.py b/tests/v4/test_index_search.py index 5d5c602b94..5c1a473074 100644 --- a/tests/v4/test_index_search.py +++ b/tests/v4/test_index_search.py @@ -26,8 +26,8 @@ def test_grid_indexing_fpoints(field_cone): for yi_expected in range(grid.ydim - 1): for xi_expected in range(grid.xdim - 1): - x = grid.lon[yi_expected, xi_expected] + 0.00001 - y = grid.lat[yi_expected, xi_expected] + 0.00001 + x = np.array([grid.lon[yi_expected, xi_expected] + 0.00001]) + y = np.array([grid.lat[yi_expected, xi_expected] + 0.00001]) yi, eta, xi, xsi = _search_indices_curvilinear_2d(grid, y, x) if eta > 0.9: diff --git a/tests/v4/test_interpolation.py b/tests/v4/test_interpolation.py index 316fc2672f..959066992a 100644 --- a/tests/v4/test_interpolation.py +++ b/tests/v4/test_interpolation.py @@ -3,13 +3,15 @@ import xarray as xr from parcels._datasets.structured.generated import simple_UV_dataset +from parcels._datasets.unstructured.generic import datasets as datasets_unstructured from parcels.application_kernels.advection import AdvectionRK4_3D -from parcels.application_kernels.interpolation import XBiLinear, XTriLinear +from parcels.application_kernels.interpolation import UXPiecewiseLinearNode, XLinear from parcels.field import Field, VectorField from parcels.fieldset import FieldSet from parcels.particle import Particle, Variable from parcels.particleset import ParticleSet from parcels.tools.statuscodes import StatusCode +from parcels.uxgrid import UxGrid from parcels.xgrid import XGrid from tests.utils import TEST_DATA @@ -19,8 +21,8 @@ def test_interpolation_mesh_type(mesh_type, npart=10): ds = simple_UV_dataset(mesh_type=mesh_type) ds["U"].data[:] = 1.0 grid = XGrid.from_dataset(ds) - U = Field("U", ds["U"], grid, mesh_type=mesh_type, interp_method=XBiLinear) - V = Field("V", ds["V"], grid, mesh_type=mesh_type, interp_method=XBiLinear) + U = Field("U", ds["U"], grid, mesh_type=mesh_type, interp_method=XLinear) + V = Field("V", ds["V"], grid, mesh_type=mesh_type, interp_method=XLinear) UV = VectorField("UV", U, V) lat = 30.0 @@ -37,8 +39,20 @@ def test_interpolation_mesh_type(mesh_type, npart=10): assert U.eval(time, 0, lat, 0, applyConversion=False) == 1 +def test_default_interpolator_set_correctly(): + ds = simple_UV_dataset() + grid = XGrid.from_dataset(ds) + U = Field("U", ds["U"], grid) + assert U.interp_method == XLinear + + ds = datasets_unstructured["stommel_gyre_delaunay"] + grid = UxGrid(grid=ds.uxgrid, z=ds.coords["nz"]) + U = Field("U", ds["U"], grid) + assert U.interp_method == UXPiecewiseLinearNode + + interp_methods = { - "linear": XTriLinear, + "linear": XLinear, } diff --git a/tests/v4/test_particleset_execute.py b/tests/v4/test_particleset_execute.py index 8204dfc7a7..7cf6aa21af 100644 --- a/tests/v4/test_particleset_execute.py +++ b/tests/v4/test_particleset_execute.py @@ -11,8 +11,10 @@ UXPiecewiseConstantFace, VectorField, ) +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.tools.statuscodes import FieldOutOfBoundError, TimeExtrapolationError from parcels.uxgrid import UxGrid from parcels.xgrid import XGrid from tests import utils @@ -28,13 +30,23 @@ def fieldset() -> FieldSet: return FieldSet([U, V]) +@pytest.fixture +def zonal_flow_fieldset() -> FieldSet: + ds = simple_UV_dataset(mesh_type="flat") + ds["U"].data[:] = 1.0 + grid = XGrid.from_dataset(ds) + U = Field("U", ds["U"], grid, mesh_type="flat") + V = Field("V", ds["V"], grid, mesh_type="flat") + UV = VectorField("UV", U, V) + return FieldSet([U, V, UV]) + + def test_pset_remove_particle_in_kernel(fieldset): npart = 100 pset = ParticleSet(fieldset, lon=np.linspace(0, 1, npart), lat=np.linspace(1, 0, npart)) def DeleteKernel(particle, fieldset, time): # pragma: no cover - if particle.lon >= 0.4 and particle.lon <= 0.6: - particle.state = StatusCode.Delete + particle.state = np.where((particle.lon >= 0.4) & (particle.lon <= 0.6), StatusCode.Delete, particle.state) pset.execute(pset.Kernel(DeleteKernel), runtime=np.timedelta64(1, "s"), dt=np.timedelta64(1, "s")) indices = [i for i in range(npart) if not (40 <= i < 60)] @@ -44,12 +56,12 @@ def DeleteKernel(particle, fieldset, time): # pragma: no cover assert pset.size == 80 -def test_pset_stop_simulation(fieldset): - pset = ParticleSet(fieldset, lon=0, lat=0, pclass=Particle) +@pytest.mark.parametrize("npart", [1, 100]) +def test_pset_stop_simulation(fieldset, npart): + pset = ParticleSet(fieldset, lon=np.zeros(npart), lat=np.zeros(npart), pclass=Particle) def Delete(particle, fieldset, time): # pragma: no cover - if time >= fieldset.time_interval.left + np.timedelta64(4, "s"): - return StatusCode.StopExecution + particle[particle.time >= fieldset.time_interval.left + np.timedelta64(4, "s")].state = StatusCode.StopExecution pset.execute(Delete, dt=np.timedelta64(1, "s"), runtime=np.timedelta64(21, "s")) assert pset[0].time == fieldset.time_interval.left + np.timedelta64(4, "s") @@ -86,34 +98,138 @@ def test_execution_endtime(fieldset, starttime, endtime, dt): assert abs(pset.time_nextloop - endtime) < np.timedelta64(1, "ms") +def test_dont_run_particles_outside_starttime(fieldset): + # Test forward in time (note third particle is outside endtime) + start_times = [fieldset.time_interval.left + np.timedelta64(t, "s") for t in [0, 2, 10]] + endtime = fieldset.time_interval.left + np.timedelta64(8, "s") + + def AddLon(particle, fieldset, time): # pragma: no cover + particle.lon += 1 + + pset = ParticleSet(fieldset, lon=np.zeros(len(start_times)), lat=np.zeros(len(start_times)), time=start_times) + pset.execute(AddLon, dt=np.timedelta64(1, "s"), endtime=endtime) + + np.testing.assert_array_equal(pset.lon, [8, 6, 0]) + assert pset.time_nextloop[0:1] == endtime + assert pset.time_nextloop[2] == start_times[2] # this particle has not been executed + + # Test backward in time (note third particle is outside endtime) + start_times = [fieldset.time_interval.right - np.timedelta64(t, "s") for t in [0, 2, 10]] + endtime = fieldset.time_interval.right - np.timedelta64(8, "s") + + pset = ParticleSet(fieldset, lon=np.zeros(len(start_times)), lat=np.zeros(len(start_times)), time=start_times) + pset.execute(AddLon, dt=-np.timedelta64(1, "s"), endtime=endtime) + + np.testing.assert_array_equal(pset.lon, [8, 6, 0]) + assert pset.time_nextloop[0:1] == endtime + assert pset.time_nextloop[2] == start_times[2] # this particle has not been executed + + +def test_some_particles_throw_outofbounds(zonal_flow_fieldset): + npart = 100 + lon = np.linspace(0, 9e5, npart) + pset = ParticleSet(zonal_flow_fieldset, lon=lon, lat=np.zeros_like(lon)) + + with pytest.raises(FieldOutOfBoundError): + pset.execute(AdvectionEE, runtime=np.timedelta64(1_000_000, "s"), dt=np.timedelta64(10_000, "s")) + + +def test_delete_on_all_errors(fieldset): + def MoveRight(particle, fieldset, time): # pragma: no cover + particle.dlon += 1 + fieldset.U[particle.time, particle.depth, particle.lat, particle.lon, particle] + + def DeleteAllErrorParticles(particle, fieldset, time): # pragma: no cover + particle[particle.state > 20].state = StatusCode.Delete + + pset = ParticleSet(fieldset, lon=[1e5, 2], lat=[0, 0]) + pset.execute([MoveRight, DeleteAllErrorParticles], runtime=np.timedelta64(10, "s"), dt=np.timedelta64(1, "s")) + assert len(pset) == 0 + + +def test_some_particles_throw_outoftime(fieldset): + time = [fieldset.time_interval.left + np.timedelta64(t, "D") for t in [0, 350]] + pset = ParticleSet(fieldset, lon=np.zeros_like(time), lat=np.zeros_like(time), time=time) + + def FieldAccessOutsideTime(particle, fieldset, time): # pragma: no cover + fieldset.U[particle.time + np.timedelta64(1, "D"), particle.depth, particle.lat, particle.lon, particle] + + with pytest.raises(TimeExtrapolationError): + pset.execute(FieldAccessOutsideTime, runtime=np.timedelta64(400, "D"), dt=np.timedelta64(10, "D")) + + +def test_execution_check_stopallexecution(fieldset): + def addoneLon(particle, fieldset, time): # pragma: no cover + particle.dlon += 1 + particle[particle.lon + particle.dlon >= 10].state = StatusCode.StopAllExecution + + pset = ParticleSet(fieldset, lon=[0, 0], lat=[0, 0]) + pset.execute(addoneLon, runtime=np.timedelta64(20, "s"), dt=np.timedelta64(1, "s")) + np.testing.assert_allclose(pset.lon, 9) + np.testing.assert_allclose(pset.time - fieldset.time_interval.left, np.timedelta64(9, "s")) + + +def test_execution_recover_out_of_bounds(fieldset): + npart = 2 + + def MoveRight(particle, fieldset, time): # pragma: no cover + fieldset.U[particle.time, particle.depth, particle.lat, particle.lon + 0.1, particle] + particle.dlon += 0.1 + + def MoveLeft(particle, fieldset, time): # pragma: no cover + inds = np.where(particle.state == StatusCode.ErrorOutOfBounds) + print(inds, particle.state) + particle[inds].dlon -= 1.0 + particle[inds].state = StatusCode.Success + + lon = np.linspace(0.05, 6.95, npart) + lat = np.linspace(1, 0, npart) + pset = ParticleSet(fieldset, lon=lon, lat=lat) + pset.execute([MoveRight, MoveLeft], runtime=np.timedelta64(61, "s"), dt=np.timedelta64(1, "s")) + assert len(pset) == npart + np.testing.assert_allclose(pset.lon, [6.05, 5.95], rtol=1e-5) + np.testing.assert_allclose(pset.lat, lat, rtol=1e-5) + + @pytest.mark.parametrize( "starttime, runtime, dt", [(0, 10, 1), (0, 10, 3), (2, 16, 3), (20, 10, -1), (20, 0, -2), (5, 15, None)], ) -def test_execution_runtime(fieldset, starttime, runtime, dt): +@pytest.mark.parametrize("npart", [1, 10]) +def test_execution_runtime(fieldset, starttime, runtime, dt, npart): starttime = fieldset.time_interval.left + np.timedelta64(starttime, "s") runtime = np.timedelta64(runtime, "s") sign_dt = 1 if dt is None else np.sign(dt) dt = np.timedelta64(dt, "s") - pset = ParticleSet(fieldset, time=starttime, lon=0, lat=0) + pset = ParticleSet(fieldset, time=starttime, lon=np.zeros(npart), lat=np.zeros(npart)) pset.execute(DoNothing, runtime=runtime, dt=dt) - assert abs(pset.time_nextloop - starttime - runtime * sign_dt) < np.timedelta64(1, "ms") + assert all([abs(p.time_nextloop - starttime - runtime * sign_dt) < np.timedelta64(1, "ms") for p in pset]) + + +def test_changing_dt_in_kernel(fieldset): + def KernelCounter(particle, fieldset, time): # pragma: no cover + particle.lon += 1 + + pset = ParticleSet(fieldset, lon=np.zeros(1), lat=np.zeros(1)) + pset.execute(KernelCounter, dt=np.timedelta64(2, "s"), runtime=np.timedelta64(5, "s")) + assert pset.lon == 3 + print(pset.dt) + assert pset.dt == np.timedelta64(2, "s") -def test_execution_fail_python_exception(fieldset, npart=10): +@pytest.mark.parametrize("npart", [1, 100]) +def test_execution_fail_python_exception(fieldset, npart): pset = ParticleSet(fieldset, lon=np.linspace(0, 1, npart), lat=np.linspace(1, 0, npart)) def PythonFail(particle, fieldset, time): # pragma: no cover - if particle.time >= fieldset.time_interval.left + np.timedelta64(10, "s"): + inds = np.argwhere(particle.time >= fieldset.time_interval.left + np.timedelta64(10, "s")) + if inds.size > 0: raise RuntimeError("Enough is enough!") - else: - pass with pytest.raises(RuntimeError): pset.execute(PythonFail, runtime=np.timedelta64(20, "s"), dt=np.timedelta64(2, "s")) assert len(pset) == npart - assert pset.time[0] == fieldset.time_interval.left + np.timedelta64(10, "s") - assert all([time == fieldset.time_interval.left + np.timedelta64(0, "s") for time in pset.time[1:]]) + assert all(pset.time == fieldset.time_interval.left + np.timedelta64(10, "s")) def test_uxstommelgyre_pset_execute(): diff --git a/tests/v4/test_xgrid.py b/tests/v4/test_xgrid.py index 30bcc7c1aa..f987cf798e 100644 --- a/tests/v4/test_xgrid.py +++ b/tests/v4/test_xgrid.py @@ -7,7 +7,13 @@ from numpy.testing import assert_allclose from parcels._datasets.structured.generic import X, Y, Z, datasets -from parcels.xgrid import XGrid, _search_1d_array, _transpose_xfield_data_to_tzyx +from parcels.xgrid import ( + LEFT_OUT_OF_BOUNDS, + RIGHT_OUT_OF_BOUNDS, + XGrid, + _search_1d_array, + _transpose_xfield_data_to_tzyx, +) from tests import utils GridTestCase = namedtuple("GridTestCase", ["ds", "attr", "expected"]) @@ -125,7 +131,7 @@ def test_xgrid_search_cpoints(ds): axis_indices = {"Z": 0, "Y": yi, "X": xi} lat, lon = lat_array[yi, xi], lon_array[yi, xi] - axis_indices_bcoords = grid.search(0, lat, lon, ei=None) + axis_indices_bcoords = grid.search(0, np.atleast_1d(lat), np.atleast_1d(lon), ei=None) axis_indices_test = {k: v[0] for k, v in axis_indices_bcoords.items()} assert axis_indices == axis_indices_test @@ -150,16 +156,40 @@ def corner_to_cell_center_points(lat, lon): @pytest.mark.parametrize( "array, x, expected_xi, expected_xsi", [ - (np.array([1, 2, 3, 4, 5]), 1.1, 0, 0.1), + (np.array([1, 2, 3, 4, 5]), (1.1, 2.1), (0, 1), (0.1, 0.1)), (np.array([1, 2, 3, 4, 5]), 2.1, 1, 0.1), (np.array([1, 2, 3, 4, 5]), 3.1, 2, 0.1), (np.array([1, 2, 3, 4, 5]), 4.5, 3, 0.5), ], ) def test_search_1d_array(array, x, expected_xi, expected_xsi): + xi, xsi = _search_1d_array(array, x) + np.testing.assert_array_equal(xi, expected_xi) + np.testing.assert_allclose(xsi, expected_xsi) + + +@pytest.mark.parametrize( + "array, x, expected_xi", + [ + (np.array([1, 2, 3, 4, 5]), -0.1, LEFT_OUT_OF_BOUNDS), + (np.array([1, 2, 3, 4, 5]), 6.5, RIGHT_OUT_OF_BOUNDS), + ], +) +def test_search_1d_array_out_of_bounds(array, x, expected_xi): xi, xsi = _search_1d_array(array, x) assert xi == expected_xi - assert np.isclose(xsi, expected_xsi) + + +@pytest.mark.parametrize( + "array, x, expected_xi", + [ + (np.array([1, 2, 3, 4, 5]), (-0.1, 2.5), (LEFT_OUT_OF_BOUNDS, 1)), + (np.array([1, 2, 3, 4, 5]), (6.5, 1), (RIGHT_OUT_OF_BOUNDS, 0)), + ], +) +def test_search_1d_array_some_out_of_bounds(array, x, expected_xi): + xi, _ = _search_1d_array(array, x) + np.testing.assert_array_equal(xi, expected_xi) @pytest.mark.parametrize( diff --git a/tests/v4/utils/test_time.py b/tests/v4/utils/test_time.py index b5105ac3f9..497a4eb29e 100644 --- a/tests/v4/utils/test_time.py +++ b/tests/v4/utils/test_time.py @@ -84,9 +84,9 @@ def test_time_interval_contains(interval): right = interval.right middle = left + (right - left) / 2 - assert left in interval - assert right in interval - assert middle in interval + assert interval.is_all_time_in_interval(left) + assert interval.is_all_time_in_interval(right) + assert interval.is_all_time_in_interval(middle) @given(time_interval_strategy(calendar="365_day"), time_interval_strategy(calendar="365_day"))