slgentil / croco

Code CROCO de base + CONFIGS qui contient les différentes configurations
0 stars 1 forks source link

get_z needs update #4

Closed apatlpo closed 4 years ago

apatlpo commented 4 years ago

it probably should look like something closer to:

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

NoeLahaye commented 4 years ago

J'ai mis une nouvelle version de get_z dans gridop. Il me semble que cela fonctionne. Si cette version vous va, on peut virer l'ancienne (get_z_old).

apatlpo commented 4 years ago

vu