Update API to take `vector_field

Also add vector_field_components to the private API

Author: VeckoTheGecko

src/parcels/_core/fieldset.py

@@ -10,10 +10,17 @@
 import xarray as xr
 
 from parcels._core.field import Field, VectorField
-from parcels._core.model import CONSTANT_FIELD_MODELS, ModelData, StructuredModelData, UnstructuredModelData
+from parcels._core.model import (
+    CONSTANT_FIELD_MODELS,
+    ModelData,
+    StructuredModelData,
+    TVectorFieldMapping,
+    UnstructuredModelData,
+)
 from parcels._core.utils.string import _assert_str_and_python_varname
 from parcels._core.utils.time import get_datetime_type_calendar
 from parcels._core.utils.time import is_compatible as datetime_is_compatible
+from parcels._python import _MISSING, _MissingType
 from parcels._typing import Mesh
 from parcels.interpolators import (
     XConstantField,
@@ -201,7 +208,12 @@ def gridset(self) -> list[BaseGrid]:
         return grids
 
     @classmethod
-    def from_ugrid_conventions(cls, ds: ux.UxDataset, mesh: str = "spherical"):
+    def from_ugrid_conventions(
+        cls,
+        ds: ux.UxDataset,
+        mesh: str = "spherical",
+        vector_fields: TVectorFieldMapping | None | _MissingType = _MISSING,
+    ):
         """Create a FieldSet from a Parcels compliant uxarray.UxDataset.
 
         This is the primary ingestion method in Parcels for structured grid datasets.
@@ -225,12 +237,15 @@ def from_ugrid_conventions(cls, ds: ux.UxDataset, mesh: str = "spherical"):
         -----
         See https://ugrid-conventions.github.io/ugrid-conventions/ for more information on the UGRID conventions.
         """
-        model = UnstructuredModelData.from_ugrid_conventions(ds, mesh)
+        model = UnstructuredModelData.from_ugrid_conventions(ds, mesh, vector_fields)
         return cls([model])
 
     @classmethod
     def from_sgrid_conventions(
-        cls, ds: xr.Dataset, mesh: Mesh | None = None
+        cls,
+        ds: xr.Dataset,
+        mesh: Mesh | None = None,
+        vector_fields: TVectorFieldMapping | None | _MissingType = _MISSING,
     ):  # TODO: Update mesh to be discovered from the dataset metadata
         """Create a FieldSet from a dataset using SGRID convention metadata.
 
@@ -259,7 +274,7 @@ def from_sgrid_conventions(
 
         See https://sgrid.github.io/sgrid/ for more information on the SGRID conventions.
         """
-        model = StructuredModelData.from_sgrid_conventions(ds, mesh)
+        model = StructuredModelData.from_sgrid_conventions(ds, mesh, vector_fields)
         return cls([model])
 
 

src/parcels/_core/model.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
+from collections.abc import Mapping, Sequence
 from typing import Any, Self
 
 import cf_xarray  # noqa: F401
@@ -18,6 +19,7 @@
     assert_all_field_dims_have_axis,  # noqa: F401, leave import for now until decision is made # TODO v4: Make decision
 )
 from parcels._logger import logger
+from parcels._python import _MissingType
 from parcels._typing import Mesh
 from parcels.convert import _ds_rename_using_standard_names
 from parcels.interpolators import (
@@ -32,11 +34,14 @@
 )
 from parcels.interpolators._base import ScalarInterpolator, VectorInterpolator
 
+TVectorFieldMapping = Mapping[str, tuple[str, str] | tuple[str, str, str]]
+
 
 class ModelData(ABC):
     data: Any
     grid: BaseGrid
     field_to_interpolator: dict[str, ScalarInterpolator | VectorInterpolator]
+    vector_field_components: TVectorFieldMapping
 
     @abstractmethod
     def construct_fields(self) -> list[Field | VectorField]: ...
@@ -79,7 +84,7 @@ def preprocess_sgrid_model_data(ds: xr.Dataset) -> xr.Dataset:
 
 
 class StructuredModelData(ModelData):
-    def __init__(self, data: xr.Dataset, mesh: Mesh):
+    def __init__(self, data: xr.Dataset, mesh: Mesh, vector_field_components: TVectorFieldMapping):
         if not isinstance(data, xr.Dataset):
             raise ValueError(f"Expected `data` to be an xarray.Dataset . Got {type(data)}")
 
@@ -88,6 +93,7 @@ def __init__(self, data: xr.Dataset, mesh: Mesh):
 
         self.data = data
         self.grid = grid
+        self.vector_field_components = vector_field_components
         self.field_to_interpolator = {}
         self._fields: list[Field | VectorField] | None = None
         self.assert_valid_model_data()
@@ -133,7 +139,9 @@ def construct_fields(self) -> list[Field | VectorField]:
         return list(fields.values())
 
     @classmethod
-    def from_sgrid_conventions(cls, ds: xr.Dataset, mesh: Mesh | None) -> Self:
+    def from_sgrid_conventions(
+        cls, ds: xr.Dataset, mesh: Mesh | None, vector_fields: TVectorFieldMapping | None | _MissingType
+    ) -> Self:
         ds = ds.copy()
         if mesh is None:
             mesh = _get_mesh_type_from_sgrid_dataset(ds)
@@ -160,7 +168,7 @@ def from_sgrid_conventions(cls, ds: xr.Dataset, mesh: Mesh | None) -> Self:
         #     ds["lon"] = ds[node_dimensions[0]]
         #     ds["lat"] = ds[node_dimensions[1]]
 
-        model = cls(ds, mesh=mesh)
+        model = cls(ds, mesh=mesh, vector_field_components=vector_fields)
         model._fields = model.construct_fields()
         for f in model._fields:
             if isinstance(f, Field):
@@ -191,13 +199,14 @@ def from_sgrid_conventions(cls, ds: xr.Dataset, mesh: Mesh | None) -> Self:
             ),
         ),
         mesh=mesh,  # type:ignore
+        vector_fields=None,
     )
     for mesh in ["flat", "spherical"]
 }
 
 
 class UnstructuredModelData(ModelData):
-    def __init__(self, data: ux.UxDataset, grid: UxGrid):
+    def __init__(self, data: ux.UxDataset, grid: UxGrid, vector_field_components: TVectorFieldMapping):
         if not isinstance(data, ux.UxDataset):
             raise ValueError(f"Expected `data` to be an uxarray.UxDataset . Got {type(data)}")
 
@@ -206,6 +215,7 @@ def __init__(self, data: ux.UxDataset, grid: UxGrid):
 
         self.data = data
         self.grid = grid
+        self.vector_field_components = vector_field_components
         self.field_to_interpolator = {}
         self._fields: list[Field | VectorField] | None = None
 
@@ -239,7 +249,9 @@ def scalar_field_names(self) -> list[str]:
         return list(self.data.data_vars)
 
     @classmethod
-    def from_ugrid_conventions(cls, ds: ux.UxDataset, mesh: Mesh):
+    def from_ugrid_conventions(
+        cls, ds: ux.UxDataset, mesh: Mesh, vector_fields: TVectorFieldMapping | None | _MissingType
+    ):
         ds_dims = list(ds.dims)
         if not all(dim in ds_dims for dim in ["time", "zf", "zc"]):
             raise ValueError(
@@ -276,6 +288,17 @@ def _get_mesh_type_from_sgrid_dataset(ds_sgrid: xr.Dataset) -> Mesh:
     return "spherical" if _is_coordinate_in_degrees(ds_sgrid[fpoint_x]) else "flat"
 
 
+def _default_vector_field_components(data_vars: Sequence[str]) -> TVectorFieldMapping:
+    vars = set(data_vars)
+    ret = {}
+
+    if {"U", "V"}.issubset(vars):
+        ret["UV"] = ("U", "V")
+    if {"U", "V", "W"}.issubset(vars):
+        ret["UVW"] = ("U", "V", "W")
+    return ret
+
+
 def _is_coordinate_in_degrees(da: xr.DataArray) -> bool:
     units = da.attrs.get("units")
     if units is None:

tests/test_fieldset.py

@@ -8,6 +8,7 @@
 
 from parcels import Field, ParticleFile, ParticleSet, XGrid, convert
 from parcels._core.fieldset import FieldSet, _datetime_to_msg
+from parcels._core.model import _default_vector_field_components
 from parcels._datasets.structured.generic import datasets as datasets_structured
 from parcels._datasets.structured.generic import datasets_sgrid
 from parcels._datasets.unstructured.generic import datasets as datasets_unstructured
@@ -126,7 +127,17 @@ def test_fieldset_vectorfield_none():
     assert "UV" not in fset1.fields
 
 
-def test_resolve_vector_field_components(): ...
+@pytest.mark.parametrize(
+    "data_vars,expected",
+    [
+        (["U", "V", "land_mask"], {"UV": ("U", "V")}),
+        (["U", "V", "W", "land_mask"], {"UV": ("U", "V"), "UVW": ("U", "V", "W")}),
+        (["field1", "field2", "field3"], {}),
+    ],
+)
+def test_default_vector_field_components(data_vars, expected):
+    got = _default_vector_field_components(data_vars)
+    assert got == expected
 
 
 # TODO restructure: use adding of fieldset notation to test this