added TELEMAC support and tests

ec-jrc · May 19, 2024 · c44a011 · c44a011
1 parent b9d977c
commit c44a011
Show file tree

Hide file tree

Showing 8 changed files with 1,654 additions and 1,541 deletions.
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -38,6 +38,7 @@ pyproj = "*"
 scipy = "*"
 shapely = "*"
 xarray = {version = "*", extras = ["io", "accel"]}
+xarray-selafin = "*"
 
 [tool.poetry.group.dev.dependencies]
 covdefaults = "*"

diff --git a/requirements/requirements-dev.txt b/requirements/requirements-dev.txt
diff --git a/requirements/requirements.txt b/requirements/requirements.txt
@@ -2,101 +2,107 @@ asciitree==0.3.3 ; python_version >= "3.9" and python_version < "4.0"
 attrs==23.2.0 ; python_version >= "3.9" and python_version < "4.0"
 beautifulsoup4==4.12.3 ; python_version < "3.10" and python_version >= "3.9"
 bleach==6.1.0 ; python_version >= "3.9" and python_version < "4.0"
-bokeh==3.3.4 ; python_version >= "3.9" and python_version < "4.0"
-bottleneck==1.3.7 ; python_version >= "3.9" and python_version < "4.0"
-cartopy==0.22.0 ; python_version >= "3.9" and python_version < "4.0"
+bokeh==3.4.1 ; python_version >= "3.9" and python_version < "4.0"
+bottleneck==1.3.8 ; python_version >= "3.9" and python_version < "4.0"
+cartopy==0.23.0 ; python_version >= "3.9" and python_version < "4.0"
 certifi==2024.2.2 ; python_version >= "3.9" and python_version < "4.0"
 cftime==1.6.3 ; python_version >= "3.9" and python_version < "4.0"
 charset-normalizer==3.3.2 ; python_version >= "3.9" and python_version < "4.0"
 click-plugins==1.1.1 ; python_version >= "3.9" and python_version < "4.0"
 click==8.1.7 ; python_version >= "3.9" and python_version < "4.0"
 cligj==0.7.2 ; python_version >= "3.9" and python_version < "4"
 cloudpickle==3.0.0 ; python_version >= "3.9" and python_version < "4.0"
-colorama==0.4.6 ; python_version >= "3.9" and python_version < "4.0" and platform_system == "Windows"
-colorcet==3.0.1 ; python_version >= "3.9" and python_version < "4.0"
-contourpy==1.2.0 ; python_version >= "3.9" and python_version < "4.0"
+colorama==0.4.6 ; python_version >= "3.9" and python_version < "4.0" and (platform_system == "Windows" or sys_platform == "win32")
+colorcet==3.1.0 ; python_version >= "3.9" and python_version < "4.0"
+contourpy==1.2.1 ; python_version >= "3.9" and python_version < "4.0"
 cycler==0.12.1 ; python_version >= "3.9" and python_version < "4.0"
-dask==2024.1.1 ; python_version >= "3.9" and python_version < "4.0"
-dask[array,complete,dataframe,diagnostics,distributed]==2024.1.1 ; python_version >= "3.9" and python_version < "4.0"
-datashader==0.16.0 ; python_version >= "3.9" and python_version < "4.0"
-distributed==2024.1.1 ; python_version >= "3.9" and python_version < "4.0"
+dask-expr==1.1.1 ; python_version >= "3.9" and python_version < "4.0"
+dask==2024.5.1 ; python_version >= "3.9" and python_version < "4.0"
+dask[array,complete,dataframe,diagnostics,distributed]==2024.5.1 ; python_version >= "3.9" and python_version < "4.0"
+datashader==0.16.1 ; python_version >= "3.9" and python_version < "4.0"
+distributed==2024.5.1 ; python_version >= "3.9" and python_version < "4.0"
 docopt==0.6.2 ; python_version < "3.10" and python_version >= "3.9"
