def get_pveddy_coeffs(geofil,
                      vgeofil,
                      fil,
                      fil2,
                      xstart,
                      xend,
                      ystart,
                      yend,
                      zs=0,
                      ze=None,
                      z=np.linspace(-3000, 0, 100),
                      percx=0,
                      percy=0):

    fhgeo = dset(geofil)
    fhvgeo = dset(vgeofil)
    fh = mfdset(fil)
    fh2 = mfdset(fil2)
    db = fhvgeo.variables['g'][:]
    fhvgeo.close()
    dbi = np.append(db, 0)
    zi = fh.variables['zi'][:]
    zl = fh.variables['zl'][:]
    dbl = np.diff(zi) * 9.8 / 1031
    dt = fh.variables['average_DT'][:]
    dt = dt[:, np.newaxis, np.newaxis, np.newaxis]

    sl, dimq = rdp1.getslice(fh,
                             xstart,
                             xend,
                             ystart,
                             yend,
                             xhxq='xq',
                             yhyq='yq')
    xs = sl[3].start
    xe = sl[3].stop
    ys = sl[2].start
    ye = sl[2].stop
    zs = sl[1].start
    ze = sl[1].stop
    slmx = np.s_[:, :, sl[2], (sl[3].start - 1):sl[3].stop]
    slpx = np.s_[:, :, sl[2], (sl[3].start + 1):sl[3].stop]
    slmxpy = np.s_[:, :, sl[2].start:(sl[2].stop + 1),
                   (sl[3].start - 1):sl[3].stop]
    slpy = np.s_[:, :, sl[2].start:(sl[2].stop + 1), sl[3]]
    sl2d = sl[2:]
    slmx2d = slmx[2:]
    slpx2d = slpx[2:]

    dxbu = fhgeo.variables['dxBu'][sl2d]
    dxcv = fhgeo.variables['dxCv'][sl2d]
    e = (fh2.variables['e'][slpy] * dt).sum(axis=0, keepdims=True) / dt.sum(
        axis=0, keepdims=True)
    e = np.concatenate((e, e[:, :, :, -1:]), axis=3)
    eatq = 0.25 * (e[:, :, :-1, :-1] + e[:, :, 1:, 1:] + e[:, :, 1:, :-1] +
                   e[:, :, :-1, 1:])

    (_, _, _, q, _, _) = pv2.extract_twapv_terms(geofil,
                                                 vgeofil,
                                                 fil,
                                                 fil2,
                                                 xs - 1,
                                                 xe,
                                                 ys,
                                                 ye,
                                                 zs,
                                                 ze,
                                                 meanax=(0, ),
                                                 fil3=None,
                                                 calledfromgeteddycoeffs=True)
    (_, _, _, _, qbud,
     _) = pv2.extract_twapv_terms(geofil,
                                  vgeofil,
                                  fil,
                                  fil2,
                                  xs,
                                  xe,
                                  ys,
                                  ye,
                                  zs,
                                  ze,
                                  meanax=(0, ),
                                  fil3=None,
                                  calledfromgeteddycoeffs=True)
    fd = qbud[:, :, :, :, 12]
    q[np.isnan(q)] = 0
    qx = np.diff(q, axis=3) / dxcv
    qx = np.concatenate((qx, qx[:, :, :, -1:]), axis=3)
    qxx = np.diff(qx, axis=3) / dxbu

    fdaz = getvaratzc(fd.astype(np.float32), z.astype(np.float32),
                      eatq.astype(np.float32))
    qxxaz = getvaratzc(qxx.astype(np.float32), z.astype(np.float32),
                       eatq.astype(np.float32))

    qxxazm = np.apply_over_axes(np.nanmean, qxxaz, (0, 2))
    fdazm = np.apply_over_axes(np.nanmean, fdaz, (0, 2))

    fig, ax = plt.subplots(1, 2, sharex=True, sharey=True, figsize=(12, 3))
    cmax = np.fabs(np.nanpercentile(qxxazm, (1, 99))).max()
    im = ax[0].pcolormesh(dimq[3],
                          z,
                          qxxazm.squeeze(),
                          vmax=cmax,
                          vmin=-cmax,
                          cmap='RdBu_r')
    cb = fig.colorbar(im, ax=ax[0])
    cb.formatter.set_powerlimits((0, 0))
    cb.update_ticks()
    cmax = np.fabs(np.nanpercentile(fdazm, (1, 99))).max()
    im = ax[1].pcolormesh(dimq[3],
                          z,
                          fdazm.squeeze(),
                          vmax=cmax,
                          vmin=-cmax,
                          cmap='RdBu_r')
    cb = fig.colorbar(im, ax=ax[1])
    cb.formatter.set_powerlimits((0, 0))
    cb.update_ticks()
    for axc in ax:
        xdegtokm(axc, 0.5 * (ystart + yend))
        axc.set_xlabel('x (km)')
        axc.grid()
    ax[0].set_ylabel('z (m)')

    fig2, ax = plt.subplots(1, 1, sharey=True)
    ax.linfit(qxx, fd, percx=percx, percy=percy)
    ax.set_xlabel('qxx')
    ax.set_ylabel('grad fd')
    ax.grid()

    return fig, fig2
