Moving more diffusion unit tests from v3 to v4

Author: erikvansebille

parcels/application_kernels/advectiondiffusion.py

@@ -24,9 +24,10 @@ def AdvectionDiffusionM1(particle, fieldset, time):  # pragma: no cover
     The Wiener increment `dW` is normally distributed with zero
     mean and a standard deviation of sqrt(dt).
     """
+    dt = particle.dt / np.timedelta64(1, "s")  # noqa 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(particle.dt)))
-    dWy = random.normalvariate(0, math.sqrt(math.fabs(particle.dt)))
+    dWx = random.normalvariate(0, math.sqrt(math.fabs(dt)))
+    dWy = random.normalvariate(0, math.sqrt(math.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]
@@ -42,8 +43,8 @@ def AdvectionDiffusionM1(particle, fieldset, time):  # pragma: no cover
     by = math.sqrt(2 * fieldset.Kh_meridional[time, particle.depth, particle.lat, particle.lon])
 
     # Particle positions are updated only after evaluating all terms.
-    particle_dlon += u * particle.dt + 0.5 * dKdx * (dWx**2 + particle.dt) + bx * dWx  # noqa
-    particle_dlat += v * particle.dt + 0.5 * dKdy * (dWy**2 + particle.dt) + by * dWy  # noqa
+    particle_dlon += u * dt + 0.5 * dKdx * (dWx**2 + dt) + bx * dWx  # noqa
+    particle_dlat += v * dt + 0.5 * dKdy * (dWy**2 + dt) + by * dWy  # noqa
 
 
 def AdvectionDiffusionEM(particle, fieldset, time):  # pragma: no cover
@@ -59,9 +60,10 @@ def AdvectionDiffusionEM(particle, fieldset, time):  # pragma: no cover
     The Wiener increment `dW` is normally distributed with zero
     mean and a standard deviation of sqrt(dt).
     """
+    dt = particle.dt / np.timedelta64(1, "s")  # noqa 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(particle.dt)))
-    dWy = random.normalvariate(0, math.sqrt(math.fabs(particle.dt)))
+    dWx = random.normalvariate(0, math.sqrt(math.fabs(dt)))
+    dWy = random.normalvariate(0, math.sqrt(math.fabs(dt)))
 
     u, v = fieldset.UV[time, particle.depth, particle.lat, particle.lon]
 
@@ -78,8 +80,8 @@ def AdvectionDiffusionEM(particle, fieldset, time):  # pragma: no cover
     by = math.sqrt(2 * fieldset.Kh_meridional[time, particle.depth, particle.lat, particle.lon])
 
     # Particle positions are updated only after evaluating all terms.
-    particle_dlon += ax * particle.dt + bx * dWx  # noqa
-    particle_dlat += ay * particle.dt + by * dWy  # noqa
+    particle_dlon += ax * dt + bx * dWx  # noqa
+    particle_dlat += ay * dt + by * dWy  # noqa
 
 
 def DiffusionUniformKh(particle, fieldset, time):  # pragma: no cover

tests/test_diffusion.py

@@ -1,56 +1,16 @@
 import random
-from datetime import timedelta
 
 import numpy as np
 import pytest
 from scipy import stats
 
 from parcels import (
-    AdvectionDiffusionEM,
-    AdvectionDiffusionM1,
-    Field,
     Particle,
     ParticleSet,
-    RectilinearZGrid,
 )
 from tests.utils import create_fieldset_zeros_conversion
 
 