-fasteners==0.19 ; python_version >= "3.9" and python_version < "4.0"
-fiona==1.9.5 ; python_version >= "3.9" and python_version < "4.0"
-flox==0.9.0 ; python_version >= "3.9" and python_version < "4.0"
-fonttools==4.48.1 ; python_version >= "3.9" and python_version < "4.0"
-fsspec==2024.2.0 ; python_version >= "3.9" and python_version < "4.0"
-future==0.18.3 ; python_version >= "3.9" and python_version < "4.0"
-geopandas==0.14.3 ; python_version >= "3.9" and python_version < "4.0"
-geoviews==1.11.0 ; python_version >= "3.9" and python_version < "4.0"
+exceptiongroup==1.2.1 ; python_version >= "3.9" and python_version < "3.11"
+fasteners==0.19 ; python_version >= "3.9" and python_version < "4.0" and sys_platform != "emscripten"
+fiona==1.9.6 ; python_version >= "3.9" and python_version < "4.0"
+flox==0.9.7 ; python_version >= "3.9" and python_version < "4.0"
+fonttools==4.51.0 ; python_version >= "3.9" and python_version < "4.0"
+fsspec==2024.5.0 ; python_version >= "3.9" and python_version < "4.0"
+future==1.0.0 ; python_version >= "3.9" and python_version < "4.0"
+geopandas==0.14.4 ; python_version >= "3.9" and python_version < "4.0"
+geoviews==1.12.0 ; python_version >= "3.9" and python_version < "4.0"
 h5netcdf==1.3.0 ; python_version >= "3.9" and python_version < "4.0"
-h5py==3.10.0 ; python_version >= "3.9" and python_version < "4.0"
-holoviews==1.18.1 ; python_version >= "3.9" and python_version < "4.0"
-idna==3.6 ; python_version >= "3.9" and python_version < "4.0"
-importlib-metadata==7.0.1 ; python_version >= "3.9" and python_version < "4.0"
-importlib-resources==6.1.1 ; python_version >= "3.9" and python_version < "3.10"
-jinja2==3.1.3 ; python_version >= "3.9" and python_version < "4.0"
+h5py==3.11.0 ; python_version >= "3.9" and python_version < "4.0"
+holoviews==1.18.3 ; python_version >= "3.9" and python_version < "4.0"
+idna==3.7 ; python_version >= "3.9" and python_version < "4.0"
+importlib-metadata==7.1.0 ; python_version >= "3.9" and python_version < "3.12"
+importlib-resources==6.4.0 ; python_version >= "3.9" and python_version < "3.10"
+iniconfig==2.0.0 ; python_version >= "3.9" and python_version < "4.0"
+jinja2==3.1.4 ; python_version >= "3.9" and python_version < "4.0"
 kiwisolver==1.4.5 ; python_version >= "3.9" and python_version < "4.0"
 linkify-it-py==2.0.3 ; python_version >= "3.9" and python_version < "4.0"
 llvmlite==0.42.0 ; python_version >= "3.9" and python_version < "4.0"
 locket==1.0.0 ; python_version >= "3.9" and python_version < "4.0"
 lz4==4.3.3 ; python_version >= "3.9" and python_version < "4.0"
 markdown-it-py==3.0.0 ; python_version >= "3.9" and python_version < "4.0"
-markdown==3.5.2 ; python_version >= "3.9" and python_version < "4.0"
+markdown==3.6 ; python_version >= "3.9" and python_version < "4.0"
 markupsafe==2.1.5 ; python_version >= "3.9" and python_version < "4.0"
-matplotlib==3.8.2 ; python_version >= "3.9" and python_version < "4.0"
-mdit-py-plugins==0.4.0 ; python_version >= "3.9" and python_version < "4.0"
+matplotlib==3.9.0 ; python_version >= "3.9" and python_version < "4.0"
+mdit-py-plugins==0.4.1 ; python_version >= "3.9" and python_version < "4.0"
 mdurl==0.1.2 ; python_version >= "3.9" and python_version < "4.0"
-msgpack==1.0.7 ; python_version >= "3.9" and python_version < "4.0"
+msgpack==1.0.8 ; python_version >= "3.9" and python_version < "4.0"
 multipledispatch==1.0.0 ; python_version >= "3.9" and python_version < "4.0"
 netcdf4==1.6.5 ; python_version >= "3.9" and python_version < "4.0"
