remove mesh2d hardcoding in layer reconstruction

dfm_tools/get_nc.py

@@ -275,6 +275,8 @@ def reconstruct_zw_zcc_fromzsigma(uds):
     """
     # TODO: center values are clipped to bedlevel, so the center values of the bottom layer are currently incorrect
 
+    gridname = uds.grid.name
+
     # temporarily decode default fillvalues
     # TODO: xarray only decodes explicitly set fillvalues: https://github.com/Deltares/dfm_tools/issues/490
     uds = xu.UgridDataset(decode_default_fillvals(uds.obj),grids=[uds.grid])
@@ -298,40 +300,53 @@ def reconstruct_zw_zcc_fromzsigma(uds):
     # z(n,k,j,i) = zlev(k)
     zw_zpart = uds_zlev_int.clip(min=-uds_depth) #added clipping of zvalues with bedlevel
     zcc_zpart = uds_zlev_lay.clip(min=-uds_depth) #added clipping of zvalues with bedlevel
-    uds['mesh2d_flowelem_zw'] = zw_sigmapart.fillna(zw_zpart)
-    uds['mesh2d_flowelem_zcc'] = zcc_sigmapart.fillna(zcc_zpart)
+    uds[f'{gridname}_flowelem_zw'] = zw_sigmapart.fillna(zw_zpart)
+    uds[f'{gridname}_flowelem_zcc'] = zcc_sigmapart.fillna(zcc_zpart)
 
-    uds = uds.set_coords(['mesh2d_flowelem_zw','mesh2d_flowelem_zcc'])
+    uds = uds.set_coords([f'{gridname}_flowelem_zw',f'{gridname}_flowelem_zcc'])
     return uds
 
 
-def reconstruct_zw_zcc(ds):
+def reconstruct_zw_zcc(uds:xu.UgridDataset):
     """
