Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
60 changes: 3 additions & 57 deletions src/parcels/_core/fieldset.py
Original file line number Diff line number Diff line change
Expand Up @@ -183,6 +183,9 @@ def gridset(self) -> list[BaseGrid]:
@classmethod
def from_ugrid_conventions(cls, ds: ux.UxDataset, mesh: str = "spherical"):
"""Create a FieldSet from a Parcels compliant uxarray.UxDataset.

This is the primary ingestion method in Parcels for structured grid datasets.

The main requirements for a uxDataset are naming conventions for vertical grid dimensions & coordinates

zf - Name for coordinate and dimension for vertical positions at layer interfaces
Expand Down Expand Up @@ -225,63 +228,6 @@ def from_ugrid_conventions(cls, ds: ux.UxDataset, mesh: str = "spherical"):

return cls(list(fields.values()))

@classmethod
def from_fesom2(cls, ds: ux.UxDataset, mesh: str = "spherical"):
"""Create a FieldSet from a FESOM2 uxarray.UxDataset.

Parameters
----------
ds : uxarray.UxDataset
uxarray.UxDataset as obtained from the uxarray package.

Returns
-------
FieldSet
FieldSet object containing the fields from the dataset that can be used for a Parcels simulation.
"""
ds = ds.copy()
ds_dims = list(ds.dims)
if not all(dim in ds_dims for dim in ["time", "nz", "nz1"]):
raise ValueError(
f"Dataset missing one of the required dimensions 'time', 'nz', or 'nz1' for FESOM data. Found dimensions {ds_dims}"
)
ds = ds.rename(
{
"nz": "zf", # Vertical Interface
"nz1": "zc", # Vertical Center
}
).set_index(zf="zf", zc="zc")

return FieldSet.from_ugrid_conventions(ds, mesh=mesh)

@classmethod
def from_icon(cls, ds: ux.UxDataset, mesh: str = "spherical"):
"""Create a FieldSet from a ICON uxarray.UxDataset.

Parameters
----------
ds : uxarray.UxDataset
uxarray.UxDataset as obtained from the uxarray package.

Returns
-------
FieldSet
FieldSet object containing the fields from the dataset that can be used for a Parcels simulation.
"""
ds = ds.copy()
ds_dims = list(ds.dims)
if not all(dim in ds_dims for dim in ["time", "depth", "depth_2"]):
raise ValueError(
f"Dataset missing one of the required dimensions 'time', 'depth', or 'depth_2' for ICON data. Found dimensions {ds_dims}"
)
ds = ds.rename(
{
"depth_2": "zf", # Vertical Interface
"depth": "zc", # Vertical Center
}
).set_index(zf="zf", zc="zc")
return FieldSet.from_ugrid_conventions(ds, mesh=mesh)

@classmethod
def from_sgrid_conventions(
cls, ds: xr.Dataset, mesh: Mesh | None = None
Expand Down
189 changes: 189 additions & 0 deletions src/parcels/convert.py
Original file line number Diff line number Diff line change
Expand Up @@ -529,3 +529,192 @@ def copernicusmarine_to_sgrid(
)

return ds


# Known vertical dimension mappings by model
_FESOM2_VERTICAL_DIMS = {"interface": "nz", "center": "nz1"}
_ICON_VERTICAL_DIMS = {"interface": "depth_2", "center": "depth"}


def _detect_vertical_coordinates(
ds: ux.UxDataset,
known_mappings: dict[str, str] | None = None,
) -> tuple[str, str]:
"""Detect vertical coordinate dimensions for faces (zf) and centers (zc).

Detection strategy (with fallback):
1. Use known_mappings if provided and dimensions exist
2. Look for CF convention axis='Z' metadata
3. Find dimension pairs where sizes differ by exactly 1

Parameters
----------
ds : ux.UxDataset
UxDataset to analyze for vertical coordinates.
known_mappings : dict[str, str] | None
Optional mapping with keys "interface" and "center" specifying
the dimension names for layer interfaces (zf) and centers (zc).

Returns
-------
tuple[str, str]
Tuple of (interface_dim_name, center_dim_name).