-numba==0.59.0 ; python_version >= "3.9" and python_version < "4.0"
-numbagg==0.8.0 ; python_version >= "3.9" and python_version < "4.0"
+numba==0.59.1 ; python_version >= "3.9" and python_version < "4.0"
+numbagg==0.8.1 ; python_version >= "3.9" and python_version < "4.0"
 numcodecs==0.12.1 ; python_version >= "3.9" and python_version < "4.0"
-numpy-groupies==0.10.2 ; python_version >= "3.9" and python_version < "4.0"
+numpy-groupies==0.11.1 ; python_version >= "3.9" and python_version < "4.0"
 numpy-indexed==0.3.7 ; python_version >= "3.9" and python_version < "4.0"
 numpy==1.26.4 ; python_version >= "3.9" and python_version < "4.0"
 opt-einsum==3.3.0 ; python_version >= "3.9" and python_version < "4.0"
-packaging==23.2 ; python_version >= "3.9" and python_version < "4.0"
-pandas==2.2.0 ; python_version >= "3.9" and python_version < "4.0"
-panel==1.3.8 ; python_version >= "3.9" and python_version < "4.0"
-param==2.0.2 ; python_version >= "3.9" and python_version < "4.0"
-partd==1.4.1 ; python_version >= "3.9" and python_version < "4.0"
-pillow==10.2.0 ; python_version >= "3.9" and python_version < "4.0"
-platformdirs==4.2.0 ; python_version >= "3.9" and python_version < "4.0"
-pooch==1.8.0 ; python_version >= "3.9" and python_version < "4.0"
+packaging==24.0 ; python_version >= "3.9" and python_version < "4.0"
+pandas==2.2.2 ; python_version >= "3.9" and python_version < "4.0"
+panel==1.4.2 ; python_version >= "3.9" and python_version < "4.0"
+param==2.1.0 ; python_version >= "3.9" and python_version < "4.0"
+partd==1.4.2 ; python_version >= "3.9" and python_version < "4.0"
+pillow==10.3.0 ; python_version >= "3.9" and python_version < "4.0"
+platformdirs==4.2.2 ; python_version >= "3.9" and python_version < "4.0"
+pluggy==1.5.0 ; python_version >= "3.9" and python_version < "4.0"
+pooch==1.8.1 ; python_version >= "3.9" and python_version < "4.0"
 psutil==5.9.8 ; python_version >= "3.9" and python_version < "4.0"
 pyarrow-hotfix==0.6 ; python_version >= "3.9" and python_version < "4.0"
-pyarrow==15.0.0 ; python_version >= "3.9" and python_version < "4.0"
+pyarrow==16.1.0 ; python_version >= "3.9" and python_version < "4.0"
 pyct==0.5.0 ; python_version >= "3.9" and python_version < "4.0"
 pydap==3.4.1 ; python_version < "3.10" and python_version >= "3.9"
-pyparsing==3.1.1 ; python_version >= "3.9" and python_version < "4.0"
+pyparsing==3.1.2 ; python_version >= "3.9" and python_version < "4.0"
 pyproj==3.6.1 ; python_version >= "3.9" and python_version < "4.0"
 pyshp==2.3.1 ; python_version >= "3.9" and python_version < "4.0"
-python-dateutil==2.8.2 ; python_version >= "3.9" and python_version < "4.0"
+pytest==8.2.0 ; python_version >= "3.9" and python_version < "4.0"
+python-dateutil==2.9.0.post0 ; python_version >= "3.9" and python_version < "4.0"
 pytz==2024.1 ; python_version >= "3.9" and python_version < "4.0"
-pyviz-comms==3.0.1 ; python_version >= "3.9" and python_version < "4.0"
+pyviz-comms==3.0.2 ; python_version >= "3.9" and python_version < "4.0"
 pyyaml==6.0.1 ; python_version >= "3.9" and python_version < "4.0"
 requests==2.31.0 ; python_version >= "3.9" and python_version < "4.0"