def get_heddy_coeffs(geofil,
                     vgeofil,
                     fil,
                     fil2,
                     xstart,
                     xend,
                     ystart,
                     yend,
                     zs=0,
                     ze=None,
                     z=np.linspace(-3000, 0, 100),
                     percx=0,
                     percy=0):

    fhgeo = dset(geofil)
    fhvgeo = dset(vgeofil)
    fh = mfdset(fil)
    fh2 = mfdset(fil2)
    db = fhvgeo.variables['g'][:]
    fhvgeo.close()
    dbi = np.append(db, 0)
    zi = fh.variables['zi'][:]
    zl = fh.variables['zl'][:]
    dbl = np.diff(zi) * 9.8 / 1031
    dt = fh.variables['average_DT'][:]
    dt = dt[:, np.newaxis, np.newaxis, np.newaxis]

    sl, dimv = rdp1.getslice(fh, xstart, xend, ystart, yend, yhyq='yq')
    slmx = np.s_[:, :, sl[2], (sl[3].start - 1):sl[3].stop]
    slpx = np.s_[:, :, sl[2], (sl[3].start + 1):sl[3].stop]
    slmxpy = np.s_[:, :, sl[2].start:(sl[2].stop + 1),
                   (sl[3].start - 1):sl[3].stop]
    slpy = np.s_[:, :, sl[2].start:(sl[2].stop + 1), sl[3]]
    sl2d = sl[2:]
    slmx2d = slmx[2:]
    slpx2d = slpx[2:]

    dxbu = fhgeo.variables['dxBu'][slmx2d]
    dxcv = fhgeo.variables['dxCv'][sl2d]

    v, _ = pv.getvtwa(fhgeo, fh, fh2, slmx)
    vx = np.diff(v, axis=3) / dxbu
    vxx = np.diff(vx, axis=3) / dxcv
    e = (fh2.variables['e'][slpy] * dt).sum(axis=0, keepdims=True) / dt.sum(
        axis=0, keepdims=True)
    eatv = 0.5 * (e[:, :, :-1, :] + e[:, :, 1:, :])
    vxx = getvaratzc(vxx.astype(np.float32), z.astype(np.float32),
                     eatv.astype(np.float32))

    e = fh2.variables['e'][slmxpy]
    eatvmx = 0.5 * (e[:, :, :-1, :] + e[:, :, 1:, :])
    vaz = getvaratzc(v.astype(np.float32), z.astype(np.float32),
                     eatvmx.astype(np.float32))
    vazx = np.diff(vaz, axis=3) / dxbu
    vazxx = np.diff(vazx, axis=3) / dxcv

    x, y, P, _, _ = py.extract_twamomy_terms(geofil, vgeofil, fil, fil2,
                                             xstart, xend, ystart, yend, zs,
                                             ze, (0, ))
    uv = P[:, :, :, :, 5]

    vvalid = vaz[:, :, :, 1:-1]

    vazm = np.apply_over_axes(np.nanmean, vvalid, (0, 2))
    vxxm = np.apply_over_axes(np.nanmean, vxx, (0, 2))
    vazxxm = np.apply_over_axes(np.nanmean, vazxx, (0, 2))
    uvm = np.apply_over_axes(np.nanmean, uv, (0, 2))

    fig, ax = plt.subplots(1, 4, sharex=True, sharey=True, figsize=(12, 3))
    cmax = np.fabs(np.percentile(vxxm, (1, 99))).max()
    im = ax[0].pcolormesh(dimv[3],
                          z,
                          vxxm.squeeze(),
                          vmax=cmax,
                          vmin=-cmax,
                          cmap='RdBu_r')
    fig.colorbar(im, ax=ax[0])
    cmax = np.fabs(np.percentile(vazxxm, (1, 99))).max()
    im = ax[1].pcolormesh(dimv[3],
                          z,
                          vazxxm.squeeze(),
                          vmax=cmax,
                          vmin=-cmax,
                          cmap='RdBu_r')
    fig.colorbar(im, ax=ax[1])
    cmax = np.fabs(np.percentile(vvalid, (0.5, 99.5))).max()
    im = ax[2].pcolormesh(dimv[3],
                          z,
                          vazm.squeeze(),
                          vmax=cmax,
                          vmin=-cmax,
                          cmap='RdBu_r')
    fig.colorbar(im, ax=ax[2])
    cmax = np.fabs(np.percentile(uvm, (0.5, 99.5))).max()
    im = ax[3].pcolormesh(dimv[3],
                          z,
                          uvm.squeeze(),
                          vmax=cmax,
                          vmin=-cmax,
                          cmap='RdBu_r')
    cb = fig.colorbar(im, ax=ax[3])
    cb.formatter.set_powerlimits((0, 0))
    cb.update_ticks()
    for axc in ax:
        xdegtokm(axc, 0.5 * (ystart + yend))
        axc.set_xlabel('x (km)')
        axc.grid()
    ax[0].set_ylabel('z (m)')

    fig2, ax = plt.subplots(1, 2, sharey=True)
    ax[0].linfit(vvalid, uv, percx=percx, percy=percy)
    ax[0].set_xlabel('v (m/s)')
    ax[0].set_ylabel('RS div (m/s2)')
    ax[0].grid()

    ax[1].linfit(vxx, uv, percx=percx, percy=percy)
    ax[1].set_xlabel('vxx (1/ms)')
    ax[1].grid()

    return fig, fig2
