Add tests for new FieldSet object

parcels/_datasets/structured/generic.py

@@ -24,6 +24,10 @@ def _rotated_curvilinear_grid():
         {
             "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)),
+            "U (A grid)": (["ZG", "YG", "XG"], np.random.rand(3 * N, 2 * N, N)),
+            "V (A grid)": (["ZG", "YG", "XG"], np.random.rand(3 * N, 2 * N, N)),
+            "U (C grid)": (["ZG", "YC", "XG"], np.random.rand(3 * N, 2 * N, N)),
+            "V (C grid)": (["ZG", "YG", "XC"], np.random.rand(3 * N, 2 * N, N)),
         },
         coords={
             "XG": (["XG"], XG, {"axis": "X", "c_grid_axis_shift": -0.5}),
@@ -95,6 +99,10 @@ def _unrolled_cone_curvilinear_grid():
         {
             "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)),
+            "U (A grid)": (["ZG", "YG", "XG"], np.random.rand(3 * N, 2 * N, N)),
+            "V (A grid)": (["ZG", "YG", "XG"], np.random.rand(3 * N, 2 * N, N)),
+            "U (C grid)": (["ZG", "YC", "XG"], np.random.rand(3 * N, 2 * N, N)),
+            "V (C grid)": (["ZG", "YG", "XC"], np.random.rand(3 * N, 2 * N, N)),
         },
         coords={
             "XG": (["XG"], XG, {"axis": "X", "c_grid_axis_shift": -0.5}),
@@ -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)": (["ZG", "YG", "XG"], np.random.rand(3 * N, 2 * N, N)),
+            "V (A grid)": (["ZG", "YG", "XG"], np.random.rand(3 * N, 2 * N, N)),
+            "U (C grid)": (["ZG", "YC", "XG"], np.random.rand(3 * N, 2 * N, N)),
+            "V (C grid)": (["ZG", "YG", "XC"], np.random.rand(3 * N, 2 * N, N)),
         },
         coords={
             "XG": (

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

@@ -11,12 +11,12 @@
 from uxarray.grid.neighbors import _barycentric_coordinates
 
 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,
@@ -185,7 +185,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

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 @@
     """
 
     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 @@
         * `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,21 +152,25 @@
                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
-                allow_time_extrapolation=True,
             )
         )
 
@@ -185,7 +204,10 @@
         `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

parcels/tools/_helpers.py

@@ -3,17 +3,12 @@
 from __future__ import annotations
 
 import functools
-import textwrap
 import warnings
 from collections.abc import Callable
 from datetime import timedelta
-from typing import TYPE_CHECKING, Any
 
 import numpy as np
 
-if TYPE_CHECKING:
-    from parcels import Field, FieldSet, ParticleSet
-
 PACKAGE = "Parcels"
 
 
@@ -68,81 +63,6 @@ def patch_docstring(obj: Callable, extra: str) -> None:
     obj.__doc__ = f"{obj.__doc__ or ''}{extra}".strip()
 
 
-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"
-
-
 def timedelta_to_float(dt: float | timedelta | np.timedelta64) -> float:
     """Convert a timedelta to a float in seconds."""
     if isinstance(dt, timedelta):

tests/v4/test_field.py

@@ -54,10 +54,9 @@ def test_field_incompatible_combination(data, grid):
     [
         pytest.param(
             structured_datasets["ds_2d_left"]["data_g"], Grid(structured_datasets["ds_2d_left"]), id="ds_2d_left"
-        ),
+        ),  # TODO: Perhaps this test should be expanded to cover more datasets?
     ],
 )
-@pytest.mark.xfail(reason="Structured grid creation is not implemented yet")
 def test_field_structured_grid_creation(data, grid):
     """Test creating a field."""
     field = Field(
@@ -66,15 +65,10 @@ def test_field_structured_grid_creation(data, grid):
         grid=grid,
     )
     assert field.name == "test_field"
-    assert field.data == data
+    assert field.data.equals(data)
     assert field.grid == grid
 
 
-def test_field_structured_grid_creation_spherical():
-    # Field(..., mesh_type="spherical")
-    ...
-
-
 def test_field_unstructured_grid_creation(): ...