Merge branch 'v4-dev' into remove-time-extrapolation-option

Author: fluidnumericsJoe

parcels/_core/utils/time.py

@@ -4,6 +4,7 @@
 from typing import TypeVar
 
 import cftime
+import numpy as np
 
 T = TypeVar("T", datetime, cftime.datetime)
 
@@ -24,10 +25,10 @@ class TimeInterval:
     """
 
     def __init__(self, left: T, right: T) -> None:
-        if not isinstance(left, (datetime, cftime.datetime)):
-            raise ValueError(f"Expected left to be a datetime or cftime.datetime, got {type(left)}.")
-        if not isinstance(right, (datetime, cftime.datetime)):
-            raise ValueError(f"Expected right to be a datetime or cftime.datetime, got {type(right)}.")
+        if not isinstance(left, (datetime, cftime.datetime, np.datetime64)):
+            raise ValueError(f"Expected right to be a datetime, cftime.datetime, or np.datetime64. Got {type(left)}.")
+        if not isinstance(right, (datetime, cftime.datetime, np.datetime64)):
+            raise ValueError(f"Expected right to be a datetime, cftime.datetime, or np.datetime64. Got {type(right)}.")
         if left >= right:
             raise ValueError(f"Expected left to be strictly less than right, got left={left} and right={right}.")
         if not is_compatible(left, right):

parcels/_datasets/structured/generic.py

@@ -6,7 +6,7 @@
 __all__ = ["N", "T", "datasets"]
 
 N = 30
-T = 10
+T = 13
 
 
 def _rotated_curvilinear_grid():
@@ -22,8 +22,12 @@ def _rotated_curvilinear_grid():
 
     return xr.Dataset(
         {
-            "data_g": (["ZG", "YG", "XG"], np.random.rand(3 * N, 2 * N, N)),
-            "data_c": (["ZC", "YC", "XC"], np.random.rand(3 * N, 2 * N, N)),
+            "data_g": (["time", "ZG", "YG", "XG"], np.random.rand(T, 3 * N, 2 * N, N)),
+            "data_c": (["time", "ZC", "YC", "XC"], np.random.rand(T, 3 * N, 2 * N, N)),
+            "U (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, 3 * N, 2 * N, N)),
+            "V (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, 3 * N, 2 * N, N)),
+            "U (C grid)": (["time", "ZG", "YC", "XG"], np.random.rand(T, 3 * N, 2 * N, N)),
+            "V (C grid)": (["time", "ZG", "YG", "XC"], np.random.rand(T, 3 * N, 2 * N, N)),
         },
         coords={
             "XG": (["XG"], XG, {"axis": "X", "c_grid_axis_shift": -0.5}),
@@ -41,7 +45,7 @@ def _rotated_curvilinear_grid():
                 {"axis": "Z"},
             ),
             "depth": (["ZG"], np.arange(3 * N), {"axis": "Z"}),
-            "time": (["time"], np.arange(T), {"axis": "T"}),
+            "time": (["time"], xr.date_range("2000", "2001", T), {"axis": "T"}),
             "lon": (
                 ["YG", "XG"],
                 LON,
@@ -93,8 +97,12 @@ def _unrolled_cone_curvilinear_grid():
 
     return xr.Dataset(
         {
-            "data_g": (["ZG", "YG", "XG"], np.random.rand(3 * N, 2 * N, N)),
-            "data_c": (["ZC", "YC", "XC"], np.random.rand(3 * N, 2 * N, N)),
+            "data_g": (["time", "ZG", "YG", "XG"], np.random.rand(T, 3 * N, 2 * N, N)),
+            "data_c": (["time", "ZC", "YC", "XC"], np.random.rand(T, 3 * N, 2 * N, N)),
+            "U (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, 3 * N, 2 * N, N)),
+            "V (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, 3 * N, 2 * N, N)),
+            "U (C grid)": (["time", "ZG", "YC", "XG"], np.random.rand(T, 3 * N, 2 * N, N)),
+            "V (C grid)": (["time", "ZG", "YG", "XC"], np.random.rand(T, 3 * N, 2 * N, N)),
         },
         coords={
             "XG": (["XG"], XG, {"axis": "X", "c_grid_axis_shift": -0.5}),
@@ -112,7 +120,7 @@ def _unrolled_cone_curvilinear_grid():
                 {"axis": "Z"},
             ),
             "depth": (["ZG"], np.arange(3 * N), {"axis": "Z"}),
-            "time": (["time"], np.arange(T), {"axis": "T"}),
+            "time": (["time"], xr.date_range("2000", "2001", T), {"axis": "T"}),
             "lon": (
                 ["YG", "XG"],
                 LON,
@@ -133,6 +141,10 @@ def _unrolled_cone_curvilinear_grid():
         {
             "data_g": (["time", "ZG", "YG", "XG"], np.random.rand(T, 3 * N, 2 * N, N)),
             "data_c": (["time", "ZC", "YC", "XC"], np.random.rand(T, 3 * N, 2 * N, N)),
+            "U (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, 3 * N, 2 * N, N)),
+            "V (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, 3 * N, 2 * N, N)),
+            "U (C grid)": (["time", "ZG", "YC", "XG"], np.random.rand(T, 3 * N, 2 * N, N)),
+            "V (C grid)": (["time", "ZG", "YG", "XC"], np.random.rand(T, 3 * N, 2 * N, N)),
         },
         coords={
             "XG": (
@@ -164,7 +176,7 @@ def _unrolled_cone_curvilinear_grid():
             "lon": (["XG"], 2 * np.pi / N * np.arange(0, N)),
             "lat": (["YG"], 2 * np.pi / (2 * N) * np.arange(0, 2 * N)),
             "depth": (["ZG"], np.arange(3 * N)),
-            "time": (["time"], np.arange(T), {"axis": "T"}),
+            "time": (["time"], xr.date_range("2000", "2001", T), {"axis": "T"}),
         },
     ),
     "2d_left_unrolled_cone": _unrolled_cone_curvilinear_grid(),

parcels/_reprs.py

@@ -0,0 +1,84 @@
+"""Parcels reprs"""
+
+from __future__ import annotations
+
+import textwrap
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from parcels import Field, FieldSet, ParticleSet
+
+
+def field_repr(field: Field) -> str:
+    """Return a pretty repr for Field"""
+    out = f"""<{type(field).__name__}>
+    name            : {field.name!r}
+    data            : {field.data!r}
+    extrapolate time: {field.allow_time_extrapolation!r}
+"""
+    return textwrap.dedent(out).strip()
+
+
+def _format_list_items_multiline(items: list[str], level: int = 1) -> str:
+    """Given a list of strings, formats them across multiple lines.
+
+    Uses indentation levels of 4 spaces provided by ``level``.
+
+    Example
+    -------
+    >>> output = _format_list_items_multiline(["item1", "item2", "item3"], 4)
+    >>> f"my_items: {output}"
+    my_items: [
+        item1,
+        item2,
+        item3,
+    ]
+    """
+    if len(items) == 0:
+        return "[]"
+
+    assert level >= 1, "Indentation level >=1 supported"
+    indentation_str = level * 4 * " "
+    indentation_str_end = (level - 1) * 4 * " "
+
+    items_str = ",\n".join([textwrap.indent(i, indentation_str) for i in items])
+    return f"[\n{items_str}\n{indentation_str_end}]"
+
+
+def particleset_repr(pset: ParticleSet) -> str:
+    """Return a pretty repr for ParticleSet"""
+    if len(pset) < 10:
+        particles = [repr(p) for p in pset]
+    else:
+        particles = [repr(pset[i]) for i in range(7)] + ["..."]
+
+    out = f"""<{type(pset).__name__}>
+    fieldset   :
+{textwrap.indent(repr(pset.fieldset), " " * 8)}
+    pclass     : {pset.pclass}
+    repeatdt   : {pset.repeatdt}
+    # particles: {len(pset)}
+    particles  : {_format_list_items_multiline(particles, level=2)}
+"""
+    return textwrap.dedent(out).strip()
+
+
+def fieldset_repr(fieldset: FieldSet) -> str:
+    """Return a pretty repr for FieldSet"""
+    fields_repr = "\n".join([repr(f) for f in fieldset.get_fields()])
+
+    out = f"""<{type(fieldset).__name__}>
+    fields:
+{textwrap.indent(fields_repr, 8 * " ")}
+"""
+    return textwrap.dedent(out).strip()
+
+
+def default_repr(obj: Any):
+    if is_builtin_object(obj):
+        return repr(obj)
+    return object.__repr__(obj)
+
+
+def is_builtin_object(obj):
+    return obj.__class__.__module__ == "builtins"

parcels/field.py

@@ -1,22 +1,24 @@
+from __future__ import annotations
+
 import inspect
 import warnings
 from collections.abc import Callable
 from datetime import datetime
 from enum import IntEnum
-from typing import TYPE_CHECKING
 
 import numpy as np
 import uxarray as ux
 import xarray as xr
 from uxarray.grid.neighbors import _barycentric_coordinates
 
+from parcels._core.utils.time import TimeInterval
 from parcels._core.utils.unstructured import get_vertical_location_from_dims
+from parcels._reprs import default_repr, field_repr
 from parcels._typing import (
     Mesh,
     VectorType,
     assert_valid_mesh,
 )
-from parcels.tools._helpers import default_repr, field_repr
 from parcels.tools.converters import (
     UnitConverter,
     unitconverters_map,
@@ -33,9 +35,6 @@
 
 from ._index_search import _search_indices_rectilinear, _search_time_index
 
-if TYPE_CHECKING:
-    pass
-
 __all__ = ["Field", "GridType", "VectorField"]
 
 
@@ -166,6 +165,7 @@ def __init__(
         self.name = name
         self.data = data
         self.grid = grid
+        self.time_interval = get_time_interval(data)
 
         # For compatibility with parts of the codebase that rely on v3 definition of Grid.
         # Should be worked to be removed in v4
@@ -184,7 +184,6 @@ def __init__(
             e.add_note(f"Error validating field {name!r}.")
             raise e
 
-        self._parent_mesh = data.attrs["mesh"]
         self._mesh_type = mesh_type
 
         # Setting the interpolation method dynamically
@@ -665,3 +664,10 @@ def _assert_compatible_combination(data: xr.DataArray | ux.UxDataArray, grid: ux
             raise ValueError(
                 f"Incompatible data-grid combination. Data is a xarray.DataArray, expected `grid` to be a parcels Grid object, got {type(grid)}."
             )
+
+
+def get_time_interval(data: xr.DataArray | ux.UxDataArray) -> TimeInterval | None:
+    if "time" not in data.dims:
+        return None
+
+    return TimeInterval(data.time.values[0], data.time.values[-1])

parcels/fieldset.py

@@ -3,9 +3,10 @@
 import numpy as np
 import xarray as xr
 
+from parcels._reprs import fieldset_repr
 from parcels._typing import Mesh
 from parcels.field import Field, VectorField
-from parcels.tools._helpers import fieldset_repr
+from parcels.v4.grid import Grid
 
 __all__ = ["FieldSet"]
 
@@ -41,10 +42,21 @@ class FieldSet:
     """
 
     def __init__(self, fields: list[Field | VectorField]):