def get_deddy_coeffs_fromflx(geofil,
                             vgeofil,
                             fil,
                             fil2,
                             xstart,
                             xend,
                             ystart,
                             yend,
                             zs=0,
                             ze=None,
                             zlim=(-1500, 0),
                             percx=0,
                             percy=0,
                             nsqdep=False,
                             fil3=None,
                             swashperc=1):

    fhgeo = dset(geofil)
    fhvgeo = dset(vgeofil)
    fh = mfdset(fil)
    fh2 = mfdset(fil2)
    db = fhvgeo.variables['g'][:]
    fhvgeo.close()
    dbi = np.append(db, 0)
    zi = fh.variables['zi'][:]
    zl = fh.variables['zl'][:]
    dbl = np.diff(zi) * 9.8 / 1031
    dt = fh.variables['average_DT'][:]
    dt = dt[:, np.newaxis, np.newaxis, np.newaxis]
    z = np.linspace(zlim[0], zlim[1], 100)

    sl, dimv = rdp1.getslice(fh, xstart, xend, ystart, yend, yhyq='yq')
    sl2d = sl[2:]
    dycv = fhgeo.variables['dyCv'][sl2d]
    slpy = np.s_[:, :, sl[2].start:(sl[2].stop + 1), sl[3]]
    f = fhgeo.variables['f'][sl2d]

    e = (fh2.variables['e'][slpy] * dt).sum(axis=0,
                                            keepdims=True) / dt.sum(axis=0)
    eatv = 0.5 * (e[:, :, :-1, :] + e[:, :, 1:, :])
    v, h = pv.getvtwa(fhgeo, fh, fh2, sl)
    v = v[:, :, :, :-1]
    h = h[:, :, :, :-1]
    sig = h / dbl[:, np.newaxis, np.newaxis]
    hi = 0.5 * (h[:, :-1] + h[:, 1:])
    sigi = hi / dbi[1:-1, np.newaxis, np.newaxis]

    vb = -np.diff(v, axis=1) / dbi[1:-1, np.newaxis, np.newaxis]
    #vb = np.concatenate((vb[:,:1,:,:],vb,vb[:,-1:,:,:]),axis=1)

    vi = np.concatenate((v[:, :1, :, :], v, -v[:, -1:, :, :]), axis=1)
    vi = 0.5 * (vi[:, :-1, :, :] + vi[:, 1:, :, :])
    vib = -np.diff(vi, axis=1) / dbl[:, np.newaxis, np.newaxis]
    if nsqdep:
        fsqvb = f**2 * vb * sigi
    else:
        fsqvb = f**2 * vb

    fsqvbaz = getvaratzc(fsqvb.astype(np.float32), z.astype(np.float32),
                         eatv.astype(np.float32))

    esq = (fh.variables['esq'][slpy] * dt).sum(axis=0,
                                               keepdims=True) / np.sum(dt)
    elmforydiff = 0.5 * (e[:, 0:-1, :, :] + e[:, 1:, :, :])
    edlsqm = (esq - elmforydiff**2)
    edlsqmy = np.diff(edlsqm, axis=2) / dycv

    hpfv = (fh.variables['twa_hpfv'][sl] * dt).sum(axis=0,
                                                   keepdims=True) / np.sum(dt)
    pfvm = (fh2.variables['PFv'][sl] * dt).sum(axis=0,
                                               keepdims=True) / np.sum(dt)
    edpfvdmb = -(-hpfv + h * pfvm -
                 0.5 * edlsqmy * dbl[:, np.newaxis, np.newaxis])

    fh.close()
    fh2.close()
    fhgeo.close()

    zdpfvd = sint.cumtrapz(edpfvdmb, dx=-dbl[0], axis=1)
    #    zdpfvd = np.concatenate((np.zeros(zdpfvd[:,:1,:,:].shape),
    #                             zdpfvd,
    #                             np.zeros(zdpfvd[:,:1,:,:].shape)),axis=1)
    #    zdpfvd = 0.5*(zdpfvd[:,:-1,:,:]+zdpfvd[:,1:,:,:])
    zdpfvdaz = getvaratzc(zdpfvd.astype(np.float32), z.astype(np.float32),
                          eatv.astype(np.float32))

    fsqvbazm = np.apply_over_axes(np.nanmean, fsqvbaz, (0, 2))
    zdpfvdazm = np.apply_over_axes(np.nanmean, zdpfvdaz, (0, 2))

    fig, ax = plt.subplots(1, 2, sharex=True, sharey=True, figsize=(12, 3))
    cmax = np.fabs(np.percentile(zdpfvdazm, (1, 99))).max()
    im = ax[0].pcolormesh(dimv[3],
                          z,
                          zdpfvdazm.squeeze(),
                          vmax=cmax,
                          vmin=-cmax,
                          cmap='RdBu_r')
    fig.colorbar(im, ax=ax[0])

    cmax = np.fabs(np.percentile(fsqvbazm, (1, 99))).max()
    im = ax[1].pcolormesh(dimv[3],
                          z,
                          fsqvbazm.squeeze(),
                          vmax=cmax,
                          vmin=-cmax,
                          cmap='RdBu_r')
    fig.colorbar(im, ax=ax[1])

    if fil3:
        fh3 = mfdset(fil3)
        slmxtn = np.s_[-1:, sl[1], sl[2], (sl[3].start - 1):sl[3].stop]
        islayerdeep0 = fh3.variables['islayerdeep'][-1:, 0, 0, 0]
        islayerdeep = fh3.variables['islayerdeep'][slmxtn]
        if np.ma.is_masked(islayerdeep):
            islayerdeep = islayerdeep.filled(np.nan)
            islayerdeep[:, :, :, -1:] = islayerdeep[:, :, :, -2:-1]

        swash = (islayerdeep0 - islayerdeep) / islayerdeep0 * 100
        swash = 0.5 * (swash[:, :, :, :-1] + swash[:, :, :, 1:])
        fh3.close()
        swash = getvaratzc(swash.astype(np.float32), z.astype(np.float32),
                           eatv.astype(np.float32))
        swash = np.apply_over_axes(np.nanmean, swash, (0, 2))
        em = np.apply_over_axes(np.nanmean, e, (0, 2))
        xx, zz = np.meshgrid(dimv[3], zi)
        for axc in ax:
            axc.contour(dimv[3],
                        z,
                        swash.squeeze(),
                        np.array([swashperc]),
                        colors='k')
            axc.contour(xx, em.squeeze(), zz, ls='k--')
            axc.set_ylim(zlim[0], zlim[1])

    fig2, ax = plt.subplots(1, 1, sharey=True, figsize=(12, 3))
    ax.linfit(fsqvbaz, zdpfvdaz, percx=percx, percy=percy)
    ax.set_xlabel(r'$f^2 v_b$')
    ax.set_ylabel(r'$\zeta^{\prime} m_y^{\prime}$')
    return fig, fig2