-scipy==1.12.0 ; python_version >= "3.9" and python_version < "4.0"
-setuptools==69.0.3 ; python_version >= "3.9" and python_version < "4.0"
-shapely==2.0.2 ; python_version >= "3.9" and python_version < "4.0"
+scipy==1.13.0 ; python_version >= "3.9" and python_version < "4.0"
+shapely==2.0.4 ; python_version >= "3.9" and python_version < "4.0"
 six==1.16.0 ; python_version >= "3.9" and python_version < "4.0"
 sortedcontainers==2.4.0 ; python_version >= "3.9" and python_version < "4.0"
 soupsieve==2.5 ; python_version < "3.10" and python_version >= "3.9"
 tblib==3.0.0 ; python_version >= "3.9" and python_version < "4.0"
+tomli==2.0.1 ; python_version >= "3.9" and python_version < "3.11"
 toolz==0.12.1 ; python_version >= "3.9" and python_version < "4.0"
 tornado==6.4 ; python_version >= "3.9" and python_version < "4.0"
-tqdm==4.66.1 ; python_version >= "3.9" and python_version < "4.0"
-typing-extensions==4.9.0 ; python_version >= "3.9" and python_version < "4.0"
-tzdata==2023.4 ; python_version >= "3.9" and python_version < "4.0"
-uc-micro-py==1.0.2 ; python_version >= "3.9" and python_version < "4.0"
-urllib3==2.2.0 ; python_version >= "3.9" and python_version < "4.0"
+tqdm==4.66.4 ; python_version >= "3.9" and python_version < "4.0"
+typing-extensions==4.11.0 ; python_version >= "3.9" and python_version < "4.0"
+tzdata==2024.1 ; python_version >= "3.9" and python_version < "4.0"
+uc-micro-py==1.0.3 ; python_version >= "3.9" and python_version < "4.0"
+urllib3==2.2.1 ; python_version >= "3.9" and python_version < "4.0"
 webencodings==0.5.1 ; python_version >= "3.9" and python_version < "4.0"
 webob==1.8.7 ; python_version >= "3.9" and python_version < "3.10"
-xarray==2024.1.1 ; python_version >= "3.9" and python_version < "4.0"
-xarray[accel,io]==2024.1.1 ; python_version >= "3.9" and python_version < "4.0"
-xyzservices==2023.10.1 ; python_version >= "3.9" and python_version < "4.0"
-zarr==2.16.1 ; python_version >= "3.9" and python_version < "4.0"
+xarray-selafin==0.1.5 ; python_version >= "3.9" and python_version < "4.0"
+xarray==2024.5.0 ; python_version >= "3.9" and python_version < "4.0"
+xarray[accel,io]==2024.5.0 ; python_version >= "3.9" and python_version < "4.0"
+xyzservices==2024.4.0 ; python_version >= "3.9" and python_version < "4.0"
+zarr==2.18.1 ; python_version >= "3.9" and python_version < "4.0"
 zict==3.0.0 ; python_version >= "3.9" and python_version < "4.0"
-zipp==3.17.0 ; python_version >= "3.9" and python_version < "4.0"
+zipp==3.18.2 ; python_version >= "3.9" and python_version < "3.12"
diff --git a/tests/api_test.py b/tests/api_test.py
@@ -9,11 +9,19 @@
 from thalassa import normalization
 
 ADCIRC_NC = DATA_DIR / "fort.63.nc"
+SELAFIN = DATA_DIR / "r3d_tidal_flats.slf"
 
+@pytest.mark.parametrize(
+    "file,variable",
+    [
+        pytest.param(ADCIRC_NC, "zeta"),
+        pytest.param(SELAFIN, "U"),
+    ],
+)
 
-def test_main_api():
-    variable = "zeta"
-    ds = api.open_dataset(ADCIRC_NC)
+
+def test_main_api(file, variable):
+    ds = api.open_dataset(file)
     assert normalization.is_generic(ds)
 
     # Create objects

diff --git a/tests/normalization_test.py b/tests/normalization_test.py
@@ -12,6 +12,7 @@
     "ds,expected_fmt",
     [
         pytest.param(api.open_dataset(DATA_DIR / "fort.63.nc", normalize=False), THALASSA_FORMATS.ADCIRC, id="ADCIRC"),
+        pytest.param(api.open_dataset(DATA_DIR / "r3d_tidal_flats.slf", normalize=False), THALASSA_FORMATS.TELEMAC, id="TELEMAC"),
         pytest.param(xr.Dataset(), THALASSA_FORMATS.UNKNOWN, id="Unknown"),
     ],
 )