-        # TODO Nick : Enforce fields to be list of Field or VectorField objects
+        for field in fields:
+            if not isinstance(field, (Field, VectorField)):
+                raise ValueError(f"Expected `field` to be a Field or VectorField object. Got {field}")
+
         self.fields = {f.name: f for f in fields}
+        self.constants = {}
 
-    # TODO : Nick : Add _getattr_ magic method to allow access to fields by name
+    def __getattr__(self, name):
+        """Get the field by name. If the field is not found, check if it's a constant."""
+        if name in self.fields:
+            return self.fields[name]
+        elif name in self.constants:
+            return self.constants[name]
+        else:
+            raise AttributeError(f"FieldSet has no attribute '{name}'")
 
     @property
     def time_interval(self):
@@ -112,6 +124,9 @@ def add_field(self, field: Field, name: str | None = None):
         * `Unit converters <../examples/tutorial_unitconverters.ipynb>`__ (Default value = None)
 
         """
+        if not isinstance(field, (Field, VectorField)):
+            raise ValueError(f"Expected `field` to be a Field or VectorField object. Got {type(field)}")
+
         name = field.name if name is None else name
 
         if name in self.fields:
@@ -137,19 +152,24 @@ def add_constant_field(self, name: str, value, mesh: Mesh = "flat"):
                correction for zonal velocity U near the poles.
             2. flat: No conversion, lat/lon are assumed to be in m.
         """