Raises
------
ValueError
If vertical coordinates cannot be detected.
"""
ds_dims = set(ds.dims)

# Strategy 1: Use known mappings if provided and dimensions exist
if known_mappings is not None:
interface_dim = known_mappings.get("interface")
center_dim = known_mappings.get("center")
if interface_dim in ds_dims and center_dim in ds_dims:
logger.info(
f"Using known vertical dimension mapping: {interface_dim!r} (interfaces) and {center_dim!r} (centers)."
)
return interface_dim, center_dim
logger.debug(f"Known mappings {known_mappings} not found in dataset dimensions {ds_dims}. Trying CF metadata.")

# Strategy 2: Look for CF convention axis='Z' metadata
z_dims = []
for dim in ds_dims:
if dim in ds.coords:
coord = ds.coords[dim]
if coord.attrs.get("axis") == "Z":
z_dims.append(dim)
elif coord.attrs.get("positive") in ("up", "down"):
z_dims.append(dim)
elif "depth" in coord.attrs.get("standard_name", "").lower():
z_dims.append(dim)

if len(z_dims) == 2:
# Sort by size - interface has n+1 values, center has n
z_dims_sorted = sorted(z_dims, key=lambda d: ds.sizes[d], reverse=True)
interface_dim, center_dim = z_dims_sorted
if ds.sizes[interface_dim] == ds.sizes[center_dim] + 1:
logger.info(
f"Detected vertical dimensions from CF metadata: {interface_dim!r} (interfaces) and {center_dim!r} (centers)."
)
return interface_dim, center_dim

# Strategy 3: Find dimension pairs where sizes differ by exactly 1
# Skip known non-vertical dimensions
skip_dims = {"time", "n_face", "n_node", "n_edge", "n_max_face_nodes"}
candidate_dims = [d for d in ds_dims if d not in skip_dims]

for dim1 in candidate_dims:
for dim2 in candidate_dims:
if dim1 != dim2:
size1, size2 = ds.sizes[dim1], ds.sizes[dim2]
if size1 == size2 + 1:
logger.info(
f"Auto-detected vertical dimensions by size difference: {dim1!r} (interfaces, size={size1}) "
f"and {dim2!r} (centers, size={size2})."
)
return dim1, dim2

raise ValueError(
f"Could not detect vertical coordinate dimensions in dataset with dims {list(ds_dims)}. "
"Please ensure the dataset has vertical layer interface and center dimensions, "
"or rename them manually to 'zf' (interfaces) and 'zc' (centers)."
)


def _rename_vertical_dims(
ds: ux.UxDataset,
interface_dim: str,
center_dim: str,
) -> ux.UxDataset:
"""Rename vertical dimensions to zf (interfaces) and zc (centers).

Parameters
----------
ds : ux.UxDataset
UxDataset with vertical dimensions to rename.
interface_dim : str
Current name of the interface dimension.
center_dim : str
Current name of the center dimension.

Returns
-------
ux.UxDataset
Dataset with renamed dimensions and indexed coordinates.
"""
rename_dict = {}
if interface_dim != "zf":
rename_dict[interface_dim] = "zf"
if center_dim != "zc":
rename_dict[center_dim] = "zc"

if rename_dict:
logger.info(f"Renaming vertical dimensions: {rename_dict}")
ds = ds.rename(rename_dict)

ds = ds.set_index(zf="zf", zc="zc")
return ds


def fesom_to_ugrid(ds: ux.UxDataset) -> ux.UxDataset:
"""Convert FESOM2 UxDataset to Parcels UGRID-compliant format.

Renames vertical dimensions:
- nz -> zf (vertical layer faces/interfaces)
- nz1 -> zc (vertical layer centers)

Parameters
----------
ds : ux.UxDataset
FESOM2 UxDataset as obtained from uxarray.

Returns
-------
ux.UxDataset
UGRID-compliant dataset ready for FieldSet.from_ugrid_conventions().

Examples
--------
>>> import uxarray as ux
>>> from parcels import FieldSet
>>> from parcels.convert import fesom_to_ugrid
>>> ds = ux.open_mfdataset(grid_path, data_path)
>>> ds_ugrid = fesom_to_ugrid(ds)
>>> fieldset = FieldSet.from_ugrid_conventions(ds_ugrid, mesh="flat")
"""
ds = ds.copy()
interface_dim, center_dim = _detect_vertical_coordinates(ds, _FESOM2_VERTICAL_DIMS)
return _rename_vertical_dims(ds, interface_dim, center_dim)


def icon_to_ugrid(ds: ux.UxDataset) -> ux.UxDataset:
"""Convert ICON UxDataset to Parcels UGRID-compliant format.

Renames vertical dimensions:
- depth_2 -> zf (vertical layer faces/interfaces)
- depth -> zc (vertical layer centers)

Parameters
----------
ds : ux.UxDataset
ICON UxDataset as obtained from uxarray.

Returns
-------
ux.UxDataset
UGRID-compliant dataset ready for FieldSet.from_ugrid_conventions().