-@pytest.mark.v4alpha
-@pytest.mark.xfail(reason="GH1946")
-@pytest.mark.parametrize("mesh", ["spherical", "flat"])
-@pytest.mark.parametrize("kernel", [AdvectionDiffusionM1, AdvectionDiffusionEM])
-def test_fieldKh_SpatiallyVaryingDiffusion(mesh, kernel):
-    """Test advection-diffusion kernels on a non-uniform diffusivity field with a linear gradient in one direction."""
-    xdim = 200
-    ydim = 100
-    mesh_conversion = 1 / 1852.0 / 60 if mesh == "spherical" else 1
-    fieldset = create_fieldset_zeros_conversion(mesh=mesh, xdim=xdim, ydim=ydim, mesh_conversion=mesh_conversion)
-
-    Kh = np.zeros((ydim, xdim), dtype=np.float32)
-    for x in range(xdim):
-        Kh[:, x] = np.tanh(fieldset.U.lon[x] / fieldset.U.lon[-1] * 10.0) * xdim / 2.0 + xdim / 2.0 + 100.0
-
-    grid = RectilinearZGrid(lon=fieldset.U.lon, lat=fieldset.U.lat, mesh=mesh)
-    fieldset.add_field(Field("Kh_zonal", Kh, grid=grid))
-    fieldset.add_field(Field("Kh_meridional", Kh, grid=grid))
-    fieldset.add_constant("dres", fieldset.U.lon[1] - fieldset.U.lon[0])
-
-    npart = 100
-    runtime = timedelta(days=1)
-
-    random.seed(1636)
-    pset = ParticleSet(fieldset=fieldset, pclass=Particle, lon=np.zeros(npart), lat=np.zeros(npart))
-    pset.execute(pset.Kernel(kernel), runtime=runtime, dt=timedelta(hours=1))
-
-    lats = pset.lat
-    lons = pset.lon
-    tol = 2000 * mesh_conversion  # effectively 2000 m errors (because of low numbers of particles)
-    assert np.allclose(np.mean(lons), 0, atol=tol)
-    assert np.allclose(np.mean(lats), 0, atol=tol)
-    assert stats.skew(lons) > stats.skew(lats)
-
-
 @pytest.mark.v4alpha
 @pytest.mark.xfail(reason="GH1946")
 @pytest.mark.parametrize("lambd", [1, 5])

tests/v4/test_diffusion.py

@@ -2,9 +2,10 @@
 
 import numpy as np
 import pytest
+from scipy import stats
 
 from parcels._datasets.structured.generic import simple_UV_dataset
-from parcels.application_kernels import DiffusionUniformKh
+from parcels.application_kernels import AdvectionDiffusionEM, AdvectionDiffusionM1, DiffusionUniformKh
 from parcels.field import Field, VectorField
 from parcels.fieldset import FieldSet
 from parcels.particleset import ParticleSet
@@ -71,3 +72,41 @@ def test_fieldKh_Brownian(mesh_type):
     assert np.allclose(np.std(pset.lon), expected_std_lon, atol=tol)
     assert np.allclose(np.mean(pset.lon), 0, atol=tol)
     assert np.allclose(np.mean(pset.lat), 0, atol=tol)
+
+
+@pytest.mark.parametrize("mesh_type", ["spherical", "flat"])
+@pytest.mark.parametrize("kernel", [AdvectionDiffusionM1, AdvectionDiffusionEM])
+def test_fieldKh_SpatiallyVaryingDiffusion(mesh_type, kernel):
+    """Test advection-diffusion kernels on a non-uniform diffusivity field with a linear gradient in one direction."""
+    ydim, xdim = 100, 200
+
+    mesh_conversion = 1 / 1852.0 / 60 if mesh_type == "spherical" else 1
+    ds = simple_UV_dataset(dims=(2, 1, ydim, xdim), mesh_type=mesh_type)
+    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=BiLinear)
+    V = Field("V", ds["V"], grid, mesh_type=mesh_type, interp_method=BiLinear)
+
+    Kh = np.zeros((ydim, xdim), dtype=np.float32)
+    for x in range(xdim):
+        Kh[:, x] = np.tanh(ds["lon"][x] / ds["lon"][-1] * 10.0) * xdim / 2.0 + xdim / 2.0 + 100.0
+
+    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=BiLinear)
+    Kh_meridional = Field("Kh_meridional", ds["Kh_meridional"], grid=grid, mesh_type=mesh_type, interp_method=BiLinear)
+    UV = VectorField("UV", U, V)
+    fieldset = FieldSet([U, V, UV, Kh_zonal, Kh_meridional])
+    fieldset.add_constant("dres", ds["lon"][1] - ds["lon"][0])
+
+    npart = 100
+
+    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)