-        time = 0.0
-        values = np.full((1, 1, 1, 1), value)
-        data = xr.DataArray(
-            data=values,
-            name=name,
-            dims="null",
-            coords=[time, [0], [0], [0]],
-            attrs=dict(description="null", units="null", location="node", mesh="constant", mesh_type=mesh),
+        da = xr.DataArray(
+            data=np.full((1, 1, 1, 1), value),
+            dims=["T", "ZG", "YG", "XG"],
+            coords={
+                "ZG": (["ZG"], np.arange(1), {"axis": "Z"}),
+                "YG": (["YG"], np.arange(1), {"axis": "Y"}),
+                "XG": (["XG"], np.arange(1), {"axis": "X"}),
+                "lon": (["XG"], np.arange(1), {"axis": "X"}),
+                "lat": (["YG"], np.arange(1), {"axis": "Y"}),
+                "depth": (["ZG"], np.arange(1), {"axis": "Z"}),
+            },
         )
+        grid = Grid(da)
         self.add_field(
             Field(
                 name,
-                data,
+                da,
+                grid,
                 interp_method=None,  # TODO : Need to define an interpolation method for constants
             )
         )
@@ -184,7 +204,10 @@ def add_constant(self, name, value):
         `Diffusion <../examples/tutorial_diffusion.ipynb>`__
         `Periodic boundaries <../examples/tutorial_periodic_boundaries.ipynb>`__
         """
-        setattr(self, name, value)
+        if name in self.constants:
+            raise ValueError(f"FieldSet already has a constant with name '{name}'")
+
+        self.constants[name] = np.float32(value)
 
     # def computeTimeChunk(self, time=0.0, dt=1):
     #     """Load a chunk of three data time steps into the FieldSet.

parcels/particlefile.py

@@ -10,7 +10,8 @@
 
 import parcels
 from parcels._compat import MPI
-from parcels.tools._helpers import default_repr, timedelta_to_float
+from parcels._reprs import default_repr
+from parcels.tools._helpers import timedelta_to_float
 from parcels.tools.warnings import FileWarning
 
 __all__ = ["ParticleFile"]

parcels/particleset.py

@@ -11,6 +11,7 @@
 from tqdm import tqdm
 
 from parcels._compat import MPI
+from parcels._reprs import particleset_repr
 from parcels.application_kernels.advection import AdvectionRK4
 from parcels.field import Field
 from parcels.grid import GridType
@@ -25,7 +26,7 @@
 from parcels.particle import Particle, Variable
 from parcels.particledata import ParticleData, ParticleDataIterator
 from parcels.particlefile import ParticleFile
-from parcels.tools._helpers import particleset_repr, timedelta_to_float
+from parcels.tools._helpers import timedelta_to_float
 from parcels.tools.converters import _get_cftime_calendars, convert_to_flat_array
 from parcels.tools.loggers import logger
 from parcels.tools.statuscodes import StatusCode