def add_vertical_spacing(ds):
grid = Grid(ds)
ds.coords['dst'] = calculate_ds(ds, dim='st')
ds.coords['dswt'] = calculate_ds(ds, dim='sw')
ds.coords['dzt'] = calculate_dz(ds['eta_t'], ds['ht'], ds['dst'])
ds.coords['dzu'] = grid.min(grid.min(ds['dzt'], 'X'), 'Y')
# # Avoids suspected computation midway during the dask graph creation...Should probably raise an xarray issue about this.
# ds['dzt'] = calculate_dz(ds['eta_t'], ds['ht'], ds['dst'])
# ds['dzu'] = grid.min(grid.min(ds['dzt'], 'X'),'Y')
# # Lets try this as a workaround...
# for vv in ['dzt', 'dzu']:
# ds.coords[vv] = ds[vv]
# the dzwt value is dependent on the model version (finite vol vs. engergetically consistent?; See MOM5 manual section 10.4.2)
return ds
),
)
# B-grid, periodic in X direction
g = Grid(
d,
periodic=["X"],
coords={
"X": {"center": "xt_ocean", "right": "xu_ocean"},
"Y": {"center": "yt_ocean", "right": "yu_ocean"},
},
)
# dzu is the minimum of the neighbouring dzt
# vanished layers just get 0 thickness, rather than missing
d["dzu"] = g.min(d.dzt, ["X", "Y"], boundary="fill").chunk({"xu_ocean": None})
d["dzt"] = d.dzt
# we needed a slightly bigger domain to do the above, so now we restrict to the proper domain
d = d.isel(yu_ocean=slice(None, -1), yt_ocean=slice(1, None))
# we don't actually need vc (it's never used by the obc code...)
# flip, sum from the bottom up, then unflip
# d["vc"] = (
# d.dzt.isel(st_ocean=slice(None, None, -1))
# .cumsum(dim="st_ocean")
# .isel(st_ocean=slice(None, None, -1))
# + dg
# )
# construct the regridders for t and q points
dg = xr.open_dataset("/g/data/x77/ahg157/inputs/mom6/panan/ocean_hgrid_0025.nc").isel(