diff --git a/parcels/__init__.py b/parcels/__init__.py index 015a135b55..3eb58c2fb1 100644 --- a/parcels/__init__.py +++ b/parcels/__init__.py @@ -2,19 +2,89 @@ __version__ = version -import warnings as _warnings +import warnings as _stdlib_warnings -from parcels.application_kernels import * -from parcels.field import * -from parcels.fieldset import * -from parcels.interaction import * -from parcels.kernel import * -from parcels.particle import * -from parcels.particlefile import * -from parcels.particleset import * -from parcels.tools import * +from parcels._core.basegrid import BaseGrid +from parcels._core.converters import ( + Geographic, + GeographicPolar, + GeographicPolarSquare, + GeographicSquare, + UnitConverter, +) +from parcels._core.field import Field, VectorField +from parcels._core.fieldset import FieldSet +from parcels._core.kernel import Kernel +from parcels._core.particle import ( + KernelParticle, # ? remove? + Particle, + ParticleClass, + Variable, +) +from parcels._core.particlefile import ParticleFile +from parcels._core.particleset import ParticleSet +from parcels._core.statuscodes import ( + AllParcelsErrorCodes, + FieldInterpolationError, + FieldOutOfBoundError, + FieldSamplingError, + KernelError, + StatusCode, + TimeExtrapolationError, +) +from parcels._core.uxgrid import UxGrid +from parcels._core.warnings import ( + FieldSetWarning, + FileWarning, + KernelWarning, + ParticleSetWarning, +) +from parcels._core.xgrid import XGrid +from parcels._logger import logger +from parcels._tutorial import download_example_dataset, list_example_datasets + +__all__ = [ # noqa: RUF022 + # Core classes + "BaseGrid", + "Field", + "VectorField", + "FieldSet", + "Kernel", + "Particle", + "ParticleClass", + "ParticleFile", + "ParticleSet", + "Variable", + "XGrid", + "UxGrid", + # Converters + "Geographic", + "GeographicPolar", + "GeographicPolarSquare", + "GeographicSquare", + "UnitConverter", + # Status codes and errors + "AllParcelsErrorCodes", + "FieldInterpolationError", + "FieldOutOfBoundError", + "FieldSamplingError", + "KernelError", + "StatusCode", + "TimeExtrapolationError", + # Warnings + "FieldSetWarning", + "FileWarning", + "KernelWarning", + "ParticleSetWarning", + # Utilities + "logger", + "download_example_dataset", + "list_example_datasets", + # (marked for potential removal) + "KernelParticle", +] -_warnings.warn( +_stdlib_warnings.warn( "This is an alpha version of Parcels v4. The API is not stable and may change without deprecation warnings.", UserWarning, stacklevel=2, diff --git a/parcels/basegrid.py b/parcels/_core/basegrid.py similarity index 99% rename from parcels/basegrid.py rename to parcels/_core/basegrid.py index 37d56f6127..a05ea6dfc3 100644 --- a/parcels/basegrid.py +++ b/parcels/_core/basegrid.py @@ -6,7 +6,7 @@ import numpy as np -from parcels.spatialhash import SpatialHash +from parcels._core.spatialhash import SpatialHash if TYPE_CHECKING: import numpy as np diff --git a/parcels/_core/utils/common.py b/parcels/_core/constants.py similarity index 100% rename from parcels/_core/utils/common.py rename to parcels/_core/constants.py diff --git a/parcels/tools/converters.py b/parcels/_core/converters.py similarity index 94% rename from parcels/tools/converters.py rename to parcels/_core/converters.py index 68c3810b6c..35fddc18bc 100644 --- a/parcels/tools/converters.py +++ b/parcels/_core/converters.py @@ -11,12 +11,12 @@ "GeographicPolarSquare", "GeographicSquare", "UnitConverter", - "convert_to_flat_array", - "unitconverters_map", + "_convert_to_flat_array", + "_unitconverters_map", ] -def convert_to_flat_array(var: npt.ArrayLike) -> npt.NDArray: +def _convert_to_flat_array(var: npt.ArrayLike) -> npt.NDArray: """Convert lists and single integers/floats to one-dimensional numpy arrays Parameters @@ -96,7 +96,7 @@ def to_source(self, value, z, y, x): return value * pow(1000.0 * 1.852 * 60.0 * np.cos(y * pi / 180), 2) -unitconverters_map = { +_unitconverters_map = { "U": GeographicPolar(), "V": Geographic(), "Kh_zonal": GeographicPolarSquare(), diff --git a/parcels/field.py b/parcels/_core/field.py similarity index 96% rename from parcels/field.py rename to parcels/_core/field.py index 212c042c0a..d582e5b3c6 100644 --- a/parcels/field.py +++ b/parcels/_core/field.py @@ -8,29 +8,28 @@ import uxarray as ux import xarray as xr +from parcels._core.converters import ( + UnitConverter, + _unitconverters_map, +) +from parcels._core.index_search import GRID_SEARCH_ERROR, LEFT_OUT_OF_BOUNDS, RIGHT_OUT_OF_BOUNDS, _search_time_index +from parcels._core.particle import KernelParticle +from parcels._core.statuscodes import ( + AllParcelsErrorCodes, + StatusCode, +) from parcels._core.utils.time import TimeInterval +from parcels._core.uxgrid import UxGrid +from parcels._core.xgrid import XGrid, _transpose_xfield_data_to_tzyx from parcels._reprs import default_repr from parcels._typing import VectorType -from parcels.application_kernels.interpolation import ( +from parcels.interpolators import ( UXPiecewiseLinearNode, XLinear, ZeroInterpolator, ZeroInterpolator_Vector, ) -from parcels.particle import KernelParticle -from parcels.tools._helpers import _assert_same_function_signature -from parcels.tools.converters import ( - UnitConverter, - unitconverters_map, -) -from parcels.tools.statuscodes import ( - AllParcelsErrorCodes, - StatusCode, -) -from parcels.uxgrid import UxGrid -from parcels.xgrid import LEFT_OUT_OF_BOUNDS, RIGHT_OUT_OF_BOUNDS, XGrid, _transpose_xfield_data_to_tzyx - -from ._index_search import GRID_SEARCH_ERROR, _search_time_index +from parcels.utils._helpers import _assert_same_function_signature __all__ = ["Field", "VectorField"] @@ -143,10 +142,10 @@ def __init__( self.igrid = -1 # Default the grid index to -1 - if self.grid._mesh == "flat" or (self.name not in unitconverters_map.keys()): + if self.grid._mesh == "flat" or (self.name not in _unitconverters_map.keys()): self.units = UnitConverter() elif self.grid._mesh == "spherical": - self.units = unitconverters_map[self.name] + self.units = _unitconverters_map[self.name] if self.data.shape[0] > 1: if "time" not in self.data.coords: diff --git a/parcels/fieldset.py b/parcels/_core/fieldset.py similarity index 98% rename from parcels/fieldset.py rename to parcels/_core/fieldset.py index 95df623493..bbc064a19e 100644 --- a/parcels/fieldset.py +++ b/parcels/_core/fieldset.py @@ -9,17 +9,17 @@ import xarray as xr import xgcm +from parcels._core.converters import Geographic, GeographicPolar +from parcels._core.field import Field, VectorField from parcels._core.utils.time import get_datetime_type_calendar from parcels._core.utils.time import is_compatible as datetime_is_compatible +from parcels._core.xgrid import _DEFAULT_XGCM_KWARGS, XGrid +from parcels._logger import logger from parcels._typing import Mesh -from parcels.field import Field, VectorField -from parcels.tools.converters import Geographic, GeographicPolar -from parcels.tools.loggers import logger -from parcels.xgrid import _DEFAULT_XGCM_KWARGS, XGrid if TYPE_CHECKING: + from parcels._core.basegrid import BaseGrid from parcels._typing import TimeLike - from parcels.basegrid import BaseGrid __all__ = ["FieldSet"] diff --git a/parcels/_index_search.py b/parcels/_core/index_search.py similarity index 84% rename from parcels/_index_search.py rename to parcels/_core/index_search.py index 08d73e560f..64689f49de 100644 --- a/parcels/_index_search.py +++ b/parcels/_core/index_search.py @@ -5,15 +5,61 @@ import numpy as np -from parcels.tools.statuscodes import _raise_time_extrapolation_error +from parcels._core.statuscodes import _raise_time_extrapolation_error if TYPE_CHECKING: + from parcels._core.field import Field from parcels.xgrid import XGrid - from .field import Field - GRID_SEARCH_ERROR = -3 +LEFT_OUT_OF_BOUNDS = -2 +RIGHT_OUT_OF_BOUNDS = -1 + + +def _search_1d_array( + arr: np.array, + x: float, +) -> tuple[int, int]: + """ + Searches for particle locations in a 1D array and returns barycentric coordinate along dimension. + + Assumptions: + - array is strictly monotonically increasing. + + Parameters + ---------- + arr : np.array + 1D array to search in. + x : float + Position in the 1D array to search for. + + Returns + ------- + array of int + Index of the element just before the position x in the array. Note that this index is -2 if the index is left out of bounds and -1 if the index is right out of bounds. + array of 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 np.zeros(shape=x.shape, dtype=np.int32), np.zeros_like(x) + index = np.searchsorted(arr, x, side="right") - 1 + # Use broadcasting to avoid repeated array access + arr_index = arr[index] + arr_next = arr[np.clip(index + 1, 1, len(arr) - 1)] # Ensure we don't go out of bounds + bcoord = (x - arr_index) / (arr_next - arr_index) + + # TODO check how we can avoid searchsorted when grid spacing is uniform + # dx = arr[1] - arr[0] + # index = ((x - arr[0]) / dx).astype(int) + # index = np.clip(index, 0, len(arr) - 2) + # bcoord = (x - arr[index]) / dx + + index = np.where(x < arr[0], LEFT_OUT_OF_BOUNDS, index) + index = np.where(x >= arr[-1], RIGHT_OUT_OF_BOUNDS, index) + + return np.atleast_1d(index), np.atleast_1d(bcoord) def _search_time_index(field: Field, time: datetime): diff --git a/parcels/kernel.py b/parcels/_core/kernel.py similarity index 97% rename from parcels/kernel.py rename to parcels/_core/kernel.py index 2f71263516..cbcc5d2cd0 100644 --- a/parcels/kernel.py +++ b/parcels/_core/kernel.py @@ -8,14 +8,8 @@ import numpy as np -from parcels.application_kernels.advection import ( - AdvectionAnalytical, - AdvectionRK4, - AdvectionRK45, -) -from parcels.basegrid import GridType -from parcels.tools._helpers import _assert_same_function_signature -from parcels.tools.statuscodes import ( +from parcels._core.basegrid import GridType +from parcels._core.statuscodes import ( StatusCode, _raise_field_interpolation_error, _raise_field_out_of_bound_error, @@ -24,7 +18,13 @@ _raise_grid_searching_error, _raise_time_extrapolation_error, ) -from parcels.tools.warnings import KernelWarning +from parcels._core.warnings import KernelWarning +from parcels.kernels import ( + AdvectionAnalytical, + AdvectionRK4, + AdvectionRK45, +) +from parcels.utils._helpers import _assert_same_function_signature if TYPE_CHECKING: from collections.abc import Callable diff --git a/parcels/particle.py b/parcels/_core/particle.py similarity index 99% rename from parcels/particle.py rename to parcels/_core/particle.py index 59c3457087..4d3ff09459 100644 --- a/parcels/particle.py +++ b/parcels/_core/particle.py @@ -8,9 +8,9 @@ import numpy as np from parcels._compat import _attrgetter_helper +from parcels._core.statuscodes import StatusCode from parcels._core.utils.time import TimeInterval from parcels._reprs import _format_list_items_multiline -from parcels.tools.statuscodes import StatusCode __all__ = ["KernelParticle", "Particle", "ParticleClass", "Variable"] _TO_WRITE_OPTIONS = [True, False, "once"] diff --git a/parcels/particlefile.py b/parcels/_core/particlefile.py similarity index 98% rename from parcels/particlefile.py rename to parcels/_core/particlefile.py index bcf865cee2..f3ae6c041e 100644 --- a/parcels/particlefile.py +++ b/parcels/_core/particlefile.py @@ -14,13 +14,13 @@ from zarr.storage import DirectoryStore import parcels -from parcels.particle import _SAME_AS_FIELDSET_TIME_INTERVAL, ParticleClass -from parcels.tools._helpers import timedelta_to_float +from parcels._core.particle import _SAME_AS_FIELDSET_TIME_INTERVAL, ParticleClass +from parcels.utils._helpers import timedelta_to_float if TYPE_CHECKING: + from parcels._core.particle import Variable + from parcels._core.particleset import ParticleSet from parcels._core.utils.time import TimeInterval - from parcels.particle import Variable - from parcels.particleset import ParticleSet __all__ = ["ParticleFile"] diff --git a/parcels/particleset.py b/parcels/_core/particleset.py similarity index 97% rename from parcels/particleset.py rename to parcels/_core/particleset.py index 0590dc0679..bdeef9ae78 100644 --- a/parcels/particleset.py +++ b/parcels/_core/particleset.py @@ -9,15 +9,15 @@ from tqdm import tqdm from zarr.storage import DirectoryStore +from parcels._core.converters import _convert_to_flat_array +from parcels._core.kernel import Kernel +from parcels._core.particle import KernelParticle, Particle, create_particle_data +from parcels._core.statuscodes import StatusCode from parcels._core.utils.time import TimeInterval, maybe_convert_python_timedelta_to_numpy +from parcels._core.warnings import ParticleSetWarning +from parcels._logger import logger from parcels._reprs import particleset_repr -from parcels.application_kernels.advection import AdvectionRK4 -from parcels.kernel import Kernel -from parcels.particle import KernelParticle, Particle, create_particle_data -from parcels.tools.converters import convert_to_flat_array -from parcels.tools.loggers import logger -from parcels.tools.statuscodes import StatusCode -from parcels.tools.warnings import ParticleSetWarning +from parcels.kernels import AdvectionRK4 __all__ = ["ParticleSet"] @@ -78,9 +78,9 @@ def __init__( self._interaction_kernel = None self.fieldset = fieldset - lon = np.empty(shape=0) if lon is None else convert_to_flat_array(lon) - lat = np.empty(shape=0) if lat is None else convert_to_flat_array(lat) - time = np.empty(shape=0) if time is None else convert_to_flat_array(time) + lon = np.empty(shape=0) if lon is None else _convert_to_flat_array(lon) + lat = np.empty(shape=0) if lat is None else _convert_to_flat_array(lat) + time = np.empty(shape=0) if time is None else _convert_to_flat_array(time) if trajectory_ids is None: trajectory_ids = np.arange(lon.size) @@ -92,7 +92,7 @@ def __init__( mindepth = min(mindepth, field.grid.depth[0]) depth = np.ones(lon.size) * mindepth else: - depth = convert_to_flat_array(depth) + depth = _convert_to_flat_array(depth) 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: @@ -114,7 +114,7 @@ def __init__( for kwvar in kwargs: if kwvar not in ["partition_function"]: - kwargs[kwvar] = convert_to_flat_array(kwargs[kwvar]) + kwargs[kwvar] = _convert_to_flat_array(kwargs[kwvar]) assert lon.size == kwargs[kwvar].size, ( f"{kwvar} and positions (lon, lat, depth) don't have the same lengths." ) diff --git a/parcels/spatialhash.py b/parcels/_core/spatialhash.py similarity index 97% rename from parcels/spatialhash.py rename to parcels/_core/spatialhash.py index 5e4aaa4aa1..bb13851672 100644 --- a/parcels/spatialhash.py +++ b/parcels/_core/spatialhash.py @@ -1,18 +1,23 @@ import numpy as np -import parcels -from parcels._index_search import GRID_SEARCH_ERROR, _latlon_rad_to_xyz, curvilinear_point_in_cell, uxgrid_point_in_cell +from parcels._core.index_search import ( + GRID_SEARCH_ERROR, + _latlon_rad_to_xyz, + curvilinear_point_in_cell, + uxgrid_point_in_cell, +) +from parcels._python import isinstance_noimport class SpatialHash: """Custom data structure that is used for performing grid searches using Spatial Hashing. This class constructs an overlying - uniformly spaced rectilinear grid, called the "hash grid" on top parcels.xgrid.XGrid. It is particularly useful for grid searching + uniformly spaced rectilinear grid, called the "hash grid" on top parcels.XGrid. It is particularly useful for grid searching on curvilinear grids. Faces in the Xgrid are related to the cells in the hash grid by determining the hash cells the bounding box of the unstructured face cells overlap with. Parameters ---------- - grid : parcels.xgrid.XGrid + grid : parcels.XGrid Source grid used to construct the hash grid and hash table Note @@ -25,17 +30,17 @@ def __init__( grid, bitwidth=1023, ): - if isinstance(grid, parcels.xgrid.XGrid): + if isinstance_noimport(grid, "XGrid"): self._point_in_cell = curvilinear_point_in_cell - elif isinstance(grid, parcels.uxgrid.UxGrid): + elif isinstance_noimport(grid, "UxGrid"): self._point_in_cell = uxgrid_point_in_cell else: - raise ValueError("Expected `grid` to be a parcels.xgrid.XGrid or parcels.uxgrid.UxGrid") + raise ValueError("Expected `grid` to be a parcels.XGrid or parcels.UxGrid") self._source_grid = grid self._bitwidth = bitwidth # Max integer to use per coordinate in quantization (10 bits = 0..1023) - if isinstance(grid, parcels.xgrid.XGrid): + if isinstance_noimport(grid, "XGrid"): if self._source_grid._mesh == "spherical": # Boundaries of the hash grid are the unit cube self._xmin = -1.0 @@ -120,7 +125,7 @@ def __init__( self._zlow = np.zeros_like(self._xlow) self._zhigh = np.zeros_like(self._xlow) - elif isinstance(grid, parcels.uxgrid.UxGrid): + elif isinstance_noimport(grid, "UxGrid"): if self._source_grid._mesh == "spherical": # Boundaries of the hash grid are the unit cube self._xmin = -1.0 diff --git a/parcels/tools/statuscodes.py b/parcels/_core/statuscodes.py similarity index 100% rename from parcels/tools/statuscodes.py rename to parcels/_core/statuscodes.py diff --git a/parcels/uxgrid.py b/parcels/_core/uxgrid.py similarity index 96% rename from parcels/uxgrid.py rename to parcels/_core/uxgrid.py index 3aa91e1777..11da6daffa 100644 --- a/parcels/uxgrid.py +++ b/parcels/_core/uxgrid.py @@ -5,11 +5,9 @@ import numpy as np import uxarray as ux -from parcels._index_search import GRID_SEARCH_ERROR, uxgrid_point_in_cell +from parcels._core.basegrid import BaseGrid +from parcels._core.index_search import GRID_SEARCH_ERROR, _search_1d_array, uxgrid_point_in_cell from parcels._typing import assert_valid_mesh -from parcels.xgrid import _search_1d_array - -from .basegrid import BaseGrid _UXGRID_AXES = Literal["Z", "FACE"] diff --git a/parcels/tools/warnings.py b/parcels/_core/warnings.py similarity index 100% rename from parcels/tools/warnings.py rename to parcels/_core/warnings.py diff --git a/parcels/xgrid.py b/parcels/_core/xgrid.py similarity index 91% rename from parcels/xgrid.py rename to parcels/_core/xgrid.py index 83bbb46f74..d20d53f8b3 100644 --- a/parcels/xgrid.py +++ b/parcels/_core/xgrid.py @@ -7,9 +7,9 @@ import xarray as xr import xgcm -from parcels._index_search import _search_indices_curvilinear_2d +from parcels._core.basegrid import BaseGrid +from parcels._core.index_search import _search_1d_array, _search_indices_curvilinear_2d from parcels._typing import assert_valid_mesh -from parcels.basegrid import BaseGrid _XGRID_AXES = Literal["X", "Y", "Z"] _XGRID_AXES_ORDERING: Sequence[_XGRID_AXES] = "ZYX" @@ -22,9 +22,6 @@ _DEFAULT_XGCM_KWARGS = {"periodic": False} -LEFT_OUT_OF_BOUNDS = -2 -RIGHT_OUT_OF_BOUNDS = -1 - def get_cell_count_along_dim(ds: xr.Dataset, axis: xgcm.Axis) -> int: first_coord = list(axis.coords.items())[0] @@ -477,51 +474,6 @@ def assert_valid_lat_lon(da_lat, da_lon, axes: _XGCM_AXES): ) -def _search_1d_array( - arr: np.array, - x: float, -) -> tuple[int, int]: - """ - Searches for particle locations in a 1D array and returns barycentric coordinate along dimension. - - Assumptions: - - array is strictly monotonically increasing. - - Parameters - ---------- - arr : np.array - 1D array to search in. - x : float - Position in the 1D array to search for. - - Returns - ------- - array of int - Index of the element just before the position x in the array. Note that this index is -2 if the index is left out of bounds and -1 if the index is right out of bounds. - array of 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 np.zeros(shape=x.shape, dtype=np.int32), np.zeros_like(x) - index = np.searchsorted(arr, x, side="right") - 1 - # Use broadcasting to avoid repeated array access - arr_index = arr[index] - arr_next = arr[np.clip(index + 1, 1, len(arr) - 1)] # Ensure we don't go out of bounds - bcoord = (x - arr_index) / (arr_next - arr_index) - - # TODO check how we can avoid searchsorted when grid spacing is uniform - # dx = arr[1] - arr[0] - # index = ((x - arr[0]) / dx).astype(int) - # index = np.clip(index, 0, len(arr) - 2) - # bcoord = (x - arr[index]) / dx - - index = np.where(x < arr[0], LEFT_OUT_OF_BOUNDS, index) - index = np.where(x >= arr[-1], RIGHT_OUT_OF_BOUNDS, index) - - return np.atleast_1d(index), np.atleast_1d(bcoord) - - def _convert_center_pos_to_fpoint( *, index: int, bcoord: float, xgcm_position: _XGCM_AXIS_POSITION, f_points_xgcm_position: _XGCM_AXIS_POSITION ) -> tuple[int, float]: diff --git a/parcels/_decorators.py b/parcels/_decorators.py new file mode 100644 index 0000000000..697696a190 --- /dev/null +++ b/parcels/_decorators.py @@ -0,0 +1,60 @@ +"""Utilities to help with deprecations.""" + +from __future__ import annotations + +import functools +import warnings +from collections.abc import Callable + +PACKAGE = "Parcels" + + +def deprecated(msg: str = "") -> Callable: + """Decorator marking a function as being deprecated + + Parameters + ---------- + msg : str, optional + Custom message to append to the deprecation warning. + + Examples + -------- + ``` + @deprecated("Please use `another_function` instead") + def some_old_function(x, y): + return x + y + + @deprecated() + def some_other_old_function(x, y): + return x + y + ``` + """ + if msg: + msg = " " + msg + + def decorator(func): + @functools.wraps(func) + def wrapper(*args, **kwargs): + msg_formatted = ( + f"`{func.__qualname__}` is deprecated and will be removed in a future release of {PACKAGE}.{msg}" + ) + + warnings.warn(msg_formatted, category=DeprecationWarning, stacklevel=3) + return func(*args, **kwargs) + + _patch_docstring(wrapper, f"\n\n.. deprecated:: {msg}") + return wrapper + + return decorator + + +def deprecated_made_private(func: Callable) -> Callable: + return deprecated( + "It has moved to the internal API as it is not expected to be directly used by " + "the end-user. If you feel that you use this code directly in your scripts, please " + "comment on our tracking issue at https://github.com/OceanParcels/Parcels/issues/1695.", + )(func) + + +def _patch_docstring(obj: Callable, extra: str) -> None: + obj.__doc__ = f"{obj.__doc__ or ''}{extra}".strip() diff --git a/parcels/_interpolation.py b/parcels/_interpolation.py index 88a2864dca..e38f0d0d0f 100644 --- a/parcels/_interpolation.py +++ b/parcels/_interpolation.py @@ -4,7 +4,7 @@ import numpy as np from parcels._typing import GridIndexingType -from parcels.tools._helpers import should_calculate_next_ti +from parcels.utils._helpers import should_calculate_next_ti @dataclass diff --git a/parcels/tools/loggers.py b/parcels/_logger.py similarity index 100% rename from parcels/tools/loggers.py rename to parcels/_logger.py diff --git a/parcels/_python.py b/parcels/_python.py new file mode 100644 index 0000000000..91f5b46458 --- /dev/null +++ b/parcels/_python.py @@ -0,0 +1,28 @@ +# Generic Python helpers + + +def isinstance_noimport(obj, class_or_tuple): + """A version of isinstance that does not require importing the class. + This is useful to avoid circular imports. + """ + return ( + type(obj).__name__ == class_or_tuple + if isinstance(class_or_tuple, str) + else type(obj).__name__ in class_or_tuple + ) + + +def test_isinstance_noimport(): + class A: + pass + + class B: + pass + + a = A() + b = B() + + assert isinstance_noimport(a, "A") + assert not isinstance_noimport(a, "B") + assert isinstance_noimport(b, ("A", "B")) + assert not isinstance_noimport(b, "C") diff --git a/parcels/tools/exampledata_utils.py b/parcels/_tutorial.py similarity index 97% rename from parcels/tools/exampledata_utils.py rename to parcels/_tutorial.py index 5513987770..00dc3bae61 100644 --- a/parcels/tools/exampledata_utils.py +++ b/parcels/_tutorial.py @@ -5,7 +5,7 @@ import pooch import xarray as xr -from parcels.tools._v3to4 import patch_dataset_v4_compat +from parcels._v3to4 import patch_dataset_v4_compat __all__ = ["download_example_dataset", "list_example_datasets"] @@ -180,12 +180,12 @@ def download_example_dataset(dataset: str, data_home=None): for file_name in odie.registry: if file_name.startswith(dataset): should_patch = dataset == "GlobCurrent_example_data" - odie.fetch(file_name, processor=v4_compat_patch if should_patch else None) + odie.fetch(file_name, processor=_v4_compat_patch if should_patch else None) return dataset_folder -def v4_compat_patch(fname, action, pup): +def _v4_compat_patch(fname, action, pup): """ Patch the GlobCurrent example dataset to be compatible with v4. diff --git a/parcels/tools/_v3to4.py b/parcels/_v3to4.py similarity index 100% rename from parcels/tools/_v3to4.py rename to parcels/_v3to4.py diff --git a/parcels/application_kernels/__init__.py b/parcels/application_kernels/__init__.py deleted file mode 100644 index cd933324ed..0000000000 --- a/parcels/application_kernels/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .advection import * -from .advectiondiffusion import * -from .interaction import * -from .interpolation import * diff --git a/parcels/application_kernels/interpolation.py b/parcels/interpolators.py similarity index 99% rename from parcels/application_kernels/interpolation.py rename to parcels/interpolators.py index 8caf5231a3..4a2f9ad293 100644 --- a/parcels/application_kernels/interpolation.py +++ b/parcels/interpolators.py @@ -8,12 +8,12 @@ import xarray as xr from dask import is_dask_collection -import parcels.tools.interpolation_utils as i_u +import parcels.utils.interpolation_utils as i_u if TYPE_CHECKING: - from parcels.field import Field, VectorField - from parcels.uxgrid import _UXGRID_AXES - from parcels.xgrid import _XGRID_AXES + from parcels._core.field import Field, VectorField + from parcels._core.uxgrid import _UXGRID_AXES + from parcels._core.xgrid import _XGRID_AXES __all__ = [ "CGrid_Tracer", diff --git a/parcels/application_kernels/EOSseawaterproperties.py b/parcels/kernels/EOSseawaterproperties.py similarity index 100% rename from parcels/application_kernels/EOSseawaterproperties.py rename to parcels/kernels/EOSseawaterproperties.py diff --git a/parcels/application_kernels/TEOSseawaterdensity.py b/parcels/kernels/TEOSseawaterdensity.py similarity index 100% rename from parcels/application_kernels/TEOSseawaterdensity.py rename to parcels/kernels/TEOSseawaterdensity.py diff --git a/parcels/kernels/__init__.py b/parcels/kernels/__init__.py new file mode 100644 index 0000000000..02e3e1f2f7 --- /dev/null +++ b/parcels/kernels/__init__.py @@ -0,0 +1,36 @@ +from .advection import ( + AdvectionAnalytical, + AdvectionEE, + AdvectionRK4, + AdvectionRK4_3D, + AdvectionRK4_3D_CROCO, + AdvectionRK45, +) +from .advectiondiffusion import ( + AdvectionDiffusionEM, + AdvectionDiffusionM1, + DiffusionUniformKh, +) +from .interaction import ( + AsymmetricAttraction, + MergeWithNearestNeighbor, + NearestNeighborWithinRange, +) + +__all__ = [ # noqa: RUF022 + # advection + "AdvectionAnalytical", + "AdvectionEE", + "AdvectionRK4_3D_CROCO", + "AdvectionRK4_3D", + "AdvectionRK4", + "AdvectionRK45", + # advectiondiffusion + "AdvectionDiffusionEM", + "AdvectionDiffusionM1", + "DiffusionUniformKh", + # interaction + "AsymmetricAttraction", + "MergeWithNearestNeighbor", + "NearestNeighborWithinRange", +] diff --git a/parcels/application_kernels/advection.py b/parcels/kernels/advection.py similarity index 99% rename from parcels/application_kernels/advection.py rename to parcels/kernels/advection.py index 1023f8d88e..8d21392903 100644 --- a/parcels/application_kernels/advection.py +++ b/parcels/kernels/advection.py @@ -4,7 +4,7 @@ import numpy as np -from parcels.tools.statuscodes import StatusCode +from parcels._core.statuscodes import StatusCode __all__ = [ "AdvectionAnalytical", @@ -193,7 +193,7 @@ def AdvectionAnalytical(particles, fieldset): # pragma: no cover """ import numpy as np - import parcels.tools.interpolation_utils as i_u + import parcels.utils.interpolation_utils as i_u tol = 1e-10 I_s = 10 # number of intermediate time steps diff --git a/parcels/application_kernels/advectiondiffusion.py b/parcels/kernels/advectiondiffusion.py similarity index 100% rename from parcels/application_kernels/advectiondiffusion.py rename to parcels/kernels/advectiondiffusion.py diff --git a/parcels/application_kernels/interaction.py b/parcels/kernels/interaction.py similarity index 98% rename from parcels/application_kernels/interaction.py rename to parcels/kernels/interaction.py index 41f4b18dd8..2eb919ef9d 100644 --- a/parcels/application_kernels/interaction.py +++ b/parcels/kernels/interaction.py @@ -2,7 +2,7 @@ import numpy as np -from parcels.tools.statuscodes import StatusCode +from parcels._core.statuscodes import StatusCode __all__ = ["AsymmetricAttraction", "MergeWithNearestNeighbor", "NearestNeighborWithinRange"] diff --git a/parcels/tools/__init__.py b/parcels/tools/__init__.py deleted file mode 100644 index 6909c6fb44..0000000000 --- a/parcels/tools/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -from .converters import * -from .exampledata_utils import * -from .interpolation_utils import * -from .loggers import * -from .statuscodes import * -from .timer import * -from .warnings import * diff --git a/parcels/_core/utils/structured.py b/parcels/utils/__init__.py similarity index 100% rename from parcels/_core/utils/structured.py rename to parcels/utils/__init__.py diff --git a/parcels/tools/_helpers.py b/parcels/utils/_helpers.py similarity index 51% rename from parcels/tools/_helpers.py rename to parcels/utils/_helpers.py index 7432572767..0b078110e0 100644 --- a/parcels/tools/_helpers.py +++ b/parcels/utils/_helpers.py @@ -2,9 +2,7 @@ from __future__ import annotations -import functools import inspect -import warnings from collections.abc import Callable from datetime import timedelta @@ -13,57 +11,6 @@ PACKAGE = "Parcels" -def deprecated(msg: str = "") -> Callable: - """Decorator marking a function as being deprecated - - Parameters - ---------- - msg : str, optional - Custom message to append to the deprecation warning. - - Examples - -------- - ``` - @deprecated("Please use `another_function` instead") - def some_old_function(x, y): - return x + y - - @deprecated() - def some_other_old_function(x, y): - return x + y - ``` - """ - if msg: - msg = " " + msg - - def decorator(func): - @functools.wraps(func) - def wrapper(*args, **kwargs): - msg_formatted = ( - f"`{func.__qualname__}` is deprecated and will be removed in a future release of {PACKAGE}.{msg}" - ) - - warnings.warn(msg_formatted, category=DeprecationWarning, stacklevel=3) - return func(*args, **kwargs) - - patch_docstring(wrapper, f"\n\n.. deprecated:: {msg}") - return wrapper - - return decorator - - -def deprecated_made_private(func: Callable) -> Callable: - return deprecated( - "It has moved to the internal API as it is not expected to be directly used by " - "the end-user. If you feel that you use this code directly in your scripts, please " - "comment on our tracking issue at https://github.com/OceanParcels/Parcels/issues/1695.", - )(func) - - -def patch_docstring(obj: Callable, extra: str) -> None: - obj.__doc__ = f"{obj.__doc__ or ''}{extra}".strip() - - def timedelta_to_float(dt: float | timedelta | np.timedelta64) -> float: """Convert a timedelta to a float in seconds.""" if isinstance(dt, timedelta): diff --git a/parcels/tools/interpolation_utils.py b/parcels/utils/interpolation_utils.py similarity index 100% rename from parcels/tools/interpolation_utils.py rename to parcels/utils/interpolation_utils.py diff --git a/parcels/tools/timer.py b/parcels/utils/timer.py similarity index 100% rename from parcels/tools/timer.py rename to parcels/utils/timer.py diff --git a/pyproject.toml b/pyproject.toml index 06436fb899..5dc6626461 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -51,8 +51,8 @@ local_scheme = "no-local-version" [tool.pytest.ini_options] addopts = ["--strict-config", "--strict-markers"] xfail_strict = true -# testpaths = ["tests", "docs/examples"] # TODO v4: Re-enable once `tests/v4` separation is gone -testpaths = ["tests/v4"] +# testpaths = ["tests", "docs/examples"] # TODO v4: Re-enable once examples are back/fixed +testpaths = ["tests"] python_files = ["test_*.py", "example_*.py", "*tutorial*"] minversion = "7" markers = [ # can be skipped by doing `pytest -m "not slow"` etc. diff --git a/tests-v3/test_advection.py b/tests-v3/test_advection.py new file mode 100644 index 0000000000..de8b004d79 --- /dev/null +++ b/tests-v3/test_advection.py @@ -0,0 +1,179 @@ +import numpy as np +import pytest +import xarray as xr + +from parcels import ( + AdvectionAnalytical, + AdvectionDiffusionEM, + AdvectionDiffusionM1, + AdvectionEE, + AdvectionRK4, + AdvectionRK4_3D, + AdvectionRK45, + FieldSet, + Particle, + ParticleSet, + Variable, +) +from tests.utils import TEST_DATA + +kernel = { + "EE": AdvectionEE, + "RK4": AdvectionRK4, + "RK45": AdvectionRK45, + "AA": AdvectionAnalytical, + "AdvDiffEM": AdvectionDiffusionEM, + "AdvDiffM1": AdvectionDiffusionM1, +} + + +@pytest.fixture +def lon(): + xdim = 200 + return np.linspace(-170, 170, xdim, dtype=np.float32) + + +@pytest.fixture +def lat(): + ydim = 100 + return np.linspace(-80, 80, ydim, dtype=np.float32) + + +@pytest.fixture +def depth(): + zdim = 2 + return np.linspace(0, 30, zdim, dtype=np.float32) + + +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="When refactoring fieldfilebuffer croco support was dropped. This will be fixed in v4.") +def test_conversion_3DCROCO(): + """Test of the (SciPy) version of the conversion from depth to sigma in CROCO + + Values below are retrieved using xroms and hardcoded in the method (to avoid dependency on xroms): + ```py + x, y = 10, 20 + s_xroms = ds.s_w.values + z_xroms = ds.z_w.isel(time=0).isel(eta_rho=y).isel(xi_rho=x).values + lat, lon = ds.y_rho.values[y, x], ds.x_rho.values[y, x] + ``` + """ + fieldset = FieldSet.from_modulefile(TEST_DATA / "fieldset_CROCO3D.py") + + lat, lon = 78000.0, 38000.0 + s_xroms = np.array([-1.0, -0.9, -0.8, -0.7, -0.6, -0.5, -0.4, -0.3, -0.2, -0.1, 0.0], dtype=np.float32) + z_xroms = np.array( + [ + -1.26000000e02, + -1.10585846e02, + -9.60985413e01, + -8.24131317e01, + -6.94126511e01, + -5.69870148e01, + -4.50318756e01, + -3.34476166e01, + -2.21383114e01, + -1.10107975e01, + 2.62768921e-02, + ], + dtype=np.float32, + ) + + sigma = np.zeros_like(z_xroms) + from parcels.field import _croco_from_z_to_sigma_scipy + + for zi, z in enumerate(z_xroms): + sigma[zi] = _croco_from_z_to_sigma_scipy(fieldset, 0, z, lat, lon, None) + + assert np.allclose(sigma, s_xroms, atol=1e-3) + + +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="CROCO 3D interpolation is not yet implemented correctly in v4. ") +def test_advection_3DCROCO(): + fieldset = FieldSet.from_modulefile(TEST_DATA / "fieldset_CROCO3D.py") + + runtime = 1e4 + X, Z = np.meshgrid([40e3, 80e3, 120e3], [-10, -130]) + Y = np.ones(X.size) * 100e3 + + pclass = Particle.add_variable(Variable("w")) + pset = ParticleSet(fieldset=fieldset, pclass=pclass, lon=X, lat=Y, depth=Z) + + def SampleW(particle, fieldset, time): # pragma: no cover + particle.w = fieldset.W[time, particle.depth, particle.lat, particle.lon] + + pset.execute([AdvectionRK4_3D, SampleW], runtime=runtime, dt=100) + assert np.allclose(pset.depth, Z.flatten(), atol=5) # TODO lower this atol + assert np.allclose(pset.lon_nextloop, [x + runtime for x in X.flatten()], atol=1e-3) + + +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="When refactoring fieldfilebuffer croco support was dropped. This will be fixed in v4.") +def test_advection_2DCROCO(): + fieldset = FieldSet.from_modulefile(TEST_DATA / "fieldset_CROCO2D.py") + + runtime = 1e4 + X = np.array([40e3, 80e3, 120e3]) + Y = np.ones(X.size) * 100e3 + Z = np.zeros(X.size) + pset = ParticleSet(fieldset=fieldset, pclass=Particle, lon=X, lat=Y, depth=Z) + + pset.execute([AdvectionRK4], runtime=runtime, dt=100) + assert np.allclose(pset.depth, Z.flatten(), atol=1e-3) + assert np.allclose(pset.lon_nextloop, [x + runtime for x in X], atol=1e-3) + + +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") +def test_analyticalAgrid(): + lon = np.arange(0, 15, dtype=np.float32) + lat = np.arange(0, 15, dtype=np.float32) + U = np.ones((lat.size, lon.size), dtype=np.float32) + V = np.ones((lat.size, lon.size), dtype=np.float32) + fieldset = FieldSet.from_data({"U": U, "V": V}, {"lon": lon, "lat": lat}, mesh="flat") + pset = ParticleSet(fieldset, pclass=Particle, lon=1, lat=1) + + with pytest.raises(NotImplementedError): + pset.execute(AdvectionAnalytical, runtime=1) + + +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1927") +@pytest.mark.parametrize("u", [1, -0.2, -0.3, 0]) +@pytest.mark.parametrize("v", [1, -0.3, 0, -1]) +@pytest.mark.parametrize("w", [None, 1, -0.3, 0, -1]) +@pytest.mark.parametrize("direction", [1, -1]) +def test_uniform_analytical(u, v, w, direction, tmp_zarrfile): + lon = np.arange(0, 15, dtype=np.float32) + lat = np.arange(0, 15, dtype=np.float32) + if w is not None: + depth = np.arange(0, 40, 2, dtype=np.float32) + U = u * np.ones((depth.size, lat.size, lon.size), dtype=np.float32) + V = v * np.ones((depth.size, lat.size, lon.size), dtype=np.float32) + W = w * np.ones((depth.size, lat.size, lon.size), dtype=np.float32) + fieldset = FieldSet.from_data({"U": U, "V": V, "W": W}, {"lon": lon, "lat": lat, "depth": depth}, mesh="flat") + fieldset.W.interp_method = "cgrid_velocity" + else: + U = u * np.ones((lat.size, lon.size), dtype=np.float32) + V = v * np.ones((lat.size, lon.size), dtype=np.float32) + fieldset = FieldSet.from_data({"U": U, "V": V}, {"lon": lon, "lat": lat}, mesh="flat") + fieldset.U.interp_method = "cgrid_velocity" + fieldset.V.interp_method = "cgrid_velocity" + + x0, y0, z0 = 6.1, 6.2, 20 + pset = ParticleSet(fieldset, pclass=Particle, lon=x0, lat=y0, depth=z0) + + outfile = pset.ParticleFile(name=tmp_zarrfile, outputdt=1, chunks=(1, 1)) + pset.execute(AdvectionAnalytical, runtime=4, dt=direction, output_file=outfile) + assert np.abs(pset.lon - x0 - pset.time * u) < 1e-6 + assert np.abs(pset.lat - y0 - pset.time * v) < 1e-6 + if w is not None: + assert np.abs(pset.depth - z0 - pset.time * w) < 1e-4 + + ds = xr.open_zarr(tmp_zarrfile) + times = (direction * ds["time"][:]).values.astype("timedelta64[s]")[0] + timeref = np.arange(1, 5).astype("timedelta64[s]") + assert np.allclose(times, timeref, atol=np.timedelta64(1, "ms")) + lons = ds["lon"][:].values + assert np.allclose(lons, x0 + direction * u * np.arange(1, 5)) diff --git a/tests/test_examples.py b/tests-v3/test_examples.py similarity index 100% rename from tests/test_examples.py rename to tests-v3/test_examples.py diff --git a/tests-v3/test_fieldset.py b/tests-v3/test_fieldset.py new file mode 100644 index 0000000000..fd83a28bef --- /dev/null +++ b/tests-v3/test_fieldset.py @@ -0,0 +1,417 @@ +from datetime import timedelta + +import numpy as np +import pytest +import xarray as xr + +from parcels import ( + AdvectionRK4, + AdvectionRK4_3D, + FieldSet, + Particle, + ParticleSet, + Variable, +) +from parcels.field import VectorField +from parcels.tools.converters import GeographicPolar, UnitConverter +from tests.utils import TEST_DATA + + +def generate_fieldset_data(xdim, ydim, zdim=1, tdim=1): + lon = np.linspace(0.0, 10.0, xdim, dtype=np.float32) + lat = np.linspace(0.0, 10.0, ydim, dtype=np.float32) + depth = np.zeros(zdim, dtype=np.float32) + time = np.zeros(tdim, dtype=np.float64) + if zdim == 1 and tdim == 1: + U, V = np.meshgrid(lon, lat) + dimensions = {"lat": lat, "lon": lon} + else: + U = np.ones((tdim, zdim, ydim, xdim)) + V = np.ones((tdim, zdim, ydim, xdim)) + dimensions = {"lat": lat, "lon": lon, "depth": depth, "time": time} + data = {"U": np.array(U, dtype=np.float32), "V": np.array(V, dtype=np.float32)} + + return (data, dimensions) + + +def to_xarray_dataset(data: dict[str, np.array], dimensions: dict[str, np.array]) -> xr.Dataset: + assert len(dimensions) in [2, 4], "this function only deals with output from generate_fieldset_data()" + + if len(dimensions) == 4: + return xr.Dataset( + { + "U": (["time", "depth", "lat", "lon"], data["U"]), + "V": (["time", "depth", "lat", "lon"], data["V"]), + }, + coords=dimensions, + ) + + return xr.Dataset( + { + "U": (["lat", "lon"], data["U"]), + "V": (["lat", "lon"], data["V"]), + }, + coords=dimensions, + ) + + +@pytest.mark.v4remove +@pytest.mark.xfail(reason="GH1946") +@pytest.fixture +def multifile_fieldset(tmp_path): + stem = "test_subsets" + + timestamps = np.arange(0, 4, 1) * 86400.0 + timestamps = np.expand_dims(timestamps, 1) + + ufiles = [] + vfiles = [] + for index, timestamp in enumerate(timestamps): + data, dimensions = generate_fieldset_data(100, 100) + path = tmp_path / f"{stem}_{index}.nc" + to_xarray_dataset(data, dimensions).pipe(assign_dataset_timestamp_dim, timestamp).to_netcdf(path) + ufiles.append(path) + vfiles.append(path) + + files = {"U": ufiles, "V": vfiles} + variables = {"U": "U", "V": "V"} + dimensions = {"lon": "lon", "lat": "lat", "time": "time"} + return FieldSet.from_netcdf(files, variables, dimensions) + + +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") +def test_fieldset_from_modulefile(): + nemo_fname = str(TEST_DATA / "fieldset_nemo.py") + nemo_error_fname = str(TEST_DATA / "fieldset_nemo_error.py") + + fieldset = FieldSet.from_modulefile(nemo_fname) + + fieldset = FieldSet.from_modulefile(nemo_fname) + assert fieldset.U.grid.lon.shape[1] == 21 + + with pytest.raises(IOError): + FieldSet.from_modulefile(nemo_error_fname) + + FieldSet.from_modulefile(nemo_error_fname, modulename="random_function_name") + + with pytest.raises(IOError): + FieldSet.from_modulefile(nemo_error_fname, modulename="none_returning_function") + + +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") +def test_field_from_netcdf_fieldtypes(): + filenames = { + "varU": { + "lon": str(TEST_DATA / "mask_nemo_cross_180lon.nc"), + "lat": str(TEST_DATA / "mask_nemo_cross_180lon.nc"), + "data": str(TEST_DATA / "Uu_eastward_nemo_cross_180lon.nc"), + }, + "varV": { + "lon": str(TEST_DATA / "mask_nemo_cross_180lon.nc"), + "lat": str(TEST_DATA / "mask_nemo_cross_180lon.nc"), + "data": str(TEST_DATA / "Vv_eastward_nemo_cross_180lon.nc"), + }, + } + variables = {"varU": "U", "varV": "V"} + dimensions = {"lon": "glamf", "lat": "gphif"} + + # first try without setting fieldtype + fset = FieldSet.from_nemo(filenames, variables, dimensions) + assert isinstance(fset.varU.units, UnitConverter) + + # now try with setting fieldtype + fset = FieldSet.from_nemo(filenames, variables, dimensions, fieldtype={"varU": "U", "varV": "V"}) + assert isinstance(fset.varU.units, GeographicPolar) + + +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") +def test_fieldset_from_agrid_dataset(): + filenames = { + "lon": str(TEST_DATA / "mask_nemo_cross_180lon.nc"), + "lat": str(TEST_DATA / "mask_nemo_cross_180lon.nc"), + "data": str(TEST_DATA / "Uu_eastward_nemo_cross_180lon.nc"), + } + variable = {"U": "U"} + dimensions = {"lon": "glamf", "lat": "gphif"} + FieldSet.from_a_grid_dataset(filenames, variable, dimensions) + + +@pytest.mark.v4remove +@pytest.mark.xfail(reason="GH1946") +def test_fieldset_from_cgrid_interpmethod(): + filenames = { + "lon": str(TEST_DATA / "mask_nemo_cross_180lon.nc"), + "lat": str(TEST_DATA / "mask_nemo_cross_180lon.nc"), + "data": str(TEST_DATA / "Uu_eastward_nemo_cross_180lon.nc"), + } + variable = "U" + dimensions = {"lon": "glamf", "lat": "gphif"} + + with pytest.raises(TypeError): + # should fail because FieldSet.from_c_grid_dataset does not support interp_method + FieldSet.from_c_grid_dataset(filenames, variable, dimensions, interp_method="partialslip") + + +@pytest.mark.v4future +@pytest.mark.xfail(reason="GH1946") +@pytest.mark.parametrize("calltype", ["from_nemo"]) +def test_illegal_dimensionsdict(calltype): + with pytest.raises(NameError): + if calltype == "from_data": + data, dimensions = generate_fieldset_data(10, 10) + dimensions["test"] = None + FieldSet.from_data(data, dimensions) + elif calltype == "from_nemo": + fname = str(TEST_DATA / "mask_nemo_cross_180lon.nc") + filenames = {"dx": fname, "mesh_mask": fname} + variables = {"dx": "e1u"} + dimensions = {"lon": "glamu", "lat": "gphiu", "test": "test"} + FieldSet.from_nemo(filenames, variables, dimensions) + + +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") +@pytest.mark.parametrize("gridtype", ["A", "C"]) +def test_fieldset_dimlength1_cgrid(gridtype): + fieldset = FieldSet.from_data({"U": 0, "V": 0}, {"lon": 0, "lat": 0}) # TODO : Remove from_data + if gridtype == "C": + fieldset.U.interp_method = "cgrid_velocity" + fieldset.V.interp_method = "cgrid_velocity" + try: + fieldset._check_complete() + success = True if gridtype == "A" else False + except NotImplementedError: + success = True if gridtype == "C" else False + assert success + + +def assign_dataset_timestamp_dim(ds, timestamp): + """Expand dim to 'time' and assign timestamp.""" + ds.expand_dims("time") + ds["time"] = timestamp + return ds + + +def addConst(particle, fieldset, time): # pragma: no cover + particle.lon = particle.lon + fieldset.movewest + fieldset.moveeast + + +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") +def test_fieldset_constant(): + data, dimensions = generate_fieldset_data(100, 100) + fieldset = FieldSet.from_data(data, dimensions) # TODO : Remove from_data + westval = -0.2 + eastval = 0.3 + fieldset.add_constant("movewest", westval) + fieldset.add_constant("moveeast", eastval) + assert fieldset.movewest == westval + + pset = ParticleSet.from_line(fieldset, size=1, pclass=Particle, start=(0.5, 0.5), finish=(0.5, 0.5)) + pset.execute(pset.Kernel(addConst), dt=1, runtime=1) + assert abs(pset.lon[0] - (0.5 + westval + eastval)) < 1e-4 + + +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") +@pytest.mark.parametrize("swapUV", [False, True]) +def test_vector_fields(swapUV): + lon = np.linspace(0.0, 10.0, 12, dtype=np.float32) + lat = np.linspace(0.0, 10.0, 10, dtype=np.float32) + U = np.ones((10, 12), dtype=np.float32) + V = np.zeros((10, 12), dtype=np.float32) + data = {"U": U, "V": V} + dimensions = {"U": {"lat": lat, "lon": lon}, "V": {"lat": lat, "lon": lon}} + fieldset = FieldSet.from_data(data, dimensions, mesh="flat") # TODO : Remove from_data + if swapUV: # we test that we can freely edit whatever UV field + UV = VectorField("UV", fieldset.V, fieldset.U) + fieldset.add_vector_field(UV) + + pset = ParticleSet.from_line(fieldset, size=1, pclass=Particle, start=(0.5, 0.5), finish=(0.5, 0.5)) + pset.execute(AdvectionRK4, dt=1, runtime=2) + if swapUV: + assert abs(pset.lon[0] - 0.5) < 1e-9 + assert abs(pset.lat[0] - 1.5) < 1e-9 + else: + assert abs(pset.lon[0] - 1.5) < 1e-9 + assert abs(pset.lat[0] - 0.5) < 1e-9 + + +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946, originated in GH938") +def test_add_second_vector_field(): + lon = np.linspace(0.0, 10.0, 12, dtype=np.float32) + lat = np.linspace(0.0, 10.0, 10, dtype=np.float32) + U = np.ones((10, 12), dtype=np.float32) + V = np.zeros((10, 12), dtype=np.float32) + data = {"U": U, "V": V} + dimensions = {"U": {"lat": lat, "lon": lon}, "V": {"lat": lat, "lon": lon}} + fieldset = FieldSet.from_data(data, dimensions, mesh="flat") # TODO : Remove from_data + + data2 = {"U2": U, "V2": V} + dimensions2 = {"lon": [ln + 0.1 for ln in lon], "lat": [lt - 0.1 for lt in lat]} + fieldset2 = FieldSet.from_data(data2, dimensions2, mesh="flat") # TODO : Remove from_data + + UV2 = VectorField("UV2", fieldset2.U2, fieldset2.V2) + fieldset.add_vector_field(UV2) + + def SampleUV2(particle, fieldset, time): # pragma: no cover + u, v = fieldset.UV2[time, particle.depth, particle.lat, particle.lon] + particle.dlon += u * particle.dt + particle.dlat += v * particle.dt + + pset = ParticleSet(fieldset, pclass=Particle, lon=0.5, lat=0.5) + pset.execute(AdvectionRK4 + pset.Kernel(SampleUV2), dt=1, runtime=2) + + assert abs(pset.lon[0] - 2.5) < 1e-9 + assert abs(pset.lat[0] - 0.5) < 1e-9 + + +@pytest.mark.v4remove +@pytest.mark.xfail(reason="time_periodic removed in v4") +@pytest.mark.parametrize("use_xarray", [True, False]) +@pytest.mark.parametrize("time_periodic", [86400.0, False]) +@pytest.mark.parametrize("dt_sign", [-1, 1]) +def test_periodic(use_xarray, time_periodic, dt_sign): + lon = np.array([0, 1], dtype=np.float32) + lat = np.array([0, 1], dtype=np.float32) + depth = np.array([0, 1], dtype=np.float32) + tsize = 24 * 60 + 1 + period = 86400 + time = np.linspace(0, period, tsize, dtype=np.float64) + + def temp_func(time): + return 20 + 2 * np.sin(time * 2 * np.pi / period) + + temp_vec = temp_func(time) + + U = np.zeros((tsize, 2, 2, 2), dtype=np.float32) + V = np.zeros((tsize, 2, 2, 2), dtype=np.float32) + V[:, 0, 0, 0] = 1e-5 + W = np.zeros((tsize, 2, 2, 2), dtype=np.float32) + temp = np.zeros((tsize, 2, 2, 2), dtype=np.float32) + temp[:, :, :, :] = temp_vec + D = np.ones((2, 2), dtype=np.float32) # adding non-timevarying field + + full_dims = {"lon": lon, "lat": lat, "depth": depth, "time": time} + dimensions = {"U": full_dims, "V": full_dims, "W": full_dims, "temp": full_dims, "D": {"lon": lon, "lat": lat}} + if use_xarray: + coords = {"lat": lat, "lon": lon, "depth": depth, "time": time} + variables = {"U": "Uxr", "V": "Vxr", "W": "Wxr", "temp": "Txr", "D": "Dxr"} + dimnames = {"lon": "lon", "lat": "lat", "depth": "depth", "time": "time"} + ds = xr.Dataset( + { + "Uxr": xr.DataArray(U, coords=coords, dims=("time", "depth", "lat", "lon")), + "Vxr": xr.DataArray(V, coords=coords, dims=("time", "depth", "lat", "lon")), + "Wxr": xr.DataArray(W, coords=coords, dims=("time", "depth", "lat", "lon")), + "Txr": xr.DataArray(temp, coords=coords, dims=("time", "depth", "lat", "lon")), + "Dxr": xr.DataArray(D, coords={"lat": lat, "lon": lon}, dims=("lat", "lon")), + } + ) + fieldset = FieldSet.from_xarray_dataset( + ds, + variables, + {"U": dimnames, "V": dimnames, "W": dimnames, "temp": dimnames, "D": {"lon": "lon", "lat": "lat"}}, + time_periodic=time_periodic, + allow_time_extrapolation=True, + ) + else: + data = {"U": U, "V": V, "W": W, "temp": temp, "D": D} + fieldset = FieldSet.from_data( + data, dimensions, mesh="flat", time_periodic=time_periodic, allow_time_extrapolation=True + ) # TODO : Remove from_data + + def sampleTemp(particle, fieldset, time): # pragma: no cover + particle.temp = fieldset.temp[time, particle.depth, particle.lat, particle.lon] + # test if we can interpolate UV and UVW together + (particle.u1, particle.v1) = fieldset.UV[time, particle.depth, particle.lat, particle.lon] + (particle.u2, particle.v2, w_) = fieldset.UVW[time, particle.depth, particle.lat, particle.lon] + # test if we can sample a non-timevarying field too + particle.d = fieldset.D[0, 0, particle.lat, particle.lon] + + MyParticle = Particle.add_variables( + [ + Variable("temp", dtype=np.float32, initial=20.0), + Variable("u1", dtype=np.float32, initial=0.0), + Variable("u2", dtype=np.float32, initial=0.0), + Variable("v1", dtype=np.float32, initial=0.0), + Variable("v2", dtype=np.float32, initial=0.0), + Variable("d", dtype=np.float32, initial=0.0), + ] + ) + + pset = ParticleSet(fieldset, pclass=MyParticle, lon=[0.5], lat=[0.5], depth=[0.5]) + pset.execute( + AdvectionRK4_3D + pset.Kernel(sampleTemp), runtime=timedelta(hours=51), dt=timedelta(hours=dt_sign * 1) + ) + + if time_periodic is not False: + t = pset.time[0] + temp_theo = temp_func(t) + elif dt_sign == 1: + temp_theo = temp_vec[-1] + elif dt_sign == -1: + temp_theo = temp_vec[0] + assert np.allclose(temp_theo, pset.temp[0], atol=1e-5) + assert np.allclose(pset.u1[0], pset.u2[0]) + assert np.allclose(pset.v1[0], pset.v2[0]) + assert np.allclose(pset.d[0], 1.0) + + +@pytest.mark.v4alpha +@pytest.mark.xfail(reason="GH1946") +def test_fieldset_from_data_gridtypes(): + """Simple test for fieldset initialisation from data.""" + xdim, ydim, zdim = 20, 10, 4 + + lon = np.linspace(0.0, 10.0, xdim, dtype=np.float32) + lat = np.linspace(0.0, 10.0, ydim, dtype=np.float32) + depth = np.linspace(0.0, 1.0, zdim, dtype=np.float32) + depth_s = np.ones((zdim, ydim, xdim)) + U = np.ones((zdim, ydim, xdim)) + V = np.ones((zdim, ydim, xdim)) + dimensions = {"lat": lat, "lon": lon, "depth": depth} + data = {"U": np.array(U, dtype=np.float32), "V": np.array(V, dtype=np.float32)} + lonm, latm = np.meshgrid(lon, lat) + for k in range(zdim): + data["U"][k, :, :] = lonm * (depth[k] + 1) + 0.1 + depth_s[k, :, :] = depth[k] + + # Rectilinear Z grid + fieldset = FieldSet.from_data(data, dimensions, mesh="flat") # TODO : Remove from_data + pset = ParticleSet(fieldset, Particle, [0, 0], [0, 0], [0, 0.4]) + pset.execute(AdvectionRK4, runtime=1.5, dt=0.5) + plon = pset.lon + plat = pset.lat + # sol of dx/dt = (init_depth+1)*x+0.1; x(0)=0 + assert np.allclose(plon, [0.17173462592827032, 0.2177736932123214]) + assert np.allclose(plat, [1, 1]) + + # Rectilinear S grid + dimensions["depth"] = depth_s + fieldset = FieldSet.from_data(data, dimensions, mesh="flat") # TODO : Remove from_data + pset = ParticleSet(fieldset, Particle, [0, 0], [0, 0], [0, 0.4]) + pset.execute(AdvectionRK4, runtime=1.5, dt=0.5) + assert np.allclose(plon, pset.lon) + assert np.allclose(plat, pset.lat) + + # Curvilinear Z grid + dimensions["lon"] = lonm + dimensions["lat"] = latm + dimensions["depth"] = depth + fieldset = FieldSet.from_data(data, dimensions, mesh="flat") # TODO : Remove from_data + pset = ParticleSet(fieldset, Particle, [0, 0], [0, 0], [0, 0.4]) + pset.execute(AdvectionRK4, runtime=1.5, dt=0.5) + assert np.allclose(plon, pset.lon) + assert np.allclose(plat, pset.lat) + + # Curvilinear S grid + dimensions["depth"] = depth_s + fieldset = FieldSet.from_data(data, dimensions, mesh="flat") # TODO : Remove from_data + pset = ParticleSet(fieldset, Particle, [0, 0], [0, 0], [0, 0.4]) + pset.execute(AdvectionRK4, runtime=1.5, dt=0.5) + assert np.allclose(plon, pset.lon) + assert np.allclose(plat, pset.lat) diff --git a/tests/test_fieldset_sampling.py b/tests-v3/test_fieldset_sampling.py similarity index 100% rename from tests/test_fieldset_sampling.py rename to tests-v3/test_fieldset_sampling.py diff --git a/tests/test_interaction.py b/tests-v3/test_interaction.py similarity index 100% rename from tests/test_interaction.py rename to tests-v3/test_interaction.py diff --git a/tests-v3/test_interpolation.py b/tests-v3/test_interpolation.py new file mode 100644 index 0000000000..67eab64204 --- /dev/null +++ b/tests-v3/test_interpolation.py @@ -0,0 +1,64 @@ +import pytest + +import parcels._interpolation as interpolation +from tests.utils import create_fieldset_zeros_3d + + +@pytest.fixture +def tmp_interpolator_registry(): + """Resets the interpolator registry after the test. Vital when testing manipulating the registry.""" + old_2d = interpolation._interpolator_registry_2d.copy() + old_3d = interpolation._interpolator_registry_3d.copy() + yield + interpolation._interpolator_registry_2d = old_2d + interpolation._interpolator_registry_3d = old_3d + + +@pytest.mark.usefixtures("tmp_interpolator_registry") +def test_interpolation_registry(): + @interpolation.register_3d_interpolator("test") + @interpolation.register_2d_interpolator("test") + def some_function(): + return "test" + + assert "test" in interpolation.get_2d_interpolator_registry() + assert "test" in interpolation.get_3d_interpolator_registry() + + f = interpolation.get_2d_interpolator_registry()["test"] + g = interpolation.get_3d_interpolator_registry()["test"] + assert f() == g() == "test" + + +@pytest.mark.v4remove +@pytest.mark.xfail(reason="GH1946") +@pytest.mark.usefixtures("tmp_interpolator_registry") +def test_interpolator_override(): + fieldset = create_fieldset_zeros_3d() + + @interpolation.register_3d_interpolator("linear") + def test_interpolator(ctx: interpolation.InterpolationContext3D): + raise NotImplementedError + + with pytest.raises(NotImplementedError): + fieldset.U[0, 0.5, 0.5, 0.5] + + +@pytest.mark.v4remove +@pytest.mark.xfail(reason="GH1946") +@pytest.mark.usefixtures("tmp_interpolator_registry") +def test_full_depth_provided_to_interpolators(): + """The full depth needs to be provided to the interpolation schemes as some interpolators + need to know whether they are at the surface or bottom of the water column. + + https://github.com/OceanParcels/Parcels/pull/1816#discussion_r1908840408 + """ + xdim, ydim, zdim = 10, 11, 12 + fieldset = create_fieldset_zeros_3d(xdim=xdim, ydim=ydim, zdim=zdim) + + @interpolation.register_3d_interpolator("linear") + def test_interpolator2(ctx: interpolation.InterpolationContext3D): + assert ctx.data.shape[1] == zdim + # The array z dimension is the same as the fieldset z dimension + return 0 + + fieldset.U[0.5, 0.5, 0.5, 0.5] diff --git a/tests/test_kernel_execution.py b/tests-v3/test_kernel_execution.py similarity index 100% rename from tests/test_kernel_execution.py rename to tests-v3/test_kernel_execution.py diff --git a/tests/test_kernel_language.py b/tests-v3/test_kernel_language.py similarity index 100% rename from tests/test_kernel_language.py rename to tests-v3/test_kernel_language.py diff --git a/tests/test_mpirun.py b/tests-v3/test_mpirun.py similarity index 100% rename from tests/test_mpirun.py rename to tests-v3/test_mpirun.py diff --git a/tests/test_particles.py b/tests-v3/test_particles.py similarity index 100% rename from tests/test_particles.py rename to tests-v3/test_particles.py diff --git a/tests/test_particlesets.py b/tests-v3/test_particlesets.py similarity index 100% rename from tests/test_particlesets.py rename to tests-v3/test_particlesets.py diff --git a/tests/test_reprs.py b/tests-v3/test_reprs.py similarity index 100% rename from tests/test_reprs.py rename to tests-v3/test_reprs.py diff --git a/tests/test_typing.py b/tests-v3/test_typing.py similarity index 100% rename from tests/test_typing.py rename to tests-v3/test_typing.py diff --git a/tests/tools/test_exampledata_utils.py b/tests-v3/tools/test_exampledata_utils.py similarity index 100% rename from tests/tools/test_exampledata_utils.py rename to tests-v3/tools/test_exampledata_utils.py diff --git a/tests/tools/test_helpers.py b/tests-v3/tools/test_helpers.py similarity index 100% rename from tests/tools/test_helpers.py rename to tests-v3/tools/test_helpers.py diff --git a/tests/tools/test_warnings.py b/tests-v3/tools/test_warnings.py similarity index 100% rename from tests/tools/test_warnings.py rename to tests-v3/tools/test_warnings.py diff --git a/tests/common_kernels.py b/tests/common_kernels.py index 58f674cf82..53e3c7310b 100644 --- a/tests/common_kernels.py +++ b/tests/common_kernels.py @@ -2,7 +2,7 @@ import numpy as np -from parcels.tools.statuscodes import StatusCode +from parcels import StatusCode def DoNothing(particles, fieldset): # pragma: no cover diff --git a/tests/test_advection.py b/tests/test_advection.py index de8b004d79..c4e8a82ce9 100644 --- a/tests/test_advection.py +++ b/tests/test_advection.py @@ -1,179 +1,564 @@ import numpy as np import pytest import xarray as xr - -from parcels import ( - AdvectionAnalytical, +import xgcm + +import parcels +from parcels import Field, FieldSet, Particle, ParticleFile, ParticleSet, StatusCode, Variable, VectorField, XGrid +from parcels._datasets.structured.generated import ( + decaying_moving_eddy_dataset, + moving_eddy_dataset, + peninsula_dataset, + radial_rotation_dataset, + simple_UV_dataset, + stommel_gyre_dataset, +) +from parcels.interpolators import CGrid_Velocity, XLinear +from parcels.kernels import ( AdvectionDiffusionEM, AdvectionDiffusionM1, AdvectionEE, AdvectionRK4, AdvectionRK4_3D, AdvectionRK45, - FieldSet, - Particle, - ParticleSet, - Variable, ) -from tests.utils import TEST_DATA +from tests.utils import round_and_hash_float_array kernel = { "EE": AdvectionEE, "RK4": AdvectionRK4, + "RK4_3D": AdvectionRK4_3D, "RK45": AdvectionRK45, - "AA": AdvectionAnalytical, + # "AA": AdvectionAnalytical, "AdvDiffEM": AdvectionDiffusionEM, "AdvDiffM1": AdvectionDiffusionM1, } -@pytest.fixture -def lon(): - xdim = 200 - return np.linspace(-170, 170, xdim, dtype=np.float32) +@pytest.mark.parametrize("mesh", ["spherical", "flat"]) +def test_advection_zonal(mesh, npart=10): + """Particles at high latitude move geographically faster due to the pole correction in `GeographicPolar`.""" + ds = simple_UV_dataset(mesh=mesh) + ds["U"].data[:] = 1.0 + grid = XGrid.from_dataset(ds, mesh=mesh) + U = Field("U", ds["U"], grid, interp_method=XLinear) + V = Field("V", ds["V"], grid, interp_method=XLinear) + 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 == "spherical": + assert (np.diff(pset.lon) > 1.0e-4).all() + else: + assert (np.diff(pset.lon) < 1.0e-4).all() + + +def test_advection_zonal_with_particlefile(tmp_store): + """Particles at high latitude move geographically faster due to the pole correction in `GeographicPolar`.""" + npart = 10 + ds = simple_UV_dataset(mesh="flat") + ds["U"].data[:] = 1.0 + grid = XGrid.from_dataset(ds, mesh="flat") + U = Field("U", ds["U"], grid, interp_method=XLinear) + V = Field("V", ds["V"], grid, interp_method=XLinear) + 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)) + pfile = ParticleFile(tmp_store, outputdt=np.timedelta64(30, "m")) + pset.execute(AdvectionRK4, runtime=np.timedelta64(2, "h"), dt=np.timedelta64(15, "m"), output_file=pfile) + + assert (np.diff(pset.lon) < 1.0e-4).all() + ds = xr.open_zarr(tmp_store) + np.testing.assert_allclose(ds.isel(obs=-1).lon.values, pset.lon) + + +def periodicBC(particles, fieldset): + particles.total_dlon += particles.dlon + particles.lon = np.fmod(particles.lon, 2) + + +def test_advection_zonal_periodic(): + ds = simple_UV_dataset(dims=(2, 2, 2, 2), mesh="flat") + ds["U"].data[:] = 0.1 + ds["lon"].data = np.array([0, 2]) + ds["lat"].data = np.array([0, 2]) + + # add a halo + halo = ds.isel(XG=0) + halo.lon.values = ds.lon.values[1] + 1 + halo.XG.values = ds.XG.values[1] + 2 + ds = xr.concat([ds, halo], dim="XG") + + grid = XGrid.from_dataset(ds) + U = Field("U", ds["U"], grid, interp_method=XLinear) + V = Field("V", ds["V"], grid, interp_method=XLinear) + UV = VectorField("UV", U, V) + fieldset = FieldSet([U, V, UV]) + + PeriodicParticle = Particle.add_variable(Variable("total_dlon", initial=0)) + startlon = np.array([0.5, 0.4]) + pset = ParticleSet(fieldset, pclass=PeriodicParticle, lon=startlon, lat=[0.5, 0.5]) + pset.execute([AdvectionEE, periodicBC], runtime=np.timedelta64(40, "s"), dt=np.timedelta64(1, "s")) + np.testing.assert_allclose(pset.total_dlon, 4, atol=1e-5) + np.testing.assert_allclose(pset.lon + pset.dlon, startlon, atol=1e-5) + np.testing.assert_allclose(pset.lat, 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="flat") + ds["U"].data[:] = 1.0 + grid = XGrid.from_dataset(ds) + U = Field("U", ds["U"], grid, interp_method=XLinear) + U.data[:, 0, :, :] = 0.0 # Set U to 0 at the surface + V = Field("V", ds["V"], grid, interp_method=XLinear) + 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") + np.testing.assert_allclose(pset.lon, expected_lon, 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="flat") + grid = XGrid.from_dataset(ds) + U = Field("U", ds["U"], grid, interp_method=XLinear) + U.data[:] = 0.01 # Set U to small value (to avoid horizontal out of bounds) + V = Field("V", ds["V"], grid, interp_method=XLinear) + W = Field("W", ds["V"], grid, interp_method=XLinear) # 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(particles, fieldset): # pragma: no cover + particles.state = np.where(particles.state == StatusCode.ErrorOutOfBounds, StatusCode.Delete, particles.state) + particles.state = np.where( + particles.state == StatusCode.ErrorThroughSurface, StatusCode.Delete, particles.state + ) + + def SubmergeParticle(particles, fieldset): # pragma: no cover + if len(particles.state) == 0: + return + inds = np.argwhere(particles.state == StatusCode.ErrorThroughSurface).flatten() + if len(inds) == 0: + return + dt = particles.dt / np.timedelta64(1, "s") + (u, v) = fieldset.UV[particles[inds]] + particles[inds].dlon = u * dt + particles[inds].dlat = v * dt + particles[inds].ddepth = 0.0 + particles[inds].depth = 0 + particles[inds].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: + np.testing.assert_allclose(pset.lon[0], 0.6, atol=1e-5) + np.testing.assert_allclose(pset.depth[0], 0, atol=1e-5) + else: + assert len(pset) == 0 -@pytest.fixture -def lat(): - ydim = 100 - return np.linspace(-80, 80, ydim, dtype=np.float32) +@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) + ) -@pytest.fixture -def depth(): - zdim = 2 - return np.linspace(0, 30, zdim, dtype=np.float32) + 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=XLinear) + V = Field("V", ds["V"], grid, interp_method=XLinear) + fields = [U, V, VectorField("UV", U, V)] + if w: + W = Field("W", ds["W"], grid, interp_method=XLinear) + 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]) + np.testing.assert_allclose(np.array([p.lon - x0 for p in pset]), 4 * u, atol=1e-6) + np.testing.assert_allclose(np.array([p.lat - y0 for p in pset]), 4 * v, atol=1e-6) + if w: + np.testing.assert_allclose(np.array([p.depth - z0 for p in pset]), 4 * w, atol=1e-6) + + +def test_radialrotation(npart=10): + ds = radial_rotation_dataset() + grid = XGrid.from_dataset(ds, mesh="flat") + U = parcels.Field("U", ds["U"], grid, interp_method=XLinear) + V = parcels.Field("V", ds["V"], grid, interp_method=XLinear) + UV = parcels.VectorField("UV", U, V) + fieldset = parcels.FieldSet([U, V, UV]) + + dt = np.timedelta64(30, "s") + lon = np.linspace(32, 50, npart) + lat = np.ones(npart) * 30 + starttime = np.arange(np.timedelta64(0, "s"), npart * dt, dt) + + pset = parcels.ParticleSet(fieldset, lon=lon, lat=lat, time=starttime) + pset.execute(parcels.kernels.AdvectionRK4, endtime=np.timedelta64(10, "m"), dt=dt) + + theta = 2 * np.pi * (pset.time - starttime) / np.timedelta64(24 * 3600, "s") + true_lon = (lon - 30.0) * np.cos(theta) + 30.0 + true_lat = -(lon - 30.0) * np.sin(theta) + 30.0 + + np.testing.assert_allclose(pset.lon, true_lon, atol=5e-2) + np.testing.assert_allclose(pset.lat, true_lat, atol=5e-2) + + +@pytest.mark.parametrize( + "method, rtol", + [ + ("EE", 1e-2), + ("AdvDiffEM", 1e-2), + ("AdvDiffM1", 1e-2), + ("RK4", 1e-5), + ("RK4_3D", 1e-5), + ("RK45", 1e-4), + ], +) +def test_moving_eddy(method, rtol): + ds = moving_eddy_dataset() + grid = XGrid.from_dataset(ds) + U = Field("U", ds["U"], grid, interp_method=XLinear) + V = Field("V", ds["V"], grid, interp_method=XLinear) + if method == "RK4_3D": + # Using W to test 3D advection (assuming same velocity as V) + W = Field("W", ds["V"], grid, interp_method=XLinear) + UVW = VectorField("UVW", U, V, W) + fieldset = FieldSet([U, V, W, UVW]) + else: + UV = VectorField("UV", U, V) + fieldset = FieldSet([U, V, UV]) + if method in ["AdvDiffEM", "AdvDiffM1"]: + # Add zero diffusivity field for diffusion kernels + ds["Kh"] = (["time", "depth", "YG", "XG"], np.full(ds["U"].shape, 0)) + fieldset.add_field(Field("Kh", ds["Kh"], grid, interp_method=XLinear), "Kh_zonal") + fieldset.add_field(Field("Kh", ds["Kh"], grid, interp_method=XLinear), "Kh_meridional") + fieldset.add_constant("dres", 0.1) + + start_lon, start_lat, start_depth = 12000, 12500, 12500 + dt = np.timedelta64(30, "m") + + if method == "RK45": + fieldset.add_constant("RK45_tol", rtol) + + pset = ParticleSet(fieldset, lon=start_lon, lat=start_lat, depth=start_depth, time=np.timedelta64(0, "s")) + pset.execute(kernel[method], dt=dt, endtime=np.timedelta64(1, "h")) + + def truth_moving(x_0, y_0, t): + t /= np.timedelta64(1, "s") + lat = y_0 - (ds.u_0 - ds.u_g) / ds.f * (1 - np.cos(ds.f * t)) + lon = x_0 + ds.u_g * t + (ds.u_0 - ds.u_g) / ds.f * np.sin(ds.f * t) + return lon, lat + + exp_lon, exp_lat = truth_moving(start_lon, start_lat, pset.time[0]) + np.testing.assert_allclose(pset.lon, exp_lon, rtol=rtol) + np.testing.assert_allclose(pset.lat, exp_lat, rtol=rtol) + if method == "RK4_3D": + np.testing.assert_allclose(pset.depth, exp_lat, rtol=rtol) + + +@pytest.mark.parametrize( + "method, rtol", + [ + ("EE", 1e-1), + ("RK4", 1e-5), + ("RK45", 1e-4), + ], +) +def test_decaying_moving_eddy(method, rtol): + ds = decaying_moving_eddy_dataset() + grid = XGrid.from_dataset(ds) + U = Field("U", ds["U"], grid, interp_method=XLinear) + V = Field("V", ds["V"], grid, interp_method=XLinear) + UV = VectorField("UV", U, V) + fieldset = FieldSet([U, V, UV]) + + start_lon, start_lat = 10000, 10000 + dt = np.timedelta64(60, "m") + + if method == "RK45": + fieldset.add_constant("RK45_tol", rtol) + fieldset.add_constant("RK45_min_dt", 10 * 60) + + pset = ParticleSet(fieldset, lon=start_lon, lat=start_lat, time=np.timedelta64(0, "s")) + pset.execute(kernel[method], dt=dt, endtime=np.timedelta64(1, "D")) + + def truth_moving(x_0, y_0, t): + t /= np.timedelta64(1, "s") + lon = ( + x_0 + + (ds.u_g / ds.gamma_g) * (1 - np.exp(-ds.gamma_g * t)) + + ds.f + * ((ds.u_0 - ds.u_g) / (ds.f**2 + ds.gamma**2)) + * ((ds.gamma / ds.f) + np.exp(-ds.gamma * t) * (np.sin(ds.f * t) - (ds.gamma / ds.f) * np.cos(ds.f * t))) + ) + lat = y_0 - ((ds.u_0 - ds.u_g) / (ds.f**2 + ds.gamma**2)) * ds.f * ( + 1 - np.exp(-ds.gamma * t) * (np.cos(ds.f * t) + (ds.gamma / ds.f) * np.sin(ds.f * t)) + ) + return lon, lat + + exp_lon, exp_lat = truth_moving(start_lon, start_lat, pset.time[0]) + np.testing.assert_allclose(pset.lon, exp_lon, rtol=rtol) + np.testing.assert_allclose(pset.lat, exp_lat, rtol=rtol) + + +@pytest.mark.parametrize( + "method, rtol", + [ + ("RK4", 0.1), + ("RK45", 0.1), + ], +) +@pytest.mark.parametrize("grid_type", ["A", "C"]) +def test_stommelgyre_fieldset(method, rtol, grid_type): + npart = 2 + ds = stommel_gyre_dataset(grid_type=grid_type) + grid = XGrid.from_dataset(ds) + vector_interp_method = None if grid_type == "A" else CGrid_Velocity + U = Field("U", ds["U"], grid) + V = Field("V", ds["V"], grid) + P = Field("P", ds["P"], grid, interp_method=XLinear) + UV = VectorField("UV", U, V, vector_interp_method=vector_interp_method) + fieldset = FieldSet([U, V, P, UV]) + + dt = np.timedelta64(30, "m") + runtime = np.timedelta64(1, "D") + start_lon = np.linspace(10e3, 100e3, npart) + start_lat = np.ones_like(start_lon) * 5000e3 + + if method == "RK45": + fieldset.add_constant("RK45_tol", rtol) + + SampleParticle = Particle.add_variable( + [Variable("p", initial=0.0, dtype=np.float32), Variable("p_start", initial=0.0, dtype=np.float32)] + ) + def UpdateP(particles, fieldset): # pragma: no cover + particles.p = fieldset.P[particles.time, particles.depth, particles.lat, particles.lon] + particles.p_start = np.where(particles.time == 0, particles.p, particles.p_start) -@pytest.mark.v4alpha -@pytest.mark.xfail(reason="When refactoring fieldfilebuffer croco support was dropped. This will be fixed in v4.") -def test_conversion_3DCROCO(): - """Test of the (SciPy) version of the conversion from depth to sigma in CROCO + pset = ParticleSet(fieldset, pclass=SampleParticle, lon=start_lon, lat=start_lat, time=np.timedelta64(0, "s")) + pset.execute([kernel[method], UpdateP], dt=dt, runtime=runtime) + np.testing.assert_allclose(pset.p, pset.p_start, rtol=rtol) - Values below are retrieved using xroms and hardcoded in the method (to avoid dependency on xroms): - ```py - x, y = 10, 20 - s_xroms = ds.s_w.values - z_xroms = ds.z_w.isel(time=0).isel(eta_rho=y).isel(xi_rho=x).values - lat, lon = ds.y_rho.values[y, x], ds.x_rho.values[y, x] - ``` - """ - fieldset = FieldSet.from_modulefile(TEST_DATA / "fieldset_CROCO3D.py") - lat, lon = 78000.0, 38000.0 - s_xroms = np.array([-1.0, -0.9, -0.8, -0.7, -0.6, -0.5, -0.4, -0.3, -0.2, -0.1, 0.0], dtype=np.float32) - z_xroms = np.array( - [ - -1.26000000e02, - -1.10585846e02, - -9.60985413e01, - -8.24131317e01, - -6.94126511e01, - -5.69870148e01, - -4.50318756e01, - -3.34476166e01, - -2.21383114e01, - -1.10107975e01, - 2.62768921e-02, - ], - dtype=np.float32, +@pytest.mark.parametrize( + "method, rtol", + [ + ("RK4", 5e-3), + ("RK45", 1e-4), + ], +) +@pytest.mark.parametrize("grid_type", ["A"]) # TODO also implement C-grid once available +def test_peninsula_fieldset(method, rtol, grid_type): + npart = 2 + ds = peninsula_dataset(grid_type=grid_type) + grid = XGrid.from_dataset(ds) + U = Field("U", ds["U"], grid, interp_method=XLinear) + V = Field("V", ds["V"], grid, interp_method=XLinear) + P = Field("P", ds["P"], grid, interp_method=XLinear) + UV = VectorField("UV", U, V) + fieldset = FieldSet([U, V, P, UV]) + + dt = np.timedelta64(30, "m") + runtime = np.timedelta64(1, "D") + start_lat = np.linspace(3e3, 47e3, npart) + start_lon = 3e3 * np.ones_like(start_lat) + + if method == "RK45": + fieldset.add_constant("RK45_tol", rtol) + + SampleParticle = Particle.add_variable( + [Variable("p", initial=0.0, dtype=np.float32), Variable("p_start", initial=0.0, dtype=np.float32)] ) - sigma = np.zeros_like(z_xroms) - from parcels.field import _croco_from_z_to_sigma_scipy + def UpdateP(particles, fieldset): # pragma: no cover + particles.p = fieldset.P[particles.time, particles.depth, particles.lat, particles.lon] + particles.p_start = np.where(particles.time == 0, particles.p, particles.p_start) - for zi, z in enumerate(z_xroms): - sigma[zi] = _croco_from_z_to_sigma_scipy(fieldset, 0, z, lat, lon, None) + pset = ParticleSet(fieldset, pclass=SampleParticle, lon=start_lon, lat=start_lat, time=np.timedelta64(0, "s")) + pset.execute([kernel[method], UpdateP], dt=dt, runtime=runtime) + np.testing.assert_allclose(pset.p, pset.p_start, rtol=rtol) - assert np.allclose(sigma, s_xroms, atol=1e-3) +def test_nemo_curvilinear_fieldset(): + data_folder = parcels.download_example_dataset("NemoCurvilinear_data") + files = data_folder.glob("*.nc4") + ds = xr.open_mfdataset(files, combine="nested", data_vars="minimal", coords="minimal", compat="override") + ds = ( + ds.isel(time_counter=0, drop=True) + .isel(time=0, drop=True) + .isel(z_a=0, drop=True) + .rename({"glamf": "lon", "gphif": "lat", "z": "depth"}) + ) -@pytest.mark.v4alpha -@pytest.mark.xfail(reason="CROCO 3D interpolation is not yet implemented correctly in v4. ") -def test_advection_3DCROCO(): - fieldset = FieldSet.from_modulefile(TEST_DATA / "fieldset_CROCO3D.py") - - runtime = 1e4 - X, Z = np.meshgrid([40e3, 80e3, 120e3], [-10, -130]) - Y = np.ones(X.size) * 100e3 + xgcm_grid = xgcm.Grid( + ds, + coords={ + "X": {"left": "x"}, + "Y": {"left": "y"}, + }, + periodic=False, + autoparse_metadata=False, + ) + grid = XGrid(xgcm_grid, mesh="spherical") - pclass = Particle.add_variable(Variable("w")) - pset = ParticleSet(fieldset=fieldset, pclass=pclass, lon=X, lat=Y, depth=Z) + U = parcels.Field("U", ds["U"], grid) + V = parcels.Field("V", ds["V"], grid) + U.units = parcels.GeographicPolar() + V.units = parcels.Geographic() + UV = parcels.VectorField("UV", U, V, vector_interp_method=CGrid_Velocity) + fieldset = parcels.FieldSet([U, V, UV]) - def SampleW(particle, fieldset, time): # pragma: no cover - particle.w = fieldset.W[time, particle.depth, particle.lat, particle.lon] + npart = 20 + lonp = 30 * np.ones(npart) + latp = np.linspace(-70, 88, npart) + runtime = np.timedelta64(12, "h") # TODO increase to 160 days - pset.execute([AdvectionRK4_3D, SampleW], runtime=runtime, dt=100) - assert np.allclose(pset.depth, Z.flatten(), atol=5) # TODO lower this atol - assert np.allclose(pset.lon_nextloop, [x + runtime for x in X.flatten()], atol=1e-3) + def periodicBC(particles, fieldset): # pragma: no cover + particles.dlon = np.where(particles.lon > 180, particles.dlon - 360, particles.dlon) + pset = parcels.ParticleSet(fieldset, lon=lonp, lat=latp) + pset.execute([AdvectionEE, periodicBC], runtime=runtime, dt=np.timedelta64(6, "h")) + np.testing.assert_allclose(pset.lat, latp, atol=1e-1) -@pytest.mark.v4alpha -@pytest.mark.xfail(reason="When refactoring fieldfilebuffer croco support was dropped. This will be fixed in v4.") -def test_advection_2DCROCO(): - fieldset = FieldSet.from_modulefile(TEST_DATA / "fieldset_CROCO2D.py") - runtime = 1e4 - X = np.array([40e3, 80e3, 120e3]) - Y = np.ones(X.size) * 100e3 - Z = np.zeros(X.size) - pset = ParticleSet(fieldset=fieldset, pclass=Particle, lon=X, lat=Y, depth=Z) +@pytest.mark.parametrize("method", ["RK4", "RK4_3D"]) +def test_nemo_3D_curvilinear_fieldset(method): + download_dir = parcels.download_example_dataset("NemoNorthSeaORCA025-N006_data") + ufiles = download_dir.glob("*U.nc") + dsu = xr.open_mfdataset(ufiles, decode_times=False, drop_variables=["nav_lat", "nav_lon"]) + dsu = dsu.rename({"time_counter": "time", "uo": "U"}) - pset.execute([AdvectionRK4], runtime=runtime, dt=100) - assert np.allclose(pset.depth, Z.flatten(), atol=1e-3) - assert np.allclose(pset.lon_nextloop, [x + runtime for x in X], atol=1e-3) + vfiles = download_dir.glob("*V.nc") + dsv = xr.open_mfdataset(vfiles, decode_times=False, drop_variables=["nav_lat", "nav_lon"]) + dsv = dsv.rename({"time_counter": "time", "vo": "V"}) + wfiles = download_dir.glob("*W.nc") + dsw = xr.open_mfdataset(wfiles, decode_times=False, drop_variables=["nav_lat", "nav_lon"]) + dsw = dsw.rename({"time_counter": "time", "depthw": "depth", "wo": "W"}) -@pytest.mark.v4alpha -@pytest.mark.xfail(reason="GH1946") -def test_analyticalAgrid(): - lon = np.arange(0, 15, dtype=np.float32) - lat = np.arange(0, 15, dtype=np.float32) - U = np.ones((lat.size, lon.size), dtype=np.float32) - V = np.ones((lat.size, lon.size), dtype=np.float32) - fieldset = FieldSet.from_data({"U": U, "V": V}, {"lon": lon, "lat": lat}, mesh="flat") - pset = ParticleSet(fieldset, pclass=Particle, lon=1, lat=1) + dsu = dsu.assign_coords(depthu=dsw.depth.values) + dsu = dsu.rename({"depthu": "depth"}) - with pytest.raises(NotImplementedError): - pset.execute(AdvectionAnalytical, runtime=1) + dsv = dsv.assign_coords(depthv=dsw.depth.values) + dsv = dsv.rename({"depthv": "depth"}) + coord_file = f"{download_dir}/coordinates.nc" + dscoord = xr.open_dataset(coord_file, decode_times=False).rename({"glamf": "lon", "gphif": "lat"}) + dscoord = dscoord.isel(time=0, drop=True) -@pytest.mark.v4alpha -@pytest.mark.xfail(reason="GH1927") -@pytest.mark.parametrize("u", [1, -0.2, -0.3, 0]) -@pytest.mark.parametrize("v", [1, -0.3, 0, -1]) -@pytest.mark.parametrize("w", [None, 1, -0.3, 0, -1]) -@pytest.mark.parametrize("direction", [1, -1]) -def test_uniform_analytical(u, v, w, direction, tmp_zarrfile): - lon = np.arange(0, 15, dtype=np.float32) - lat = np.arange(0, 15, dtype=np.float32) - if w is not None: - depth = np.arange(0, 40, 2, dtype=np.float32) - U = u * np.ones((depth.size, lat.size, lon.size), dtype=np.float32) - V = v * np.ones((depth.size, lat.size, lon.size), dtype=np.float32) - W = w * np.ones((depth.size, lat.size, lon.size), dtype=np.float32) - fieldset = FieldSet.from_data({"U": U, "V": V, "W": W}, {"lon": lon, "lat": lat, "depth": depth}, mesh="flat") - fieldset.W.interp_method = "cgrid_velocity" - else: - U = u * np.ones((lat.size, lon.size), dtype=np.float32) - V = v * np.ones((lat.size, lon.size), dtype=np.float32) - fieldset = FieldSet.from_data({"U": U, "V": V}, {"lon": lon, "lat": lat}, mesh="flat") - fieldset.U.interp_method = "cgrid_velocity" - fieldset.V.interp_method = "cgrid_velocity" - - x0, y0, z0 = 6.1, 6.2, 20 - pset = ParticleSet(fieldset, pclass=Particle, lon=x0, lat=y0, depth=z0) - - outfile = pset.ParticleFile(name=tmp_zarrfile, outputdt=1, chunks=(1, 1)) - pset.execute(AdvectionAnalytical, runtime=4, dt=direction, output_file=outfile) - assert np.abs(pset.lon - x0 - pset.time * u) < 1e-6 - assert np.abs(pset.lat - y0 - pset.time * v) < 1e-6 - if w is not None: - assert np.abs(pset.depth - z0 - pset.time * w) < 1e-4 - - ds = xr.open_zarr(tmp_zarrfile) - times = (direction * ds["time"][:]).values.astype("timedelta64[s]")[0] - timeref = np.arange(1, 5).astype("timedelta64[s]") - assert np.allclose(times, timeref, atol=np.timedelta64(1, "ms")) - lons = ds["lon"][:].values - assert np.allclose(lons, x0 + direction * u * np.arange(1, 5)) + ds = xr.merge([dsu, dsv, dsw, dscoord]) + ds = ds.drop_vars( + [ + "uos", + "vos", + "nav_lev", + "nav_lon", + "nav_lat", + "tauvo", + "tauuo", + "time_steps", + "gphiu", + "gphiv", + "gphit", + "glamu", + "glamv", + "glamt", + "time_centered_bounds", + "time_counter_bounds", + "time_centered", + ] + ) + ds = ds.drop_vars(["e1f", "e1t", "e1u", "e1v", "e2f", "e2t", "e2u", "e2v"]) + ds["time"] = [np.timedelta64(int(t), "s") + np.datetime64("1900-01-01") for t in ds["time"]] + + ds["W"] *= -1 # Invert W velocity + + xgcm_grid = xgcm.Grid( + ds, + coords={ + "X": {"left": "x"}, + "Y": {"left": "y"}, + "Z": {"left": "depth"}, + "T": {"center": "time"}, + }, + periodic=False, + autoparse_metadata=False, + ) + grid = XGrid(xgcm_grid, mesh="spherical") + + U = parcels.Field("U", ds["U"], grid) + V = parcels.Field("V", ds["V"], grid) + W = parcels.Field("W", ds["W"], grid) + U.units = parcels.GeographicPolar() + V.units = parcels.Geographic() + UV = parcels.VectorField("UV", U, V, vector_interp_method=CGrid_Velocity) + UVW = parcels.VectorField("UVW", U, V, W, vector_interp_method=CGrid_Velocity) + fieldset = parcels.FieldSet([U, V, W, UV, UVW]) + + npart = 10 + lons = np.linspace(1.9, 3.4, npart) + lats = np.linspace(52.5, 51.6, npart) + pset = parcels.ParticleSet(fieldset, lon=lons, lat=lats, depth=np.ones_like(lons)) + + pset.execute(kernel[method], runtime=np.timedelta64(4, "D"), dt=np.timedelta64(6, "h")) + + if method == "RK4": + np.testing.assert_equal(round_and_hash_float_array([p.lon for p in pset], decimals=5), 29977383852960156017546) + elif method == "RK4_3D": + # TODO check why decimals needs to be so low in RK4_3D (compare to v3) + np.testing.assert_equal( + round_and_hash_float_array([p.depth for p in pset], decimals=1), 29747210774230389239432 + ) diff --git a/tests/v4/test_basegrid.py b/tests/test_basegrid.py similarity index 97% rename from tests/v4/test_basegrid.py rename to tests/test_basegrid.py index 5f18b162af..de11f00e23 100644 --- a/tests/v4/test_basegrid.py +++ b/tests/test_basegrid.py @@ -5,7 +5,7 @@ import numpy as np import pytest -from parcels.basegrid import BaseGrid +from parcels._core.basegrid import BaseGrid class MockGrid(BaseGrid): diff --git a/tests/test_data/README.md b/tests/test_data/README.md new file mode 100644 index 0000000000..083506fb69 --- /dev/null +++ b/tests/test_data/README.md @@ -0,0 +1 @@ +Test data that was used primary in v3 (or in migrating from v3 to v4). In v4 this subpackage is superceded by the `parcels._datasets` package as well as `pooch` for fetching realistic data. diff --git a/tests/v4/test_datasets.py b/tests/test_datasets.py similarity index 93% rename from tests/v4/test_datasets.py rename to tests/test_datasets.py index e6a550d960..2750583f58 100644 --- a/tests/v4/test_datasets.py +++ b/tests/test_datasets.py @@ -1,7 +1,7 @@ import xgcm +from parcels._core.xgrid import _DEFAULT_XGCM_KWARGS from parcels._datasets.structured.generic import datasets -from parcels.xgrid import _DEFAULT_XGCM_KWARGS def test_left_indexed_dataset(): diff --git a/tests/v4/test_diffusion.py b/tests/test_diffusion.py similarity index 94% rename from tests/v4/test_diffusion.py rename to tests/test_diffusion.py index 65c9aad008..c7ecaf0057 100644 --- a/tests/v4/test_diffusion.py +++ b/tests/test_diffusion.py @@ -4,14 +4,10 @@ import pytest from scipy import stats +from parcels import Field, FieldSet, Particle, ParticleSet, Variable, VectorField, XGrid from parcels._datasets.structured.generated import simple_UV_dataset -from parcels.application_kernels import AdvectionDiffusionEM, AdvectionDiffusionM1, DiffusionUniformKh -from parcels.application_kernels.interpolation import XLinear -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 parcels.interpolators import XLinear +from parcels.kernels import AdvectionDiffusionEM, AdvectionDiffusionM1, DiffusionUniformKh from tests.utils import create_fieldset_zeros_conversion diff --git a/tests/v4/test_field.py b/tests/test_field.py similarity index 98% rename from tests/v4/test_field.py rename to tests/test_field.py index 75c6f6dc98..fda23f4798 100644 --- a/tests/v4/test_field.py +++ b/tests/test_field.py @@ -5,12 +5,11 @@ import uxarray as ux import xarray as xr -from parcels import Field, UXPiecewiseConstantFace, UXPiecewiseLinearNode, VectorField +from parcels import Field, UxGrid, VectorField, XGrid from parcels._datasets.structured.generic import T as T_structured from parcels._datasets.structured.generic import datasets as datasets_structured from parcels._datasets.unstructured.generic import datasets as datasets_unstructured -from parcels.uxgrid import UxGrid -from parcels.xgrid import XGrid +from parcels.interpolators import UXPiecewiseConstantFace, UXPiecewiseLinearNode def test_field_init_param_types(): diff --git a/tests/test_fieldset.py b/tests/test_fieldset.py index fd83a28bef..c96682ed0c 100644 --- a/tests/test_fieldset.py +++ b/tests/test_fieldset.py @@ -1,417 +1,272 @@ from datetime import timedelta +import cf_xarray # noqa: F401 +import cftime import numpy as np import pytest import xarray as xr -from parcels import ( - AdvectionRK4, - AdvectionRK4_3D, - FieldSet, - Particle, - ParticleSet, - Variable, -) -from parcels.field import VectorField -from parcels.tools.converters import GeographicPolar, UnitConverter -from tests.utils import TEST_DATA - - -def generate_fieldset_data(xdim, ydim, zdim=1, tdim=1): - lon = np.linspace(0.0, 10.0, xdim, dtype=np.float32) - lat = np.linspace(0.0, 10.0, ydim, dtype=np.float32) - depth = np.zeros(zdim, dtype=np.float32) - time = np.zeros(tdim, dtype=np.float64) - if zdim == 1 and tdim == 1: - U, V = np.meshgrid(lon, lat) - dimensions = {"lat": lat, "lon": lon} - else: - U = np.ones((tdim, zdim, ydim, xdim)) - V = np.ones((tdim, zdim, ydim, xdim)) - dimensions = {"lat": lat, "lon": lon, "depth": depth, "time": time} - data = {"U": np.array(U, dtype=np.float32), "V": np.array(V, dtype=np.float32)} - - return (data, dimensions) - - -def to_xarray_dataset(data: dict[str, np.array], dimensions: dict[str, np.array]) -> xr.Dataset: - assert len(dimensions) in [2, 4], "this function only deals with output from generate_fieldset_data()" - - if len(dimensions) == 4: - return xr.Dataset( - { - "U": (["time", "depth", "lat", "lon"], data["U"]), - "V": (["time", "depth", "lat", "lon"], data["V"]), - }, - coords=dimensions, - ) - - return xr.Dataset( - { - "U": (["lat", "lon"], data["U"]), - "V": (["lat", "lon"], data["V"]), - }, - coords=dimensions, - ) +from parcels import Field, VectorField, XGrid +from parcels._core.fieldset import CalendarError, FieldSet, _datetime_to_msg +from parcels._datasets.structured.circulation_models import datasets as datasets_circulation_models +from parcels._datasets.structured.generic import T as T_structured +from parcels._datasets.structured.generic import datasets as datasets_structured +from tests import utils + +ds = datasets_structured["ds_2d_left"] -@pytest.mark.v4remove -@pytest.mark.xfail(reason="GH1946") @pytest.fixture -def multifile_fieldset(tmp_path): - stem = "test_subsets" - - timestamps = np.arange(0, 4, 1) * 86400.0 - timestamps = np.expand_dims(timestamps, 1) - - ufiles = [] - vfiles = [] - for index, timestamp in enumerate(timestamps): - data, dimensions = generate_fieldset_data(100, 100) - path = tmp_path / f"{stem}_{index}.nc" - to_xarray_dataset(data, dimensions).pipe(assign_dataset_timestamp_dim, timestamp).to_netcdf(path) - ufiles.append(path) - vfiles.append(path) - - files = {"U": ufiles, "V": vfiles} - variables = {"U": "U", "V": "V"} - dimensions = {"lon": "lon", "lat": "lat", "time": "time"} - return FieldSet.from_netcdf(files, variables, dimensions) - - -@pytest.mark.v4alpha -@pytest.mark.xfail(reason="GH1946") -def test_fieldset_from_modulefile(): - nemo_fname = str(TEST_DATA / "fieldset_nemo.py") - nemo_error_fname = str(TEST_DATA / "fieldset_nemo_error.py") - - fieldset = FieldSet.from_modulefile(nemo_fname) - - fieldset = FieldSet.from_modulefile(nemo_fname) - assert fieldset.U.grid.lon.shape[1] == 21 - - with pytest.raises(IOError): - FieldSet.from_modulefile(nemo_error_fname) - - FieldSet.from_modulefile(nemo_error_fname, modulename="random_function_name") - - with pytest.raises(IOError): - FieldSet.from_modulefile(nemo_error_fname, modulename="none_returning_function") - - -@pytest.mark.v4alpha -@pytest.mark.xfail(reason="GH1946") -def test_field_from_netcdf_fieldtypes(): - filenames = { - "varU": { - "lon": str(TEST_DATA / "mask_nemo_cross_180lon.nc"), - "lat": str(TEST_DATA / "mask_nemo_cross_180lon.nc"), - "data": str(TEST_DATA / "Uu_eastward_nemo_cross_180lon.nc"), - }, - "varV": { - "lon": str(TEST_DATA / "mask_nemo_cross_180lon.nc"), - "lat": str(TEST_DATA / "mask_nemo_cross_180lon.nc"), - "data": str(TEST_DATA / "Vv_eastward_nemo_cross_180lon.nc"), - }, - } - variables = {"varU": "U", "varV": "V"} - dimensions = {"lon": "glamf", "lat": "gphif"} - - # first try without setting fieldtype - fset = FieldSet.from_nemo(filenames, variables, dimensions) - assert isinstance(fset.varU.units, UnitConverter) - - # now try with setting fieldtype - fset = FieldSet.from_nemo(filenames, variables, dimensions, fieldtype={"varU": "U", "varV": "V"}) - assert isinstance(fset.varU.units, GeographicPolar) - - -@pytest.mark.v4alpha -@pytest.mark.xfail(reason="GH1946") -def test_fieldset_from_agrid_dataset(): - filenames = { - "lon": str(TEST_DATA / "mask_nemo_cross_180lon.nc"), - "lat": str(TEST_DATA / "mask_nemo_cross_180lon.nc"), - "data": str(TEST_DATA / "Uu_eastward_nemo_cross_180lon.nc"), - } - variable = {"U": "U"} - dimensions = {"lon": "glamf", "lat": "gphif"} - FieldSet.from_a_grid_dataset(filenames, variable, dimensions) - - -@pytest.mark.v4remove -@pytest.mark.xfail(reason="GH1946") -def test_fieldset_from_cgrid_interpmethod(): - filenames = { - "lon": str(TEST_DATA / "mask_nemo_cross_180lon.nc"), - "lat": str(TEST_DATA / "mask_nemo_cross_180lon.nc"), - "data": str(TEST_DATA / "Uu_eastward_nemo_cross_180lon.nc"), - } - variable = "U" - dimensions = {"lon": "glamf", "lat": "gphif"} - - with pytest.raises(TypeError): - # should fail because FieldSet.from_c_grid_dataset does not support interp_method - FieldSet.from_c_grid_dataset(filenames, variable, dimensions, interp_method="partialslip") - - -@pytest.mark.v4future -@pytest.mark.xfail(reason="GH1946") -@pytest.mark.parametrize("calltype", ["from_nemo"]) -def test_illegal_dimensionsdict(calltype): - with pytest.raises(NameError): - if calltype == "from_data": - data, dimensions = generate_fieldset_data(10, 10) - dimensions["test"] = None - FieldSet.from_data(data, dimensions) - elif calltype == "from_nemo": - fname = str(TEST_DATA / "mask_nemo_cross_180lon.nc") - filenames = {"dx": fname, "mesh_mask": fname} - variables = {"dx": "e1u"} - dimensions = {"lon": "glamu", "lat": "gphiu", "test": "test"} - FieldSet.from_nemo(filenames, variables, dimensions) - - -@pytest.mark.v4alpha -@pytest.mark.xfail(reason="GH1946") -@pytest.mark.parametrize("gridtype", ["A", "C"]) -def test_fieldset_dimlength1_cgrid(gridtype): - fieldset = FieldSet.from_data({"U": 0, "V": 0}, {"lon": 0, "lat": 0}) # TODO : Remove from_data - if gridtype == "C": - fieldset.U.interp_method = "cgrid_velocity" - fieldset.V.interp_method = "cgrid_velocity" - try: - fieldset._check_complete() - success = True if gridtype == "A" else False - except NotImplementedError: - success = True if gridtype == "C" else False - assert success - - -def assign_dataset_timestamp_dim(ds, timestamp): - """Expand dim to 'time' and assign timestamp.""" - ds.expand_dims("time") - ds["time"] = timestamp - return ds - - -def addConst(particle, fieldset, time): # pragma: no cover - particle.lon = particle.lon + fieldset.movewest + fieldset.moveeast - - -@pytest.mark.v4alpha -@pytest.mark.xfail(reason="GH1946") -def test_fieldset_constant(): - data, dimensions = generate_fieldset_data(100, 100) - fieldset = FieldSet.from_data(data, dimensions) # TODO : Remove from_data - westval = -0.2 - eastval = 0.3 - fieldset.add_constant("movewest", westval) - fieldset.add_constant("moveeast", eastval) - assert fieldset.movewest == westval - - pset = ParticleSet.from_line(fieldset, size=1, pclass=Particle, start=(0.5, 0.5), finish=(0.5, 0.5)) - pset.execute(pset.Kernel(addConst), dt=1, runtime=1) - assert abs(pset.lon[0] - (0.5 + westval + eastval)) < 1e-4 - - -@pytest.mark.v4alpha -@pytest.mark.xfail(reason="GH1946") -@pytest.mark.parametrize("swapUV", [False, True]) -def test_vector_fields(swapUV): - lon = np.linspace(0.0, 10.0, 12, dtype=np.float32) - lat = np.linspace(0.0, 10.0, 10, dtype=np.float32) - U = np.ones((10, 12), dtype=np.float32) - V = np.zeros((10, 12), dtype=np.float32) - data = {"U": U, "V": V} - dimensions = {"U": {"lat": lat, "lon": lon}, "V": {"lat": lat, "lon": lon}} - fieldset = FieldSet.from_data(data, dimensions, mesh="flat") # TODO : Remove from_data - if swapUV: # we test that we can freely edit whatever UV field - UV = VectorField("UV", fieldset.V, fieldset.U) - fieldset.add_vector_field(UV) - - pset = ParticleSet.from_line(fieldset, size=1, pclass=Particle, start=(0.5, 0.5), finish=(0.5, 0.5)) - pset.execute(AdvectionRK4, dt=1, runtime=2) - if swapUV: - assert abs(pset.lon[0] - 0.5) < 1e-9 - assert abs(pset.lat[0] - 1.5) < 1e-9 - else: - assert abs(pset.lon[0] - 1.5) < 1e-9 - assert abs(pset.lat[0] - 0.5) < 1e-9 - - -@pytest.mark.v4alpha -@pytest.mark.xfail(reason="GH1946, originated in GH938") -def test_add_second_vector_field(): - lon = np.linspace(0.0, 10.0, 12, dtype=np.float32) - lat = np.linspace(0.0, 10.0, 10, dtype=np.float32) - U = np.ones((10, 12), dtype=np.float32) - V = np.zeros((10, 12), dtype=np.float32) - data = {"U": U, "V": V} - dimensions = {"U": {"lat": lat, "lon": lon}, "V": {"lat": lat, "lon": lon}} - fieldset = FieldSet.from_data(data, dimensions, mesh="flat") # TODO : Remove from_data - - data2 = {"U2": U, "V2": V} - dimensions2 = {"lon": [ln + 0.1 for ln in lon], "lat": [lt - 0.1 for lt in lat]} - fieldset2 = FieldSet.from_data(data2, dimensions2, mesh="flat") # TODO : Remove from_data - - UV2 = VectorField("UV2", fieldset2.U2, fieldset2.V2) - fieldset.add_vector_field(UV2) - - def SampleUV2(particle, fieldset, time): # pragma: no cover - u, v = fieldset.UV2[time, particle.depth, particle.lat, particle.lon] - particle.dlon += u * particle.dt - particle.dlat += v * particle.dt - - pset = ParticleSet(fieldset, pclass=Particle, lon=0.5, lat=0.5) - pset.execute(AdvectionRK4 + pset.Kernel(SampleUV2), dt=1, runtime=2) - - assert abs(pset.lon[0] - 2.5) < 1e-9 - assert abs(pset.lat[0] - 0.5) < 1e-9 - - -@pytest.mark.v4remove -@pytest.mark.xfail(reason="time_periodic removed in v4") -@pytest.mark.parametrize("use_xarray", [True, False]) -@pytest.mark.parametrize("time_periodic", [86400.0, False]) -@pytest.mark.parametrize("dt_sign", [-1, 1]) -def test_periodic(use_xarray, time_periodic, dt_sign): - lon = np.array([0, 1], dtype=np.float32) - lat = np.array([0, 1], dtype=np.float32) - depth = np.array([0, 1], dtype=np.float32) - tsize = 24 * 60 + 1 - period = 86400 - time = np.linspace(0, period, tsize, dtype=np.float64) - - def temp_func(time): - return 20 + 2 * np.sin(time * 2 * np.pi / period) - - temp_vec = temp_func(time) - - U = np.zeros((tsize, 2, 2, 2), dtype=np.float32) - V = np.zeros((tsize, 2, 2, 2), dtype=np.float32) - V[:, 0, 0, 0] = 1e-5 - W = np.zeros((tsize, 2, 2, 2), dtype=np.float32) - temp = np.zeros((tsize, 2, 2, 2), dtype=np.float32) - temp[:, :, :, :] = temp_vec - D = np.ones((2, 2), dtype=np.float32) # adding non-timevarying field - - full_dims = {"lon": lon, "lat": lat, "depth": depth, "time": time} - dimensions = {"U": full_dims, "V": full_dims, "W": full_dims, "temp": full_dims, "D": {"lon": lon, "lat": lat}} - if use_xarray: - coords = {"lat": lat, "lon": lon, "depth": depth, "time": time} - variables = {"U": "Uxr", "V": "Vxr", "W": "Wxr", "temp": "Txr", "D": "Dxr"} - dimnames = {"lon": "lon", "lat": "lat", "depth": "depth", "time": "time"} - ds = xr.Dataset( - { - "Uxr": xr.DataArray(U, coords=coords, dims=("time", "depth", "lat", "lon")), - "Vxr": xr.DataArray(V, coords=coords, dims=("time", "depth", "lat", "lon")), - "Wxr": xr.DataArray(W, coords=coords, dims=("time", "depth", "lat", "lon")), - "Txr": xr.DataArray(temp, coords=coords, dims=("time", "depth", "lat", "lon")), - "Dxr": xr.DataArray(D, coords={"lat": lat, "lon": lon}, dims=("lat", "lon")), - } - ) - fieldset = FieldSet.from_xarray_dataset( - ds, - variables, - {"U": dimnames, "V": dimnames, "W": dimnames, "temp": dimnames, "D": {"lon": "lon", "lat": "lat"}}, - time_periodic=time_periodic, - allow_time_extrapolation=True, - ) - else: - data = {"U": U, "V": V, "W": W, "temp": temp, "D": D} - fieldset = FieldSet.from_data( - data, dimensions, mesh="flat", time_periodic=time_periodic, allow_time_extrapolation=True - ) # TODO : Remove from_data - - def sampleTemp(particle, fieldset, time): # pragma: no cover - particle.temp = fieldset.temp[time, particle.depth, particle.lat, particle.lon] - # test if we can interpolate UV and UVW together - (particle.u1, particle.v1) = fieldset.UV[time, particle.depth, particle.lat, particle.lon] - (particle.u2, particle.v2, w_) = fieldset.UVW[time, particle.depth, particle.lat, particle.lon] - # test if we can sample a non-timevarying field too - particle.d = fieldset.D[0, 0, particle.lat, particle.lon] - - MyParticle = Particle.add_variables( - [ - Variable("temp", dtype=np.float32, initial=20.0), - Variable("u1", dtype=np.float32, initial=0.0), - Variable("u2", dtype=np.float32, initial=0.0), - Variable("v1", dtype=np.float32, initial=0.0), - Variable("v2", dtype=np.float32, initial=0.0), - Variable("d", dtype=np.float32, initial=0.0), - ] +def fieldset() -> FieldSet: + """Fixture to create a FieldSet object for testing.""" + grid = XGrid.from_dataset(ds, mesh="flat") + U = Field("U", ds["U (A grid)"], grid) + V = Field("V", ds["V (A grid)"], grid) + UV = VectorField("UV", U, V) + + return FieldSet( + [U, V, UV], + ) + + +def test_fieldset_init_wrong_types(): + with pytest.raises(ValueError, match="Expected `field` to be a Field or VectorField object. Got .*"): + FieldSet([1.0, 2.0, 3.0]) + + +def test_fieldset_add_constant(fieldset): + fieldset.add_constant("test_constant", 1.0) + assert fieldset.test_constant == 1.0 + + +def test_fieldset_add_constant_field(fieldset): + fieldset.add_constant_field("test_constant_field", 1.0) + + # Get a point in the domain + time = ds["time"].mean() + depth = ds["depth"].mean() + lat = ds["lat"].mean() + lon = ds["lon"].mean() + + pytest.xfail(reason="Not yet implemented interpolation.") + assert fieldset.test_constant_field[time, depth, lat, lon] == 1.0 + + +def test_fieldset_add_field(fieldset): + grid = XGrid.from_dataset(ds, mesh="flat") + field = Field("test_field", ds["U (A grid)"], grid) + fieldset.add_field(field) + assert fieldset.test_field == field + + +def test_fieldset_add_field_wrong_type(fieldset): + not_a_field = 1.0 + with pytest.raises(ValueError, match="Expected `field` to be a Field or VectorField object. Got .*"): + fieldset.add_field(not_a_field, "test_field") + + +def test_fieldset_add_field_already_exists(fieldset): + grid = XGrid.from_dataset(ds, mesh="flat") + field = Field("test_field", ds["U (A grid)"], grid) + fieldset.add_field(field, "test_field") + with pytest.raises(ValueError, match="FieldSet already has a Field with name 'test_field'"): + fieldset.add_field(field, "test_field") + + +def test_fieldset_gridset(fieldset): + assert fieldset.fields["U"].grid in fieldset.gridset + assert fieldset.fields["V"].grid in fieldset.gridset + assert fieldset.fields["UV"].grid in fieldset.gridset + assert len(fieldset.gridset) == 1 + + fieldset.add_constant_field("constant_field", 1.0) + assert len(fieldset.gridset) == 2 + + +@pytest.mark.parametrize("ds", [pytest.param(ds, id=k) for k, ds in datasets_structured.items()]) +def test_fieldset_from_structured_generic_datasets(ds): + grid = XGrid.from_dataset(ds, mesh="flat") + fields = [] + for var in ds.data_vars: + fields.append(Field(var, ds[var], grid)) + + fieldset = FieldSet(fields) + + assert len(fieldset.fields) == len(ds.data_vars) + for field in fieldset.fields.values(): + utils.assert_valid_field_data(field.data, field.grid) + + assert len(fieldset.gridset) == 1 + + +def test_fieldset_gridset_multiple_grids(): ... + + +def test_fieldset_time_interval(): + grid1 = XGrid.from_dataset(ds, mesh="flat") + field1 = Field("field1", ds["U (A grid)"], grid1) + + ds2 = ds.copy() + ds2["time"] = (ds2["time"].dims, ds2["time"].data + np.timedelta64(timedelta(days=1)), ds2["time"].attrs) + grid2 = XGrid.from_dataset(ds2, mesh="flat") + field2 = Field("field2", ds2["U (A grid)"], grid2) + + fieldset = FieldSet([field1, field2]) + fieldset.add_constant_field("constant_field", 1.0) + + assert fieldset.time_interval.left == np.datetime64("2000-01-02") + assert fieldset.time_interval.right == np.datetime64("2001-01-01") + + +def test_fieldset_time_interval_constant_fields(): + fieldset = FieldSet([]) + fieldset.add_constant_field("constant_field", 1.0) + fieldset.add_constant_field("constant_field2", 2.0) + + assert fieldset.time_interval is None + + +def test_fieldset_init_incompatible_calendars(): + ds1 = ds.copy() + ds1["time"] = ( + ds1["time"].dims, + xr.date_range("2000", "2001", T_structured, calendar="365_day", use_cftime=True), + ds1["time"].attrs, + ) + + grid = XGrid.from_dataset(ds1, mesh="flat") + U = Field("U", ds1["U (A grid)"], grid) + V = Field("V", ds1["V (A grid)"], grid) + UV = VectorField("UV", U, V) + + ds2 = ds.copy() + ds2["time"] = ( + ds2["time"].dims, + xr.date_range("2000", "2001", T_structured, calendar="360_day", use_cftime=True), + ds2["time"].attrs, + ) + grid2 = XGrid.from_dataset(ds2, mesh="flat") + incompatible_calendar = Field("test", ds2["data_g"], grid2) + + with pytest.raises(CalendarError, match="Expected field '.*' to have calendar compatible with datetime object"): + FieldSet([U, V, UV, incompatible_calendar]) + + +def test_fieldset_add_field_incompatible_calendars(fieldset): + ds_test = ds.copy() + ds_test["time"] = ( + ds_test["time"].dims, + xr.date_range("2000", "2001", T_structured, calendar="360_day", use_cftime=True), + ds_test["time"].attrs, ) + grid = XGrid.from_dataset(ds_test, mesh="flat") + field = Field("test_field", ds_test["data_g"], grid) - pset = ParticleSet(fieldset, pclass=MyParticle, lon=[0.5], lat=[0.5], depth=[0.5]) - pset.execute( - AdvectionRK4_3D + pset.Kernel(sampleTemp), runtime=timedelta(hours=51), dt=timedelta(hours=dt_sign * 1) + with pytest.raises(CalendarError, match="Expected field '.*' to have calendar compatible with datetime object"): + fieldset.add_field(field, "test_field") + + ds_test = ds.copy() + ds_test["time"] = ( + ds_test["time"].dims, + np.linspace(0, 100, T_structured, dtype="timedelta64[s]"), + ds_test["time"].attrs, ) + grid = XGrid.from_dataset(ds_test, mesh="flat") + field = Field("test_field", ds_test["data_g"], grid) + + with pytest.raises(CalendarError, match="Expected field '.*' to have calendar compatible with datetime object"): + fieldset.add_field(field, "test_field") + + +@pytest.mark.parametrize( + "input_, expected", + [ + (cftime.DatetimeNoLeap(2000, 1, 1), " with cftime calendar noleap'"), + (cftime.Datetime360Day(2000, 1, 1), " with cftime calendar 360_day'"), + (cftime.DatetimeJulian(2000, 1, 1), " with cftime calendar julian'"), + ( + cftime.DatetimeGregorian(2000, 1, 1), + " with cftime calendar standard'", + ), + (np.datetime64("2000-01-01"), ""), + (cftime.datetime(2000, 1, 1), " with cftime calendar standard'"), + ], +) +def test_datetime_to_msg(input_, expected): + assert _datetime_to_msg(input_) == expected + + +def test_fieldset_samegrids_UV(): + """Test that if a simple fieldset with U and V is created, that only one grid object is defined.""" + ... + + +def test_fieldset_grid_deduplication(): + """Tests that for a full fieldset that the number of grid objects is as expected + (sharing of grid objects so that the particle location is not duplicated). + + When grid deduplication is actually implemented, this might need to be refactored + into multiple tests (/more might be needed). + """ + ... + - if time_periodic is not False: - t = pset.time[0] - temp_theo = temp_func(t) - elif dt_sign == 1: - temp_theo = temp_vec[-1] - elif dt_sign == -1: - temp_theo = temp_vec[0] - assert np.allclose(temp_theo, pset.temp[0], atol=1e-5) - assert np.allclose(pset.u1[0], pset.u2[0]) - assert np.allclose(pset.v1[0], pset.v2[0]) - assert np.allclose(pset.d[0], 1.0) - - -@pytest.mark.v4alpha -@pytest.mark.xfail(reason="GH1946") -def test_fieldset_from_data_gridtypes(): - """Simple test for fieldset initialisation from data.""" - xdim, ydim, zdim = 20, 10, 4 - - lon = np.linspace(0.0, 10.0, xdim, dtype=np.float32) - lat = np.linspace(0.0, 10.0, ydim, dtype=np.float32) - depth = np.linspace(0.0, 1.0, zdim, dtype=np.float32) - depth_s = np.ones((zdim, ydim, xdim)) - U = np.ones((zdim, ydim, xdim)) - V = np.ones((zdim, ydim, xdim)) - dimensions = {"lat": lat, "lon": lon, "depth": depth} - data = {"U": np.array(U, dtype=np.float32), "V": np.array(V, dtype=np.float32)} - lonm, latm = np.meshgrid(lon, lat) - for k in range(zdim): - data["U"][k, :, :] = lonm * (depth[k] + 1) + 0.1 - depth_s[k, :, :] = depth[k] - - # Rectilinear Z grid - fieldset = FieldSet.from_data(data, dimensions, mesh="flat") # TODO : Remove from_data - pset = ParticleSet(fieldset, Particle, [0, 0], [0, 0], [0, 0.4]) - pset.execute(AdvectionRK4, runtime=1.5, dt=0.5) - plon = pset.lon - plat = pset.lat - # sol of dx/dt = (init_depth+1)*x+0.1; x(0)=0 - assert np.allclose(plon, [0.17173462592827032, 0.2177736932123214]) - assert np.allclose(plat, [1, 1]) - - # Rectilinear S grid - dimensions["depth"] = depth_s - fieldset = FieldSet.from_data(data, dimensions, mesh="flat") # TODO : Remove from_data - pset = ParticleSet(fieldset, Particle, [0, 0], [0, 0], [0, 0.4]) - pset.execute(AdvectionRK4, runtime=1.5, dt=0.5) - assert np.allclose(plon, pset.lon) - assert np.allclose(plat, pset.lat) - - # Curvilinear Z grid - dimensions["lon"] = lonm - dimensions["lat"] = latm - dimensions["depth"] = depth - fieldset = FieldSet.from_data(data, dimensions, mesh="flat") # TODO : Remove from_data - pset = ParticleSet(fieldset, Particle, [0, 0], [0, 0], [0, 0.4]) - pset.execute(AdvectionRK4, runtime=1.5, dt=0.5) - assert np.allclose(plon, pset.lon) - assert np.allclose(plat, pset.lat) - - # Curvilinear S grid - dimensions["depth"] = depth_s - fieldset = FieldSet.from_data(data, dimensions, mesh="flat") # TODO : Remove from_data - pset = ParticleSet(fieldset, Particle, [0, 0], [0, 0], [0, 0.4]) - pset.execute(AdvectionRK4, runtime=1.5, dt=0.5) - assert np.allclose(plon, pset.lon) - assert np.allclose(plat, pset.lat) +def test_fieldset_add_field_after_pset(): + # ? Should it be allowed to add fields (normal or vector) after a ParticleSet has been initialized? + ... + + +_COPERNICUS_DATASETS = [ + datasets_circulation_models["ds_copernicusmarine"], + datasets_circulation_models["ds_copernicusmarine_waves"], +] + + +@pytest.mark.parametrize("ds", _COPERNICUS_DATASETS) +def test_fieldset_from_copernicusmarine(ds, caplog): + fieldset = FieldSet.from_copernicusmarine(ds) + assert "U" in fieldset.fields + assert "V" in fieldset.fields + assert "UV" in fieldset.fields + assert "renamed it to 'U'" in caplog.text + assert "renamed it to 'V'" in caplog.text + + +def test_fieldset_from_copernicusmarine_no_currents(caplog): + ds = datasets_circulation_models["ds_copernicusmarine"].cf.drop_vars( + ["eastward_sea_water_velocity", "northward_sea_water_velocity"] + ) + fieldset = FieldSet.from_copernicusmarine(ds) + assert "U" not in fieldset.fields + assert "V" not in fieldset.fields + assert "UV" not in fieldset.fields + assert caplog.text == "" + + +@pytest.mark.parametrize("ds", _COPERNICUS_DATASETS) +def test_fieldset_from_copernicusmarine_no_logs(ds, caplog): + ds = ds.copy() + zeros = xr.zeros_like(list(ds.data_vars.values())[0]) + ds["U"] = zeros + ds["V"] = zeros + + fieldset = FieldSet.from_copernicusmarine(ds) + assert "U" in fieldset.fields + assert "V" in fieldset.fields + assert "UV" in fieldset.fields + assert caplog.text == "" + + +def test_fieldset_from_copernicusmarine_with_W(caplog): + ds = datasets_circulation_models["ds_copernicusmarine"] + ds = ds.copy() + ds["wo"] = ds["uo"] + ds["wo"].attrs["standard_name"] = "vertical_sea_water_velocity" + + fieldset = FieldSet.from_copernicusmarine(ds) + assert "U" in fieldset.fields + assert "V" in fieldset.fields + assert "W" in fieldset.fields + assert "UV" not in fieldset.fields + assert "UVW" in fieldset.fields + assert "renamed it to 'W'" in caplog.text diff --git a/tests/v4/test_index_search.py b/tests/test_index_search.py similarity index 93% rename from tests/v4/test_index_search.py rename to tests/test_index_search.py index b664ae42f9..ef27bcd65f 100644 --- a/tests/v4/test_index_search.py +++ b/tests/test_index_search.py @@ -3,11 +3,10 @@ import xarray as xr import xgcm +from parcels import Field, XGrid +from parcels._core.index_search import _search_indices_curvilinear_2d from parcels._datasets.structured.generic import datasets -from parcels._index_search import _search_indices_curvilinear_2d -from parcels.field import Field -from parcels.tools.exampledata_utils import download_example_dataset -from parcels.xgrid import XGrid +from parcels._tutorial import download_example_dataset @pytest.fixture diff --git a/tests/test_interpolation.py b/tests/test_interpolation.py index 67eab64204..3b9478d3cb 100644 --- a/tests/test_interpolation.py +++ b/tests/test_interpolation.py @@ -1,64 +1,229 @@ +import numpy as np import pytest - -import parcels._interpolation as interpolation -from tests.utils import create_fieldset_zeros_3d +import xarray as xr + +from parcels import ( + Field, + FieldSet, + Particle, + ParticleFile, + ParticleSet, + StatusCode, + UxGrid, + Variable, + VectorField, + XGrid, +) +from parcels._core.index_search import _search_time_index +from parcels._datasets.structured.generated import simple_UV_dataset +from parcels._datasets.unstructured.generic import datasets as datasets_unstructured +from parcels.interpolators import ( + UXPiecewiseLinearNode, + XFreeslip, + XLinear, + XNearest, + XPartialslip, + ZeroInterpolator, +) +from parcels.kernels import AdvectionRK4_3D +from tests.utils import TEST_DATA @pytest.fixture -def tmp_interpolator_registry(): - """Resets the interpolator registry after the test. Vital when testing manipulating the registry.""" - old_2d = interpolation._interpolator_registry_2d.copy() - old_3d = interpolation._interpolator_registry_3d.copy() - yield - interpolation._interpolator_registry_2d = old_2d - interpolation._interpolator_registry_3d = old_3d - - -@pytest.mark.usefixtures("tmp_interpolator_registry") -def test_interpolation_registry(): - @interpolation.register_3d_interpolator("test") - @interpolation.register_2d_interpolator("test") - def some_function(): - return "test" - - assert "test" in interpolation.get_2d_interpolator_registry() - assert "test" in interpolation.get_3d_interpolator_registry() - - f = interpolation.get_2d_interpolator_registry()["test"] - g = interpolation.get_3d_interpolator_registry()["test"] - assert f() == g() == "test" - - -@pytest.mark.v4remove -@pytest.mark.xfail(reason="GH1946") -@pytest.mark.usefixtures("tmp_interpolator_registry") -def test_interpolator_override(): - fieldset = create_fieldset_zeros_3d() - - @interpolation.register_3d_interpolator("linear") - def test_interpolator(ctx: interpolation.InterpolationContext3D): - raise NotImplementedError - - with pytest.raises(NotImplementedError): - fieldset.U[0, 0.5, 0.5, 0.5] - - -@pytest.mark.v4remove -@pytest.mark.xfail(reason="GH1946") -@pytest.mark.usefixtures("tmp_interpolator_registry") -def test_full_depth_provided_to_interpolators(): - """The full depth needs to be provided to the interpolation schemes as some interpolators - need to know whether they are at the surface or bottom of the water column. - - https://github.com/OceanParcels/Parcels/pull/1816#discussion_r1908840408 - """ - xdim, ydim, zdim = 10, 11, 12 - fieldset = create_fieldset_zeros_3d(xdim=xdim, ydim=ydim, zdim=zdim) - - @interpolation.register_3d_interpolator("linear") - def test_interpolator2(ctx: interpolation.InterpolationContext3D): - assert ctx.data.shape[1] == zdim - # The array z dimension is the same as the fieldset z dimension - return 0 - - fieldset.U[0.5, 0.5, 0.5, 0.5] +def field(): + """Reference data used for testing interpolation.""" + z0 = np.array( # each x is +1 from the previous, each y is +2 from the previous + [ + [0.0, 1.0, 2.0, 3.0], + [2.0, 3.0, 4.0, 5.0], + [4.0, 5.0, 6.0, 7.0], + [6.0, 7.0, 8.0, 9.0], + ] + ) + spatial_data = np.array([z0, z0 + 3, z0 + 6, z0 + 9]) # each z is +3 from the previous + temporal_data = np.array([spatial_data, spatial_data + 10, spatial_data + 20]) # each t is +10 from the previous + + ds = xr.Dataset( + {"U": (["time", "depth", "lat", "lon"], temporal_data)}, + coords={ + "time": (["time"], [np.timedelta64(t, "s") for t in [0, 2, 4]], {"axis": "T"}), + "depth": (["depth"], [0, 1, 2, 3], {"axis": "Z"}), + "lat": (["lat"], [0, 1, 2, 3], {"axis": "Y", "c_grid_axis_shift": -0.5}), + "lon": (["lon"], [0, 1, 2, 3], {"axis": "X", "c_grid_axis_shift": -0.5}), + "x": (["x"], [0.5, 1.5, 2.5, 3.5], {"axis": "X"}), + "y": (["y"], [0.5, 1.5, 2.5, 3.5], {"axis": "Y"}), + }, + ) + return Field("U", ds["U"], XGrid.from_dataset(ds)) + + +@pytest.mark.parametrize( + "func, t, z, y, x, expected", + [ + pytest.param(ZeroInterpolator, np.timedelta64(1, "s"), 2.5, 0.49, 0.51, 0, id="Zero"), + pytest.param( + XLinear, + [np.timedelta64(0, "s"), np.timedelta64(1, "s")], + [0, 0], + [0.49, 0.49], + [0.51, 0.51], + [1.49, 6.49], + id="Linear", + ), + pytest.param(XLinear, np.timedelta64(1, "s"), 2.5, 0.49, 0.51, 13.99, id="Linear-2"), + pytest.param( + XNearest, + [np.timedelta64(0, "s"), np.timedelta64(3, "s")], + [0.2, 0.2], + [0.2, 0.2], + [0.51, 0.51], + [1.0, 16.0], + id="Nearest", + ), + ], +) +def test_raw_2d_interpolation(field, func, t, z, y, x, expected): + """Test the interpolation functions on the Field.""" + tau, ti = _search_time_index(field, t) + position = field.grid.search(z, y, x) + + value = func(field, ti, position, tau, 0, 0, y, x) + np.testing.assert_equal(value, expected) + + +@pytest.mark.parametrize( + "func, t, z, y, x, expected", + [ + (XPartialslip, np.timedelta64(1, "s"), 0, 0, 0.0, [[1], [1]]), + (XFreeslip, np.timedelta64(1, "s"), 0, 0.5, 1.5, [[1], [0.5]]), + (XPartialslip, np.timedelta64(1, "s"), 0, 2.5, 1.5, [[0.75], [0.5]]), + (XFreeslip, np.timedelta64(1, "s"), 0, 2.5, 1.5, [[1], [0.5]]), + (XPartialslip, np.timedelta64(1, "s"), 0, 1.5, 0.5, [[0.5], [0.75]]), + (XFreeslip, np.timedelta64(1, "s"), 0, 1.5, 0.5, [[0.5], [1]]), + ( + XFreeslip, + [np.timedelta64(1, "s"), np.timedelta64(0, "s")], + [0, 2], + [1.5, 1.5], + [2.5, 0.5], + [[0.5, 0.5], [1, 1]], + ), + ], +) +def test_spatial_slip_interpolation(field, func, t, z, y, x, expected): + field.data[:] = 1.0 + field.data[:, :, 1:3, 1:3] = 0.0 # Set zero land value to test spatial slip + U = field + V = field + UV = VectorField("UV", U, V, vector_interp_method=func) + + velocities = UV[t, z, y, x] + np.testing.assert_array_almost_equal(velocities, expected) + + +@pytest.mark.parametrize("mesh", ["spherical", "flat"]) +def test_interpolation_mesh_type(mesh, npart=10): + ds = simple_UV_dataset(mesh=mesh) + ds["U"].data[:] = 1.0 + grid = XGrid.from_dataset(ds, mesh=mesh) + U = Field("U", ds["U"], grid, interp_method=XLinear) + V = Field("V", ds["V"], grid, interp_method=XLinear) + UV = VectorField("UV", U, V) + + lat = 30.0 + time = U.time_interval.left + u_expected = 1.0 if mesh == "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 + + +def test_default_interpolator_set_correctly(): + ds = simple_UV_dataset() + grid = XGrid.from_dataset(ds) + U = Field("U", ds["U"], grid) + assert U.interp_method == XLinear + + ds = datasets_unstructured["stommel_gyre_delaunay"] + grid = UxGrid(grid=ds.uxgrid, z=ds.coords["nz"]) + U = Field("U", ds["U"], grid) + assert U.interp_method == UXPiecewiseLinearNode + + +interp_methods = { + "linear": XLinear, +} + + +@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, mesh="flat") + U = Field("U", ds["U"], grid, interp_method=interp_methods[interp_name]) + V = Field("V", ds["V"], grid, interp_method=interp_methods[interp_name]) + W = Field("W", ds["W"], grid, 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.state = StatusCode.Delete + + outfile = 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) diff --git a/tests/v4/test_kernel.py b/tests/test_kernel.py similarity index 97% rename from tests/v4/test_kernel.py rename to tests/test_kernel.py index 3ece81a744..023a1cada1 100644 --- a/tests/v4/test_kernel.py +++ b/tests/test_kernel.py @@ -2,15 +2,15 @@ import pytest from parcels import ( - AdvectionRK4, Field, FieldSet, + Kernel, + Particle, ParticleSet, + XGrid, ) from parcels._datasets.structured.generic import datasets as datasets_structured -from parcels.kernel import Kernel -from parcels.particle import Particle -from parcels.xgrid import XGrid +from parcels.kernels import AdvectionRK4 from tests.common_kernels import MoveEast, MoveNorth diff --git a/tests/v4/test_particle.py b/tests/test_particle.py similarity index 97% rename from tests/v4/test_particle.py rename to tests/test_particle.py index 3ca0a6c131..fdc05ea147 100644 --- a/tests/v4/test_particle.py +++ b/tests/test_particle.py @@ -1,9 +1,15 @@ import numpy as np import pytest +from parcels._core.particle import ( + _SAME_AS_FIELDSET_TIME_INTERVAL, + Particle, + ParticleClass, + Variable, + create_particle_data, +) from parcels._core.utils.time import TimeInterval from parcels._datasets.structured.generic import TIME -from parcels.particle import _SAME_AS_FIELDSET_TIME_INTERVAL, Particle, ParticleClass, Variable, create_particle_data def test_variable_init(): diff --git a/tests/v4/test_particlefile.py b/tests/test_particlefile.py similarity index 98% rename from tests/v4/test_particlefile.py rename to tests/test_particlefile.py index 5937f59f7f..fbe222f95b 100755 --- a/tests/v4/test_particlefile.py +++ b/tests/test_particlefile.py @@ -8,13 +8,11 @@ from zarr.storage import MemoryStore import parcels -from parcels import AdvectionRK4, Field, FieldSet, Particle, ParticleSet, Variable, VectorField +from parcels import Field, FieldSet, Particle, ParticleFile, ParticleSet, StatusCode, Variable, VectorField, XGrid +from parcels._core.particle import Particle, create_particle_data, get_default_particle from parcels._core.utils.time import TimeInterval from parcels._datasets.structured.generic import datasets -from parcels.particle import Particle, create_particle_data, get_default_particle -from parcels.particlefile import ParticleFile -from parcels.tools.statuscodes import StatusCode -from parcels.xgrid import XGrid +from parcels.kernels import AdvectionRK4 from tests.common_kernels import DoNothing diff --git a/tests/v4/test_particleset.py b/tests/test_particleset.py similarity index 99% rename from tests/v4/test_particleset.py rename to tests/test_particleset.py index b99b708b25..3f18b02943 100644 --- a/tests/v4/test_particleset.py +++ b/tests/test_particleset.py @@ -13,9 +13,9 @@ ParticleSet, ParticleSetWarning, Variable, + XGrid, ) from parcels._datasets.structured.generic import datasets as datasets_structured -from parcels.xgrid import XGrid from tests.common_kernels import DoNothing from tests.utils import round_and_hash_float_array diff --git a/tests/v4/test_particleset_execute.py b/tests/test_particleset_execute.py similarity index 98% rename from tests/v4/test_particleset_execute.py rename to tests/test_particleset_execute.py index fbad3d5ece..0aacb1a2d5 100644 --- a/tests/v4/test_particleset_execute.py +++ b/tests/test_particleset_execute.py @@ -5,23 +5,25 @@ import pytest from parcels import ( - AdvectionEE, Field, + FieldInterpolationError, + FieldOutOfBoundError, FieldSet, Particle, + ParticleFile, ParticleSet, StatusCode, - UXPiecewiseConstantFace, + TimeExtrapolationError, + UxGrid, Variable, VectorField, + XGrid, ) from parcels._datasets.structured.generated import simple_UV_dataset from parcels._datasets.structured.generic import datasets as datasets_structured from parcels._datasets.unstructured.generic import datasets as datasets_unstructured -from parcels.particlefile import ParticleFile -from parcels.tools.statuscodes import FieldInterpolationError, FieldOutOfBoundError, TimeExtrapolationError -from parcels.uxgrid import UxGrid -from parcels.xgrid import XGrid +from parcels.interpolators import UXPiecewiseConstantFace +from parcels.kernels import AdvectionEE from tests import utils from tests.common_kernels import DoNothing diff --git a/tests/v4/test_spatialhash.py b/tests/test_spatialhash.py similarity index 96% rename from tests/v4/test_spatialhash.py rename to tests/test_spatialhash.py index bdec911735..500ce562cd 100644 --- a/tests/v4/test_spatialhash.py +++ b/tests/test_spatialhash.py @@ -1,7 +1,7 @@ import numpy as np +from parcels import XGrid from parcels._datasets.structured.generic import datasets -from parcels.xgrid import XGrid def test_spatialhash_init(): diff --git a/tests/v4/test_structured_gcm.py b/tests/test_structured_gcm.py similarity index 100% rename from tests/v4/test_structured_gcm.py rename to tests/test_structured_gcm.py diff --git a/tests/v4/test_utils.py b/tests/test_utils.py similarity index 100% rename from tests/v4/test_utils.py rename to tests/test_utils.py diff --git a/tests/v4/test_uxarray_fieldset.py b/tests/test_uxarray_fieldset.py similarity index 99% rename from tests/v4/test_uxarray_fieldset.py rename to tests/test_uxarray_fieldset.py index 0ab13f8cae..d279dac4f2 100644 --- a/tests/v4/test_uxarray_fieldset.py +++ b/tests/test_uxarray_fieldset.py @@ -7,13 +7,15 @@ FieldSet, Particle, ParticleSet, - UXPiecewiseConstantFace, - UXPiecewiseLinearNode, + UxGrid, VectorField, download_example_dataset, ) from parcels._datasets.unstructured.generic import datasets as datasets_unstructured -from parcels.uxgrid import UxGrid +from parcels.interpolators import ( + UXPiecewiseConstantFace, + UXPiecewiseLinearNode, +) @pytest.fixture diff --git a/tests/v4/test_uxgrid.py b/tests/test_uxgrid.py similarity index 95% rename from tests/v4/test_uxgrid.py rename to tests/test_uxgrid.py index 9c675b278b..91d33b8f0d 100644 --- a/tests/v4/test_uxgrid.py +++ b/tests/test_uxgrid.py @@ -1,7 +1,7 @@ import pytest +from parcels import UxGrid from parcels._datasets.unstructured.generic import datasets as uxdatasets -from parcels.uxgrid import UxGrid @pytest.mark.parametrize("uxds", [pytest.param(uxds, id=key) for key, uxds in uxdatasets.items()]) diff --git a/tests/v4/test_xarray_fieldset.py b/tests/test_xarray_fieldset.py similarity index 100% rename from tests/v4/test_xarray_fieldset.py rename to tests/test_xarray_fieldset.py diff --git a/tests/v4/test_xgrid.py b/tests/test_xgrid.py similarity index 98% rename from tests/v4/test_xgrid.py rename to tests/test_xgrid.py index eb9e8524e8..d93e4689ad 100644 --- a/tests/v4/test_xgrid.py +++ b/tests/test_xgrid.py @@ -6,14 +6,12 @@ import xarray as xr from numpy.testing import assert_allclose -from parcels._datasets.structured.generic import X, Y, Z, datasets -from parcels.xgrid import ( - LEFT_OUT_OF_BOUNDS, - RIGHT_OUT_OF_BOUNDS, +from parcels._core.index_search import LEFT_OUT_OF_BOUNDS, RIGHT_OUT_OF_BOUNDS, _search_1d_array +from parcels._core.xgrid import ( XGrid, - _search_1d_array, _transpose_xfield_data_to_tzyx, ) +from parcels._datasets.structured.generic import X, Y, Z, datasets from tests import utils GridTestCase = namedtuple("GridTestCase", ["ds", "attr", "expected"]) diff --git a/tests/utils.py b/tests/utils.py index c89cbea45b..4653c27dc3 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -9,11 +9,10 @@ import xarray as xr import parcels +from parcels import Field, FieldSet, VectorField +from parcels._core.xgrid import _FIELD_DATA_ORDERING, XGrid, get_axis_from_dim_name from parcels._datasets.structured.generated import simple_UV_dataset -from parcels.application_kernels.interpolation import XLinear -from parcels.field import Field, VectorField -from parcels.fieldset import FieldSet -from parcels.xgrid import _FIELD_DATA_ORDERING, XGrid, get_axis_from_dim_name +from parcels.interpolators import XLinear PROJECT_ROOT = Path(__file__).resolve().parents[1] TEST_ROOT = PROJECT_ROOT / "tests" diff --git a/tests/v4/utils/test_time.py b/tests/utils/test_time.py similarity index 100% rename from tests/v4/utils/test_time.py rename to tests/utils/test_time.py diff --git a/tests/v4/utils/test_unstructured.py b/tests/utils/test_unstructured.py similarity index 100% rename from tests/v4/utils/test_unstructured.py rename to tests/utils/test_unstructured.py diff --git a/tests/v4/test_advection.py b/tests/v4/test_advection.py deleted file mode 100644 index 906c55a979..0000000000 --- a/tests/v4/test_advection.py +++ /dev/null @@ -1,564 +0,0 @@ -import numpy as np -import pytest -import xarray as xr -import xgcm - -import parcels -from parcels._datasets.structured.generated import ( - decaying_moving_eddy_dataset, - moving_eddy_dataset, - peninsula_dataset, - radial_rotation_dataset, - simple_UV_dataset, - stommel_gyre_dataset, -) -from parcels.application_kernels.advection import AdvectionEE, AdvectionRK4, AdvectionRK4_3D, AdvectionRK45 -from parcels.application_kernels.advectiondiffusion import AdvectionDiffusionEM, AdvectionDiffusionM1 -from parcels.application_kernels.interpolation import CGrid_Velocity, XLinear -from parcels.field import Field, VectorField -from parcels.fieldset import FieldSet -from parcels.particle import Particle, Variable -from parcels.particlefile import ParticleFile -from parcels.particleset import ParticleSet -from parcels.tools.statuscodes import StatusCode -from parcels.xgrid import XGrid -from tests.utils import round_and_hash_float_array - -kernel = { - "EE": AdvectionEE, - "RK4": AdvectionRK4, - "RK4_3D": AdvectionRK4_3D, - "RK45": AdvectionRK45, - # "AA": AdvectionAnalytical, - "AdvDiffEM": AdvectionDiffusionEM, - "AdvDiffM1": AdvectionDiffusionM1, -} - - -@pytest.mark.parametrize("mesh", ["spherical", "flat"]) -def test_advection_zonal(mesh, npart=10): - """Particles at high latitude move geographically faster due to the pole correction in `GeographicPolar`.""" - ds = simple_UV_dataset(mesh=mesh) - ds["U"].data[:] = 1.0 - grid = XGrid.from_dataset(ds, mesh=mesh) - U = Field("U", ds["U"], grid, interp_method=XLinear) - V = Field("V", ds["V"], grid, interp_method=XLinear) - 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 == "spherical": - assert (np.diff(pset.lon) > 1.0e-4).all() - else: - assert (np.diff(pset.lon) < 1.0e-4).all() - - -def test_advection_zonal_with_particlefile(tmp_store): - """Particles at high latitude move geographically faster due to the pole correction in `GeographicPolar`.""" - npart = 10 - ds = simple_UV_dataset(mesh="flat") - ds["U"].data[:] = 1.0 - grid = XGrid.from_dataset(ds, mesh="flat") - U = Field("U", ds["U"], grid, interp_method=XLinear) - V = Field("V", ds["V"], grid, interp_method=XLinear) - 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)) - pfile = ParticleFile(tmp_store, outputdt=np.timedelta64(30, "m")) - pset.execute(AdvectionRK4, runtime=np.timedelta64(2, "h"), dt=np.timedelta64(15, "m"), output_file=pfile) - - assert (np.diff(pset.lon) < 1.0e-4).all() - ds = xr.open_zarr(tmp_store) - np.testing.assert_allclose(ds.isel(obs=-1).lon.values, pset.lon) - - -def periodicBC(particles, fieldset): - particles.total_dlon += particles.dlon - particles.lon = np.fmod(particles.lon, 2) - - -def test_advection_zonal_periodic(): - ds = simple_UV_dataset(dims=(2, 2, 2, 2), mesh="flat") - ds["U"].data[:] = 0.1 - ds["lon"].data = np.array([0, 2]) - ds["lat"].data = np.array([0, 2]) - - # add a halo - halo = ds.isel(XG=0) - halo.lon.values = ds.lon.values[1] + 1 - halo.XG.values = ds.XG.values[1] + 2 - ds = xr.concat([ds, halo], dim="XG") - - grid = XGrid.from_dataset(ds) - U = Field("U", ds["U"], grid, interp_method=XLinear) - V = Field("V", ds["V"], grid, interp_method=XLinear) - UV = VectorField("UV", U, V) - fieldset = FieldSet([U, V, UV]) - - PeriodicParticle = Particle.add_variable(Variable("total_dlon", initial=0)) - startlon = np.array([0.5, 0.4]) - pset = ParticleSet(fieldset, pclass=PeriodicParticle, lon=startlon, lat=[0.5, 0.5]) - pset.execute([AdvectionEE, periodicBC], runtime=np.timedelta64(40, "s"), dt=np.timedelta64(1, "s")) - np.testing.assert_allclose(pset.total_dlon, 4, atol=1e-5) - np.testing.assert_allclose(pset.lon + pset.dlon, startlon, atol=1e-5) - np.testing.assert_allclose(pset.lat, 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="flat") - ds["U"].data[:] = 1.0 - grid = XGrid.from_dataset(ds) - U = Field("U", ds["U"], grid, interp_method=XLinear) - U.data[:, 0, :, :] = 0.0 # Set U to 0 at the surface - V = Field("V", ds["V"], grid, interp_method=XLinear) - 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") - np.testing.assert_allclose(pset.lon, expected_lon, 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="flat") - grid = XGrid.from_dataset(ds) - U = Field("U", ds["U"], grid, interp_method=XLinear) - U.data[:] = 0.01 # Set U to small value (to avoid horizontal out of bounds) - V = Field("V", ds["V"], grid, interp_method=XLinear) - W = Field("W", ds["V"], grid, interp_method=XLinear) # 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(particles, fieldset): # pragma: no cover - particles.state = np.where(particles.state == StatusCode.ErrorOutOfBounds, StatusCode.Delete, particles.state) - particles.state = np.where( - particles.state == StatusCode.ErrorThroughSurface, StatusCode.Delete, particles.state - ) - - def SubmergeParticle(particles, fieldset): # pragma: no cover - if len(particles.state) == 0: - return - inds = np.argwhere(particles.state == StatusCode.ErrorThroughSurface).flatten() - if len(inds) == 0: - return - dt = particles.dt / np.timedelta64(1, "s") - (u, v) = fieldset.UV[particles[inds]] - particles[inds].dlon = u * dt - particles[inds].dlat = v * dt - particles[inds].ddepth = 0.0 - particles[inds].depth = 0 - particles[inds].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: - np.testing.assert_allclose(pset.lon[0], 0.6, atol=1e-5) - np.testing.assert_allclose(pset.depth[0], 0, atol=1e-5) - else: - assert len(pset) == 0 - - -@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=XLinear) - V = Field("V", ds["V"], grid, interp_method=XLinear) - fields = [U, V, VectorField("UV", U, V)] - if w: - W = Field("W", ds["W"], grid, interp_method=XLinear) - 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]) - np.testing.assert_allclose(np.array([p.lon - x0 for p in pset]), 4 * u, atol=1e-6) - np.testing.assert_allclose(np.array([p.lat - y0 for p in pset]), 4 * v, atol=1e-6) - if w: - np.testing.assert_allclose(np.array([p.depth - z0 for p in pset]), 4 * w, atol=1e-6) - - -def test_radialrotation(npart=10): - ds = radial_rotation_dataset() - grid = XGrid.from_dataset(ds, mesh="flat") - U = parcels.Field("U", ds["U"], grid, interp_method=XLinear) - V = parcels.Field("V", ds["V"], grid, interp_method=XLinear) - UV = parcels.VectorField("UV", U, V) - fieldset = parcels.FieldSet([U, V, UV]) - - dt = np.timedelta64(30, "s") - lon = np.linspace(32, 50, npart) - lat = np.ones(npart) * 30 - starttime = np.arange(np.timedelta64(0, "s"), npart * dt, dt) - - pset = parcels.ParticleSet(fieldset, lon=lon, lat=lat, time=starttime) - pset.execute(parcels.AdvectionRK4, endtime=np.timedelta64(10, "m"), dt=dt) - - theta = 2 * np.pi * (pset.time - starttime) / np.timedelta64(24 * 3600, "s") - true_lon = (lon - 30.0) * np.cos(theta) + 30.0 - true_lat = -(lon - 30.0) * np.sin(theta) + 30.0 - - np.testing.assert_allclose(pset.lon, true_lon, atol=5e-2) - np.testing.assert_allclose(pset.lat, true_lat, atol=5e-2) - - -@pytest.mark.parametrize( - "method, rtol", - [ - ("EE", 1e-2), - ("AdvDiffEM", 1e-2), - ("AdvDiffM1", 1e-2), - ("RK4", 1e-5), - ("RK4_3D", 1e-5), - ("RK45", 1e-4), - ], -) -def test_moving_eddy(method, rtol): - ds = moving_eddy_dataset() - grid = XGrid.from_dataset(ds) - U = Field("U", ds["U"], grid, interp_method=XLinear) - V = Field("V", ds["V"], grid, interp_method=XLinear) - if method == "RK4_3D": - # Using W to test 3D advection (assuming same velocity as V) - W = Field("W", ds["V"], grid, interp_method=XLinear) - UVW = VectorField("UVW", U, V, W) - fieldset = FieldSet([U, V, W, UVW]) - else: - UV = VectorField("UV", U, V) - fieldset = FieldSet([U, V, UV]) - if method in ["AdvDiffEM", "AdvDiffM1"]: - # Add zero diffusivity field for diffusion kernels - ds["Kh"] = (["time", "depth", "YG", "XG"], np.full(ds["U"].shape, 0)) - fieldset.add_field(Field("Kh", ds["Kh"], grid, interp_method=XLinear), "Kh_zonal") - fieldset.add_field(Field("Kh", ds["Kh"], grid, interp_method=XLinear), "Kh_meridional") - fieldset.add_constant("dres", 0.1) - - start_lon, start_lat, start_depth = 12000, 12500, 12500 - dt = np.timedelta64(30, "m") - - if method == "RK45": - fieldset.add_constant("RK45_tol", rtol) - - pset = ParticleSet(fieldset, lon=start_lon, lat=start_lat, depth=start_depth, time=np.timedelta64(0, "s")) - pset.execute(kernel[method], dt=dt, endtime=np.timedelta64(1, "h")) - - def truth_moving(x_0, y_0, t): - t /= np.timedelta64(1, "s") - lat = y_0 - (ds.u_0 - ds.u_g) / ds.f * (1 - np.cos(ds.f * t)) - lon = x_0 + ds.u_g * t + (ds.u_0 - ds.u_g) / ds.f * np.sin(ds.f * t) - return lon, lat - - exp_lon, exp_lat = truth_moving(start_lon, start_lat, pset.time[0]) - np.testing.assert_allclose(pset.lon, exp_lon, rtol=rtol) - np.testing.assert_allclose(pset.lat, exp_lat, rtol=rtol) - if method == "RK4_3D": - np.testing.assert_allclose(pset.depth, exp_lat, rtol=rtol) - - -@pytest.mark.parametrize( - "method, rtol", - [ - ("EE", 1e-1), - ("RK4", 1e-5), - ("RK45", 1e-4), - ], -) -def test_decaying_moving_eddy(method, rtol): - ds = decaying_moving_eddy_dataset() - grid = XGrid.from_dataset(ds) - U = Field("U", ds["U"], grid, interp_method=XLinear) - V = Field("V", ds["V"], grid, interp_method=XLinear) - UV = VectorField("UV", U, V) - fieldset = FieldSet([U, V, UV]) - - start_lon, start_lat = 10000, 10000 - dt = np.timedelta64(60, "m") - - if method == "RK45": - fieldset.add_constant("RK45_tol", rtol) - fieldset.add_constant("RK45_min_dt", 10 * 60) - - pset = ParticleSet(fieldset, lon=start_lon, lat=start_lat, time=np.timedelta64(0, "s")) - pset.execute(kernel[method], dt=dt, endtime=np.timedelta64(1, "D")) - - def truth_moving(x_0, y_0, t): - t /= np.timedelta64(1, "s") - lon = ( - x_0 - + (ds.u_g / ds.gamma_g) * (1 - np.exp(-ds.gamma_g * t)) - + ds.f - * ((ds.u_0 - ds.u_g) / (ds.f**2 + ds.gamma**2)) - * ((ds.gamma / ds.f) + np.exp(-ds.gamma * t) * (np.sin(ds.f * t) - (ds.gamma / ds.f) * np.cos(ds.f * t))) - ) - lat = y_0 - ((ds.u_0 - ds.u_g) / (ds.f**2 + ds.gamma**2)) * ds.f * ( - 1 - np.exp(-ds.gamma * t) * (np.cos(ds.f * t) + (ds.gamma / ds.f) * np.sin(ds.f * t)) - ) - return lon, lat - - exp_lon, exp_lat = truth_moving(start_lon, start_lat, pset.time[0]) - np.testing.assert_allclose(pset.lon, exp_lon, rtol=rtol) - np.testing.assert_allclose(pset.lat, exp_lat, rtol=rtol) - - -@pytest.mark.parametrize( - "method, rtol", - [ - ("RK4", 0.1), - ("RK45", 0.1), - ], -) -@pytest.mark.parametrize("grid_type", ["A", "C"]) -def test_stommelgyre_fieldset(method, rtol, grid_type): - npart = 2 - ds = stommel_gyre_dataset(grid_type=grid_type) - grid = XGrid.from_dataset(ds) - vector_interp_method = None if grid_type == "A" else CGrid_Velocity - U = Field("U", ds["U"], grid) - V = Field("V", ds["V"], grid) - P = Field("P", ds["P"], grid, interp_method=XLinear) - UV = VectorField("UV", U, V, vector_interp_method=vector_interp_method) - fieldset = FieldSet([U, V, P, UV]) - - dt = np.timedelta64(30, "m") - runtime = np.timedelta64(1, "D") - start_lon = np.linspace(10e3, 100e3, npart) - start_lat = np.ones_like(start_lon) * 5000e3 - - if method == "RK45": - fieldset.add_constant("RK45_tol", rtol) - - SampleParticle = Particle.add_variable( - [Variable("p", initial=0.0, dtype=np.float32), Variable("p_start", initial=0.0, dtype=np.float32)] - ) - - def UpdateP(particles, fieldset): # pragma: no cover - particles.p = fieldset.P[particles.time, particles.depth, particles.lat, particles.lon] - particles.p_start = np.where(particles.time == 0, particles.p, particles.p_start) - - pset = ParticleSet(fieldset, pclass=SampleParticle, lon=start_lon, lat=start_lat, time=np.timedelta64(0, "s")) - pset.execute([kernel[method], UpdateP], dt=dt, runtime=runtime) - np.testing.assert_allclose(pset.p, pset.p_start, rtol=rtol) - - -@pytest.mark.parametrize( - "method, rtol", - [ - ("RK4", 5e-3), - ("RK45", 1e-4), - ], -) -@pytest.mark.parametrize("grid_type", ["A"]) # TODO also implement C-grid once available -def test_peninsula_fieldset(method, rtol, grid_type): - npart = 2 - ds = peninsula_dataset(grid_type=grid_type) - grid = XGrid.from_dataset(ds) - U = Field("U", ds["U"], grid, interp_method=XLinear) - V = Field("V", ds["V"], grid, interp_method=XLinear) - P = Field("P", ds["P"], grid, interp_method=XLinear) - UV = VectorField("UV", U, V) - fieldset = FieldSet([U, V, P, UV]) - - dt = np.timedelta64(30, "m") - runtime = np.timedelta64(1, "D") - start_lat = np.linspace(3e3, 47e3, npart) - start_lon = 3e3 * np.ones_like(start_lat) - - if method == "RK45": - fieldset.add_constant("RK45_tol", rtol) - - SampleParticle = Particle.add_variable( - [Variable("p", initial=0.0, dtype=np.float32), Variable("p_start", initial=0.0, dtype=np.float32)] - ) - - def UpdateP(particles, fieldset): # pragma: no cover - particles.p = fieldset.P[particles.time, particles.depth, particles.lat, particles.lon] - particles.p_start = np.where(particles.time == 0, particles.p, particles.p_start) - - pset = ParticleSet(fieldset, pclass=SampleParticle, lon=start_lon, lat=start_lat, time=np.timedelta64(0, "s")) - pset.execute([kernel[method], UpdateP], dt=dt, runtime=runtime) - np.testing.assert_allclose(pset.p, pset.p_start, rtol=rtol) - - -def test_nemo_curvilinear_fieldset(): - data_folder = parcels.download_example_dataset("NemoCurvilinear_data") - files = data_folder.glob("*.nc4") - ds = xr.open_mfdataset(files, combine="nested", data_vars="minimal", coords="minimal", compat="override") - ds = ( - ds.isel(time_counter=0, drop=True) - .isel(time=0, drop=True) - .isel(z_a=0, drop=True) - .rename({"glamf": "lon", "gphif": "lat", "z": "depth"}) - ) - - xgcm_grid = xgcm.Grid( - ds, - coords={ - "X": {"left": "x"}, - "Y": {"left": "y"}, - }, - periodic=False, - autoparse_metadata=False, - ) - grid = XGrid(xgcm_grid, mesh="spherical") - - U = parcels.Field("U", ds["U"], grid) - V = parcels.Field("V", ds["V"], grid) - U.units = parcels.GeographicPolar() - V.units = parcels.Geographic() - UV = parcels.VectorField("UV", U, V, vector_interp_method=CGrid_Velocity) - fieldset = parcels.FieldSet([U, V, UV]) - - npart = 20 - lonp = 30 * np.ones(npart) - latp = np.linspace(-70, 88, npart) - runtime = np.timedelta64(12, "h") # TODO increase to 160 days - - def periodicBC(particles, fieldset): # pragma: no cover - particles.dlon = np.where(particles.lon > 180, particles.dlon - 360, particles.dlon) - - pset = parcels.ParticleSet(fieldset, lon=lonp, lat=latp) - pset.execute([AdvectionEE, periodicBC], runtime=runtime, dt=np.timedelta64(6, "h")) - np.testing.assert_allclose(pset.lat, latp, atol=1e-1) - - -@pytest.mark.parametrize("method", ["RK4", "RK4_3D"]) -def test_nemo_3D_curvilinear_fieldset(method): - download_dir = parcels.download_example_dataset("NemoNorthSeaORCA025-N006_data") - ufiles = download_dir.glob("*U.nc") - dsu = xr.open_mfdataset(ufiles, decode_times=False, drop_variables=["nav_lat", "nav_lon"]) - dsu = dsu.rename({"time_counter": "time", "uo": "U"}) - - vfiles = download_dir.glob("*V.nc") - dsv = xr.open_mfdataset(vfiles, decode_times=False, drop_variables=["nav_lat", "nav_lon"]) - dsv = dsv.rename({"time_counter": "time", "vo": "V"}) - - wfiles = download_dir.glob("*W.nc") - dsw = xr.open_mfdataset(wfiles, decode_times=False, drop_variables=["nav_lat", "nav_lon"]) - dsw = dsw.rename({"time_counter": "time", "depthw": "depth", "wo": "W"}) - - dsu = dsu.assign_coords(depthu=dsw.depth.values) - dsu = dsu.rename({"depthu": "depth"}) - - dsv = dsv.assign_coords(depthv=dsw.depth.values) - dsv = dsv.rename({"depthv": "depth"}) - - coord_file = f"{download_dir}/coordinates.nc" - dscoord = xr.open_dataset(coord_file, decode_times=False).rename({"glamf": "lon", "gphif": "lat"}) - dscoord = dscoord.isel(time=0, drop=True) - - ds = xr.merge([dsu, dsv, dsw, dscoord]) - ds = ds.drop_vars( - [ - "uos", - "vos", - "nav_lev", - "nav_lon", - "nav_lat", - "tauvo", - "tauuo", - "time_steps", - "gphiu", - "gphiv", - "gphit", - "glamu", - "glamv", - "glamt", - "time_centered_bounds", - "time_counter_bounds", - "time_centered", - ] - ) - ds = ds.drop_vars(["e1f", "e1t", "e1u", "e1v", "e2f", "e2t", "e2u", "e2v"]) - ds["time"] = [np.timedelta64(int(t), "s") + np.datetime64("1900-01-01") for t in ds["time"]] - - ds["W"] *= -1 # Invert W velocity - - xgcm_grid = xgcm.Grid( - ds, - coords={ - "X": {"left": "x"}, - "Y": {"left": "y"}, - "Z": {"left": "depth"}, - "T": {"center": "time"}, - }, - periodic=False, - autoparse_metadata=False, - ) - grid = XGrid(xgcm_grid, mesh="spherical") - - U = parcels.Field("U", ds["U"], grid) - V = parcels.Field("V", ds["V"], grid) - W = parcels.Field("W", ds["W"], grid) - U.units = parcels.GeographicPolar() - V.units = parcels.Geographic() - UV = parcels.VectorField("UV", U, V, vector_interp_method=CGrid_Velocity) - UVW = parcels.VectorField("UVW", U, V, W, vector_interp_method=CGrid_Velocity) - fieldset = parcels.FieldSet([U, V, W, UV, UVW]) - - npart = 10 - lons = np.linspace(1.9, 3.4, npart) - lats = np.linspace(52.5, 51.6, npart) - pset = parcels.ParticleSet(fieldset, lon=lons, lat=lats, depth=np.ones_like(lons)) - - pset.execute(kernel[method], runtime=np.timedelta64(4, "D"), dt=np.timedelta64(6, "h")) - - if method == "RK4": - np.testing.assert_equal(round_and_hash_float_array([p.lon for p in pset], decimals=5), 29977383852960156017546) - elif method == "RK4_3D": - # TODO check why decimals needs to be so low in RK4_3D (compare to v3) - np.testing.assert_equal( - round_and_hash_float_array([p.depth for p in pset], decimals=1), 29747210774230389239432 - ) diff --git a/tests/v4/test_fieldset.py b/tests/v4/test_fieldset.py deleted file mode 100644 index 0787886be1..0000000000 --- a/tests/v4/test_fieldset.py +++ /dev/null @@ -1,273 +0,0 @@ -from datetime import timedelta - -import cf_xarray # noqa: F401 -import cftime -import numpy as np -import pytest -import xarray as xr - -from parcels._datasets.structured.circulation_models import datasets as datasets_circulation_models -from parcels._datasets.structured.generic import T as T_structured -from parcels._datasets.structured.generic import datasets as datasets_structured -from parcels.field import Field, VectorField -from parcels.fieldset import CalendarError, FieldSet, _datetime_to_msg -from parcels.xgrid import XGrid -from tests import utils - -ds = datasets_structured["ds_2d_left"] - - -@pytest.fixture -def fieldset() -> FieldSet: - """Fixture to create a FieldSet object for testing.""" - grid = XGrid.from_dataset(ds, mesh="flat") - U = Field("U", ds["U (A grid)"], grid) - V = Field("V", ds["V (A grid)"], grid) - UV = VectorField("UV", U, V) - - return FieldSet( - [U, V, UV], - ) - - -def test_fieldset_init_wrong_types(): - with pytest.raises(ValueError, match="Expected `field` to be a Field or VectorField object. Got .*"): - FieldSet([1.0, 2.0, 3.0]) - - -def test_fieldset_add_constant(fieldset): - fieldset.add_constant("test_constant", 1.0) - assert fieldset.test_constant == 1.0 - - -def test_fieldset_add_constant_field(fieldset): - fieldset.add_constant_field("test_constant_field", 1.0) - - # Get a point in the domain - time = ds["time"].mean() - depth = ds["depth"].mean() - lat = ds["lat"].mean() - lon = ds["lon"].mean() - - pytest.xfail(reason="Not yet implemented interpolation.") - assert fieldset.test_constant_field[time, depth, lat, lon] == 1.0 - - -def test_fieldset_add_field(fieldset): - grid = XGrid.from_dataset(ds, mesh="flat") - field = Field("test_field", ds["U (A grid)"], grid) - fieldset.add_field(field) - assert fieldset.test_field == field - - -def test_fieldset_add_field_wrong_type(fieldset): - not_a_field = 1.0 - with pytest.raises(ValueError, match="Expected `field` to be a Field or VectorField object. Got .*"): - fieldset.add_field(not_a_field, "test_field") - - -def test_fieldset_add_field_already_exists(fieldset): - grid = XGrid.from_dataset(ds, mesh="flat") - field = Field("test_field", ds["U (A grid)"], grid) - fieldset.add_field(field, "test_field") - with pytest.raises(ValueError, match="FieldSet already has a Field with name 'test_field'"): - fieldset.add_field(field, "test_field") - - -def test_fieldset_gridset(fieldset): - assert fieldset.fields["U"].grid in fieldset.gridset - assert fieldset.fields["V"].grid in fieldset.gridset - assert fieldset.fields["UV"].grid in fieldset.gridset - assert len(fieldset.gridset) == 1 - - fieldset.add_constant_field("constant_field", 1.0) - assert len(fieldset.gridset) == 2 - - -@pytest.mark.parametrize("ds", [pytest.param(ds, id=k) for k, ds in datasets_structured.items()]) -def test_fieldset_from_structured_generic_datasets(ds): - grid = XGrid.from_dataset(ds, mesh="flat") - fields = [] - for var in ds.data_vars: - fields.append(Field(var, ds[var], grid)) - - fieldset = FieldSet(fields) - - assert len(fieldset.fields) == len(ds.data_vars) - for field in fieldset.fields.values(): - utils.assert_valid_field_data(field.data, field.grid) - - assert len(fieldset.gridset) == 1 - - -def test_fieldset_gridset_multiple_grids(): ... - - -def test_fieldset_time_interval(): - grid1 = XGrid.from_dataset(ds, mesh="flat") - field1 = Field("field1", ds["U (A grid)"], grid1) - - ds2 = ds.copy() - ds2["time"] = (ds2["time"].dims, ds2["time"].data + np.timedelta64(timedelta(days=1)), ds2["time"].attrs) - grid2 = XGrid.from_dataset(ds2, mesh="flat") - field2 = Field("field2", ds2["U (A grid)"], grid2) - - fieldset = FieldSet([field1, field2]) - fieldset.add_constant_field("constant_field", 1.0) - - assert fieldset.time_interval.left == np.datetime64("2000-01-02") - assert fieldset.time_interval.right == np.datetime64("2001-01-01") - - -def test_fieldset_time_interval_constant_fields(): - fieldset = FieldSet([]) - fieldset.add_constant_field("constant_field", 1.0) - fieldset.add_constant_field("constant_field2", 2.0) - - assert fieldset.time_interval is None - - -def test_fieldset_init_incompatible_calendars(): - ds1 = ds.copy() - ds1["time"] = ( - ds1["time"].dims, - xr.date_range("2000", "2001", T_structured, calendar="365_day", use_cftime=True), - ds1["time"].attrs, - ) - - grid = XGrid.from_dataset(ds1, mesh="flat") - U = Field("U", ds1["U (A grid)"], grid) - V = Field("V", ds1["V (A grid)"], grid) - UV = VectorField("UV", U, V) - - ds2 = ds.copy() - ds2["time"] = ( - ds2["time"].dims, - xr.date_range("2000", "2001", T_structured, calendar="360_day", use_cftime=True), - ds2["time"].attrs, - ) - grid2 = XGrid.from_dataset(ds2, mesh="flat") - incompatible_calendar = Field("test", ds2["data_g"], grid2) - - with pytest.raises(CalendarError, match="Expected field '.*' to have calendar compatible with datetime object"): - FieldSet([U, V, UV, incompatible_calendar]) - - -def test_fieldset_add_field_incompatible_calendars(fieldset): - ds_test = ds.copy() - ds_test["time"] = ( - ds_test["time"].dims, - xr.date_range("2000", "2001", T_structured, calendar="360_day", use_cftime=True), - ds_test["time"].attrs, - ) - grid = XGrid.from_dataset(ds_test, mesh="flat") - field = Field("test_field", ds_test["data_g"], grid) - - with pytest.raises(CalendarError, match="Expected field '.*' to have calendar compatible with datetime object"): - fieldset.add_field(field, "test_field") - - ds_test = ds.copy() - ds_test["time"] = ( - ds_test["time"].dims, - np.linspace(0, 100, T_structured, dtype="timedelta64[s]"), - ds_test["time"].attrs, - ) - grid = XGrid.from_dataset(ds_test, mesh="flat") - field = Field("test_field", ds_test["data_g"], grid) - - with pytest.raises(CalendarError, match="Expected field '.*' to have calendar compatible with datetime object"): - fieldset.add_field(field, "test_field") - - -@pytest.mark.parametrize( - "input_, expected", - [ - (cftime.DatetimeNoLeap(2000, 1, 1), " with cftime calendar noleap'"), - (cftime.Datetime360Day(2000, 1, 1), " with cftime calendar 360_day'"), - (cftime.DatetimeJulian(2000, 1, 1), " with cftime calendar julian'"), - ( - cftime.DatetimeGregorian(2000, 1, 1), - " with cftime calendar standard'", - ), - (np.datetime64("2000-01-01"), ""), - (cftime.datetime(2000, 1, 1), " with cftime calendar standard'"), - ], -) -def test_datetime_to_msg(input_, expected): - assert _datetime_to_msg(input_) == expected - - -def test_fieldset_samegrids_UV(): - """Test that if a simple fieldset with U and V is created, that only one grid object is defined.""" - ... - - -def test_fieldset_grid_deduplication(): - """Tests that for a full fieldset that the number of grid objects is as expected - (sharing of grid objects so that the particle location is not duplicated). - - When grid deduplication is actually implemented, this might need to be refactored - into multiple tests (/more might be needed). - """ - ... - - -def test_fieldset_add_field_after_pset(): - # ? Should it be allowed to add fields (normal or vector) after a ParticleSet has been initialized? - ... - - -_COPERNICUS_DATASETS = [ - datasets_circulation_models["ds_copernicusmarine"], - datasets_circulation_models["ds_copernicusmarine_waves"], -] - - -@pytest.mark.parametrize("ds", _COPERNICUS_DATASETS) -def test_fieldset_from_copernicusmarine(ds, caplog): - fieldset = FieldSet.from_copernicusmarine(ds) - assert "U" in fieldset.fields - assert "V" in fieldset.fields - assert "UV" in fieldset.fields - assert "renamed it to 'U'" in caplog.text - assert "renamed it to 'V'" in caplog.text - - -def test_fieldset_from_copernicusmarine_no_currents(caplog): - ds = datasets_circulation_models["ds_copernicusmarine"].cf.drop_vars( - ["eastward_sea_water_velocity", "northward_sea_water_velocity"] - ) - fieldset = FieldSet.from_copernicusmarine(ds) - assert "U" not in fieldset.fields - assert "V" not in fieldset.fields - assert "UV" not in fieldset.fields - assert caplog.text == "" - - -@pytest.mark.parametrize("ds", _COPERNICUS_DATASETS) -def test_fieldset_from_copernicusmarine_no_logs(ds, caplog): - ds = ds.copy() - zeros = xr.zeros_like(list(ds.data_vars.values())[0]) - ds["U"] = zeros - ds["V"] = zeros - - fieldset = FieldSet.from_copernicusmarine(ds) - assert "U" in fieldset.fields - assert "V" in fieldset.fields - assert "UV" in fieldset.fields - assert caplog.text == "" - - -def test_fieldset_from_copernicusmarine_with_W(caplog): - ds = datasets_circulation_models["ds_copernicusmarine"] - ds = ds.copy() - ds["wo"] = ds["uo"] - ds["wo"].attrs["standard_name"] = "vertical_sea_water_velocity" - - fieldset = FieldSet.from_copernicusmarine(ds) - assert "U" in fieldset.fields - assert "V" in fieldset.fields - assert "W" in fieldset.fields - assert "UV" not in fieldset.fields - assert "UVW" in fieldset.fields - assert "renamed it to 'W'" in caplog.text diff --git a/tests/v4/test_interpolation.py b/tests/v4/test_interpolation.py deleted file mode 100644 index a1e18e2489..0000000000 --- a/tests/v4/test_interpolation.py +++ /dev/null @@ -1,225 +0,0 @@ -import numpy as np -import pytest -import xarray as xr - -from parcels._datasets.structured.generated import simple_UV_dataset -from parcels._datasets.unstructured.generic import datasets as datasets_unstructured -from parcels._index_search import _search_time_index -from parcels.application_kernels.advection import AdvectionRK4_3D -from parcels.application_kernels.interpolation import ( - UXPiecewiseLinearNode, - XFreeslip, - XLinear, - XNearest, - XPartialslip, - ZeroInterpolator, -) -from parcels.field import Field, VectorField -from parcels.fieldset import FieldSet -from parcels.particle import Particle, Variable -from parcels.particlefile import ParticleFile -from parcels.particleset import ParticleSet -from parcels.tools.statuscodes import StatusCode -from parcels.uxgrid import UxGrid -from parcels.xgrid import XGrid -from tests.utils import TEST_DATA - - -@pytest.fixture -def field(): - """Reference data used for testing interpolation.""" - z0 = np.array( # each x is +1 from the previous, each y is +2 from the previous - [ - [0.0, 1.0, 2.0, 3.0], - [2.0, 3.0, 4.0, 5.0], - [4.0, 5.0, 6.0, 7.0], - [6.0, 7.0, 8.0, 9.0], - ] - ) - spatial_data = np.array([z0, z0 + 3, z0 + 6, z0 + 9]) # each z is +3 from the previous - temporal_data = np.array([spatial_data, spatial_data + 10, spatial_data + 20]) # each t is +10 from the previous - - ds = xr.Dataset( - {"U": (["time", "depth", "lat", "lon"], temporal_data)}, - coords={ - "time": (["time"], [np.timedelta64(t, "s") for t in [0, 2, 4]], {"axis": "T"}), - "depth": (["depth"], [0, 1, 2, 3], {"axis": "Z"}), - "lat": (["lat"], [0, 1, 2, 3], {"axis": "Y", "c_grid_axis_shift": -0.5}), - "lon": (["lon"], [0, 1, 2, 3], {"axis": "X", "c_grid_axis_shift": -0.5}), - "x": (["x"], [0.5, 1.5, 2.5, 3.5], {"axis": "X"}), - "y": (["y"], [0.5, 1.5, 2.5, 3.5], {"axis": "Y"}), - }, - ) - return Field("U", ds["U"], XGrid.from_dataset(ds)) - - -@pytest.mark.parametrize( - "func, t, z, y, x, expected", - [ - pytest.param(ZeroInterpolator, np.timedelta64(1, "s"), 2.5, 0.49, 0.51, 0, id="Zero"), - pytest.param( - XLinear, - [np.timedelta64(0, "s"), np.timedelta64(1, "s")], - [0, 0], - [0.49, 0.49], - [0.51, 0.51], - [1.49, 6.49], - id="Linear", - ), - pytest.param(XLinear, np.timedelta64(1, "s"), 2.5, 0.49, 0.51, 13.99, id="Linear-2"), - pytest.param( - XNearest, - [np.timedelta64(0, "s"), np.timedelta64(3, "s")], - [0.2, 0.2], - [0.2, 0.2], - [0.51, 0.51], - [1.0, 16.0], - id="Nearest", - ), - ], -) -def test_raw_2d_interpolation(field, func, t, z, y, x, expected): - """Test the interpolation functions on the Field.""" - tau, ti = _search_time_index(field, t) - position = field.grid.search(z, y, x) - - value = func(field, ti, position, tau, 0, 0, y, x) - np.testing.assert_equal(value, expected) - - -@pytest.mark.parametrize( - "func, t, z, y, x, expected", - [ - (XPartialslip, np.timedelta64(1, "s"), 0, 0, 0.0, [[1], [1]]), - (XFreeslip, np.timedelta64(1, "s"), 0, 0.5, 1.5, [[1], [0.5]]), - (XPartialslip, np.timedelta64(1, "s"), 0, 2.5, 1.5, [[0.75], [0.5]]), - (XFreeslip, np.timedelta64(1, "s"), 0, 2.5, 1.5, [[1], [0.5]]), - (XPartialslip, np.timedelta64(1, "s"), 0, 1.5, 0.5, [[0.5], [0.75]]), - (XFreeslip, np.timedelta64(1, "s"), 0, 1.5, 0.5, [[0.5], [1]]), - ( - XFreeslip, - [np.timedelta64(1, "s"), np.timedelta64(0, "s")], - [0, 2], - [1.5, 1.5], - [2.5, 0.5], - [[0.5, 0.5], [1, 1]], - ), - ], -) -def test_spatial_slip_interpolation(field, func, t, z, y, x, expected): - field.data[:] = 1.0 - field.data[:, :, 1:3, 1:3] = 0.0 # Set zero land value to test spatial slip - U = field - V = field - UV = VectorField("UV", U, V, vector_interp_method=func) - - velocities = UV[t, z, y, x] - np.testing.assert_array_almost_equal(velocities, expected) - - -@pytest.mark.parametrize("mesh", ["spherical", "flat"]) -def test_interpolation_mesh_type(mesh, npart=10): - ds = simple_UV_dataset(mesh=mesh) - ds["U"].data[:] = 1.0 - grid = XGrid.from_dataset(ds, mesh=mesh) - U = Field("U", ds["U"], grid, interp_method=XLinear) - V = Field("V", ds["V"], grid, interp_method=XLinear) - UV = VectorField("UV", U, V) - - lat = 30.0 - time = U.time_interval.left - u_expected = 1.0 if mesh == "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 - - -def test_default_interpolator_set_correctly(): - ds = simple_UV_dataset() - grid = XGrid.from_dataset(ds) - U = Field("U", ds["U"], grid) - assert U.interp_method == XLinear - - ds = datasets_unstructured["stommel_gyre_delaunay"] - grid = UxGrid(grid=ds.uxgrid, z=ds.coords["nz"]) - U = Field("U", ds["U"], grid) - assert U.interp_method == UXPiecewiseLinearNode - - -interp_methods = { - "linear": XLinear, -} - - -@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, mesh="flat") - U = Field("U", ds["U"], grid, interp_method=interp_methods[interp_name]) - V = Field("V", ds["V"], grid, interp_method=interp_methods[interp_name]) - W = Field("W", ds["W"], grid, 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.state = StatusCode.Delete - - outfile = 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)