diff --git a/parcels/_datasets/structured/generated.py b/parcels/_datasets/structured/generated.py index 10f09e527..0a91f6864 100644 --- a/parcels/_datasets/structured/generated.py +++ b/parcels/_datasets/structured/generated.py @@ -1,3 +1,5 @@ +import math + import numpy as np import xarray as xr @@ -18,3 +20,267 @@ def simple_UV_dataset(dims=(360, 2, 30, 4), maxdepth=1, mesh_type="spherical"): "lon": (["XG"], np.linspace(-max_lon, max_lon, dims[3]), {"axis": "X", "c_grid_axis_shift": -0.5}), }, ) + + +def radial_rotation_dataset(xdim=200, ydim=200): # 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 + + U = np.zeros((2, 1, ydim, xdim), dtype=np.float32) + V = np.zeros((2, 1, ydim, xdim), dtype=np.float32) + + omega = 2 * np.pi / 86400.0 # Define the rotational period as 1 day. + + 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 + + U[:, :, j, i] = r * np.sin(theta) * omega + V[:, :, j, i] = -r * np.cos(theta) * omega + + return xr.Dataset( + {"U": (["time", "depth", "YG", "XG"], U), "V": (["time", "depth", "YG", "XG"], V)}, + 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}), + }, + ) + + +def moving_eddy_dataset(xdim=2, ydim=2): # TODO check if this also works with xdim=1, ydim=1 + """Create a dataset with an eddy moving in time. Note that there is no spatial variation in the flow.""" + f, u_0, u_g = 1.0e-4, 0.3, 0.04 # Some constants + + lon = np.linspace(0, 25000, xdim, dtype=np.float32) + lat = np.linspace(0, 25000, ydim, dtype=np.float32) + + time = np.arange(np.timedelta64(0, "s"), np.timedelta64(7, "h"), np.timedelta64(1, "m")) + + U = np.zeros((len(time), 1, ydim, xdim), dtype=np.float32) + V = np.zeros((len(time), 1, ydim, xdim), dtype=np.float32) + + for t in range(len(time)): + U[t, :, :, :] = u_g + (u_0 - u_g) * np.cos(f * (time[t] / np.timedelta64(1, "s"))) + V[t, :, :, :] = -(u_0 - u_g) * np.sin(f * (time[t] / np.timedelta64(1, "s"))) + + return xr.Dataset( + {"U": (["time", "depth", "YG", "XG"], U), "V": (["time", "depth", "YG", "XG"], V)}, + coords={ + "time": (["time"], time, {"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}), + }, + attrs={ + "u_0": u_0, + "u_g": u_g, + "f": f, + }, + ) + + +def decaying_moving_eddy_dataset(xdim=2, ydim=2): + """Simulate an ocean that accelerates subject to Coriolis force + and dissipative effects, upon which a geostrophic current is + superimposed. + + The original test description can be found in: N. Fabbroni, 2009, + Numerical Simulation of Passive tracers dispersion in the sea, + Ph.D. dissertation, University of Bologna + http://amsdottorato.unibo.it/1733/1/Fabbroni_Nicoletta_Tesi.pdf + """ + u_g = 0.04 # Geostrophic current + u_0 = 0.3 # Initial speed in x dirrection. v_0 = 0 + gamma = 1.0 / (2.89 * 86400) # Dissipitave effects due to viscousity. + gamma_g = 1.0 / (28.9 * 86400) + f = 1.0e-4 # Coriolis parameter. + + time = np.arange(np.timedelta64(0, "s"), np.timedelta64(1, "D") + np.timedelta64(1, "h"), np.timedelta64(2, "m")) + lon = np.linspace(0, 20000, xdim, dtype=np.float32) + lat = np.linspace(5000, 12000, ydim, dtype=np.float32) + + U = np.zeros((time.size, 1, lat.size, lon.size), dtype=np.float32) + V = np.zeros((time.size, 1, lat.size, lon.size), dtype=np.float32) + + for t in range(time.size): + t_float = time[t] / np.timedelta64(1, "s") + U[t, :, :, :] = u_g * np.exp(-gamma_g * t_float) + (u_0 - u_g) * np.exp(-gamma * t_float) * np.cos(f * t_float) + V[t, :, :, :] = -(u_0 - u_g) * np.exp(-gamma * t_float) * np.sin(f * t_float) + + return xr.Dataset( + {"U": (["time", "depth", "YG", "XG"], U), "V": (["time", "depth", "YG", "XG"], V)}, + coords={ + "time": (["time"], time, {"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}), + }, + attrs={ + "u_0": u_0, + "u_g": u_g, + "f": f, + "gamma": gamma, + "gamma_g": gamma_g, + }, + ) + + +def peninsula_dataset(xdim=100, ydim=50, mesh="flat", grid_type="A"): + """Construct a fieldset encapsulating the flow field around an idealised peninsula. + + Parameters + ---------- + xdim : + Horizontal dimension of the generated fieldset + ydim : + Vertical dimension of the generated fieldset + mesh : str + String indicating the type of mesh coordinates and + units used during velocity interpolation: + + 1. spherical: Lat and lon in degree, with a + correction for zonal velocity U near the poles. + 2. flat (default): No conversion, lat/lon are assumed to be in m. + grid_type : + Option whether grid is either Arakawa A (default) or C + + The original test description can be found in Fig. 2.2.3 in: + North, E. W., Gallego, A., Petitgas, P. (Eds). 2009. Manual of + recommended practices for modelling physical - biological + interactions during fish early life. + ICES Cooperative Research Report No. 295. 111 pp. + http://archimer.ifremer.fr/doc/00157/26792/24888.pdf + """ + domainsizeX, domainsizeY = (1.0e5, 5.0e4) + La = np.linspace(1e3, domainsizeX, xdim, dtype=np.float32) + Wa = np.linspace(1e3, domainsizeY, ydim, dtype=np.float32) + + u0 = 1 + x0 = domainsizeX / 2 + R = 0.32 * domainsizeX / 2 + + # Create the fields + P = np.zeros((ydim, xdim), dtype=np.float32) + U = np.zeros_like(P) + V = np.zeros_like(P) + x, y = np.meshgrid(La, Wa, sparse=True, indexing="xy") + P[:, :] = u0 * R**2 * y / ((x - x0) ** 2 + y**2) - u0 * y + + # Set land points to zero + landpoints = P >= 0.0 + P[landpoints] = 0.0 + + if grid_type == "A": + U[:, :] = u0 - u0 * R**2 * ((x - x0) ** 2 - y**2) / (((x - x0) ** 2 + y**2) ** 2) + V[:, :] = -2 * u0 * R**2 * ((x - x0) * y) / (((x - x0) ** 2 + y**2) ** 2) + U[landpoints] = 0.0 + V[landpoints] = 0.0 + Udims = ["YC", "XG"] + Vdims = ["YG", "XC"] + elif grid_type == "C": + U = np.zeros(P.shape) + V = np.zeros(P.shape) + V[:, 1:] = (P[:, 1:] - P[:, :-1]) / (La[1] - La[0]) + U[1:, :] = -(P[1:, :] - P[:-1, :]) / (Wa[1] - Wa[0]) + Udims = ["YG", "XG"] + Vdims = ["YG", "XG"] + else: + raise RuntimeError(f"Grid_type {grid_type} is not a valid option") + + # Convert from m to lat/lon for spherical meshes + lon = La / 1852.0 / 60.0 if mesh == "spherical" else La + lat = Wa / 1852.0 / 60.0 if mesh == "spherical" else Wa + + return xr.Dataset( + { + "U": (Udims, U), + "V": (Vdims, V), + "P": (["YG", "XG"], P), + }, + coords={ + "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}), + }, + ) + + +def stommel_gyre_dataset(xdim=200, ydim=200, grid_type="A"): + """Simulate a periodic current along a western boundary, with significantly + larger velocities along the western edge than the rest of the region + + The original test description can be found in: N. Fabbroni, 2009, + Numerical Simulation of Passive tracers dispersion in the sea, + Ph.D. dissertation, University of Bologna + http://amsdottorato.unibo.it/1733/1/Fabbroni_Nicoletta_Tesi.pdf + """ + a = b = 10000 * 1e3 + scalefac = 0.05 # to scale for physically meaningful velocities + dx, dy = a / xdim, b / ydim + + # Coordinates of the test fieldset (on A-grid in deg) + lon = np.linspace(0, a, xdim, dtype=np.float32) + lat = np.linspace(0, b, ydim, dtype=np.float32) + + # Define arrays U (zonal), V (meridional) and P (sea surface height) + U = np.zeros((lat.size, lon.size), dtype=np.float32) + V = np.zeros((lat.size, lon.size), dtype=np.float32) + P = np.zeros((lat.size, lon.size), dtype=np.float32) + + beta = 2e-11 + r = 1 / (11.6 * 86400) + es = r / (beta * a) + + for j in range(lat.size): + for i in range(lon.size): + xi = lon[i] / a + yi = lat[j] / b + P[j, i] = (1 - math.exp(-xi / es) - xi) * math.pi * np.sin(math.pi * yi) * scalefac + if grid_type == "A": + U[j, i] = -(1 - math.exp(-xi / es) - xi) * math.pi**2 * np.cos(math.pi * yi) * scalefac + V[j, i] = (math.exp(-xi / es) / es - 1) * math.pi * np.sin(math.pi * yi) * scalefac + if grid_type == "C": + V[:, 1:] = (P[:, 1:] - P[:, 0:-1]) / dx * a + U[1:, :] = -(P[1:, :] - P[0:-1, :]) / dy * b + Udims = ["YC", "XG"] + Vdims = ["YG", "XC"] + else: + Udims = ["YG", "XG"] + Vdims = ["YG", "XG"] + + return xr.Dataset( + {"U": (Udims, U), "V": (Vdims, V), "P": (["YG", "XG"], P)}, + coords={ + "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}), + }, + ) diff --git a/parcels/application_kernels/advection.py b/parcels/application_kernels/advection.py index d9773926e..d21111d98 100644 --- a/parcels/application_kernels/advection.py +++ b/parcels/application_kernels/advection.py @@ -115,9 +115,15 @@ 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 = min( - particle.next_dt / np.timedelta64(1, "s"), fieldset.RK45_max_dt - ) # TODO: improve API for converting dt to seconds + 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 + c = [1.0 / 4.0, 3.0 / 8.0, 12.0 / 13.0, 1.0, 1.0 / 2.0] A = [ [1.0 / 4.0, 0.0, 0.0, 0.0, 0.0], @@ -162,7 +168,7 @@ def AdvectionRK45(particle, fieldset, time): # pragma: no cover 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)): + 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 diff --git a/parcels/application_kernels/interpolation.py b/parcels/application_kernels/interpolation.py index f47411652..a41b76afc 100644 --- a/parcels/application_kernels/interpolation.py +++ b/parcels/application_kernels/interpolation.py @@ -40,7 +40,10 @@ def XBiLinear( zi, _ = position["Z"] data = field.data.data[:, zi, yi : yi + 2, xi : xi + 2] - data = (1 - tau) * data[ti, :, :] + tau * data[ti + 1, :, :] + if tau > 0: + data = (1 - tau) * data[ti, :, :] + tau * data[ti + 1, :, :] + else: + data = data[ti, :, :] return ( (1 - xsi) * (1 - eta) * data[0, 0] diff --git a/parcels/fieldset.py b/parcels/fieldset.py index 5937c7f57..113ef637d 100644 --- a/parcels/fieldset.py +++ b/parcels/fieldset.py @@ -164,7 +164,8 @@ def add_constant(self, name, value): """ if name in self.constants: raise ValueError(f"FieldSet already has a constant with name '{name}'") - + if not isinstance(value, (float, np.floating, int, np.integer)): + raise ValueError(f"FieldSet constants have to be of type float or int, got a {type(value)}") self.constants[name] = np.float32(value) @property diff --git a/parcels/kernel.py b/parcels/kernel.py index 2a4c95726..73a91adb8 100644 --- a/parcels/kernel.py +++ b/parcels/kernel.py @@ -303,7 +303,7 @@ def evaluate_particle(self, p, endtime): 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 == StatusCode.StopExecution: + if res in [StatusCode.StopExecution, StatusCode.Repeat]: break if res is None: diff --git a/parcels/xgrid.py b/parcels/xgrid.py index 3c8b9a3b8..95d72eb2e 100644 --- a/parcels/xgrid.py +++ b/parcels/xgrid.py @@ -271,7 +271,10 @@ def _gtype(self): def search(self, z, y, x, ei=None): ds = self.xgcm_grid._ds - zi, zeta = _search_1d_array(ds.depth.values, z) + if "Z" in self.axes: + zi, zeta = _search_1d_array(ds.depth.values, z) + else: + zi, zeta = 0, 0.0 if zi == -1: if zeta < 0: raise FieldOutOfBoundError( diff --git a/tests/v4/test_advection.py b/tests/v4/test_advection.py index 7d0905751..b53946c3a 100644 --- a/tests/v4/test_advection.py +++ b/tests/v4/test_advection.py @@ -2,7 +2,15 @@ import pytest import xarray as xr -from parcels._datasets.structured.generated import simple_UV_dataset +import parcels +from parcels._datasets.structured.generated import ( + decaying_moving_eddy_dataset, + moving_eddy_dataset, + peninsula_dataset, + radial_rotation_dataset, + simple_UV_dataset, + stommel_gyre_dataset, +) 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 @@ -65,8 +73,8 @@ def test_advection_zonal_periodic(): PeriodicParticle = Particle.add_variable(Variable("total_dlon", initial=0)) pset = ParticleSet(fieldset, pclass=PeriodicParticle, lon=[0.5], lat=[0.5]) pset.execute([AdvectionEE, periodicBC], runtime=np.timedelta64(40, "s"), dt=np.timedelta64(1, "s")) - assert np.isclose(pset.total_dlon[0], 4, atol=1e-5) - assert np.isclose(pset.lon_nextloop[0], 0.5, atol=1e-5) + 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) def test_horizontal_advection_in_3D_flow(npart=10): @@ -84,7 +92,7 @@ def test_horizontal_advection_in_3D_flow(npart=10): pset.execute(AdvectionRK4, runtime=np.timedelta64(2, "h"), dt=np.timedelta64(15, "m")) expected_lon = pset.depth * (pset.time - fieldset.time_interval.left) / np.timedelta64(1, "s") - assert np.allclose(expected_lon, pset.lon_nextloop, atol=1.0e-1) + np.testing.assert_allclose(expected_lon, pset.lon_nextloop, atol=1.0e-1) @pytest.mark.parametrize("direction", ["up", "down"]) @@ -124,8 +132,8 @@ def SubmergeParticle(particle, fieldset, time): # pragma: no cover pset.execute(kernels, runtime=np.timedelta64(11, "s"), dt=np.timedelta64(1, "s")) if direction == "up" and wErrorThroughSurface: - assert np.allclose(pset.lon[0], 0.6) - assert np.allclose(pset.depth[0], 0) + np.testing.assert_allclose(pset.lon[0], 0.6, atol=1e-5) + np.testing.assert_allclose(pset.depth[0], 0, atol=1e-5) else: assert len(pset) == 0 @@ -181,10 +189,34 @@ def test_length1dimensions(u, v, w): # TODO: Refactor this test to be more read pset.execute(kernel, runtime=np.timedelta64(5, "s"), dt=np.timedelta64(1, "s")) assert len(pset.lon) == len([p.lon for p in pset]) - assert ((np.array([p.lon - x0 for p in pset]) - 4 * u) < 1e-6).all() - assert ((np.array([p.lat - y0 for p in pset]) - 4 * v) < 1e-6).all() + np.testing.assert_allclose(np.array([p.lon - x0 for p in pset]), 4 * u, atol=1e-6) + np.testing.assert_allclose(np.array([p.lat - y0 for p in pset]), 4 * v, atol=1e-6) if w: - assert ((np.array([p.depth - z0 for p in pset]) - 4 * w) < 1e-6).all() + np.testing.assert_allclose(np.array([p.depth - z0 for p in pset]), 4 * w, atol=1e-6) + + +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) + UV = parcels.VectorField("UV", U, V) + fieldset = parcels.FieldSet([U, V, UV]) + + dt = np.timedelta64(30, "s") + lon = np.linspace(32, 50, npart) + lat = np.ones(npart) * 30 + starttime = np.arange(np.timedelta64(0, "s"), npart * dt, dt) + + pset = parcels.ParticleSet(fieldset, lon=lon, lat=lat, time=starttime) + pset.execute(parcels.AdvectionRK4, endtime=np.timedelta64(10, "m"), dt=dt) + + theta = 2 * np.pi * (pset.time_nextloop - starttime) / np.timedelta64(24 * 3600, "s") + true_lon = (lon - 30.0) * np.cos(theta) + 30.0 + true_lat = -(lon - 30.0) * np.sin(theta) + 30.0 + + np.testing.assert_allclose(pset.lon, true_lon, atol=5e-2) + np.testing.assert_allclose(pset.lat, true_lat, atol=5e-2) @pytest.mark.parametrize( @@ -195,29 +227,12 @@ def test_length1dimensions(u, v, w): # TODO: Refactor this test to be more read ("AdvDiffM1", 1e-2), ("RK4", 1e-5), ("RK4_3D", 1e-5), - pytest.param("RK45", 1e-5, marks=pytest.mark.xfail(reason="Started failing in GH2123 - not sure why")), + ("RK45", 1e-5), ], ) -def test_moving_eddy(method, rtol): # TODO: Refactor this test to be more readable - f, u_0, u_g = 1.0e-4, 0.3, 0.04 # Some constants - start_lon, start_lat = 12000, 12500 - - def truth_moving(x_0, y_0, t): - t /= np.timedelta64(1, "s") - lat = y_0 - (u_0 - u_g) / f * (1 - np.cos(f * t)) - lon = x_0 + u_g * t + (u_0 - u_g) / f * np.sin(f * t) - return lon, lat - - dt = np.timedelta64(3, "m") - time = np.arange(np.timedelta64(0, "s"), np.timedelta64(7, "h"), np.timedelta64(1, "m")) - ds = simple_UV_dataset(dims=(len(time), 2, 2, 2), mesh_type="flat", maxdepth=25000) +def test_moving_eddy(method, rtol): + ds = moving_eddy_dataset() grid = XGrid.from_dataset(ds) - for t in range(len(time)): - ds["U"].data[t, :, :, :] = u_g + (u_0 - u_g) * np.cos(f * (time[t] / np.timedelta64(1, "s"))) - ds["V"].data[t, :, :, :] = -(u_0 - u_g) * np.sin(f * (time[t] / np.timedelta64(1, "s"))) - ds["lon"].data = np.array([0, 25000]) - ds["lat"].data = np.array([0, 25000]) - ds = ds.assign_coords(time=time) U = Field("U", ds["U"], grid, interp_method=XBiLinear) V = Field("V", ds["V"], grid, interp_method=XBiLinear) if method == "RK4_3D": @@ -225,22 +240,23 @@ def truth_moving(x_0, y_0, t): W = Field("W", ds["V"], grid, interp_method=XTriLinear) UVW = VectorField("UVW", U, V, W) fieldset = FieldSet([U, V, W, UVW]) - start_depth = start_lat else: UV = VectorField("UV", U, V) fieldset = FieldSet([U, V, UV]) - start_depth = 0 + 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_constant("dres", 0.1) + + start_lon, start_lat, start_depth = 12000, 12500, 12500 + dt = np.timedelta64(3, "m") if method == "RK45": # Use RK45Particles to set next_dt RK45Particles = Particle.add_variable(Variable("next_dt", initial=dt, dtype=np.timedelta64)) fieldset.add_constant("RK45_tol", 1e-6) - elif method in ["AdvDiffEM", "AdvDiffM1"]: - # Add zero diffusivity field for diffusion kernels - ds["Kh"] = (["time", "depth", "YG", "XG"], np.full((len(time), 2, 2, 2), 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_constant("dres", 0.1) pclass = RK45Particles if method == "RK45" else Particle pset = ParticleSet( @@ -248,8 +264,118 @@ def truth_moving(x_0, y_0, t): ) pset.execute(kernel[method], dt=dt, endtime=np.timedelta64(6, "h")) - exp_lon, exp_lat = truth_moving(start_lon, start_lat, pset.time[0]) - assert np.allclose(pset.lon_nextloop, exp_lon, rtol=rtol) - assert np.allclose(pset.lat_nextloop, exp_lat, rtol=rtol) + def truth_moving(x_0, y_0, t): + t /= np.timedelta64(1, "s") + lat = y_0 - (ds.u_0 - ds.u_g) / ds.f * (1 - np.cos(ds.f * t)) + lon = x_0 + ds.u_g * t + (ds.u_0 - ds.u_g) / ds.f * np.sin(ds.f * t) + return lon, lat + + exp_lon, exp_lat = truth_moving(start_lon, start_lat, pset.time_nextloop[0]) + np.testing.assert_allclose(pset.lon_nextloop, exp_lon, rtol=rtol) + np.testing.assert_allclose(pset.lat_nextloop, exp_lat, rtol=rtol) if method == "RK4_3D": - assert np.allclose(pset.depth_nextloop, exp_lat, rtol=rtol) + np.testing.assert_allclose(pset.depth_nextloop, exp_lat, rtol=rtol) + + +@pytest.mark.parametrize( + "method, rtol", + [ + ("EE", 1e-2), + ("RK4", 1e-5), + ("RK45", 1e-5), + ], +) +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) + UV = VectorField("UV", U, V) + fieldset = FieldSet([U, V, UV]) + + start_lon, start_lat = 10000, 10000 + dt = np.timedelta64(2, "m") + + if method == "RK45": + # Use RK45Particles to set next_dt + RK45Particles = Particle.add_variable(Variable("next_dt", initial=dt, dtype=np.timedelta64)) + fieldset.add_constant("RK45_tol", 1e-6) + + pclass = RK45Particles if method == "RK45" else Particle + + pset = ParticleSet(fieldset, pclass=pclass, 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): + t /= np.timedelta64(1, "s") + lon = ( + x_0 + + (ds.u_g / ds.gamma_g) * (1 - np.exp(-ds.gamma_g * t)) + + ds.f + * ((ds.u_0 - ds.u_g) / (ds.f**2 + ds.gamma**2)) + * ((ds.gamma / ds.f) + np.exp(-ds.gamma * t) * (np.sin(ds.f * t) - (ds.gamma / ds.f) * np.cos(ds.f * t))) + ) + lat = y_0 - ((ds.u_0 - ds.u_g) / (ds.f**2 + ds.gamma**2)) * ds.f * ( + 1 - np.exp(-ds.gamma * t) * (np.cos(ds.f * t) + (ds.gamma / ds.f) * np.sin(ds.f * t)) + ) + return lon, lat + + exp_lon, exp_lat = truth_moving(start_lon, start_lat, pset.time_nextloop[0]) + np.testing.assert_allclose(pset.lon_nextloop, exp_lon, rtol=rtol) + 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", + [ + ("RK4", 1), + ("RK45", 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): + 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") + 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) + UV = VectorField("UV", U, V) + fieldset = FieldSet([U, V, P, UV]) + + dt = np.timedelta64(1, "m") # TODO check these settings (and possibly increase) + + 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=np.timedelta64), + ] + ) + fieldset.add_constant("RK45_tol", 1e-6) + else: + 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] + + 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) diff --git a/tests/v4/test_diffusion.py b/tests/v4/test_diffusion.py index 31f984504..3519af4d0 100644 --- a/tests/v4/test_diffusion.py +++ b/tests/v4/test_diffusion.py @@ -75,7 +75,7 @@ def test_fieldKh_SpatiallyVaryingDiffusion(mesh_type, kernel): Kh_meridional = Field("Kh_meridional", ds["Kh_meridional"], grid=grid, mesh_type=mesh_type, interp_method=XBiLinear) UV = VectorField("UV", U, V) fieldset = FieldSet([U, V, UV, Kh_zonal, Kh_meridional]) - fieldset.add_constant("dres", ds["lon"][1] - ds["lon"][0]) + fieldset.add_constant("dres", float(ds["lon"][1] - ds["lon"][0])) npart = 100