@@ -24,6 +25,7 @@ def test_infer_format(ds, expected_fmt):
     "path,expected",
     [
         pytest.param(DATA_DIR / "fort.63.nc", True, id="ADCIRC"),
+        pytest.param(DATA_DIR / "r3d_tidal_flats.slf", True, id="TELEMAC"),
         pytest.param(__file__, False, id="Unknown"),
     ],
 )

diff --git a/thalassa/api.py b/thalassa/api.py
@@ -74,7 +74,7 @@ def open_dataset(
     import xarray as xr
 
     default_kwargs: dict[str, T.Any] = dict(
-        mask_and_scale=True,
+        # mask_and_scale=True,
         cache=False,
         drop_variables=ADCIRC_VARIABLES_TO_BE_DROPPED,
     )

diff --git a/thalassa/normalization.py b/thalassa/normalization.py
@@ -19,19 +19,29 @@ class THALASSA_FORMATS(enum.Enum):
     UNKNOWN = "UNKNOWN"
     ADCIRC = "ADCIRC"
     SCHISM = "SCHISM"
+    TELEMAC = "TELEMAC"
     GENERIC = "GENERIC"
     PYPOSEIDON = "PYPOSEIDON"
 
 
 # fmt: off
+EDGE_DIM = "edge"
+FACE_DIM = "face"
+NODE_DIM = "node"
+VERTICAL_DIM = "layer"
+CONNECTIVITY = "face_nodes"
+VERTICE_DIM = "max_no_vertices"
+X_DIM = "lon"
+Y_DIM = "lat"
+
 _GENERIC_DIMS = {
-    "node",
+    NODE_DIM,
     "triface",
     "three",
 }
 _GENERIC_VARS = {
-    "lon",
-    "lat",
+    X_DIM,
+    Y_DIM,
     "triface_nodes",
 }
 _SCHISM_DIMS = {
@@ -45,6 +55,14 @@ class THALASSA_FORMATS(enum.Enum):
     "SCHISM_hgrid_node_y",
     "SCHISM_hgrid_face_nodes",
 }
+_TELEMAC_DIMS = {
+    "node",
+}
+_TELEMAC_VARS = {
+    "x",
+    "y",
+    "ikle2",
+}
 _PYPOSEIDON_DIMS = {
     "nSCHISM_hgrid_face",
     "nSCHISM_hgrid_node",
@@ -76,6 +94,11 @@ def is_schism(ds: xarray.Dataset) -> bool:
     return _SCHISM_DIMS.issubset(ds.dims) and _SCHISM_VARS.issubset(total_vars)
 
 
+def is_telemac(ds: xarray.Dataset) -> bool:
+    total_vars = list(ds.data_vars.keys()) + list(ds.coords.keys()) + list(ds.attrs.keys())
+    return _TELEMAC_DIMS.issubset(ds.dims) and _TELEMAC_VARS.issubset(total_vars)
+
+
 def is_pyposeidon(ds: xarray.Dataset) -> bool:
     return _PYPOSEIDON_DIMS.issubset(ds.dims) and _PYPOSEIDON_VARS.issubset(ds.data_vars)
 
@@ -87,6 +110,8 @@ def is_adcirc(ds: xarray.Dataset) -> bool:
 def infer_format(ds: xarray.Dataset) -> THALASSA_FORMATS:
     if is_schism(ds):
         fmt = THALASSA_FORMATS.SCHISM
+    if is_telemac(ds):
+        fmt = THALASSA_FORMATS.TELEMAC
     elif is_pyposeidon(ds):
         fmt = THALASSA_FORMATS.PYPOSEIDON
     elif is_generic(ds):
@@ -120,37 +145,59 @@ def normalize_generic(ds: xarray.Dataset) -> xarray.Dataset:
 def normalize_schism(ds: xarray.Dataset) -> xarray.Dataset:
     ds = ds.rename(
         {
-            "nSCHISM_hgrid_edge": "edge",
-            "nSCHISM_hgrid_face": "face",
-            "nSCHISM_hgrid_node": "node",
-            "SCHISM_hgrid_face_nodes": "face_nodes",
-            "nMaxSCHISM_hgrid_face_nodes": "max_no_vertices",
-            "SCHISM_hgrid_node_x": "lon",
-            "SCHISM_hgrid_node_y": "lat",
+            "nSCHISM_hgrid_edge": EDGE_DIM,
+            "nSCHISM_hgrid_face": FACE_DIM,
+            "nSCHISM_hgrid_node": NODE_DIM,
+            "SCHISM_hgrid_face_nodes": CONNECTIVITY,
+            "nMaxSCHISM_hgrid_face_nodes": VERTICE_DIM,
+            "SCHISM_hgrid_node_x": X_DIM,
+            "SCHISM_hgrid_node_y": Y_DIM,
         },
     )
     if "nSCHISM_vgrid_layers" in ds.dims:
         # I.e. OLD Schism IO or "merged" new IO
         ds = ds.rename(
             {
-                "nSCHISM_vgrid_layers": "layer",
+                "nSCHISM_vgrid_layers": VERTICAL_DIM,
             },
         )
     # SCHISM output uses one-based indices for `face_nodes`
     # Let's ensure that we use zero-based indices everywhere.
-    ds["face_nodes"] -= 1
+    ds[CONNECTIVITY] -= 1
+    return ds
+
+
+def normalize_telemac(ds: xarray.Dataset) -> xarray.Dataset:
+    ds = ds.rename(
+        {
+            "node": NODE_DIM,
+            "x": X_DIM,
+            "y": Y_DIM,
+        },
+    )
+    if "plan" in ds.dims:
+        # I.e. OLD Schism IO or "merged" new IO
+        ds = ds.rename(
+            {
+                "plan": VERTICAL_DIM,
+            },
+        )
+
+    # TELEMAC output uses one-based indices for `face_nodes`
+    # Let's ensure that we use zero-based indices everywhere.
+    ds[CONNECTIVITY] = ((FACE_DIM, VERTICE_DIM), ds.attrs['ikle2'] - 1)
     return ds
 
 
 def normalize_pyposeidon(ds: xarray.Dataset) -> xarray.Dataset:
     ds = ds.rename(
         {
-            "nSCHISM_hgrid_face": "face",
-            "nSCHISM_hgrid_node": "node",
-            "SCHISM_hgrid_face_nodes": "face_nodes",
-            "nMaxSCHISM_hgrid_face_nodes": "max_no_vertices",
-            "SCHISM_hgrid_node_x": "lon",
-            "SCHISM_hgrid_node_y": "lat",
+            "nSCHISM_hgrid_face": FACE_DIM,
+            "nSCHISM_hgrid_node": NODE_DIM,
+            "SCHISM_hgrid_face_nodes": CONNECTIVITY,
+            "nMaxSCHISM_hgrid_face_nodes": VERTICE_DIM,
+            "SCHISM_hgrid_node_x": X_DIM,
+            "SCHISM_hgrid_node_y": Y_DIM,
         },
     )
     return ds
@@ -159,23 +206,24 @@ def normalize_pyposeidon(ds: xarray.Dataset) -> xarray.Dataset:
 def normalize_adcirc(ds: xarray.Dataset) -> xarray.Dataset:
     ds = ds.rename(
         {
-            "x": "lon",
-            "y": "lat",
-            "element": "face_nodes",
-            "nvertex": "max_no_vertices",
-            "nele": "face",
+            "x": X_DIM,
+            "y": Y_DIM,
+            "element": CONNECTIVITY,
+            "nvertex": VERTICE_DIM,
+            "nele": FACE_DIM,
         },
     )
     # ADCIRC output uses one-based indices for `face_nodes`
     # Let's ensure that we use zero-based indices everywhere.
-    ds["face_nodes"] -= 1
+    ds[CONNECTIVITY] -= 1
     return ds
 
 
 NORMALIZE_DISPATCHER = {
     THALASSA_FORMATS.ADCIRC: normalize_adcirc,
     THALASSA_FORMATS.GENERIC: normalize_generic,
     THALASSA_FORMATS.SCHISM: normalize_schism,
+    THALASSA_FORMATS.TELEMAC: normalize_telemac,
     THALASSA_FORMATS.PYPOSEIDON: normalize_pyposeidon,
 }