def get_deddy_coeffs(geofil,
                     vgeofil,
                     fil,
                     fil2,
                     xstart,
                     xend,
                     ystart,
                     yend,
                     zs=0,
                     ze=None,
                     z=np.linspace(-1500, 0, 100),
                     perc=5):

    fhgeo = dset(geofil)
    fhvgeo = dset(vgeofil)
    fh = mfdset(fil)
    fh2 = mfdset(fil2)
    db = fhvgeo.variables['g'][:]
    fhvgeo.close()
    dbi = np.append(db, 0)
    zi = fh.variables['zi'][:]
    zl = fh.variables['zl'][:]
    dbl = np.diff(zi) * 9.8 / 1031

    sl, dimv = rdp1.getslice(fh, xstart, xend, ystart, yend, yhyq='yq')
    sl2d = sl[2:]
    slpy = np.s_[:, :, sl[2].start:(sl[2].stop + 1), sl[3]]
    f = fhgeo.variables['f'][sl2d]

    e = fh2.variables['e'][slpy]
    eatv = 0.5 * (e[:, :, :-1, :] + e[:, :, 1:, :])
    v, h = pv.getvtwa(fhgeo, fh, fh2, sl)
    v = v[:, :, :, :-1]
    h = h[:, :, :, :-1]
    sig = h / dbl[:, np.newaxis, np.newaxis]
    sigaz = getvaratzc(sig.astype(np.float32), z.astype(np.float32),
                       eatv.astype(np.float32))

    vb = -np.diff(v, axis=1) / dbi[1:-1, np.newaxis, np.newaxis]
    vb = np.concatenate((vb[:, :1, :, :], vb, vb[:, -1:, :, :]), axis=1)
    vbb = -np.diff(vb, axis=1) / dbl[:, np.newaxis, np.newaxis]
    vbb = getvaratzc(vbb.astype(np.float32), z.astype(np.float32),
                     eatv.astype(np.float32))

    vi = np.concatenate((v[:, :1, :, :], v, -v[:, -1:, :, :]), axis=1)
    vi = 0.5 * (vi[:, :-1, :, :] + vi[:, 1:, :, :])
    vib = -np.diff(vi, axis=1) / dbl[:, np.newaxis, np.newaxis]
    svib = sig * vib
    svibb = -np.diff(svib, axis=1) / dbi[1:-1, np.newaxis, np.newaxis]
    svibb = np.concatenate((svibb[:, :1, :, :], svibb, svibb[:, -1:, :, :]),
                           axis=1)
    svibb = 0.5 * (svibb[:, :-1, :, :] + svibb[:, 1:, :, :])
    svibb = getvaratzc(svibb.astype(np.float32), z.astype(np.float32),
                       eatv.astype(np.float32))

    sv = sig * v
    svb = -np.diff(sv, axis=1) / dbi[1:-1, np.newaxis, np.newaxis]
    svb = np.concatenate((svb[:, :1, :, :], svb, svb[:, -1:, :, :]), axis=1)
    svbb = -np.diff(svb, axis=1) / dbl[:, np.newaxis, np.newaxis]
    svbb = getvaratzc(svbb.astype(np.float32), z.astype(np.float32),
                      eatv.astype(np.float32))

    x, y, P, _, _ = py.extract_twamomy_terms(geofil,
                                             vgeofil,
                                             fil,
                                             fil2,
                                             xstart,
                                             xend,
                                             ystart,
                                             yend,
                                             zs,
                                             ze, (0, ),
                                             z=z)
    fdrag = P[:, :, :, :, 7]

    vbbm = np.apply_over_axes(np.nanmean, vbb, (0, 2))
    svibbm = np.apply_over_axes(np.nanmean, svibb, (0, 2))
    svbbm = np.apply_over_axes(np.nanmean, svbb, (0, 2))
    fdragm = np.apply_over_axes(np.nanmean, fdrag, (0, 2))

    fig, ax = plt.subplots(1, 4, sharex=True, sharey=True, figsize=(12, 3))
    cmax = np.fabs(np.percentile(vbbm, (1, 99))).max()
    im = ax[0].pcolormesh(dimv[3],
                          z,
                          vbbm.squeeze(),
                          vmax=cmax,
                          vmin=-cmax,
                          cmap='RdBu_r')
    fig.colorbar(im, ax=ax[0])

    cmax = np.fabs(np.percentile(svibbm, (1, 99))).max()
    im = ax[1].pcolormesh(dimv[3],
                          z,
                          svibbm.squeeze(),
                          vmax=cmax,
                          vmin=-cmax,
                          cmap='RdBu_r')
    fig.colorbar(im, ax=ax[1])

    cmax = np.fabs(np.percentile(svbbm, (1, 99))).max()
    im = ax[2].pcolormesh(dimv[3],
                          z,
                          svbbm.squeeze(),
                          vmax=cmax,
                          vmin=-cmax,
                          cmap='RdBu_r')
    fig.colorbar(im, ax=ax[2])

    cmax = np.fabs(np.percentile(fdragm, (1, 99))).max()
    im = ax[3].pcolormesh(dimv[3],
                          z,
                          fdragm.squeeze(),
                          vmax=cmax,
                          vmin=-cmax,
                          cmap='RdBu_r')
    fig.colorbar(im, ax=ax[3])

    fig2, ax = plt.subplots(1, 3, sharey=True, figsize=(12, 3))
    ax[0].linfit(vbb, sigaz * fdrag / f**2)
    ax[0].set_xlabel(r'$v_{bb}$')
    ax[1].linfit(svibb, sigaz * fdrag / f**2)
    ax[1].set_xlabel(r'$(\sigma v_b)_b$')
    ax[2].linfit(svbb, sigaz * fdrag / f**2)
    ax[2].set_xlabel(r'$(\sigma v)_{bb}$')

    ax[0].set_ylabel('Form Drag')
    for axc in ax:
        axc.grid()

    return fig, fig2
