diff --git a/parcels/_datasets/structured/generated.py b/parcels/_datasets/structured/generated.py new file mode 100644 index 0000000000..10f09e5277 --- /dev/null +++ b/parcels/_datasets/structured/generated.py @@ -0,0 +1,20 @@ +import numpy as np +import xarray as xr + + +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, 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"}), + "XG": (["XG"], np.arange(dims[3]), {"axis": "X", "c_grid_axis_shift": -0.5}), + "lat": (["YG"], np.linspace(-90, 90, dims[2]), {"axis": "Y", "c_grid_axis_shift": 0.5}), + "lon": (["XG"], np.linspace(-max_lon, max_lon, dims[3]), {"axis": "X", "c_grid_axis_shift": -0.5}), + }, + ) diff --git a/parcels/_index_search.py b/parcels/_index_search.py index 4638b18486..397d372164 100644 --- a/parcels/_index_search.py +++ b/parcels/_index_search.py @@ -63,8 +63,8 @@ def _search_time_index(field: Field, time: datetime): tau = 1 else: tau = ( - (time - field.data.time[ti]).dt.total_seconds() - / (field.data.time[ti + 1] - field.data.time[ti]).dt.total_seconds() + (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 ) diff --git a/parcels/_reprs.py b/parcels/_reprs.py index c888945caf..6ea42992ac 100644 --- a/parcels/_reprs.py +++ b/parcels/_reprs.py @@ -55,8 +55,7 @@ def particleset_repr(pset: ParticleSet) -> str: out = f"""<{type(pset).__name__}> fieldset : {textwrap.indent(repr(pset.fieldset), " " * 8)} - pclass : {pset.pclass} - repeatdt : {pset.repeatdt} + ptype : {pset._ptype} # particles: {len(pset)} particles : {_format_list_items_multiline(particles, level=2)} """ diff --git a/parcels/application_kernels/advection.py b/parcels/application_kernels/advection.py index dc5b801dcc..82cb6c970a 100644 --- a/parcels/application_kernels/advection.py +++ b/parcels/application_kernels/advection.py @@ -113,7 +113,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 = min(particle.next_dt, fieldset.RK45_max_dt) / np.timedelta64(1, "s") # noqa TODO improve API for converting dt to seconds + dt = min(particle.next_dt / np.timedelta64(1, "s"), fieldset.RK45_max_dt) # noqa 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 = [ [1.0 / 4.0, 0.0, 0.0, 0.0, 0.0], diff --git a/parcels/application_kernels/advectiondiffusion.py b/parcels/application_kernels/advectiondiffusion.py index 7f120b54db..8f83d26aeb 100644 --- a/parcels/application_kernels/advectiondiffusion.py +++ b/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 @@ -100,9 +102,10 @@ def DiffusionUniformKh(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))) bx = math.sqrt(2 * fieldset.Kh_zonal[particle]) by = math.sqrt(2 * fieldset.Kh_meridional[particle]) diff --git a/parcels/application_kernels/interpolation.py b/parcels/application_kernels/interpolation.py index 1622ffcac7..f474116529 100644 --- a/parcels/application_kernels/interpolation.py +++ b/parcels/application_kernels/interpolation.py @@ -7,16 +7,135 @@ import numpy as np from parcels.field import Field +from parcels.tools.statuscodes import ( + FieldOutOfBoundError, +) if TYPE_CHECKING: from parcels.uxgrid import _UXGRID_AXES + from parcels.xgrid import _XGRID_AXES __all__ = [ "UXPiecewiseConstantFace", "UXPiecewiseLinearNode", + "XBiLinear", + "XBiLinearPeriodic", + "XTriLinear", ] +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] + 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 XBiLinearPeriodic( + 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 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] + + +def XTriLinear( + 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, +): + """Trilinear interpolation on a regular grid.""" + xi, xsi = position["X"] + yi, eta = position["Y"] + zi, zeta = position["Z"] + + if zi < 0 or xi < 0 or yi < 0: + raise FieldOutOfBoundError + + 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, :, :] + 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] + else: + return data[0, 0] + + def UXPiecewiseConstantFace( field: Field, ti: int, diff --git a/parcels/kernel.py b/parcels/kernel.py index e80b6e03c9..56ae8b2f98 100644 --- a/parcels/kernel.py +++ b/parcels/kernel.py @@ -12,8 +12,6 @@ from parcels.application_kernels.advection import ( AdvectionAnalytical, - AdvectionRK4_3D, - AdvectionRK4_3D_CROCO, AdvectionRK45, ) from parcels.basegrid import GridType @@ -124,9 +122,10 @@ def __init__( # Derive meta information from pyfunc, if not given self.check_fieldsets_in_kernels(pyfunc) - if (pyfunc is AdvectionRK4_3D) and fieldset.U.gridindexingtype == "croco": - pyfunc = AdvectionRK4_3D_CROCO - self.funcname = "AdvectionRK4_3D_CROCO" + # # TODO will be implemented when we support CROCO again + # if (pyfunc is AdvectionRK4_3D) and fieldset.U.gridindexingtype == "croco": + # pyfunc = AdvectionRK4_3D_CROCO + # self.funcname = "AdvectionRK4_3D_CROCO" if funcvars is not None: # TODO v4: check if needed from here onwards self.funcvars = funcvars @@ -385,13 +384,12 @@ def evaluate_particle(self, p, endtime): return p pre_dt = p.dt - # TODO implement below later again - # try: # Use next_dt from AdvectionRK45 if it is set - # if abs(endtime - p.time_nextloop) < abs(p.next_dt) - 1e-6: - # p.next_dt = abs(endtime - p.time_nextloop) * sign_dt - # except AttributeError: - if sign_dt * (endtime - p.time_nextloop) <= p.dt: - p.dt = sign_dt * (endtime - p.time_nextloop) + 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 = self._pyfunc(p, self._fieldset, p.time_nextloop) if res is None: diff --git a/parcels/particleset.py b/parcels/particleset.py index f018a356da..8448a4b2d8 100644 --- a/parcels/particleset.py +++ b/parcels/particleset.py @@ -109,7 +109,9 @@ def __init__( assert lon.size == lat.size and lon.size == depth.size, "lon, lat, depth don't all have the same lenghts" if time is None or len(time) == 0: - time = np.datetime64("NaT", "ns") # do not set a time yet (because sign_dt not known) + time = type(fieldset.U.data.time[0].values)( + "NaT", "ns" + ) # do not set a time yet (because sign_dt not known) elif type(time[0]) in [np.datetime64, np.timedelta64]: pass # already in the right format else: @@ -156,7 +158,7 @@ def __init__( if isinstance(v.initial, attrgetter): initial = v.initial(self) else: - initial = v.initial * np.ones(len(trajectory_ids), dtype=v.dtype) + initial = [np.array(v.initial, dtype=v.dtype)] * len(trajectory_ids) self._data[v.name] = initial # update initial values provided on ParticleSet creation @@ -843,16 +845,20 @@ def _warn_outputdt_release_desync(outputdt: float, starttime: float, release_tim def _warn_particle_times_outside_fieldset_time_bounds(release_times: np.ndarray, time: TimeInterval): - if np.any(release_times): - if np.any(release_times < time.left): - warnings.warn( - "Some particles are set to be released outside the FieldSet's executable time domain.", - ParticleSetWarning, - stacklevel=2, - ) - if np.any(release_times > time.right): - warnings.warn( - "Some particles are set to be released after the fieldset's last time and the fields are not constant in time.", - ParticleSetWarning, - stacklevel=2, - ) + if np.isnat(release_times).all(): + return + + if isinstance(time.left, np.datetime64) and isinstance(release_times[0], np.timedelta64): + release_times = np.array([t + time.left for t in release_times]) + if np.any(release_times < time.left): + warnings.warn( + "Some particles are set to be released outside the FieldSet's executable time domain.", + ParticleSetWarning, + stacklevel=2, + ) + if np.any(release_times > time.right): + warnings.warn( + "Some particles are set to be released after the fieldset's last time and the fields are not constant in time.", + ParticleSetWarning, + stacklevel=2, + ) diff --git a/parcels/xgrid.py b/parcels/xgrid.py index 81391ee818..3c8b9a3b86 100644 --- a/parcels/xgrid.py +++ b/parcels/xgrid.py @@ -9,6 +9,7 @@ from parcels import xgcm from parcels._index_search import _search_indices_curvilinear_2d from parcels.basegrid import BaseGrid +from parcels.tools.statuscodes import FieldOutOfBoundError, FieldOutOfBoundSurfaceError _XGRID_AXES = Literal["X", "Y", "Z"] _XGRID_AXES_ORDERING: Sequence[_XGRID_AXES] = "ZYX" @@ -271,6 +272,15 @@ def search(self, z, y, x, ei=None): ds = self.xgcm_grid._ds zi, zeta = _search_1d_array(ds.depth.values, z) + 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) @@ -453,6 +463,9 @@ def _search_1d_array( 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 diff --git a/tests/test_advection.py b/tests/test_advection.py index f0f710637e..de8b004d79 100644 --- a/tests/test_advection.py +++ b/tests/test_advection.py @@ -1,6 +1,3 @@ -import math -from datetime import timedelta - import numpy as np import pytest import xarray as xr @@ -13,11 +10,9 @@ AdvectionRK4, AdvectionRK4_3D, AdvectionRK45, - Field, FieldSet, Particle, ParticleSet, - StatusCode, Variable, ) from tests.utils import TEST_DATA @@ -31,13 +26,6 @@ "AdvDiffM1": AdvectionDiffusionM1, } -# Some constants -f = 1.0e-4 -u_0 = 0.3 -u_g = 0.04 -gamma = 1 / (86400.0 * 2.89) -gamma_g = 1 / (86400.0 * 28.9) - @pytest.fixture def lon(): @@ -57,146 +45,6 @@ def depth(): return np.linspace(0, 30, zdim, dtype=np.float32) -@pytest.mark.v4alpha -@pytest.mark.xfail(reason="GH1946") -def test_advection_zonal(lon, lat, depth): - """Particles at high latitude move geographically faster due to the pole correction in `GeographicPolar`.""" - npart = 10 - data2D = { - "U": np.ones((lat.size, lon.size), dtype=np.float32), - "V": np.zeros((lat.size, lon.size), dtype=np.float32), - } - data3D = { - "U": np.ones((depth.size, lat.size, lon.size), dtype=np.float32), - "V": np.zeros((depth.size, lat.size, lon.size), dtype=np.float32), - } - dimensions = {"lon": lon, "lat": lat} - fieldset2D = FieldSet.from_data(data2D, dimensions, mesh="spherical") - - pset2D = ParticleSet(fieldset2D, pclass=Particle, lon=np.zeros(npart) + 20.0, lat=np.linspace(0, 80, npart)) - pset2D.execute(AdvectionRK4, runtime=timedelta(hours=2), dt=timedelta(seconds=30)) - assert (np.diff(pset2D.lon) > 1.0e-4).all() - - dimensions["depth"] = depth - fieldset3D = FieldSet.from_data(data3D, dimensions, mesh="spherical") - pset3D = ParticleSet( - fieldset3D, - pclass=Particle, - lon=np.zeros(npart) + 20.0, - lat=np.linspace(0, 80, npart), - depth=np.zeros(npart) + 10.0, - ) - pset3D.execute(AdvectionRK4, runtime=timedelta(hours=2), dt=timedelta(seconds=30)) - assert (np.diff(pset3D.lon) > 1.0e-4).all() - - -@pytest.mark.v4alpha -@pytest.mark.xfail(reason="GH1946") -def test_advection_meridional(lon, lat): - """Particles at high latitude move geographically faster due to the pole correction in `GeographicPolar`.""" - npart = 10 - data = {"U": np.zeros((lat.size, lon.size), dtype=np.float32), "V": np.ones((lat.size, lon.size), dtype=np.float32)} - dimensions = {"lon": lon, "lat": lat} - fieldset = FieldSet.from_data(data, dimensions, mesh="spherical") - - pset = ParticleSet(fieldset, pclass=Particle, lon=np.linspace(-60, 60, npart), lat=np.linspace(0, 30, npart)) - delta_lat = np.diff(pset.lat) - pset.execute(AdvectionRK4, runtime=timedelta(hours=2), dt=timedelta(seconds=30)) - assert np.allclose(np.diff(pset.lat), delta_lat, rtol=1.0e-4) - - -@pytest.mark.v4alpha -@pytest.mark.xfail(reason="GH1946") -def test_advection_3D(): - """Flat 2D zonal flow that increases linearly with depth from 0 m/s to 1 m/s.""" - xdim = ydim = zdim = 2 - npart = 11 - dimensions = { - "lon": np.linspace(0.0, 1e4, xdim, dtype=np.float32), - "lat": np.linspace(0.0, 1e4, ydim, dtype=np.float32), - "depth": np.linspace(0.0, 1.0, zdim, dtype=np.float32), - } - data = {"U": np.ones((zdim, ydim, xdim), dtype=np.float32), "V": np.zeros((zdim, ydim, xdim), dtype=np.float32)} - data["U"][0, :, :] = 0.0 - fieldset = FieldSet.from_data(data, dimensions, mesh="flat") - - pset = ParticleSet( - fieldset, pclass=Particle, lon=np.zeros(npart), lat=np.zeros(npart) + 1e2, depth=np.linspace(0, 1, npart) - ) - time = timedelta(hours=2).total_seconds() - pset.execute(AdvectionRK4, runtime=time, dt=timedelta(seconds=30)) - assert np.allclose(pset.depth * pset.time, pset.lon, atol=1.0e-1) - - -@pytest.mark.v4alpha -@pytest.mark.xfail(reason="GH1946") -@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") - - def DeleteParticle(particle, fieldset, time): # pragma: no cover - if particle.state == StatusCode.ErrorOutOfBounds or particle.state == StatusCode.ErrorThroughSurface: - particle.delete() - - def SubmergeParticle(particle, fieldset, time): # pragma: no cover - if particle.state == StatusCode.ErrorThroughSurface: - (u, v) = fieldset.UV[particle] - particle_dlon = u * particle.dt # noqa - particle_dlat = v * particle.dt # noqa - particle_ddepth = 0.0 # noqa - particle.depth = 0 - particle.state = StatusCode.Evaluate - - kernels = [AdvectionRK4_3D] - if wErrorThroughSurface: - kernels.append(SubmergeParticle) - kernels.append(DeleteParticle) - - pset = ParticleSet(fieldset=fieldset, pclass=Particle, lon=0.5, lat=0.5, depth=0.9) - pset.execute(kernels, runtime=11.0, dt=1) - - if direction == "up" and wErrorThroughSurface: - assert np.allclose(pset.lon[0], 0.6) - assert np.allclose(pset.depth[0], 0) - else: - assert len(pset) == 0 - - -@pytest.mark.v4alpha -@pytest.mark.xfail(reason="GH1946") -@pytest.mark.parametrize("rk45_tol", [10, 100]) -def test_advection_RK45(lon, lat, rk45_tol): - npart = 10 - data2D = { - "U": np.ones((lat.size, lon.size), dtype=np.float32), - "V": np.zeros((lat.size, lon.size), dtype=np.float32), - } - dimensions = {"lon": lon, "lat": lat} - fieldset = FieldSet.from_data(data2D, dimensions, mesh="spherical") - fieldset.add_constant("RK45_tol", rk45_tol) - - dt = timedelta(seconds=30).total_seconds() - RK45Particles = Particle.add_variable("next_dt", dtype=np.float32, initial=dt) - pset = ParticleSet(fieldset, pclass=RK45Particles, lon=np.zeros(npart) + 20.0, lat=np.linspace(0, 80, npart)) - pset.execute(AdvectionRK45, runtime=timedelta(hours=2), dt=dt) - assert (np.diff(pset.lon) > 1.0e-4).all() - assert np.isclose(fieldset.RK45_tol, rk45_tol / (1852 * 60)) - print(fieldset.RK45_tol) - - @pytest.mark.v4alpha @pytest.mark.xfail(reason="When refactoring fieldfilebuffer croco support was dropped. This will be fixed in v4.") def test_conversion_3DCROCO(): @@ -276,372 +124,6 @@ def test_advection_2DCROCO(): assert np.allclose(pset.lon_nextloop, [x + runtime for x in X], atol=1e-3) -def create_periodic_fieldset(xdim, ydim, uvel, vvel): - dimensions = { - "lon": np.linspace(0.0, 1.0, xdim + 1, dtype=np.float32)[1:], # don't include both 0 and 1, for periodic b.c. - "lat": np.linspace(0.0, 1.0, ydim + 1, dtype=np.float32)[1:], - } - - data = {"U": uvel * np.ones((ydim, xdim), dtype=np.float32), "V": vvel * np.ones((ydim, xdim), dtype=np.float32)} - return FieldSet.from_data(data, dimensions, mesh="spherical") - - -def periodicBC(particle, fieldset, time): # pragma: no cover - particle.lon = math.fmod(particle.lon, 1) - particle.lat = math.fmod(particle.lat, 1) - - -@pytest.mark.v4alpha -@pytest.mark.xfail(reason="Calls fieldset.add_periodic_halo(). In v4, interpolation should work without adding halo.") -def test_advection_periodic_zonal(): - xdim, ydim, halosize = 100, 100, 3 - fieldset = create_periodic_fieldset(xdim, ydim, uvel=1.0, vvel=0.0) - fieldset.add_periodic_halo(zonal=True, halosize=halosize) - assert len(fieldset.U.lon) == xdim + 2 * halosize - - pset = ParticleSet(fieldset, pclass=Particle, lon=[0.5], lat=[0.5]) - pset.execute(AdvectionRK4 + pset.Kernel(periodicBC), runtime=timedelta(hours=20), dt=timedelta(seconds=30)) - assert abs(pset.lon[0] - 0.15) < 0.1 - - -@pytest.mark.v4alpha -@pytest.mark.xfail(reason="Calls fieldset.add_periodic_halo(). In v4, interpolation should work without adding halo.") -def test_advection_periodic_meridional(): - xdim, ydim = 100, 100 - fieldset = create_periodic_fieldset(xdim, ydim, uvel=0.0, vvel=1.0) - fieldset.add_periodic_halo(meridional=True) - assert len(fieldset.U.lat) == ydim + 10 # default halo size is 5 grid points - - pset = ParticleSet(fieldset, pclass=Particle, lon=[0.5], lat=[0.5]) - pset.execute(AdvectionRK4 + pset.Kernel(periodicBC), runtime=timedelta(hours=20), dt=timedelta(seconds=30)) - assert abs(pset.lat[0] - 0.15) < 0.1 - - -@pytest.mark.v4alpha -@pytest.mark.xfail(reason="Calls fieldset.add_periodic_halo(). In v4, interpolation should work without adding halo.") -def test_advection_periodic_zonal_meridional(): - xdim, ydim = 100, 100 - fieldset = create_periodic_fieldset(xdim, ydim, uvel=1.0, vvel=1.0) - fieldset.add_periodic_halo(zonal=True, meridional=True) - assert len(fieldset.U.lat) == ydim + 10 # default halo size is 5 grid points - assert len(fieldset.U.lon) == xdim + 10 # default halo size is 5 grid points - assert np.allclose(np.diff(fieldset.U.lat), fieldset.U.lat[1] - fieldset.U.lat[0], rtol=0.001) - assert np.allclose(np.diff(fieldset.U.lon), fieldset.U.lon[1] - fieldset.U.lon[0], rtol=0.001) - - pset = ParticleSet(fieldset, pclass=Particle, lon=[0.4], lat=[0.5]) - pset.execute(AdvectionRK4 + pset.Kernel(periodicBC), runtime=timedelta(hours=20), dt=timedelta(seconds=30)) - assert abs(pset.lon[0] - 0.05) < 0.1 - assert abs(pset.lat[0] - 0.15) < 0.1 - - -@pytest.mark.v4alpha -@pytest.mark.xfail(reason="GH1946") -@pytest.mark.parametrize("u", [-0.3, np.array(0.2)]) -@pytest.mark.parametrize("v", [0.2, np.array(1)]) -@pytest.mark.parametrize("w", [None, -0.2, np.array(0.7)]) -def test_length1dimensions(u, v, w): - (lon, xdim) = (np.linspace(-10, 10, 21), 21) if isinstance(u, np.ndarray) else (0, 1) - (lat, ydim) = (np.linspace(-15, 15, 31), 31) if isinstance(v, np.ndarray) else (-4, 1) - (depth, zdim) = (np.linspace(-5, 5, 11), 11) if (isinstance(w, np.ndarray) and w is not None) else (3, 1) - dimensions = {"lon": lon, "lat": lat, "depth": depth} - - dims = [] - if zdim > 1: - dims.append(zdim) - if ydim > 1: - dims.append(ydim) - if xdim > 1: - dims.append(xdim) - if len(dims) > 0: - U = u * np.ones(dims, dtype=np.float32) - V = v * np.ones(dims, dtype=np.float32) - if w is not None: - W = w * np.ones(dims, dtype=np.float32) - else: - U, V, W = u, v, w - - data = {"U": U, "V": V} - if w is not None: - data["W"] = W - fieldset = FieldSet.from_data(data, dimensions, mesh="flat") - - x0, y0, z0 = 2, 8, -4 - pset = ParticleSet(fieldset, pclass=Particle, lon=x0, lat=y0, depth=z0) - pfunc = AdvectionRK4 if w is None else AdvectionRK4_3D - kernel = pset.Kernel(pfunc) - pset.execute(kernel, runtime=5, dt=1) - - 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() - if 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) - - -def truth_moving(x_0, y_0, t): - lat = y_0 - (u_0 - u_g) / f * (1 - math.cos(f * t)) - lon = x_0 + u_g * t + (u_0 - u_g) / f * math.sin(f * t) - return lon, lat - - -def create_fieldset_moving(xdim=100, ydim=100, maxtime=timedelta(hours=6)): - """Generate a FieldSet encapsulating the flow field of a moving 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_g + (u_0 - u_g) * np.cos(f * time)), - "V": np.transpose(np.ones((xdim, ydim, 1), dtype=np.float32) * -(u_0 - u_g) * np.sin(f * time)), - } - return FieldSet.from_data(data, dimensions, mesh="flat") - - -@pytest.fixture -def fieldset_moving(): - return create_fieldset_moving() - - -@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_moving_eddy(fieldset_moving, method, rtol, diffField): - npart = 1 - fieldset = fieldset_moving - 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_moving(x, y, t)[0] for x, y, t in zip(lon, lat, pset.time, strict=True)] - exp_lat = [truth_moving(x, y, t)[1] for x, y, t in zip(lon, lat, pset.time, strict=True)] - assert np.allclose(pset.lon, exp_lon, rtol=rtol) - assert np.allclose(pset.lat, exp_lat, rtol=rtol) - - -def truth_decaying(x_0, y_0, t): - lat = y_0 - ( - (u_0 - u_g) * f / (f**2 + gamma**2) * (1 - np.exp(-gamma * t) * (np.cos(f * t) + gamma / f * np.sin(f * t))) - ) - lon = x_0 + ( - u_g / gamma_g * (1 - np.exp(-gamma_g * t)) - + (u_0 - u_g) - * f - / (f**2 + gamma**2) - * (gamma / f + np.exp(-gamma * t) * (math.sin(f * t) - gamma / f * math.cos(f * t))) - ) - return lon, lat - - -def create_fieldset_decaying(xdim=100, ydim=100, maxtime=timedelta(hours=6)): - """Generate a FieldSet encapsulating the flow field of a decaying 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_g * np.exp(-gamma_g * time) - + (u_0 - u_g) * np.exp(-gamma * time) * np.cos(f * time) - ), - "V": np.transpose( - np.ones((xdim, ydim, 1), dtype=np.float32) * -(u_0 - u_g) * np.exp(-gamma * time) * np.sin(f * time) - ), - } - return FieldSet.from_data(data, dimensions, mesh="flat") - - -@pytest.fixture -def fieldset_decaying(): - return create_fieldset_decaying() - - -@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), - ("AA", 1e-3, False), - ], -) -def test_decaying_eddy(fieldset_decaying, method, rtol, diffField): - npart = 1 - fieldset = fieldset_decaying - if method == "AA": - # needed for AnalyticalAdvection to work, but comes at expense of accuracy - fieldset.U.interp_method = "cgrid_velocity" - fieldset.V.interp_method = "cgrid_velocity" - - 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_decaying(x, y, t)[0] for x, y, t in zip(lon, lat, pset.time, strict=True)] - exp_lat = [truth_decaying(x, y, t)[1] for x, y, t in zip(lon, lat, pset.time, 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_analyticalAgrid(): diff --git a/tests/test_diffusion.py b/tests/test_diffusion.py deleted file mode 100644 index 1cfb71e018..0000000000 --- a/tests/test_diffusion.py +++ /dev/null @@ -1,149 +0,0 @@ -import random -from datetime import timedelta - -import numpy as np -import pytest -from scipy import stats - -from parcels import ( - AdvectionDiffusionEM, - AdvectionDiffusionM1, - DiffusionUniformKh, - 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"]) -def test_fieldKh_Brownian(mesh): - xdim = 200 - ydim = 100 - kh_zonal = 100 - kh_meridional = 50 - - 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) - - fieldset.add_constant_field("Kh_zonal", kh_zonal, mesh=mesh) - fieldset.add_constant_field("Kh_meridional", kh_meridional, mesh=mesh) - - npart = 1000 - runtime = timedelta(days=1) - - random.seed(1234) - pset = ParticleSet(fieldset=fieldset, pclass=Particle, lon=np.zeros(npart), lat=np.zeros(npart)) - pset.execute(pset.Kernel(DiffusionUniformKh), runtime=runtime, dt=timedelta(hours=1)) - - expected_std_lon = np.sqrt(2 * kh_zonal * mesh_conversion**2 * runtime.total_seconds()) - expected_std_lat = np.sqrt(2 * kh_meridional * mesh_conversion**2 * runtime.total_seconds()) - - lats = pset.lat - lons = pset.lon - - tol = 500 * mesh_conversion # effectively 500 m errors - assert np.allclose(np.std(lats), expected_std_lat, atol=tol) - assert np.allclose(np.std(lons), expected_std_lon, atol=tol) - assert np.allclose(np.mean(lons), 0, atol=tol) - assert np.allclose(np.mean(lats), 0, atol=tol) - - -@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]) -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, pclass=Particle, 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=1, dt=1) - - depth = pset.depth - expected_mean = 1.0 / fieldset.lambd - assert np.allclose(np.mean(depth), expected_mean, rtol=0.1) - - -@pytest.mark.v4alpha -@pytest.mark.xfail(reason="GH1946") -@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("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=1, dt=1) - - 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) diff --git a/tests/test_interpolation.py b/tests/test_interpolation.py index cac7abc893..6dd9ff6914 100644 --- a/tests/test_interpolation.py +++ b/tests/test_interpolation.py @@ -3,8 +3,7 @@ import xarray as xr import parcels._interpolation as interpolation -from parcels import AdvectionRK4_3D, FieldSet, Particle, ParticleSet -from tests.utils import TEST_DATA, create_fieldset_zeros_3d +from tests.utils import create_fieldset_zeros_3d @pytest.fixture @@ -109,50 +108,3 @@ def test_interpolator2(ctx: interpolation.InterpolationContext3D): return 0 fieldset.U[0.5, 0.5, 0.5, 0.5] - - -@pytest.mark.v4remove -@pytest.mark.xfail(reason="GH1946") -@pytest.mark.parametrize( - "interp_method", - [ - "linear", - "freeslip", - "nearest", - "cgrid_velocity", - ], -) -def test_interp_regression_v3(interp_method): - """Test that the v4 versions of the interpolation are the same as the v3 versions.""" - variables = {"U": "U", "V": "V", "W": "W"} - dimensions = {"time": "time", "lon": "lon", "lat": "lat", "depth": "depth"} - ds = xr.open_dataset(str(TEST_DATA / f"test_interpolation_data_random_{interp_method}.nc")) - fieldset = FieldSet.from_xarray_dataset( - ds, - variables, - dimensions, - mesh="flat", - ) - - for field in [fieldset.U, fieldset.V, fieldset.W]: # Set a land point (for testing freeslip) - field.interp_method = interp_method - - x, y, z = np.meshgrid(np.linspace(0, 1, 7), np.linspace(0, 1, 13), np.linspace(0, 1, 5)) - - TestP = Particle.add_variable("pid", dtype=np.int32, initial=0) - pset = ParticleSet(fieldset, pclass=TestP, lon=x, lat=y, depth=z, pid=np.arange(x.size)) - - def DeleteParticle(particle, fieldset, time): - if particle.state >= 50: - particle.delete() - - outfile = pset.ParticleFile(f"test_interpolation_v4_{interp_method}", outputdt=1) - pset.execute([AdvectionRK4_3D, DeleteParticle], runtime=4, dt=1, output_file=outfile) - - ds_v3 = xr.open_zarr(str(TEST_DATA / f"test_interpolation_jit_{interp_method}.zarr")) - ds_v4 = xr.open_zarr(f"test_interpolation_v4_{interp_method}.zarr") - - tol = 1e-6 - np.testing.assert_allclose(ds_v3.lon, ds_v4.lon, atol=tol) - np.testing.assert_allclose(ds_v3.lat, ds_v4.lat, atol=tol) - np.testing.assert_allclose(ds_v3.z, ds_v4.z, atol=tol) diff --git a/tests/utils.py b/tests/utils.py index 2298955179..a496caf2e1 100644 --- a/tests/utils.py +++ b/tests/utils.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.generated 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): diff --git a/tests/v4/test_advection.py b/tests/v4/test_advection.py new file mode 100644 index 0000000000..faa146619a --- /dev/null +++ b/tests/v4/test_advection.py @@ -0,0 +1,255 @@ +import numpy as np +import pytest +import xarray as xr + +from parcels._datasets.structured.generated import simple_UV_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 +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.xgrid import XGrid + +kernel = { + "EE": AdvectionEE, + "RK4": AdvectionRK4, + "RK4_3D": AdvectionRK4_3D, + "RK45": AdvectionRK45, + # "AA": AdvectionAnalytical, + "AdvDiffEM": AdvectionDiffusionEM, + "AdvDiffM1": AdvectionDiffusionM1, +} + + +@pytest.mark.parametrize("mesh_type", ["spherical", "flat"]) +def test_advection_zonal(mesh_type, npart=10): + """Particles at high latitude move geographically faster due to the pole correction in `GeographicPolar`.""" + 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) + UV = VectorField("UV", U, V) + fieldset = FieldSet([U, V, UV]) + + pset = ParticleSet(fieldset, lon=np.zeros(npart) + 20.0, lat=np.linspace(0, 80, npart)) + pset.execute(AdvectionRK4, runtime=np.timedelta64(2, "h"), dt=np.timedelta64(15, "m")) + + if mesh_type == "spherical": + assert (np.diff(pset.lon) > 1.0e-4).all() + else: + assert (np.diff(pset.lon) < 1.0e-4).all() + + +def periodicBC(particle, fieldset, time): + particle.total_dlon += particle_dlon # noqa + particle.lon = np.fmod(particle.lon, fieldset.U.grid.lon[-1]) + particle.lat = np.fmod(particle.lat, fieldset.U.grid.lat[-1]) + + +def test_advection_zonal_periodic(): + ds = simple_UV_dataset(dims=(2, 2, 2, 2), mesh_type="flat") + ds["U"].data[:] = 0.1 + ds["lon"].data = np.array([0, 2]) + ds["lat"].data = np.array([0, 2]) + + grid = XGrid.from_dataset(ds) + U = Field("U", ds["U"], grid, interp_method=XBiLinearPeriodic) + V = Field("V", ds["V"], grid, interp_method=XBiLinearPeriodic) + 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]) + 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) + + +def test_horizontal_advection_in_3D_flow(npart=10): + """Flat 2D zonal flow that increases linearly with depth from 0 m/s to 1 m/s.""" + 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.data[:, 0, :, :] = 0.0 # Set U to 0 at the surface + V = Field("V", ds["V"], grid, interp_method=XTriLinear) + UV = VectorField("UV", U, V) + fieldset = FieldSet([U, V, UV]) + + pset = ParticleSet(fieldset, lon=np.zeros(npart), lat=np.zeros(npart), depth=np.linspace(0.1, 0.9, npart)) + 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) + + +@pytest.mark.parametrize("direction", ["up", "down"]) +@pytest.mark.parametrize("wErrorThroughSurface", [True, False]) +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 + 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.delete() + + 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 # noqa + particle_dlat = v * dt # noqa + particle_ddepth = 0.0 # noqa + particle.depth = 0 + particle.state = StatusCode.Evaluate + + kernels = [AdvectionRK4_3D] + if wErrorThroughSurface: + kernels.append(SubmergeParticle) + kernels.append(DeleteParticle) + + pset = ParticleSet(fieldset=fieldset, lon=0.5, lat=0.5, depth=0.9) + 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) + else: + assert len(pset) == 0 + + +@pytest.mark.parametrize("u", [-0.3, np.array(0.2)]) +@pytest.mark.parametrize("v", [0.2, np.array(1)]) +@pytest.mark.parametrize("w", [None, -0.2, np.array(0.7)]) +def test_length1dimensions(u, v, w): # TODO: Refactor this test to be more readable (and isolate test setup) + (lon, xdim) = (np.linspace(-10, 10, 21), 21) if isinstance(u, np.ndarray) else (np.array([0]), 1) + (lat, ydim) = (np.linspace(-15, 15, 31), 31) if isinstance(v, np.ndarray) else (np.array([-4]), 1) + (depth, zdim) = ( + (np.linspace(-5, 5, 11), 11) if (isinstance(w, np.ndarray) and w is not None) else (np.array([3]), 1) + ) + + tdim = 2 # TODO make this also work for length-1 time dimensions + dims = (tdim, zdim, ydim, xdim) + U = u * np.ones(dims, dtype=np.float32) + V = v * np.ones(dims, dtype=np.float32) + if w is not None: + W = w * np.ones(dims, dtype=np.float32) + + ds = xr.Dataset( + { + "U": (["time", "depth", "YG", "XG"], U), + "V": (["time", "depth", "YG", "XG"], V), + }, + coords={ + "time": (["time"], [np.timedelta64(0, "s"), np.timedelta64(10, "s")], {"axis": "T"}), + "depth": (["depth"], depth, {"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}), + }, + ) + if w: + 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) + fields = [U, V, VectorField("UV", U, V)] + if w: + W = Field("W", ds["W"], grid, interp_method=XTriLinear) + fields.append(VectorField("UVW", U, V, W)) + fieldset = FieldSet(fields) + + x0, y0, z0 = 2, 8, -4 + pset = ParticleSet(fieldset, lon=x0, lat=y0, depth=z0) + kernel = AdvectionRK4 if w is None else AdvectionRK4_3D + 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() + if w: + assert ((np.array([p.depth - z0 for p in pset]) - 4 * w) < 1e-6).all() + + +@pytest.mark.parametrize( + "method, rtol", + [ + ("EE", 1e-2), + ("AdvDiffEM", 1e-2), + ("AdvDiffM1", 1e-2), + ("RK4", 1e-5), + ("RK4_3D", 1e-5), + ("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) + 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": + # Using W to test 3D advection (assuming same velocity as V) + 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 == "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( + 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) diff --git a/tests/v4/test_diffusion.py b/tests/v4/test_diffusion.py new file mode 100644 index 0000000000..31f984504e --- /dev/null +++ b/tests/v4/test_diffusion.py @@ -0,0 +1,144 @@ +import random + +import numpy as np +import pytest +from scipy import stats + +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.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 +from tests.utils import create_fieldset_zeros_conversion + + +@pytest.mark.parametrize("mesh_type", ["spherical", "flat"]) +def test_fieldKh_Brownian(mesh_type): + kh_zonal = 100 + kh_meridional = 50 + mesh_conversion = 1 / 1852.0 / 60 if mesh_type == "spherical" else 1 + + ds = simple_UV_dataset(dims=(2, 1, 2, 2), mesh_type=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) + 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) + UV = VectorField("UV", U, V) + fieldset = FieldSet([U, V, UV, Kh_zonal, Kh_meridional]) + + npart = 100 + runtime = np.timedelta64(2, "h") + + 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")) + + expected_std_lon = np.sqrt(2 * kh_zonal * mesh_conversion**2 * (runtime / np.timedelta64(1, "s"))) + expected_std_lat = np.sqrt(2 * kh_meridional * mesh_conversion**2 * (runtime / np.timedelta64(1, "s"))) + + tol = 500 * mesh_conversion # effectively 500 m errors + assert np.allclose(np.std(pset.lat), expected_std_lat, atol=tol) + 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=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=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]) + + 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) + + +@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) diff --git a/tests/v4/test_interpolation.py b/tests/v4/test_interpolation.py new file mode 100644 index 0000000000..82642a6b10 --- /dev/null +++ b/tests/v4/test_interpolation.py @@ -0,0 +1,109 @@ +import numpy as np +import pytest +import xarray as xr + +from parcels._datasets.structured.generated import simple_UV_dataset +from parcels.application_kernels.advection import AdvectionRK4_3D +from parcels.application_kernels.interpolation import XBiLinear, XTriLinear +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 +from tests.utils import TEST_DATA + + +@pytest.mark.parametrize("mesh_type", ["spherical", "flat"]) +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) + UV = VectorField("UV", U, V) + + lat = 30.0 + time = U.time_interval.left + u_expected = 1.0 if mesh_type == "flat" else 1.0 / (1852 * 60 * np.cos(np.radians(lat))) + + assert np.isclose(U[time, 0, lat, 0], u_expected, atol=1e-7) + assert V[time, 0, lat, 0] == 0.0 + + u, v = UV[time, 0, lat, 0] + assert np.isclose(u, u_expected, atol=1e-7) + assert v == 0.0 + + assert U.eval(time, 0, lat, 0, applyConversion=False) == 1 + + +interp_methods = { + "linear": XTriLinear, +} + + +@pytest.mark.xfail(reason="ParticleFile not implemented yet") +@pytest.mark.parametrize( + "interp_name", + [ + "linear", + # "freeslip", + # "nearest", + # "cgrid_velocity", + ], +) +def test_interp_regression_v3(interp_name): + """Test that the v4 versions of the interpolation are the same as the v3 versions.""" + ds_input = xr.open_dataset(str(TEST_DATA / f"test_interpolation_data_random_{interp_name}.nc")) + ydim = ds_input["U"].shape[2] + xdim = ds_input["U"].shape[3] + time = [np.timedelta64(int(t), "s") for t in ds_input["time"].values] + + ds = xr.Dataset( + { + "U": (["time", "depth", "YG", "XG"], ds_input["U"].values), + "V": (["time", "depth", "YG", "XG"], ds_input["V"].values), + "W": (["time", "depth", "YG", "XG"], ds_input["W"].values), + }, + coords={ + "time": (["time"], time, {"axis": "T"}), + "depth": (["depth"], ds_input["depth"].values, {"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"], ds_input["lat"].values, {"axis": "Y", "c_grid_axis_shift": 0.5}), + "lon": (["XG"], ds_input["lon"].values, {"axis": "X", "c_grid_axis_shift": -0.5}), + }, + ) + + grid = XGrid.from_dataset(ds) + U = Field("U", ds["U"], grid, mesh_type="flat", interp_method=interp_methods[interp_name]) + V = Field("V", ds["V"], grid, mesh_type="flat", interp_method=interp_methods[interp_name]) + W = Field("W", ds["W"], grid, mesh_type="flat", interp_method=interp_methods[interp_name]) + fieldset = FieldSet([U, V, W, VectorField("UVW", U, V, W)]) + + x, y, z = np.meshgrid(np.linspace(0, 1, 7), np.linspace(0, 1, 13), np.linspace(0, 1, 5)) + + TestP = Particle.add_variable(Variable("pid", dtype=np.int32, initial=0)) + pset = ParticleSet(fieldset, pclass=TestP, lon=x, lat=y, depth=z, pid=np.arange(x.size)) + + def DeleteParticle(particle, fieldset, time): + if particle.state >= 50: + particle.delete() + + outfile = pset.ParticleFile(f"test_interpolation_v4_{interp_name}", outputdt=np.timedelta64(1, "s")) + pset.execute( + [AdvectionRK4_3D, DeleteParticle], + runtime=np.timedelta64(4, "s"), + dt=np.timedelta64(1, "s"), + output_file=outfile, + ) + + print(str(TEST_DATA / f"test_interpolation_jit_{interp_name}.zarr")) + ds_v3 = xr.open_zarr(str(TEST_DATA / f"test_interpolation_jit_{interp_name}.zarr")) + ds_v4 = xr.open_zarr(f"test_interpolation_v4_{interp_name}.zarr") + + tol = 1e-6 + np.testing.assert_allclose(ds_v3.lon, ds_v4.lon, atol=tol) + np.testing.assert_allclose(ds_v3.lat, ds_v4.lat, atol=tol) + np.testing.assert_allclose(ds_v3.z, ds_v4.z, atol=tol)