Moving last diffusion tests from v3 to v4

And also moving BiLinear interpolation to interpolation.py

Author: erikvansebille

parcels/application_kernels/interpolation.py

@@ -10,13 +10,41 @@
 
 if TYPE_CHECKING:
     from parcels.uxgrid import _UXGRID_AXES
+    from parcels.xgrid import _XGRID_AXES
 
 __all__ = [
     "UXPiecewiseConstantFace",
     "UXPiecewiseLinearNode",
+    "XBiLinear",
 ]
 
 
+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, zeta = position["Z"]
+
+    data = field.data.data[:, zi, yi : yi + 2, xi : xi + 2]
+    data = (1 - tau) * data[ti, :, :] + tau * data[ti + 1, :, :]
+
+    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 UXPiecewiseConstantFace(
     field: Field,
     ti: int,

tests/test_diffusion.py

@@ -4,17 +4,16 @@
 
 import struct
 from pathlib import Path
-from typing import TYPE_CHECKING
 
 import numpy as np
 import xarray as xr
 
 import parcels
-from parcels import FieldSet
-from parcels.xgrid import _FIELD_DATA_ORDERING, get_axis_from_dim_name
-
-if TYPE_CHECKING:
-    from parcels.xgrid import XGrid
+from parcels._datasets.structured.generic import simple_UV_dataset
+from parcels.application_kernels.interpolation import XBiLinear
+from parcels.field import Field, VectorField
+from parcels.fieldset import FieldSet
+from parcels.xgrid import _FIELD_DATA_ORDERING, XGrid, get_axis_from_dim_name
 
 PROJECT_ROOT = Path(__file__).resolve().parents[1]
 TEST_ROOT = PROJECT_ROOT / "tests"
@@ -69,13 +68,18 @@ def create_fieldset_global(xdim=200, ydim=100):
     return FieldSet.from_data(data, dimensions, mesh="flat")
 
 
-def create_fieldset_zeros_conversion(mesh="spherical", xdim=200, ydim=100, mesh_conversion=1) -> FieldSet:
+def create_fieldset_zeros_conversion(mesh_type="spherical", xdim=200, ydim=100) -> FieldSet:
     """Zero velocity field with lat and lon determined by a conversion factor."""
-    lon = np.linspace(-1e5 * mesh_conversion, 1e5 * mesh_conversion, xdim, dtype=np.float32)
-    lat = np.linspace(-1e5 * mesh_conversion, 1e5 * mesh_conversion, ydim, dtype=np.float32)
-    dimensions = {"lon": lon, "lat": lat}
-    data = {"U": np.zeros((ydim, xdim), dtype=np.float32), "V": np.zeros((ydim, xdim), dtype=np.float32)}
-    return FieldSet.from_data(data, dimensions, mesh=mesh)
+    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 * 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)
+
+    UV = VectorField("UV", U, V)
+    return FieldSet([U, V, UV])
 
 
 def create_simple_pset(n=1):

tests/utils.py

@@ -6,36 +6,13 @@
 
 from parcels._datasets.structured.generic import simple_UV_dataset
 from parcels.application_kernels import AdvectionDiffusionEM, AdvectionDiffusionM1, DiffusionUniformKh
+from parcels.application_kernels.interpolation import XBiLinear
 from parcels.field import Field, VectorField
 from parcels.fieldset import FieldSet
+from parcels.particle import Particle, Variable
 from parcels.particleset import ParticleSet
-from parcels.xgrid import _XGRID_AXES, XGrid
-
-
-def BiLinear(  # TODO move to interpolation file
-    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, zeta = position["Z"]
-
-    data = field.data.data[:, zi, yi : yi + 2, xi : xi + 2]
-    data = (1 - tau) * data[ti, :, :] + tau * data[ti + 1, :, :]
-
-    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]
-    )
+from parcels.xgrid import XGrid
+from tests.utils import create_fieldset_zeros_conversion
 
 
 @pytest.mark.parametrize("mesh_type", ["spherical", "flat"])
@@ -48,12 +25,12 @@ 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=BiLinear)
-    V = Field("V", ds["V"], grid, mesh_type=mesh_type, interp_method=BiLinear)
+    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)
     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=BiLinear)
-    Kh_meridional = Field("Kh_meridional", ds["Kh_meridional"], grid=grid, mesh_type=mesh_type, interp_method=BiLinear)
+    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)
     UV = VectorField("UV", U, V)
     fieldset = FieldSet([U, V, UV, Kh_zonal, Kh_meridional])
 
@@ -85,17 +62,17 @@ 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=BiLinear)
-    V = Field("V", ds["V"], grid, mesh_type=mesh_type, interp_method=BiLinear)
+    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)
 
     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)
+    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)
     UV = VectorField("UV", U, V)
     fieldset = FieldSet([U, V, UV, Kh_zonal, Kh_meridional])
     fieldset.add_constant("dres", ds["lon"][1] - ds["lon"][0])
@@ -110,3 +87,58 @@ def test_fieldKh_SpatiallyVaryingDiffusion(mesh_type, kernel):
     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)
+
+
+@pytest.mark.parametrize("lambd", [1, 5])
+def test_randomexponential(lambd):
+    fieldset = create_fieldset_zeros_conversion()
+    npart = 1000
+
+    # Rate parameter for random.expovariate
+    fieldset.lambd = lambd
+
+    # Set random seed
+    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)
+
+    pset.execute(vertical_randomexponential, runtime=np.timedelta64(1, "s"), dt=np.timedelta64(1, "s"))
+
+    expected_mean = 1.0 / fieldset.lambd
+    assert np.allclose(np.mean(pset.depth), expected_mean, rtol=0.1)
+
+
+@pytest.mark.parametrize("mu", [0.8 * np.pi, np.pi])
+@pytest.mark.parametrize("kappa", [2, 4])
+def test_randomvonmises(mu, kappa):
+    npart = 10000
+    fieldset = create_fieldset_zeros_conversion()
+
+    # Parameters for random.vonmisesvariate
+    fieldset.mu = mu
+    fieldset.kappa = kappa
+
+    # Set random seed
+    random.seed(1234)
+
+    AngleParticle = Particle.add_variable(Variable("angle"))
+    pset = ParticleSet(
+        fieldset=fieldset, pclass=AngleParticle, lon=np.zeros(npart), lat=np.zeros(npart), depth=np.zeros(npart)
+    )
+
+    def vonmises(particle, fieldset, time):  # pragma: no cover
+        particle.angle = random.vonmisesvariate(fieldset.mu, fieldset.kappa)
+
+    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)
+    vonmises_mean = stats.vonmises.mean(kappa=kappa, loc=mu)
+    assert np.allclose(np.mean(angles), 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)