-    reconstruct full grid output (time/face-varying z-values) for all layertypes, passes on to respective reconstruction function
+    reconstruct full grid output (time/face-varying z-values) for all layertypes,
+    calls the respective reconstruction function
+
+    Parameters
+    ----------
+    uds : xu.UgridDataset
+        DESCRIPTION.
+
+    Raises
+    ------
+    KeyError
+        DESCRIPTION.
+
+    Returns
+    -------
+    uds : TYPE
+        DESCRIPTION.
+
     """
-    dimn_layer, dimn_interfaces = get_vertical_dimensions(ds)
 
-    if dimn_layer is not None: #D-FlowFM mapfile
-        gridname = ds.grid.name
-        varname_zint = f'{gridname}_flowelem_zw'
-    elif 'laydim' in ds.dims: #D-FlowFM hisfile
-        varname_zint = 'zcoordinate_w'
+    gridname = uds.grid.name
+    varname_zint = f'{gridname}_flowelem_zw'
 
     #reconstruct zw/zcc variables (if not in file) and treat as fullgrid mapfile from here
-    if varname_zint in ds.variables: #fullgrid info already available, so continuing
+    if varname_zint in uds.variables: #fullgrid info already available, so continuing
         print(f'zw/zcc (fullgrid) values already present in Dataset in variable {varname_zint}')
-    elif len(ds.filter_by_attrs(standard_name='ocean_sigma_z_coordinate')) != 0:
+    elif len(uds.filter_by_attrs(standard_name='ocean_sigma_z_coordinate')) != 0:
         print('zsigma-layer model, computing zw/zcc (fullgrid) values and treat as fullgrid model from here')
-        ds = reconstruct_zw_zcc_fromzsigma(ds)
-    elif 'mesh2d_layer_sigma' in ds.variables: #TODO: var with standard_name='ocean_sigma_coordinate' available?
+        uds = reconstruct_zw_zcc_fromzsigma(uds)
+    elif f'{gridname}_layer_sigma' in uds.variables: #TODO: var with standard_name='ocean_sigma_coordinate' available?
         print('sigma-layer model, computing zw/zcc (fullgrid) values and treat as fullgrid model from here')
-        ds = reconstruct_zw_zcc_fromsigma(ds)
-    elif 'mesh2d_layer_z' in ds.variables:
+        uds = reconstruct_zw_zcc_fromsigma(uds)
+    elif f'{gridname}_layer_z' in uds.variables:
         print('z-layer model, computing zw/zcc (fullgrid) values and treat as fullgrid model from here')
-        ds = reconstruct_zw_zcc_fromz(ds)
+        uds = reconstruct_zw_zcc_fromz(uds)
     else:
-        raise KeyError('layers present, but unknown layertype, expected one of variables: mesh2d_flowelem_zw, mesh2d_layer_sigma, mesh2d_layer_z')
-    return ds
+        raise KeyError(f'layers present, but unknown layertype, expected one of variables: {gridname}_flowelem_zw, {gridname}_layer_sigma, {gridname}_layer_z')
+    return uds
 
 
 def get_Dataset_atdepths(data_xr:xu.UgridDataset, depths, reference:str ='z0'):    
@@ -400,7 +415,8 @@ def get_Dataset_atdepths(data_xr:xu.UgridDataset, depths, reference:str ='z0'):
                                                            'positive':'up'}) #TODO: make more in line with CMEMS etc
 
     #potentially construct fullgrid info (zcc/zw)
-    data_xr = reconstruct_zw_zcc(data_xr)
+    if dimn_layer is not None: #D-FlowFM mapfile
+        data_xr = reconstruct_zw_zcc(data_xr)
 
     #correct reference level
     if reference=='z0':

docs/whats-new.md

@@ -12,6 +12,7 @@
 - moved gtsmv4.1 tide interpolation from extrapolated to rasterized dataset by [@veenstrajelmer](https://github.com/veenstrajelmer) in [#731](https://github.com/Deltares/dfm_tools/pull/731)
 - limited `dfmt.download_CMEMS()` freq argument to M/D only by [@veenstrajelmer](https://github.com/veenstrajelmer) in [#733](https://github.com/Deltares/dfm_tools/pull/733)
 - suppres xarray chunking warning by [@veenstrajelmer](https://github.com/veenstrajelmer) in [#742](https://github.com/Deltares/dfm_tools/pull/742)
+- remove mesh2d hardcoding from layer reconstruction by [@veenstrajelmer](https://github.com/veenstrajelmer) in [#747](https://github.com/Deltares/dfm_tools/pull/747)
 
 
 ## 0.18.0 (2023-12-08)

tests/test_get_nc.py

@@ -9,6 +9,7 @@
 import pytest
 import numpy as np
 import dfm_tools as dfmt
+import xarray as xr
 
 
 @pytest.mark.unittest
@@ -25,3 +26,22 @@ def test_polyline_mapslice():
     assert np.isclose(uds_crs.grid.node_x.max(), 1737.0864257281437)
     assert np.isclose(uds_crs.grid.node_y.min(), -5)
     assert np.isclose(uds_crs.grid.node_y.max(), 0.9261683648147339)
+
+
+def test_get_dataset_atdepths_hisfile():
+
+    file_nc = dfmt.data.fm_grevelingen_his(return_filepath=True)
+    ds = xr.open_dataset(file_nc)#, preprocess=dfmt.preprocess_hisnc)
+
+    depths = [-1,-4,0,-6]
+    data_fromhis_atdepths = dfmt.get_Dataset_atdepths(data_xr=ds, depths=depths, reference='z0')
+    data_xr_selzt = data_fromhis_atdepths.isel(stations=2).isel(time=slice(40,100))
+
+    tem_values = data_xr_selzt.temperature.isel(time=-1).to_numpy()
+    exp_values = np.array([5.81065253, 5.43289777, 5.36916911, 5.36916911])
+
+    assert data_xr_selzt.temperature.shape == (60,4)
+    assert np.allclose(tem_values, exp_values)
+    assert np.isclose(data_xr_selzt.temperature.sum(), 1295.56826688)
+
+