Package restructure

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,

parcels/basegrid.py → 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

parcels/tools/converters.py → 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(),

parcels/field.py → 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:

parcels/fieldset.py → 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"]
 
 

parcels/_index_search.py → 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):

parcels/kernel.py → 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

parcels/particle.py → 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"]

parcels/particlefile.py → 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"]
 

parcels/particleset.py → 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."
                 )