Examples
--------
>>> import uxarray as ux
>>> from parcels import FieldSet
>>> from parcels.convert import icon_to_ugrid
>>> ds = ux.open_mfdataset(grid_path, data_path)
>>> ds_ugrid = icon_to_ugrid(ds)
>>> fieldset = FieldSet.from_ugrid_conventions(ds_ugrid, mesh="flat")
"""
ds = ds.copy()
interface_dim, center_dim = _detect_vertical_coordinates(ds, _ICON_VERTICAL_DIMS)
return _rename_vertical_dims(ds, interface_dim, center_dim)
16 changes: 8 additions & 8 deletions tests/test_fieldset.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
import pytest
import xarray as xr

from parcels import Field, ParticleFile, ParticleSet, VectorField, XGrid
from parcels import Field, ParticleFile, ParticleSet, VectorField, XGrid, convert
from parcels._core.fieldset import CalendarError, FieldSet, _datetime_to_msg
from parcels._datasets.structured.generic import T as T_structured
from parcels._datasets.structured.generic import datasets as datasets_structured
Expand Down Expand Up @@ -243,33 +243,33 @@ def test_fieldset_add_field_after_pset():


def test_fieldset_from_icon():
ds = datasets_unstructured["icon_square_delaunay_uniform_z_coordinate"]
fieldset = FieldSet.from_icon(ds)
ds = convert.icon_to_ugrid(datasets_unstructured["icon_square_delaunay_uniform_z_coordinate"])
fieldset = FieldSet.from_ugrid_conventions(ds)
assert "U" in fieldset.fields
assert "V" in fieldset.fields
assert "UVW" in fieldset.fields


def test_fieldset_from_fesom2():
ds = datasets_unstructured["fesom2_square_delaunay_uniform_z_coordinate"]
fieldset = FieldSet.from_fesom2(ds)
ds = convert.fesom_to_ugrid(datasets_unstructured["fesom2_square_delaunay_uniform_z_coordinate"])
fieldset = FieldSet.from_ugrid_conventions(ds)
assert "U" in fieldset.fields
assert "V" in fieldset.fields
assert "UVW" in fieldset.fields


def test_fieldset_from_fesom2_missingUV():
ds = datasets_unstructured["fesom2_square_delaunay_uniform_z_coordinate"]
ds = convert.fesom_to_ugrid(datasets_unstructured["fesom2_square_delaunay_uniform_z_coordinate"])
# Intentionally create a dataset that is missing the U field
localds = ds.rename({"U": "notU"})
with pytest.raises(ValueError) as info:
_ = FieldSet.from_fesom2(localds)
_ = FieldSet.from_ugrid_conventions(localds)
assert "Dataset has only one of the two variables 'U' and 'V'" in str(info)

# Intentionally create a dataset that is missing the V field
localds = ds.rename({"V": "notV"})
with pytest.raises(ValueError) as info:
_ = FieldSet.from_fesom2(localds)
_ = FieldSet.from_ugrid_conventions(localds)
assert "Dataset has only one of the two variables 'U' and 'V'" in str(info)


Expand Down
27 changes: 7 additions & 20 deletions tests/test_uxarray_fieldset.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
download_example_dataset,
)
from parcels._datasets.unstructured.generic import datasets as datasets_unstructured
from parcels.convert import fesom_to_ugrid, icon_to_ugrid
from parcels.interpolators import (
Ux_Velocity,
UxConstantFaceConstantZC,
Expand All @@ -29,12 +30,7 @@ def ds_fesom_channel() -> ux.UxDataset:
f"{fesom_path}/w.fesom_channel.nc",
]
ds = ux.open_mfdataset(grid_path, data_path).rename_vars({"u": "U", "v": "V", "w": "W"})
ds = ds.rename(
{
"nz": "zf", # Vertical Interface
"nz1": "zc", # Vertical Center
}
).set_index(zf="zf", zc="zc")
ds = fesom_to_ugrid(ds)
return ds


Expand Down Expand Up @@ -121,12 +117,7 @@ def test_fesom2_square_delaunay_uniform_z_coordinate_eval():
Since the underlying data is constant, we can check that the values are as expected.
"""
ds = datasets_unstructured["fesom2_square_delaunay_uniform_z_coordinate"]
ds = ds.rename(
{
"nz": "zf", # Vertical Interface
"nz1": "zc", # Vertical Center
}
).set_index(zf="zf", zc="zc")
ds = fesom_to_ugrid(ds)
grid = UxGrid(ds.uxgrid, z=ds.coords["zf"], mesh="flat")
UVW = VectorField(
name="UVW",
Expand Down Expand Up @@ -174,12 +165,7 @@ def test_fesom2_square_delaunay_antimeridian_eval():
Since the underlying data is constant, we can check that the values are as expected.
"""
ds = datasets_unstructured["fesom2_square_delaunay_antimeridian"]
ds = ds.rename(
{
"nz": "zf", # Vertical Interface
"nz1": "zc", # Vertical Center
}
).set_index(zf="zf", zc="zc")
ds = fesom_to_ugrid(ds)
P = Field(
name="p",
data=ds.p,
Expand All @@ -196,15 +182,16 @@ def test_fesom2_square_delaunay_antimeridian_eval():

def test_icon_evals():
ds = datasets_unstructured["icon_square_delaunay_uniform_z_coordinate"].copy(deep=True)
fieldset = FieldSet.from_icon(ds, mesh="flat")
ds = icon_to_ugrid(ds)
fieldset = FieldSet.from_ugrid_conventions(ds, mesh="flat")

# Query points, are chosen to be just a fraction off from the center of a cell for testing
# This generic dataset has an effective lateral grid-spacing of 3 degrees and vertical grid
# spacing of 100m - shifting by 1/10 of a degree laterally and 10m vertically should keep us
# within the cell and make for easy exactness checking of constant and linear interpolation
xc = ds.uxgrid.face_lon.values
yc = ds.uxgrid.face_lat.values
zc = 0.0 * xc + ds.depth.values[1] # Make zc the same length as xc
zc = 0.0 * xc + ds.zc.values[1] # Make zc the same length as xc

tq = 0.0 * xc
xq = xc + 0.1
Expand Down
Loading