def get_z(run, zeta=None, h=None, vgrid='r', hgrid='r'):
''' compute vertical coordinates
zeta should have the size of the final output
vertical coordinate should be first
'''
ds = run['his']
N = run.N
hc = run.params['Hc']
if zeta is not None:
_zeta = zeta
else:
_zeta = ds.ssh
if h is not None:
_h = h
else:
_h = ds.h
#
if hgrid is 'u':
_zeta = rho2u(_zeta, ds)
_h = rho2u(_h, ds)
elif hgrid is 'v':
_zeta = rho2v(_zeta, ds)
_h = rho2v(_h, ds)
#
sc=run['his']['sc_'+vgrid]
cs=run['his']['Cs_'+vgrid]
#
z0 = (hc * sc + _h * cs) / (hc + _h)
z = _zeta + (_zeta + _h) * z0
#z = np.squeeze(_zeta + (_zeta + _h) * z0) # this should not be here
# manually swap dims, could also perform align with T,S
#if z.ndim ==4:
# z = z.transpose(z.dims[0], z.dims[3], z.dims[1], z.dims[2])
#elif z.ndim == 3:
# z = z.transpose(sc.dims[0], _zeta.dims[0], _zeta.dims[1])
#elif z.ndim == 2:
# z = z.transpose(sc.dims[0], _zeta.dims[0])
return z.rename('z_'+vgrid)
I am not sure transposition is necessary, to discuss
it probably should look like something closer to:
I am not sure transposition is necessary, to discuss