def get_heddy_coeffs_fromflx(geofil,
                             vgeofil,
                             fil,
                             fil2,
                             xstart,
                             xend,
                             ystart,
                             yend,
                             zs=0,
                             ze=None,
                             z=np.linspace(-3000, 0, 100),
                             perc=5,
                             htol=1e-3):

    fhgeo = dset(geofil)
    fhvgeo = dset(vgeofil)
    fh = mfdset(fil)
    fh2 = mfdset(fil2)
    db = fhvgeo.variables['g'][:]
    fhvgeo.close()
    dbi = np.append(db, 0)
    zi = fh.variables['zi'][:]
    zl = fh.variables['zl'][:]
    dbl = np.diff(zi) * 9.8 / 1031
    dt = fh.variables['average_DT'][:]
    dt = dt[:, np.newaxis, np.newaxis, np.newaxis]

    sl, dimv = rdp1.getslice(fh, xstart, xend, ystart, yend, yhyq='yq')
    slmx = np.s_[:, :, sl[2], (sl[3].start - 1):sl[3].stop]
    slpx = np.s_[:, :, sl[2], (sl[3].start + 1):sl[3].stop]
    slmxpy = np.s_[:, :, sl[2].start:(sl[2].stop + 1),
                   (sl[3].start - 1):sl[3].stop]
    slpy = np.s_[:, :, sl[2].start:(sl[2].stop + 1), sl[3]]
    slmpy = np.s_[:, :, (sl[2].start - 1):(sl[2].stop + 1), sl[3]]
    sl2d = sl[2:]
    slmx2d = slmx[2:]
    slpx2d = slpx[2:]
    slpy2d = slpy[2:]

    dxbu = fhgeo.variables['dxBu'][slmx2d]
    dybu = fhgeo.variables['dyBu'][slmx2d]
    dxcv = fhgeo.variables['dxCv'][sl2d]
    dycv = fhgeo.variables['dyCv'][sl2d]
    dxt = fhgeo.variables['dxT'][slpy2d]
    dyt = fhgeo.variables['dyT'][slpy2d]

    v, _ = pv.getvtwa(fhgeo, fh, fh2, slmx)
    vx = np.diff(v, axis=3) / dxbu
    vx = vx[:, :, :, 1:]
    e = (fh2.variables['e'][slpy] * dt).sum(axis=0, keepdims=True) / dt.sum(
        axis=0, keepdims=True)
    eatv = 0.5 * (e[:, :, :-1, :] + e[:, :, 1:, :])

    e = fh2.variables['e'][slmxpy]
    eatvmx = 0.5 * (e[:, :, :-1, :] + e[:, :, 1:, :])

    vh = (fh.variables['vh_masked'][sl] * dt).sum(axis=0,
                                                  keepdims=True) / np.sum(dt)
    h_cv = (fh.variables['h_Cv'][sl] * dt).sum(axis=0,
                                               keepdims=True) / np.sum(dt)
    h_cv[h_cv < htol] = np.nan
    sig = h_cv / dbl[:, np.newaxis, np.newaxis]
    h_vm = h_cv
    vtwa = vh / h_cv / dxcv
    vhforxdiff = (fh.variables['vh_masked'][slmx] * dt).sum(
        axis=0, keepdims=True) / np.sum(dt)
    h_cvforxdiff = (fh.variables['h_Cv'][slmx] * dt).sum(
        axis=0, keepdims=True) / np.sum(dt)
    h_cvforxdiff[h_cvforxdiff < htol] = np.nan
    vtwaforxdiff = vhforxdiff / h_cvforxdiff
    vtwaforxdiff = np.concatenate((vtwaforxdiff, vtwaforxdiff[:, :, :, -1:]),
                                  axis=3)
    vtwax = np.diff(vtwaforxdiff, axis=3) / dxbu / dybu
    vtwax = 0.5 * (vtwax[:, :, :, :-1] + vtwax[:, :, :, 1:])
    uh = (fh.variables['uh_masked'][slmxpy] * dt).filled(0).sum(
        axis=0, keepdims=True) / np.sum(dt)
    hum = 0.25 * (uh[:, :, :-1, :-1] + uh[:, :, :-1, 1:] + uh[:, :, 1:, :-1] +
                  uh[:, :, 1:, 1:]) / dycv
    huvxphvvym = (fh.variables['twa_huvxpt'][sl] * dt +
                  fh.variables['twa_hvvymt'][sl] * dt).sum(
                      axis=0, keepdims=True) / np.sum(dt)
    hvv = (fh.variables['hvv_Cv'][slmpy] * dt).sum(axis=0,
                                                   keepdims=True) / np.sum(dt)
    hvvym = np.diff(hvv, axis=2) / dxt / dyt
    hvvym = 0.5 * (hvvym[:, :, :-1, :] + hvvym[:, :, 1:, :])
    huvxm = -(huvxphvvym + hvvym)
    advx = hum * vtwax / h_vm
    humx = np.diff(np.nan_to_num(uh), axis=3) / dxt / dyt
    humx = 0.5 * (humx[:, :, :-1, :] + humx[:, :, 1:, :])
    xdivep1 = -huvxm / h_vm
    xdivep2 = advx
    xdivep3 = vtwa * humx / h_vm
    xdivep = (xdivep1 + xdivep2 + xdivep3)

    xdivep *= sig
    uv = sint.cumtrapz(
        xdivep[:, :, :, ::-1], dx=-dxbu[:, 1:-1], axis=3,
        initial=0)[:, :, :, ::-1] / sig

    uvm = np.apply_over_axes(np.nanmean, uv, (0, 2))
    vxm = np.apply_over_axes(np.nanmean, vx, (0, 2))
    uvm = getvaratzc(uvm.astype(np.float32), z.astype(np.float32),
                     eatv.astype(np.float32))
    vxm = getvaratzc(vxm.astype(np.float32), z.astype(np.float32),
                     eatv.astype(np.float32))

    fig, ax = plt.subplots(1, 2, sharex=True, sharey=True, figsize=(12, 3))
    cmax = np.fabs(np.percentile(uvm, (1, 99))).max()
    im = ax[0].pcolormesh(dimv[3],
                          z,
                          uvm.squeeze(),
                          vmax=cmax,
                          vmin=-cmax,
                          cmap='RdBu_r')
    cmax = np.fabs(np.percentile(vxm, (1, 99))).max()
    im = ax[1].pcolormesh(dimv[3],
                          z,
                          vxm.squeeze(),
                          vmax=cmax,
                          vmin=-cmax,
                          cmap='RdBu_r')
    for axc in ax:
        xdegtokm(axc, 0.5 * (ystart + yend))
        axc.set_xlabel('x (km)')
        axc.grid()
    ax[0].set_ylabel('z (m)')

    fig2, ax = plt.subplots(1, 1, sharey=True)
    ax.linfit(vx, uv)
    ax.set_xlabel('v_x (m/s)')
    ax.set_ylabel('RS (m2/s2)')
    ax.grid()

    return fig, fig2
