use xugrid.UgridDataset.from_structured2d in dfm_tools

[veenstrajelmer]

Todo:

• wait for fix of from_structured does not support non-monotonic indexes xugrid#316
• wait for xugrid release (after 13 feb 2025)
• set new minimal version in toml
• simplify open_dataset_delft3d4 with code below (also transfer other variables with MC/NC dimensions) >> or drop them just like in current function
• simplify open_dataset_curvilinear
• maybe fix layers with improve getting layer/interface dimensions? #1128, at least "vertical_dimensions" (in case of delft3d4) is not useful to add

Delft3D4:

import dfm_tools as dfmt
import xarray as xr
import xugrid as xu
import matplotlib.pyplot as plt
plt.close("all")
import numpy as np

file_nc = r'p:\archivedprojects\1220688-lake-kivu\3_modelling\1_FLOW\7_heatfluxinhis\062_netcdf\trim-thiery_002_coarse.nc'
ds = xr.open_dataset(file_nc, chunks={"time":1})

def get_delft3d4_nanmask(x,y):
    # -999.999 in kivu and 0.0 in curvedbend, both in westernscheldt
    bool_0 = (x==0) & (y==0)
    bool_1 = (x==-999) & (y==-999)
    bool_2 = (x==-999.999) & (y==-999.999)
    bool_mask = bool_0 | bool_1 | bool_2
    return bool_mask

def stack_shifted_coords(da):
    shift = 1
    np_stacked = np.stack([
        da, #ll
        da.shift(MC=shift), #lr
        da.shift(MC=shift, NC=shift), #ur
        da.shift(NC=shift), #ul
        ],axis=-1)
    da_stacked = xr.DataArray(np_stacked, dims=("M","N","four"))
    return da_stacked

XCOR_stacked = stack_shifted_coords(ds.XCOR)
YCOR_stacked = stack_shifted_coords(ds.YCOR)
mask_xy = get_delft3d4_nanmask(XCOR_stacked,YCOR_stacked)
ds['XCOR_stacked'] = XCOR_stacked.where(~mask_xy)
ds['YCOR_stacked'] = YCOR_stacked.where(~mask_xy)

# replace invalid values not with nan but with zero
# otherwise the spatial coverage is affected
mask_u1 = (ds.U1==-999) | (ds.U1==-999.999)
mask_v1 = (ds.V1==-999) | (ds.V1==-999.999)
u1_mn = ds.U1.where(~mask_u1, 0)
v1_mn = ds.V1.where(~mask_v1, 0)

# minus 0.5 since padding=low so corner value is representative for previous face
# according to that logic, method=nearest might be better (or just rename the dims)
u1_mn_cen = u1_mn.interp(MC=u1_mn.MC-0.5, method='linear').rename({'MC':'M'})
v1_mn_cen = v1_mn.interp(NC=v1_mn.NC-0.5, method='linear').rename({'NC':'N'})
# since padding=low, just renaming the dims might even be better?
# u1_mn_cen = u1_mn.rename({'MC':'M'})
# v1_mn_cen = v1_mn.rename({'NC':'N'})
# >> could also be done with `ds = ds.swap_dims({"M":"MC","N":"NC"})`

# create combined uv mask (have to rename dimensions)
mask_u1_mn = mask_u1.rename({'MC':'M'})
mask_v1_mn = mask_v1.rename({'NC':'N'})
mask_uv1_mn = mask_u1_mn & mask_v1_mn
# drop all actual missing cells
u1_mn_cen = u1_mn_cen.where(~mask_uv1_mn)
v1_mn_cen = v1_mn_cen.where(~mask_uv1_mn)

# to avoid creating large chunks, alternative is to overwrite the vars with the MN-averaged vars, but it requires passing and updating of attrs
ds = ds.drop_vars(['U1','V1'])

