Commit b6a76ceb authored by Jan Klaus Rieck's avatar Jan Klaus Rieck

handle grid W in getbox()

parent b6081fc3
Pipeline #3636 failed with stage
in 1 minute and 36 seconds
......@@ -310,8 +310,9 @@ def getbox(dataset, time=None, depth=None, y=None, x=None, surf=False):
# define possible grids
grids = possible_grids(ds)
if 'W' in grids:
del grids['W'] # do not need W for getbox...
if ('W' in grids) and ('T' not in grids) and ('wmask' not in ds):
raise ValueError('no mask for grid W in dataset'
+ 'and no other grid to create mask from')
# define empty dataset to write into
ds_sel = xr.Dataset()
......@@ -327,6 +328,22 @@ def getbox(dataset, time=None, depth=None, y=None, x=None, surf=False):
Xroll[grid['x']] = int(Xroll_argmax.max())
# apply the mask (defined on the current grid) only to
# variables that share its grid
if gr == 'W':
m = ds[grids['T']['mask']].isel(z_c=0)
m3 = ds[grids['T']['mask']]
# need to construct wmask, because its not in output
m = xr.DataArray(m.data, dims=['y_c', 'x_c'])
diff_z = int(len(ds['z_l']) - len(ds['z_c']))
if diff_z > 0:
m3 = xr.DataArray(np.vstack((m3.data[0:1, ...], m3.data)),
dims=['z_l', 'y_c', 'x_c'])
elif diff_z < 0:
m3 = xr.DataArray(m3.data[0:diff_z],
dims=['z_l', 'y_c', 'x_c'])
else:
m3 = xr.DataArray(m3.data,
dims=['z_l', 'y_c', 'x_c'])
ds = ds.update({'wmask': m3}).set_coords('wmask')
if grid['box_mask'] in ds.coords:
ds[grid['box_mask']] = ds.coords[grid['box_mask']]
else:
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment