Adding 3D-advection test

Author: erikvansebille

parcels/_datasets/structured/generic.py

@@ -134,14 +134,14 @@ def _unrolled_cone_curvilinear_grid():
     )
 
 
-def simple_UV_dataset(dims=(360, 2, 30, 4), mesh_type="spherical"):
+def simple_UV_dataset(dims=(360, 2, 30, 4), maxdepth=1, mesh_type="spherical"):
     max_lon = 180.0 if mesh_type == "spherical" else 1e6
 
     return xr.Dataset(
         {"U": (["time", "depth", "YG", "XG"], np.zeros(dims)), "V": (["time", "depth", "YG", "XG"], np.zeros(dims))},
         coords={
             "time": (["time"], xr.date_range("2000", "2001", dims[0]), {"axis": "T"}),
-            "depth": (["depth"], np.linspace(0, 1, dims[1]), {"axis": "Z"}),
+            "depth": (["depth"], np.linspace(0, maxdepth, dims[1]), {"axis": "Z"}),
             "YC": (["YC"], np.arange(dims[2]) + 0.5, {"axis": "Y"}),
             "YG": (["YG"], np.arange(dims[2]), {"axis": "Y", "c_grid_axis_shift": -0.5}),
             "XC": (["XC"], np.arange(dims[3]) + 0.5, {"axis": "X"}),

tests/test_advection.py

@@ -13,7 +13,6 @@
     AdvectionRK4,
     AdvectionRK4_3D,
     AdvectionRK45,
-    Field,
     FieldSet,
     Particle,
     ParticleSet,
@@ -274,119 +273,6 @@ def test_length1dimensions(u, v, w):
         assert ((np.array([p.depth - y0 for p in pset]) - 4 * w) < 1e-6).all()
 
 
-def truth_stationary(x_0, y_0, t):
-    lat = y_0 - u_0 / f * (1 - math.cos(f * t))
-    lon = x_0 + u_0 / f * math.sin(f * t)
-    return lon, lat
-
-
-def create_fieldset_stationary(xdim=100, ydim=100, maxtime=timedelta(hours=6)):
-    """Generate a FieldSet encapsulating the flow field of a stationary eddy.
-
-    Reference: N. Fabbroni, 2009, "Numerical simulations of passive
-    tracers dispersion in the sea"
-    """
-    time = np.arange(0.0, maxtime.total_seconds() + 1e-5, 60.0, dtype=np.float64)
-    dimensions = {
-        "lon": np.linspace(0, 25000, xdim, dtype=np.float32),
-        "lat": np.linspace(0, 25000, ydim, dtype=np.float32),
-        "time": time,
-    }
-    data = {
-        "U": np.transpose(np.ones((xdim, ydim, 1), dtype=np.float32) * u_0 * np.cos(f * time)),
-        "V": np.transpose(np.ones((xdim, ydim, 1), dtype=np.float32) * -u_0 * np.sin(f * time)),
-    }
-    fieldset = FieldSet.from_data(data, dimensions, mesh="flat")
-    # setting some constants for AdvectionRK45 kernel
-    fieldset.RK45_min_dt = 1e-3
-    fieldset.RK45_max_dt = 1e2
-    fieldset.RK45_tol = 1e-5
-    return fieldset
-
-
-@pytest.fixture
-def fieldset_stationary():
-    return create_fieldset_stationary()
-
-
-@pytest.mark.v4alpha
-@pytest.mark.xfail(reason="GH1946")
-@pytest.mark.parametrize(
-    "method, rtol, diffField",
-    [
-        ("EE", 1e-2, False),
-        ("AdvDiffEM", 1e-2, True),
-        ("AdvDiffM1", 1e-2, True),
-        ("RK4", 1e-5, False),
-        ("RK45", 1e-5, False),
-    ],
-)
-def test_stationary_eddy(fieldset_stationary, method, rtol, diffField):
-    npart = 1
-    fieldset = fieldset_stationary
-    if diffField:
-        fieldset.add_field(Field("Kh_zonal", np.zeros(fieldset.U.data.shape), grid=fieldset.U.grid))
-        fieldset.add_field(Field("Kh_meridional", np.zeros(fieldset.V.data.shape), grid=fieldset.V.grid))
-        fieldset.add_constant("dres", 0.1)
-    lon = np.linspace(12000, 21000, npart)
-    lat = np.linspace(12500, 12500, npart)
-    dt = timedelta(minutes=3).total_seconds()
-    endtime = timedelta(hours=6).total_seconds()
-
-    RK45Particles = Particle.add_variable("next_dt", dtype=np.float32, initial=dt)
-
-    pclass = RK45Particles if method == "RK45" else Particle
-    pset = ParticleSet(fieldset, pclass=pclass, lon=lon, lat=lat)
-    pset.execute(kernel[method], dt=dt, endtime=endtime)
-
-    exp_lon = [truth_stationary(x, y, pset[0].time)[0] for x, y in zip(lon, lat, strict=True)]
-    exp_lat = [truth_stationary(x, y, pset[0].time)[1] for x, y in zip(lon, lat, strict=True)]
-    assert np.allclose(pset.lon, exp_lon, rtol=rtol)
-    assert np.allclose(pset.lat, exp_lat, rtol=rtol)
-
-
-@pytest.mark.v4alpha
-@pytest.mark.xfail(reason="GH1946")
-def test_stationary_eddy_vertical():
-    npart = 1
-    lon = np.linspace(12000, 21000, npart)
-    lat = np.linspace(10000, 20000, npart)
-    depth = np.linspace(12500, 12500, npart)
-    endtime = timedelta(hours=6).total_seconds()
-    dt = timedelta(minutes=3).total_seconds()
-
-    xdim = ydim = 100
-    lon_data = np.linspace(0, 25000, xdim, dtype=np.float32)
-    lat_data = np.linspace(0, 25000, ydim, dtype=np.float32)
-    time_data = np.arange(0.0, 6 * 3600 + 1e-5, 60.0, dtype=np.float64)
-    fld1 = np.transpose(np.ones((xdim, ydim, 1), dtype=np.float32) * u_0 * np.cos(f * time_data))
-    fld2 = np.transpose(np.ones((xdim, ydim, 1), dtype=np.float32) * -u_0 * np.sin(f * time_data))
-    fldzero = np.transpose(np.zeros((xdim, ydim, 1), dtype=np.float32) * time_data)
-
-    dimensions = {"lon": lon_data, "lat": lat_data, "time": time_data}
-    data = {"U": fld1, "V": fldzero, "W": fld2}
-    fieldset = FieldSet.from_data(data, dimensions, mesh="flat")
-
-    pset = ParticleSet(fieldset, pclass=Particle, lon=lon, lat=lat, depth=depth)
-    pset.execute(AdvectionRK4_3D, dt=dt, endtime=endtime)
-    exp_lon = [truth_stationary(x, z, pset[0].time)[0] for x, z in zip(lon, depth, strict=True)]
-    exp_depth = [truth_stationary(x, z, pset[0].time)[1] for x, z in zip(lon, depth, strict=True)]
-    assert np.allclose(pset.lon, exp_lon, rtol=1e-5)
-    assert np.allclose(pset.lat, lat, rtol=1e-5)
-    assert np.allclose(pset.depth, exp_depth, rtol=1e-5)
-
-    data = {"U": fldzero, "V": fld2, "W": fld1}
-    fieldset = FieldSet.from_data(data, dimensions, mesh="flat")
-
-    pset = ParticleSet(fieldset, pclass=Particle, lon=lon, lat=lat, depth=depth)
-    pset.execute(AdvectionRK4_3D, dt=dt, endtime=endtime)
-    exp_depth = [truth_stationary(z, y, pset[0].time)[0] for z, y in zip(depth, lat, strict=True)]
-    exp_lat = [truth_stationary(z, y, pset[0].time)[1] for z, y in zip(depth, lat, strict=True)]
-    assert np.allclose(pset.lon, lon, rtol=1e-5)
-    assert np.allclose(pset.lat, exp_lat, rtol=1e-5)
-    assert np.allclose(pset.depth, exp_depth, rtol=1e-5)
-
-
 @pytest.mark.v4alpha
 @pytest.mark.xfail(reason="GH1946")
 def test_analyticalAgrid():

tests/v4/test_advection.py

@@ -116,20 +116,6 @@ def test_horizontal_advection_in_3D_flow(npart=10):
 @pytest.mark.parametrize("direction", ["up", "down"])
 @pytest.mark.parametrize("wErrorThroughSurface", [True, False])
 def test_advection_3D_outofbounds(direction, wErrorThroughSurface):
-    # xdim = ydim = zdim = 2
-    # dimensions = {
-    #     "lon": np.linspace(0.0, 1, xdim, dtype=np.float32),
-    #     "lat": np.linspace(0.0, 1, ydim, dtype=np.float32),
-    #     "depth": np.linspace(0.0, 1, zdim, dtype=np.float32),
-    # }
-    # wfac = -1.0 if direction == "up" else 1.0
-    # data = {
-    #     "U": 0.01 * np.ones((xdim, ydim, zdim), dtype=np.float32),
-    #     "V": np.zeros((xdim, ydim, zdim), dtype=np.float32),
-    #     "W": wfac * np.ones((xdim, ydim, zdim), dtype=np.float32),
-    # }
-    # fieldset = FieldSet.from_data(data, dimensions, mesh="flat")
-
     ds = simple_UV_dataset(mesh_type="flat")
     grid = XGrid.from_dataset(ds)
     U = Field("U", ds["U"], grid, interp_method=TriLinear)
@@ -177,7 +163,7 @@ def SubmergeParticle(particle, fieldset, time):  # pragma: no cover
         # ("AdvDiffEM", 1e-2),
         # ("AdvDiffM1", 1e-2),
         ("RK4", 1e-5),
-        # ('RK4_3D', 1e-5),
+        ("RK4_3D", 1e-5),
         # ("RK45", 1e-5),
     ],
 )
@@ -193,7 +179,7 @@ def truth_moving(x_0, y_0, t):
 
     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")
+    ds = simple_UV_dataset(dims=(len(time), 2, 2, 2), mesh_type="flat", maxdepth=25000)
     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")))
@@ -203,15 +189,27 @@ def truth_moving(x_0, y_0, t):
     ds = ds.assign_coords(time=time)
     U = Field("U", ds["U"], grid, interp_method=BiLinear)
     V = Field("V", ds["V"], grid, interp_method=BiLinear)
-    UV = VectorField("UV", U, V)
-    fieldset = FieldSet([U, V, UV])
+    if method == "RK4_3D":
+        # Using W to test 3D advection (assuming same velocity as V)
+        W = Field("W", ds["V"], grid, interp_method=TriLinear)
+        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
 
     RK45Particles = Particle.add_variable(Variable("next_dt", initial=dt))
 
     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, 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"))
 
     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)
+    if method == "RK4_3D":
+        assert np.allclose(pset.depth_nextloop, exp_lat, rtol=rtol)