Add open_raw_zarr helper

changes.md

@@ -0,0 +1,262 @@
+# Refactoring Summary: `field.py`, `fieldset.py`, `model.py`
+
+This document describes the refactoring introduced in commit `69338d87a89763efbb1e3886b470e09992812978` relative to `main`.
+
+---
+
+## Overview
+
+The central change is the introduction of a new `Model` abstraction layer between raw xarray/uxarray data and the `Field`/`FieldSet` objects. Previously, `Field` owned its data and grid directly. Now, `Field` is a thin view over a `Model`, and `FieldSet` is a container of `Model` objects rather than `Field` objects.
+
+---
+
+## New file: `src/parcels/_core/model.py`
+
+### `Model` (abstract base class)
+
+Abstract class with three required attributes:
+
+- `data: Any` — the underlying dataset
+- `grid: BaseGrid` — the grid object
+- `field_to_interpolator: dict[str, ScalarInterpolator | VectorInterpolator]` — maps field names to interpolator instances
+
+Abstract methods:
+
+- `construct_fields() -> list[Field | VectorField]` — build field objects from this model
+- `scalar_field_names -> list[str]` — names of scalar fields in the data
+- `assert_valid_field_data(field_data)` — validate a single field's data
+
+Concrete methods on `Model`:
+
+- `assert_valid_model_data()` — iterates `scalar_field_names` and calls `assert_valid_field_data` on each
+- `time_interval -> TimeInterval | None` — computed from `self.data`
+
+### `StructuredModel(Model)`
+
+For structured (SGRID) grid data backed by `xr.Dataset`.
+
+Constructor: `StructuredModel(data: xr.Dataset, mesh: Mesh)`
+
+- Calls `preprocess_sgrid_model_data(data)` to transpose fields to `(t, z, y, x)` order
+- Creates an `XGrid(data, mesh)` grid
+- Initializes `field_to_interpolator = {}`
+- Calls `assert_valid_model_data()` on construction
+
+`from_sgrid_conventions(cls, ds, mesh=None)` classmethod:
+
+- Copied/moved from `FieldSet.from_sgrid_conventions` — handles time axis renaming, mesh type inference
+- Sets default interpolator `XLinear()` on all scalar fields after construction
+- Returns a `StructuredModel` instance
+
+`construct_fields()`:
+
+- Creates `Field("U", self)`, `Field("V", self)` etc., then wraps them in `VectorField("UV", ...)` if U+V present
+- Uses `XLinear_Velocity()` for A-grids, `CGrid_Velocity()` for C-grids
+
+### `UnstructuredModel(Model)`
+
+For unstructured (UGRID) grid data backed by `ux.UxDataset`.
+
+Constructor: `UnstructuredModel(data: ux.UxDataset, grid: UxGrid)`
+
+`from_ugrid_conventions(cls, ds, mesh="spherical")` classmethod:
+
+- Validates required dimensions (`time`, `zf`, `zc`)
+- Creates `UxGrid`, calls `_discover_ux_U_and_V`, returns instance
+
+`construct_fields()`:
+
+- Uses `_select_uxinterpolator(da)` to pick the appropriate interpolator per field
+- Note: interpolator is passed as 3rd arg to `Field(name, model, interp)` — see Field changes below
+
+### Helper functions moved from `fieldset.py` to `model.py`
+
+- `_discover_ux_U_and_V(ds)` — unchanged logic
+- `_select_uxinterpolator(da)` — unchanged logic
+- `_get_mesh_type_from_sgrid_dataset(ds)` — unchanged logic
+- `_is_coordinate_in_degrees(da)` — unchanged logic
+- `_get_time_interval(data)` — logic adjusted: checks `"time" not in data or data["time"].size == 1` (previously checked `data.shape[0] == 1`)
+- `_assert_valid_uxdataarray(data)` — unchanged logic
+- `_assert_has_time_coordinate(da)` — new helper extracted from old `Field.__init__`
+
+### New helper in `model.py`
+
+- `preprocess_sgrid_model_data(ds)` — transposes all non-grid-topology data vars to `(t, z, y, x)` using `_transpose_xfield_data_to_tzyx`
+
+---
+
+## Changes to `src/parcels/_core/field.py`
+
+### `Field.__init__` signature change
+
+**Before:**
+
+```python
+Field(name: str, data: xr.DataArray | ux.UxDataArray, grid: UxGrid | XGrid, interp_method: Callable)
+```
+
+**After:**
+
+```python
+Field(name: str, model: Model)
+```
+
+- `data`, `grid`, and `interp_method` are no longer constructor arguments
+- The constructor only sets `self.name`, `self.model`, and `self.igrid = -1`
+- Validation (data type checks, axis checks, time interval extraction) removed from `__init__`
+
+### `Field` properties (delegating to model)
+
+Three new properties proxy into the model:
+
+```python
+@property
+def data(self):
+    return self.model.data[self.name]
+
+@property
+def grid(self):
+    return self.model.grid
+
+@property
+def time_interval(self):
+    return self.model.time_interval
+```
+
+These preserve backward compatibility for code that reads `field.data`, `field.grid`, `field.time_interval`.
+
+### `Field.interp_method` property/setter
+
+**Before:** stored as `self._interp_method`; validated via `assert_same_function_signature` against `ZeroInterpolator`
+
+**After:** stored in `self.model.field_to_interpolator[self.name]`
+
+- Getter raises `AttributeError` (not `KeyError`) if no interpolator is set for this field
+- Setter validates `isinstance(value, ScalarInterpolator)` instead of checking function signature
+
+### Interpolator call convention change
+
+**Before:** `self._interp_method(particle_positions, grid_positions, self)`
+
+**After:** `self.interp_method.interp(particle_positions, grid_positions, self)`
+
+Interpolators are now objects with an `.interp(...)` method, not plain callables.
+
+### `VectorField` changes
+
+- `interp_method` parameter type annotation changed from `Callable | None` to `VectorInterpolator | None`
+- Validation changed from `assert_same_function_signature(...)` to `isinstance(interp_method, VectorInterpolator)`
+- Setter similarly validates `isinstance(method, VectorInterpolator)`
+- Call site: `self._interp_method.interp(...)` instead of `self._interp_method(...)`
+
+### Removed from `field.py`
+
+- `_assert_valid_uxdataarray` — moved to `model.py`
+- `_assert_compatible_combination` — removed (validation now handled per-model)
+- `_get_time_interval` — moved to `model.py`
+- Imports: `uxarray`, `xarray`, `Callable`, `TimeInterval`, `ZeroInterpolator`, `ZeroInterpolator_Vector`, `assert_same_function_signature`, `_transpose_xfield_data_to_tzyx`, `assert_all_field_dims_have_axis`
+
+---
+
+## Changes to `src/parcels/_core/fieldset.py`
+
+### `FieldSet.__init__` signature change
+
+**Before:** `FieldSet(fields: list[Field | VectorField])`
+
+**After:** `FieldSet(models: list[Model])`
+
+- Now stores `self.models: list[Model]`
+- Calls `self.reconstruct_fields()` on init to build `self._fields`
+- `assert_compatible_calendars(fields)` call commented out (TODO)
+
+### New `FieldSet.fields` property
+
+`_fields` is now the backing store; `fields` is a lazy property that calls `reconstruct_fields()` if `_fields` is `None`.
+
+### New `FieldSet.reconstruct_fields()` method
+
+Iterates `self.models`, calls `model.construct_fields()` on each, flattens into `self._fields` dict.
+
+### `context` renamed to `constants`
+
+- `self.context` → `self.constants`
+- `add_context(name, value)` → `add_constant(name, value)`
+- `add_constant` now validates that `value` is `float | np.floating | int | np.integer`
+- `__setattr__` override that guarded `context` keys has been **removed**
+
+### `__getattr__` updated
+
+Now checks `self._fields` and `self.constants` (was `self.fields` and `self.context`).
+
+### New `FieldSet.__add__` operator
+
+```python
+def __add__(self, other: FieldSet) -> FieldSet:
+    assert_compatible_fieldsets(self, other)
+    combined = FieldSet(self.models + other.models)
+    combined.constants = {**self.constants, **other.constants}
+    return combined
+```
+
+### `from_ugrid_conventions` simplified
+
+**Before:** ~15 lines building grid, discovering U/V, creating Field objects, returning `cls(list(fields.values()))`
+
+**After:**
+
+```python
+model = UnstructuredModel.from_ugrid_conventions(ds, mesh)
+return cls([model])
+```
+
+### `from_sgrid_conventions` simplified
+
+**Before:** ~50 lines handling time axis, xgcm grid creation, field creation
+
+**After:**
+
+```python
+model = StructuredModel.from_sgrid_conventions(ds, mesh)
+return cls([model])
+```
+
+### `add_field` constant field creation updated
+
+The inline `xgcm.Grid(...)` call when adding a constant scalar field is replaced with constructing `XGrid(ds, mesh=mesh)` directly (after attaching SGRID metadata via `sgrid._attach_sgrid_metadata`).
+
+### New module-level function: `assert_compatible_fieldsets`
+
+```python
+def assert_compatible_fieldsets(left: FieldSet, right: FieldSet) -> None
+```
+
+Raises `ValueError` if the two fieldsets share any field names or constant names.
+
+### Removed from `fieldset.py`
+
+- `xgcm` import
+- `UxGrid` import
+- `_DEFAULT_XGCM_KWARGS` import
+- `logger` import
+- `_ds_rename_using_standard_names` import
+- Most interpolator imports (only `XConstantField` remains)
+- `_discover_ux_U_and_V` — moved to `model.py`
+- `_select_uxinterpolator` — moved to `model.py`
+- `_get_mesh_type_from_sgrid_dataset` — moved to `model.py`
+- `_is_coordinate_in_degrees` — moved to `model.py`
+
+---
+
+## Summary of architectural intent
+
+| Concern                 | Before                                             | After                                                    |
+| ----------------------- | -------------------------------------------------- | -------------------------------------------------------- |
+| Data ownership          | `Field` (held `self.data`, `self.grid`)            | `Model` (holds `self.data`, `self.grid`)                 |
+| Interpolator storage    | `Field._interp_method` (per-field callable)        | `Model.field_to_interpolator` (dict of objects)          |
+| Interpolator type       | Any callable matching `ZeroInterpolator` signature | Instance of `ScalarInterpolator` / `VectorInterpolator`  |
+| Interpolator invocation | `interp_method(positions, grid_positions, field)`  | `interp_method.interp(positions, grid_positions, field)` |
+| `FieldSet` contents     | `list[Field \| VectorField]`                       | `list[Model]`                                            |
+| Field construction      | Done in `FieldSet.from_*` classmethods             | Delegated to `Model.construct_fields()`                  |
+| `context` / `constants` | `fieldset.context` (any type)                      | `fieldset.constants` (float/int only)                    |
+| `FieldSet` combination  | Not supported                                      | `fieldset_a + fieldset_b` via `__add__`                  |