Ejemplo n.º 6
0
def getuv(geofil,
          vgeofil,
          fil,
          fil2,
          xstart,
          xend,
          ystart,
          yend,
          zs,
          ze,
          meanax,
          xyasindices=False):
    keepax = ()
    for i in range(4):
        if i not in meanax:
            keepax += (i, )

    fhvgeo = dset(vgeofil)
    db = -fhvgeo.variables['g'][:]
    dbi = np.append(db, 0)
    fhvgeo.close()
    fhgeo = dset(geofil)
    fh = mfdset(fil)
    fh2 = mfdset(fil2)
    zi = rdp1.getdims(fh)[2][0]
    dbl = np.diff(zi) * 9.8 / 1031

    slh, dimh = rdp1.getslice(fh,
                              wlon=xstart,
                              elon=xend,
                              slat=ystart,
                              nlat=yend,
                              zs=zs,
                              ze=ze)
    ys, ye = slh[2].start, slh[2].stop
    xs, xe = slh[3].start, slh[3].stop
    slhmx = np.s_[:, zs:ze, ys:ye, xs - 1:xe]
    slhmy = np.s_[:, zs:ze, ys - 1:ye, xs:xe]
    slhmpy = np.s_[:, zs:ze, ys - 1:ye + 1, xs:xe]
    slu, dimu = rdp1.getslice(fh,
                              xstart,
                              xend,
                              ystart,
                              yend,
                              zs=zs,
                              ze=ze,
                              ts=0,
                              te=None,
                              xhxq='xq',
                              yhyq='yh',
                              zlzi='zl')
    slv, dimv = rdp1.getslice(fh,
                              xstart,
                              xend,
                              ystart,
                              yend,
                              zs=zs,
                              ze=ze,
                              ts=0,
                              te=None,
                              xhxq='xh',
                              yhyq='yq',
                              zlzi='zl')
    ys, ye = slv[2].start, slv[2].stop
    xs, xe = slv[3].start, slv[3].stop
    slvpy = np.s_[:, zs:ze, ys:ye + 1, xs:xe]

    uh = getvaravg(fh2, 'uh', slu)
    u = getvaravg(fh2, 'u', slu).filled(np.nan)
    h_cu = getvaravg(fh, 'h_Cu', slu)
    h_cu = np.ma.masked_array(h_cu, mask=(h_cu < 1e-3))
    dycu = fhgeo.variables['dyCu'][slu[2:]]
    utwa = uh / h_cu / dycu

    vh = getvaravg(fh2, 'vh', slv)
    v = getvaravg(fh2, 'v', slv)
    h_cv = getvaravg(fh, 'h_Cv', slv)
    h_cv = np.ma.masked_array(h_cv, mask=(h_cv < 1e-3))
    dxcv = fhgeo.variables['dxCv'][slv[2:]]
    vtwa = vh / dxcv / h_cv

    emforxdiff = getvaravg(fh2, 'e', slhmx)
    elmforxdiff = 0.5 * (emforxdiff[:, 0:-1, :, :] + emforxdiff[:, 1:, :, :])
    elmforxdiff = np.concatenate((elmforxdiff, elmforxdiff[:, :, :, -1:]),
                                 axis=3)
    dxcuforxdiff = fhgeo.variables['dxCu'][slhmx[2:]]
    ex = np.diff(elmforxdiff, axis=3) / dxcuforxdiff

    uh = getvaravg(fh2, 'uh', slhmx)
    h_cu = getvaravg(fh, 'h_Cu', slhmx)
    h_cu = np.ma.masked_array(h_cu, mask=(h_cu < 1e-3))
    dycu = fhgeo.variables['dyCu'][slhmx[2:]]
    uzx = (uh / h_cu / dycu).filled(0) * ex
    uzx = 0.5 * (uzx[:, :, :, 1:] + uzx[:, :, :, :-1])

    emforydiff = getvaravg(fh2, 'e', slhmpy)
    elmforydiff = 0.5 * (emforydiff[:, 0:-1, :, :] + emforydiff[:, 1:, :, :])
    dycv = fhgeo.variables['dyCv'][slhmy[2:]]
    ey = np.diff(elmforydiff, axis=2) / dycv

    vh = getvaravg(fh2, 'vh', slhmy)
    h_cv = getvaravg(fh, 'h_Cv', slhmy)
    h_cv = np.ma.masked_array(h_cv, mask=(h_cv < 1e-3))
    dxcv = fhgeo.variables['dxCv'][slhmy[2:]]
    vtwa = vh / dxcv / h_cv
    vzy = vtwa * ey
    vtwa = 0.5 * (vtwa[:, :, 1:, :] + vtwa[:, :, :-1, :])
    vzy = 0.5 * (vzy[:, :, 1:, :] + vzy[:, :, :-1, :])

    wd = getvaravg(fh2, 'wd', slh)
    hw = wd * dbi[:, np.newaxis, np.newaxis]
    hw = 0.5 * (hw[:, 1:, :, :] + hw[:, :-1, :, :])
    wzb = hw / dbl[:, np.newaxis, np.newaxis]
    whash = uzx + vzy + wzb

    terms = [utwa, vtwa, whash, u, v]
    slices = [slu, slvpy, slh, slu, slvpy]
    X = [
        dimu[keepax[1]], dimv[keepax[1]], dimh[keepax[1]], dimu[keepax[1]],
        dimv[keepax[1]]
    ]
    Y = [
        dimu[keepax[0]], dimv[keepax[0]], dimh[keepax[0]], dimu[keepax[1]],
        dimv[keepax[1]]
    ]
    termsm = []
    for item in terms:
        try:
            item = item.filled(np.nan)
        except AttributeError:
            item = item
        termsm.append(np.ma.apply_over_axes(np.nanmean, item, meanax))

    if 1 in keepax:
        Y = []
        for i in range(len(terms)):
            z = np.linspace(-3000, 0, 100)
            em = fh2.variables['e'][slices[i]]
            if i == 0 or i == 3:
                em = np.concatenate((em, em[:, :, :, -1:]), axis=3)
                em = 0.5 * (em[:, :, :, :-1] + em[:, :, :, 1:])
            elif i == 1 or i == 4:
                em = 0.5 * (em[:, :, :-1, :] + em[:, :, 1:, :])
            em = np.ma.apply_over_axes(np.mean, em, meanax)
            termsm[i] = getvaratzc(termsm[i].astype(np.float32),
                                   z.astype(np.float32), em.astype(np.float32))
            Y.append(z)

    fh2.close()
    fh.close()
    fhgeo.close()

    P = []
    for i, item in enumerate(termsm):
        P.append(item)
        X[i] = np.ma.filled(X[i].astype(float), np.nan)
        Y[i] = np.ma.filled(Y[i].astype(float), np.nan)

    return X, Y, P