# compute ux/uy/umag/udir #TODO: add attrs to variables
alfas_rad = np.deg2rad(ds.ALFAS)
vel_x = u1_mn_cen*np.cos(alfas_rad) - v1_mn_cen*np.sin(alfas_rad)
vel_y = u1_mn_cen*np.sin(alfas_rad) + v1_mn_cen*np.cos(alfas_rad)
ds['ux'] = vel_x
ds['uy'] = vel_y
ds['umag'] = np.sqrt(vel_x**2 + vel_y**2)
ds['udir'] = np.rad2deg(np.arctan2(vel_y, vel_x))%360

# use same dimensions for variables on cell corners and cells faces
# ds = ds.swap_dims({"M":"MC","N":"NC"})
uds_centers = xu.UgridDataset.from_structured2d(ds, topology={"mesh2d":{"x":"M",
                                                                "y":"N",
                                                                "x_bounds":"XCOR_stacked",
                                                                "y_bounds":"YCOR_stacked",
                                                                }})

# fig, ax = plt.subplots(figsize=(6,8))
# (-uds_centers.DPS0).ugrid.plot(ax=ax, cmap='jet', vmax=0)
# uds_centers.grid.plot(ax=ax, linewidth=0.5, color='grey')
# dfmt.plot_coastlines(ax=ax, res="f")

fig, ax = plt.subplots(figsize=(6,7))
uds_centers.umag.isel(KMAXOUT_RESTR=-2, time=10).ugrid.plot(ax=ax, cmap='jet')
dfmt.plot_coastlines(ax=ax, res="f")

# uds = dfmt.open_dataset_delft3d4(file_nc)
# fig, ax = plt.subplots(figsize=(6,7))
# uds.umag.isel(KMAXOUT_RESTR=-2, time=10).ugrid.plot(ax=ax, cmap='jet')
# dfmt.plot_coastlines(ax=ax, res="f")

# save gshhg as shapefile for quickplot reference image
xlim = ax.get_xlim()
ylim = ax.get_ylim()
bbox = (xlim[0], ylim[0], xlim[1], ylim[1])
coastlines_gdb = dfmt.coastlines.get_coastlines_gdb(bbox=bbox, res="f", min_area=0, crs=None)
coastlines_gdb.to_file(r"c:\Users\veenstra\Downloads\shp_gshhg\gshhg_kivu.shp")

Some others:

# WAQUA
file_nc = r'p:\archivedprojects\11205258-006-kpp2020_rmm-g6\C_Work\07_WAQUAresultaten\j15\SDS-riv_tba_map.nc'
ds = xr.open_dataset(file_nc, chunks={'time':1})
uds = xu.UgridDataset.from_structured2d(ds, topology={"mesh2d":{"x":"N",
                                                                "y":"M",
                                                                "x_bounds":"grid_x", # XZETA
                                                                "y_bounds":"grid_y", # YZETA
                                                                }})
fig, ax = plt.subplots(figsize=(10,5))
uds.grid.plot()


# CMCC
file_nc = 'p:\\archivedprojects\\11206304-futuremares-rawdata-preps\\data\\CMIP6_BC\\CMCC-ESM2\\thetao_Omon_CMCC-ESM2_ssp126_r1i1p1f1_gn_201501-203412.nc'
ds = xr.open_dataset(file_nc, chunks={'time':1}) # without time=1 this generates >1GB of traffic before even finishing the load
uds = xu.UgridDataset.from_structured2d(ds, topology={"mesh2d":{"x":"i",
                                                                "y":"j",
                                                                "x_bounds":"vertices_longitude",
                                                                "y_bounds":"vertices_latitude",
                                                                }})
uds = uds.ugrid.to_nonperiodic(xmax=360.0)
fig, ax = plt.subplots(figsize=(10,5))
uds.grid.plot()


# CMEMS without vertices
file_nc = r"p:\11209810-cmems-nws-data-nrt\IBI\nrt\IBI_nrt_2021.nc"
ds = xr.open_dataset(file_nc, chunks={'time':1}) # without time chunks: MemoryError: Unable to allocate 84.8 GiB for an array with shape (725, 50, 755, 832) and data type float32
uds = xu.UgridDataset.from_structured2d(ds, topology={"mesh2d":{"x":"longitude",
                                                                "y":"latitude",
                                                                }})
fig, ax = plt.subplots(figsize=(10,5))
uds.grid.plot()