Skip to content

Commit 9f5623c

Browse files
VeckoTheGeckoVeckoTheGecko
authored
Fix convert.fesom_to_ugrid dimension renaming (#2579)
1 parent 11b0458 commit 9f5623c

3 files changed

Lines changed: 65 additions & 41 deletions

File tree

‎src/parcels/_datasets/remote.py‎

Lines changed: 46 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -20,9 +20,10 @@
2020
_DATA_HOME = os.environ.get("PARCELS_EXAMPLE_DATA")
2121
if _DATA_HOME is None:
2222
_DATA_HOME = pooch.os_cache("parcels")
23+
_DATA_HOME = Path(_DATA_HOME)
2324

2425
# See instructions at https://github.com/Parcels-code/parcels-data for adding new datasets
25-
_POOCH_REGISTRY_FILES: list[str] = (
26+
_ODIE_REGISTRY_FILES: list[str] = (
2627
# These datasets are from v3 and before of Parcels, where we just used netcdf files
2728
[
2829
"data/MovingEddies_data/moving_eddiesP.nc",
@@ -101,16 +102,17 @@
101102
+ [f"data/WOA_data/woa18_decav_t{m:02d}_04.nc" for m in range(1, 13)]
102103
+ ["data/CROCOidealized_data/CROCO_idealized.nc"]
103104
# These datasets are from v4 of Parcels where we're opting for Zipped zarr datasets
104-
# ...
105+
+ [
106+
"data-zarr/Benchmarks_FESOM2-baroclinic-gyre/data.zip",
107+
"data-zarr/Benchmarks_FESOM2-baroclinic-gyre/grid.zip",
108+
]
109+
+ []
105110
)
106111

107-
_POOCH_REGISTRY = {k: None for k in _POOCH_REGISTRY_FILES}
108-
109-
110112
_ODIE = pooch.create(
111113
path=_DATA_HOME,
112114
base_url=_DATA_URL,
113-
registry=_POOCH_REGISTRY,
115+
registry={k: None for k in _ODIE_REGISTRY_FILES},
114116
)
115117

116118

@@ -120,13 +122,13 @@ def open_dataset(self) -> xr.Dataset: ...
120122

121123

122124
class _V3Dataset(_ParcelsDataset):
123-
def __init__(self, path_relative_to_root: str, pre_decode_cf_callable=None):
124-
self.path_relative_to_root = path_relative_to_root # glob is allowed
125+
def __init__(self, pup: pooch.Pooch, path_relative_to_pup: str, pre_decode_cf_callable=None):
126+
self.path_relative_to_root = path_relative_to_pup # glob is allowed
125127

126128
# Function to apply to the dataset before the decoding the CF variables
127-
self.pup = _ODIE
129+
self.pup = pup
128130
self.pre_decode_cf_callable: None | Callable[[xr.Dataset], xr.Dataset] = pre_decode_cf_callable
129-
self.v3_dataset_name = path_relative_to_root.split("/")[0]
131+
self.v3_dataset_name = path_relative_to_pup.split("/")[0]
130132

131133
def open_dataset(self) -> xr.Dataset:
132134
self.download_relevant_files()
@@ -154,9 +156,9 @@ def download_relevant_files(self) -> None:
154156

155157

156158
class _ZarrZipDataset(_ParcelsDataset):
157-
def __init__(self, path_relative_to_root):
158-
self.pup = _ODIE
159-
self.path_relative_to_root = path_relative_to_root
159+
def __init__(self, pup, path_relative_to_pup):
160+
self.pup = pup
161+
self.path_relative_to_root = path_relative_to_pup
160162

161163
def open_dataset(self) -> xr.Dataset:
162164
self.pup.fetch(self.path_relative_to_root)
@@ -190,35 +192,38 @@ class _Purpose(enum.Enum):
190192
# The first here is a human readable key used to open datasets, with an object to open the datasets
191193
# fmt: off
192194
_DATASET_KEYS_AND_CONFIGS: dict[str, tuple[_V3Dataset, _Purpose]] = dict([
193-
("MovingEddies_data/P", (_V3Dataset("data/MovingEddies_data/moving_eddiesP.nc"), _Purpose.TUTORIAL)),
194-
("MovingEddies_data/U", (_V3Dataset("data/MovingEddies_data/moving_eddiesU.nc"), _Purpose.TUTORIAL)),
195-
("MovingEddies_data/V", (_V3Dataset("data/MovingEddies_data/moving_eddiesV.nc"), _Purpose.TUTORIAL)),
196-
("MITgcm_example_data/mitgcm_UV_surface_zonally_reentrant", (_V3Dataset("data/MITgcm_example_data/mitgcm_UV_surface_zonally_reentrant.nc"), _Purpose.TUTORIAL)),
197-
("OFAM_example_data/U", (_V3Dataset("data/OFAM_example_data/OFAM_simple_U.nc"), _Purpose.TUTORIAL)),
198-
("OFAM_example_data/V", (_V3Dataset("data/OFAM_example_data/OFAM_simple_V.nc"), _Purpose.TUTORIAL)),
199-
("Peninsula_data/U", (_V3Dataset("data/Peninsula_data/peninsulaU.nc"), _Purpose.TUTORIAL)),
200-
("Peninsula_data/V", (_V3Dataset("data/Peninsula_data/peninsulaV.nc"), _Purpose.TUTORIAL)),
201-
("Peninsula_data/P", (_V3Dataset("data/Peninsula_data/peninsulaP.nc"), _Purpose.TUTORIAL)),
202-
("Peninsula_data/T", (_V3Dataset("data/Peninsula_data/peninsulaT.nc"), _Purpose.TUTORIAL)),
203-
("GlobCurrent_example_data/data", (_V3Dataset("data/GlobCurrent_example_data/*000000-GLOBCURRENT-L4-CUReul_hs-ALT_SUM-v02.0-fv01.0.nc", pre_decode_cf_callable=patch_dataset_v4_compat), _Purpose.TUTORIAL)),
204-
("CopernicusMarine_data_for_Argo_tutorial/data", (_V3Dataset("data/CopernicusMarine_data_for_Argo_tutorial/cmems_mod_glo_phy-*.nc"), _Purpose.TUTORIAL)),
205-
("DecayingMovingEddy_data/U", (_V3Dataset("data/DecayingMovingEddy_data/decaying_moving_eddyU.nc"), _Purpose.TUTORIAL)),
206-
("DecayingMovingEddy_data/V", (_V3Dataset("data/DecayingMovingEddy_data/decaying_moving_eddyV.nc"), _Purpose.TUTORIAL)),
207-
("FESOM_periodic_channel/fesom_channel", (_V3Dataset("data/FESOM_periodic_channel/fesom_channel.nc"), _Purpose.TUTORIAL)),
208-
("FESOM_periodic_channel/u.fesom_channel", (_V3Dataset("data/FESOM_periodic_channel/u.fesom_channel.nc"), _Purpose.TUTORIAL)),
209-
("FESOM_periodic_channel/v.fesom_channel", (_V3Dataset("data/FESOM_periodic_channel/v.fesom_channel.nc"), _Purpose.TUTORIAL)),
210-
("FESOM_periodic_channel/w.fesom_channel", (_V3Dataset("data/FESOM_periodic_channel/w.fesom_channel.nc"), _Purpose.TUTORIAL)),
211-
("NemoCurvilinear_data_zonal/U", (_V3Dataset("data/NemoCurvilinear_data/U_purely_zonal-ORCA025_grid_U.nc4"), _Purpose.TUTORIAL)),
212-
("NemoCurvilinear_data_zonal/V", (_V3Dataset("data/NemoCurvilinear_data/V_purely_zonal-ORCA025_grid_V.nc4"), _Purpose.TUTORIAL)),
213-
("NemoCurvilinear_data_zonal/mesh_mask", (_V3Dataset("data/NemoCurvilinear_data/mesh_mask.nc4", _preprocess_drop_time_from_mesh2), _Purpose.TUTORIAL)),
214-
("NemoNorthSeaORCA025-N006_data/U", (_V3Dataset("data/NemoNorthSeaORCA025-N006_data/ORCA025-N06_200001*05U.nc"), _Purpose.TUTORIAL)),
215-
("NemoNorthSeaORCA025-N006_data/V", (_V3Dataset("data/NemoNorthSeaORCA025-N006_data/ORCA025-N06_200001*05V.nc"), _Purpose.TUTORIAL)),
216-
("NemoNorthSeaORCA025-N006_data/W", (_V3Dataset("data/NemoNorthSeaORCA025-N006_data/ORCA025-N06_200001*05W.nc"), _Purpose.TUTORIAL)),
217-
("NemoNorthSeaORCA025-N006_data/mesh_mask", (_V3Dataset("data/NemoNorthSeaORCA025-N006_data/coordinates.nc", _preprocess_drop_time_from_mesh1), _Purpose.TUTORIAL)),
195+
("MovingEddies_data/P", (_V3Dataset(_ODIE,"data/MovingEddies_data/moving_eddiesP.nc"), _Purpose.TUTORIAL)),
196+
("MovingEddies_data/U", (_V3Dataset(_ODIE,"data/MovingEddies_data/moving_eddiesU.nc"), _Purpose.TUTORIAL)),
197+
("MovingEddies_data/V", (_V3Dataset(_ODIE,"data/MovingEddies_data/moving_eddiesV.nc"), _Purpose.TUTORIAL)),
198+
("MITgcm_example_data/mitgcm_UV_surface_zonally_reentrant", (_V3Dataset(_ODIE,"data/MITgcm_example_data/mitgcm_UV_surface_zonally_reentrant.nc"), _Purpose.TUTORIAL)),
199+
("OFAM_example_data/U", (_V3Dataset(_ODIE,"data/OFAM_example_data/OFAM_simple_U.nc"), _Purpose.TUTORIAL)),
200+
("OFAM_example_data/V", (_V3Dataset(_ODIE,"data/OFAM_example_data/OFAM_simple_V.nc"), _Purpose.TUTORIAL)),
201+
("Peninsula_data/U", (_V3Dataset(_ODIE,"data/Peninsula_data/peninsulaU.nc"), _Purpose.TUTORIAL)),
202+
("Peninsula_data/V", (_V3Dataset(_ODIE,"data/Peninsula_data/peninsulaV.nc"), _Purpose.TUTORIAL)),
203+
("Peninsula_data/P", (_V3Dataset(_ODIE,"data/Peninsula_data/peninsulaP.nc"), _Purpose.TUTORIAL)),
204+
("Peninsula_data/T", (_V3Dataset(_ODIE,"data/Peninsula_data/peninsulaT.nc"), _Purpose.TUTORIAL)),
205+
("GlobCurrent_example_data/data", (_V3Dataset(_ODIE,"data/GlobCurrent_example_data/*000000-GLOBCURRENT-L4-CUReul_hs-ALT_SUM-v02.0-fv01.0.nc", pre_decode_cf_callable=patch_dataset_v4_compat), _Purpose.TUTORIAL)),
206+
("CopernicusMarine_data_for_Argo_tutorial/data", (_V3Dataset(_ODIE,"data/CopernicusMarine_data_for_Argo_tutorial/cmems_mod_glo_phy-*.nc"), _Purpose.TUTORIAL)),
207+
("DecayingMovingEddy_data/U", (_V3Dataset(_ODIE,"data/DecayingMovingEddy_data/decaying_moving_eddyU.nc"), _Purpose.TUTORIAL)),
208+
("DecayingMovingEddy_data/V", (_V3Dataset(_ODIE,"data/DecayingMovingEddy_data/decaying_moving_eddyV.nc"), _Purpose.TUTORIAL)),
209+
("FESOM_periodic_channel/fesom_channel", (_V3Dataset(_ODIE,"data/FESOM_periodic_channel/fesom_channel.nc"), _Purpose.TUTORIAL)),
210+
("FESOM_periodic_channel/u.fesom_channel", (_V3Dataset(_ODIE,"data/FESOM_periodic_channel/u.fesom_channel.nc"), _Purpose.TUTORIAL)),
211+
("FESOM_periodic_channel/v.fesom_channel", (_V3Dataset(_ODIE,"data/FESOM_periodic_channel/v.fesom_channel.nc"), _Purpose.TUTORIAL)),
212+
("FESOM_periodic_channel/w.fesom_channel", (_V3Dataset(_ODIE,"data/FESOM_periodic_channel/w.fesom_channel.nc"), _Purpose.TUTORIAL)),
213+
("NemoCurvilinear_data_zonal/U", (_V3Dataset(_ODIE,"data/NemoCurvilinear_data/U_purely_zonal-ORCA025_grid_U.nc4"), _Purpose.TUTORIAL)),
214+
("NemoCurvilinear_data_zonal/V", (_V3Dataset(_ODIE,"data/NemoCurvilinear_data/V_purely_zonal-ORCA025_grid_V.nc4"), _Purpose.TUTORIAL)),
215+
("NemoCurvilinear_data_zonal/mesh_mask", (_V3Dataset(_ODIE,"data/NemoCurvilinear_data/mesh_mask.nc4", _preprocess_drop_time_from_mesh2), _Purpose.TUTORIAL)),
216+
("NemoNorthSeaORCA025-N006_data/U", (_V3Dataset(_ODIE,"data/NemoNorthSeaORCA025-N006_data/ORCA025-N06_200001*05U.nc"), _Purpose.TUTORIAL)),
217+
("NemoNorthSeaORCA025-N006_data/V", (_V3Dataset(_ODIE,"data/NemoNorthSeaORCA025-N006_data/ORCA025-N06_200001*05V.nc"), _Purpose.TUTORIAL)),
218+
("NemoNorthSeaORCA025-N006_data/W", (_V3Dataset(_ODIE,"data/NemoNorthSeaORCA025-N006_data/ORCA025-N06_200001*05W.nc"), _Purpose.TUTORIAL)),
219+
("NemoNorthSeaORCA025-N006_data/mesh_mask", (_V3Dataset(_ODIE,"data/NemoNorthSeaORCA025-N006_data/coordinates.nc", _preprocess_drop_time_from_mesh1), _Purpose.TUTORIAL)),
218220
# "POPSouthernOcean_data/t.x1_SAMOC_flux.16900*.nc", # TODO v4: In v3 but should not be in v4 https://github.com/Parcels-code/Parcels/issues/2571#issuecomment-4214476973
219-
("SWASH_data/data", (_V3Dataset("data/SWASH_data/field_00655*.nc"), _Purpose.TUTORIAL)),
220-
("WOA_data/data", (_V3Dataset("data/WOA_data/woa18_decav_t*_04.nc", _preprocess_set_cf_calendar_360_day), _Purpose.TUTORIAL)),
221-
("CROCOidealized_data/data", (_V3Dataset("data/CROCOidealized_data/CROCO_idealized.nc"), _Purpose.TUTORIAL)),
221+
("SWASH_data/data", (_V3Dataset(_ODIE,"data/SWASH_data/field_00655*.nc"), _Purpose.TUTORIAL)),
222+
("WOA_data/data", (_V3Dataset(_ODIE,"data/WOA_data/woa18_decav_t*_04.nc", _preprocess_set_cf_calendar_360_day), _Purpose.TUTORIAL)),
223+
("CROCOidealized_data/data", (_V3Dataset(_ODIE,"data/CROCOidealized_data/CROCO_idealized.nc"), _Purpose.TUTORIAL)),
224+
] + [
225+
("Benchmarks_FESOM2-baroclinic-gyre/data", (_ZarrZipDataset(_ODIE, 'data-zarr/Benchmarks_FESOM2-baroclinic-gyre/data.zip'), _Purpose.TESTING)),
226+
("Benchmarks_FESOM2-baroclinic-gyre/grid", (_ZarrZipDataset(_ODIE, 'data-zarr/Benchmarks_FESOM2-baroclinic-gyre/grid.zip'),_Purpose.TESTING)),
222227
])
223228
# fmt: on
224229

‎src/parcels/convert.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -685,6 +685,8 @@ def fesom_to_ugrid(ds: ux.UxDataset) -> ux.UxDataset:
685685
Renames vertical dimensions:
686686
- nz -> zf (vertical layer faces/interfaces)
687687
- nz1 -> zc (vertical layer centers)
688+
- nod2 -> n_face (face)
689+
- elem -> n_node (node)
688690
689691
Parameters
690692
----------
@@ -706,6 +708,11 @@ def fesom_to_ugrid(ds: ux.UxDataset) -> ux.UxDataset:
706708
>>> fieldset = FieldSet.from_ugrid_conventions(ds_ugrid, mesh="flat")
707709
"""
708710
ds = ds.copy()
711+
712+
for try_dim, target in [("nod2", "n_face"), ("elem", "n_node")]:
713+
if try_dim in ds.dims:
714+
ds = ds.rename_dims({try_dim: target})
715+
709716
interface_dim, center_dim = _detect_vertical_coordinates(ds, _FESOM2_VERTICAL_DIMS)
710717
return _rename_vertical_dims(ds, interface_dim, center_dim)
711718

‎tests/test_convert.py‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,13 @@
11
import pytest
2+
import uxarray as ux
23
import xarray as xr
34

45
import parcels
56
import parcels.convert as convert
67
import parcels.tutorial
78
from parcels import FieldSet
89
from parcels._core.utils import sgrid
10+
from parcels._datasets.remote import open_remote_dataset
911
from parcels._datasets.structured.circulation_models import datasets as datasets_circulation_models
1012
from parcels.interpolators._xinterpolators import _get_offsets_dictionary
1113

@@ -122,3 +124,13 @@ def test_convert_copernicusmarine_no_logs(ds, caplog):
122124
assert "V" in fieldset.fields
123125
assert "UV" in fieldset.fields
124126
assert caplog.text == ""
127+
128+
129+
def test_convert_fesom_to_ugrid():
130+
grid_file = open_remote_dataset("Benchmarks_FESOM2-baroclinic-gyre/grid")
131+
data_files = open_remote_dataset("Benchmarks_FESOM2-baroclinic-gyre/data")
132+
133+
grid = ux.open_grid(grid_file)
134+
uxds = ux.UxDataset(data_files, uxgrid=grid)
135+
uxds = convert.fesom_to_ugrid(uxds)
136+
FieldSet.from_ugrid_conventions(uxds)

0 commit comments

Comments
 (0)