Add fieldset.describe

Author: VeckoTheGecko

src/parcels/_core/fieldset.py

@@ -21,6 +21,7 @@
 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 NOTSET, NotSetType
+from parcels._reprs import fieldset_describe
 from parcels.interpolators import (
     XConstantField,
 )
@@ -284,6 +285,10 @@ def from_sgrid_conventions(
         model = StructuredModelData.from_sgrid_conventions(ds, mesh, vector_fields)
         return cls([model])
 
+    def describe(self):
+        """Return a table description of a FieldSet, which fields it has and their interpolation methods."""
+        return fieldset_describe(self)
+
 
 def assert_compatible_fieldsets(left: FieldSet, right: FieldSet) -> None:
     """Assert that two FieldSets can be combined without name conflicts.

src/parcels/_reprs.py

@@ -3,14 +3,18 @@
 from __future__ import annotations
 
 import textwrap
-from typing import TYPE_CHECKING, Any, cast
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Literal, cast
 
 import numpy as np
 import xarray as xr
 
+from parcels._python import isinstance_noimport
+
 if TYPE_CHECKING:
     from parcels import Field, FieldSet, ParticleSet
     from parcels._core.field import VectorField
+    from parcels._core.model import ModelData
 
 
 def fieldset_repr(fieldset: FieldSet) -> str:
@@ -177,3 +181,69 @@ def _format_list_items_multiline(items: list[str] | dict, level: int = 1, with_b
 
 def is_builtin_object(obj):
     return obj.__class__.__module__ == "builtins"
+
+
+@dataclass
+class _FieldSetDescriptionRow:
+    type_: Literal["Field", "VectorField", "Context value"]
+    model_id: int | None
+    name: str
+    interp_method: str | None
+
+    def to_dict(self) -> dict[str, str]:
+        return {
+            "Type": self.type_,
+            "Dataset origin": str(self.model_id) if self.model_id is not None else "-",
+            "Name": self.name,
+            "Interp method": str(self.interp_method) if self.interp_method is not None else "-",
+        }
+
+
+def _print_table(rows: list[_FieldSetDescriptionRow]) -> str:
+    import pandas as pd
+
+    dicts = [r.to_dict() for r in rows]
+    return pd.DataFrame(dicts).sort_values(["Dataset origin", "Type", "Name"]).to_markdown(index=False)
+
+
+def fieldset_describe(fieldset: FieldSet) -> str:
+    rows: list[_FieldSetDescriptionRow] = []
+    models: dict[int, int] = {}  # mapping of memory ID to a human readable ID
+
+    assert fieldset._fields is not None
+
+    for field in fieldset._fields.values():
+        model_id: int
+
+        # Set human readable model ID
+        parent_id = id(_get_parent_model(field))
+        models[parent_id] = models.get(parent_id, len(models))
+        model_id = models[parent_id]
+
+        type_ = field.__class__.__name__
+        assert type_ in ("Field", "VectorField")
+
+        rows.append(
+            _FieldSetDescriptionRow(
+                type_=type_,
+                model_id=model_id,
+                name=field.name,
+                interp_method=repr(field.interp_method),
+            )
+        )
+    for k, v in fieldset.context.items():
+        rows.append(
+            _FieldSetDescriptionRow(
+                type_="Context value",
+                model_id=None,
+                name=k,
+                interp_method=repr(v),
+            )
+        )
+    return _print_table(rows)
+
+
+def _get_parent_model(field: Field | VectorField) -> ModelData:
+    if isinstance_noimport(field, "Field"):
+        return field.model
+    return field.U.model

tests/test_fieldset.py

@@ -18,6 +18,20 @@
 ds = datasets_structured["ds_2d_left"]
 
 
+@pytest.fixture
+def fieldset_two_models():
+    ds1 = datasets_structured["ds_2d_left"][["U_A_grid", "V_A_grid", "grid"]].rename({"U_A_grid": "U", "V_A_grid": "V"})
+    ds2 = datasets_structured["ds_2d_left"][["U_A_grid", "V_A_grid", "grid"]].rename(
+        {"U_A_grid": "U_wind", "V_A_grid": "V_wind"}
+    )
+
+    fset1 = FieldSet.from_sgrid_conventions(ds1, mesh="flat")
+    fset2 = FieldSet.from_sgrid_conventions(ds2, mesh="flat", vector_fields={"UV_wind": ("U_wind", "V_wind")})
+    fset2.add_context("my_value", 2.0)
+    fset2.add_context("my_list", [1, 2, "hello"])
+    return fset1 + fset2
+
+
 def test_fieldset_init_wrong_types():
     with pytest.raises(ValueError, match="Expected `model` to be a ModelData object. Got .*"):
         FieldSet([1.0, 2.0, 3.0])
@@ -373,3 +387,21 @@ def test_fieldset_add_context_values():
 
     assert fset.context["c1"] == 1.0
     assert fset.context["c2"] == 2.0
+
+
+def test_fieldset_describe(fieldset_two_models):
+    fieldset = fieldset_two_models
+    expected = """\
+| Type          | Dataset origin   | Name     | Interp method         |
+|:--------------|:-----------------|:---------|:----------------------|
+| Context value | -                | my_list  | [1, 2, 'hello']       |
+| Context value | -                | my_value | 2.0                   |
+| Field         | 0                | U        | XLinear(...)          |
+| Field         | 0                | V        | XLinear(...)          |
+| VectorField   | 0                | UV       | XLinear_Velocity(...) |
+| Field         | 1                | U_wind   | XLinear(...)          |
+| Field         | 1                | V_wind   | XLinear(...)          |
+| VectorField   | 1                | UV_wind  | XLinear_Velocity(...) |"""
+    actual = fieldset.describe()
+    # breakpoint()
+    assert actual == expected