def qgsw(Hi=None,
         c=None,
         lon=None,
         lat=None,
         tint=None,
         dtout=None,
         dt=None,
         obsspace=None,
         Hm=None,
         rappel=None,
         snu=None):
    """ QG Shallow Water model

    Args:
        Hi (2D array): Initial SSH field.
        c (same size as Hi): Rossby first baroclinic phase speed
        lon (2D array): longitudes
        lat (2D array): latitudes
        tint (scalar): Time of propagator integration in seconds. Can be positive (forward integration) or negative (backward integration)
        dtout (scalar): Time period of outputs
        dt (scalar): Propagator time step

    Returns:
        SSH: 3D array with dimensions (timesteps, height, width), SSH forecast  
    """
    way = np.sign(tint)

    ##############
    # Setups
    ##############

    grd = modgrid.grid(Hi, c, snu, lon, lat)
    #plt.figure()
    #plt.pcolor(grd.mask)
    #plt.show()
    time_abs = 0.
    index_time = 0
    if obsspace is not None:
        hg = np.empty((np.shape(obsspace)[0]))
        hg[:] = np.NAN
        iobs = np.where((way * obsspace[:, 2] >= time_abs - dt / 2)
                        & (way * obsspace[:, 2] < time_abs + dt / 2))
        if np.size(iobs) > 0:
            hg[iobs] = griddata(
                (lon.ravel(), lat.ravel()), h.ravel(),
                (obsspace[iobs, 0].squeeze(), obsspace[iobs, 1].squeeze()))
    else:
        hg = None
    nindex_time = np.int(abs(tint) / dtout + 1)
    grdnx = np.int(grd.nx)
    grdny = np.int(grd.ny)
    SSH = np.empty((nindex_time, grdny, grdnx))
    SSH[index_time, :, :] = Hi

    nstep = int(abs(tint) / dt)
    stepout = int(dtout / dt)

    ############################
    # Passive variable initializations
    ############################

    if rappel is not None:
        Qm, = modelliptic.h2pv(Hm, grd)
        #Qm[np.where(grd.mask==1)]=q[np.where(grd.mask==1)]

############################
# Active variable initializations
############################

    h = +Hi
    q, = modelliptic.h2pv(h, grd)

    hb = +h  # just for hguess

    ############################
    # Time loop
    ############################

    for step in range(nstep):
        #print step
        time_abs = (step + 1) * dt
        if (np.mod(step + 1, stepout) == 0):
            index_time += 1

        ############################
        #Initialization of previous fields
        ############################

        hguess = 2 * h - hb
        hb = +h
        qb = +q

        ########################
        # Main routines
        ########################

        # 1/
        u, v, = moddyn.h2uv(h, grd)

        # 2/
        rq, = moddyn.qrhs(u, v, qb, grd, way)

        # 3/
        if rappel is not None:
            q = qb + dt * (rq - rappel * (qb - Qm))
        else:
            q = qb + dt * rq

        # 4/
        h, = modelliptic.pv2h(q, hguess, grd)

        ############################
        #Saving outputs
        ############################

        if (np.mod(step + 1, stepout) == 0):
            SSH[index_time, :, :] = h

        if obsspace is not None:
            iobs = np.where((way * obsspace[:, 2] >= time_abs - dt / 2)
                            & (way * obsspace[:, 2] < time_abs + dt / 2))
            if np.size(iobs) > 0:
                hg[iobs] = griddata(
                    (lon.ravel(), lat.ravel()), h.ravel(),
                    (obsspace[iobs, 0].squeeze(), obsspace[iobs, 1].squeeze()))

    return SSH, hg
Example #2
0
def qgsw_stochuv(Hi=None,
                 c=None,
                 lon=None,
                 lat=None,
                 tint=None,
                 dtout=None,
                 dt=None,
                 obsspace=None,
                 snu=None,
                 qgiter=1):
    """ QG Shallow Water model

    Args:
        Hi (2D array): Initial SSH field.
        c (same size as Hi): Rossby first baroclinic phase speed
        lon (2D array): longitudes
        lat (2D array): latitudes
        tint (scalar): Time of propagator integration in seconds. Can be positive (forward integration) or negative (backward integration)
        dtout (scalar): Time period of outputs
        dt (scalar): Propagator time step

    Returns:
        SSH: 3D array with dimensions (timesteps, height, width), SSH forecast  
    """
    way = np.sign(tint)

    ##############
    # Setups
    ##############
    grd = modgrid.grid(Hi, c, snu, lon, lat)
    #plt.figure()
    #plt.pcolor(grd.mask)
    #plt.show()
    time_abs = 0.
    index_time = 0
    if obsspace is not None:
        hg = np.empty((np.shape(obsspace)[0]))
        hg[:] = np.NAN
        iobs = np.where((way * obsspace[:, 2] >= time_abs - dt / 2)
                        & (way * obsspace[:, 2] < time_abs + dt / 2))
        if np.size(iobs) > 0:
            hg[iobs] = griddata(
                (lon.ravel(), lat.ravel()), h.ravel(),
                (obsspace[iobs, 0].squeeze(), obsspace[iobs, 1].squeeze()))
    else:
        hg = None
    nindex_time = np.int(abs(tint) / dtout + 1)
    grdnx = np.int(grd.nx)
    grdny = np.int(grd.ny)
    SSH = np.empty((nindex_time, grdny, grdnx))
    SSH[index_time, :, :] = Hi

    nstep = int(abs(tint) / dt)
    stepout = int(dtout / dt)

    ############################
    # Passive variable initializations
    ############################

    ############################
    # Active variable initializations
    ############################

    h = +Hi

    q, = modelliptic.h2pv(h, grd)

    hb = +h  # just for hguess

    ############################
    # Time loop
    ############################

    for step in range(nstep):
        #print step
        time_abs = (step + 1) * dt
        if (np.mod(step + 1, stepout) == 0):
            index_time += 1

        ############################
        #Initialization of previous fields
        ############################

        hguess = 2 * h - hb
        hb = +h
        qb = +q

        ########################
        # Main routines
        ########################

        # 1/
        u, v, = moddyn.h2uv(h, grd)

        # 1b/
        # Generate spatially correlated noise # Check http://www2.geog.ucl.ac.uk/~plewis/geogg122-2011-12/dem1.html
        sizefilter = 30
        max_noise = 0.02
        nx = np.shape(u)[0] + 2 * sizefilter
        ny = np.shape(u)[1] + 2 * sizefilter
        u_uniform = np.random.rand(nx, ny)
        v_uniform = np.random.rand(nx, ny)
        x, y = np.mgrid[-sizefilter:sizefilter + 1, -sizefilter:sizefilter + 1]
        gauss = np.exp(-0.333 *
                       (x**2 / float(sizefilter) + y**2 / float(sizefilter)))
        gauss_filter = gauss / gauss.sum()
        u_corr_noise = signal.convolve(u_uniform, gauss_filter, mode='valid')
        v_corr_noise = signal.convolve(v_uniform, gauss_filter, mode='valid')
        # rescale so it lies between -max_noise and max_noise
        u_corr_noise = (
            (u_corr_noise - u_corr_noise.min()) /
            (u_corr_noise.max() - u_corr_noise.min()) - 0.5) * max_noise * 2
        v_corr_noise = (
            (v_corr_noise - v_corr_noise.min()) /
            (v_corr_noise.max() - v_corr_noise.min()) - 0.5) * max_noise * 2

        # Apply noises to u and v
        u = u + u_corr_noise
        v = v + v_corr_noise

        # 2/
        rq, = moddyn.qrhs(u, v, qb, grd, way)

        # 3/
        q = qb + dt * rq

        # 4/
        h, = modelliptic.pv2h(q, hguess, grd, qgiter)

        ############################
        #Saving outputs
        ############################

        if (np.mod(step + 1, stepout) == 0):
            SSH[index_time, :, :] = h

        if obsspace is not None:
            iobs = np.where((way * obsspace[:, 2] >= time_abs - dt / 2)
                            & (way * obsspace[:, 2] < time_abs + dt / 2))
            if np.size(iobs) > 0:
                hg[iobs] = griddata(
                    (lon.ravel(), lat.ravel()), h.ravel(),
                    (obsspace[iobs, 0].squeeze(), obsspace[iobs, 1].squeeze()))

    return SSH, hg
Example #3
0
def qgsw_tgl(Htraj=None,
             dHi=None,
             c=None,
             lon=None,
             lat=None,
             tint=None,
             dtout=None,
             dt=None,
             obsspace=None,
             rappel=None,
             snu=None):

    way = np.sign(tint)

    ##############
    # Setups
    ##############
    dHi = dHi + Htraj[0, :, :] * 0.

    grd = modgrid.grid(dHi, c, snu, lon, lat)
    time_abs = 0.
    index_time = 0
    if obsspace is not None:
        dhg = np.empty((np.shape(obsspace)[0]))
        dhg[:] = np.NAN
        iobs = np.where((way * obsspace[:, 2] >= time_abs - dt / 2)
                        & (way * obsspace[:, 2] < time_abs + dt / 2))
        if np.size(iobs) > 0:
            dhg[iobs] = griddata(
                (lon.ravel(), lat.ravel()), h.ravel(),
                (obsspace[iobs, 0].squeeze(), obsspace[iobs, 1].squeeze()))
    else:
        dhg = None

    nindex_time = np.abs(tint) / dtout + 1
    dSSH = np.empty((nindex_time, grd.ny, grd.nx))
    dSSH[index_time, :, :] = dHi

    nstep = int(abs(tint) / dt)
    stepout = int(dtout / dt)

    ############################
    # Active variable initializations
    ############################

    dh = +dHi
    dq, = modelliptic.h2pv(dh, grd)

    dhb = +dh  # just for hguess

    ############################
    # Time loop
    ############################

    for step in range(nstep):
        #print step
        time_abs = (step + 1) * dt
        if (np.mod(step + 1, stepout) == 0):
            index_time += 1

        ############################
        #Tangent update on current trajectory
        ############################

        h = Htraj[index_time, :, :].squeeze()
        q, = modelliptic.h2pv(h, grd)
        u, v, = moddyn.h2uv(h, grd)

        ############################
        # Initializations
        ############################

        dhguess = 2 * dh - dhb
        dhb = +dh
        dqb = +dq

        ########################
        # Main routines
        ########################

        # 1/
        du, dv, = moddyn.h2uv(dhb, grd)

        # 2/
        drq, = moddyn.qrhs_tgl(du, dv, dqb, u, v, q, grd, way)

        # 3/
        if rappel is not None:
            dq = dqb + dt * (drq - rappel * (dqb))
        else:
            dq = dqb + dt * drq

        # 4/ From new q, we update h
        dh, = modelliptic.pv2h(dq, dhguess, grd)

        ############################
        #Saving outputs
        ############################

        if (np.mod(step + 1, stepout) == 0):
            dSSH[index_time, :, :] = dh

        if obsspace is not None:
            iobs = np.where((way * obsspace[:, 2] >= time_abs - dt / 2)
                            & (way * obsspace[:, 2] < time_abs + dt / 2))
            if np.size(iobs) > 0:
                dhg[iobs] = griddata(
                    (lon.ravel(), lat.ravel()), dh.ravel(),
                    (obsspace[iobs, 0].squeeze(), obsspace[iobs, 1].squeeze()))

    return dSSH, dhg
Example #4
0
def qgsw(Hi=None,
         PVi=None,
         c=None,
         lon=None,
         lat=None,
         tint=None,
         dtout=None,
         dt=None,
         obsspace=None,
         scheme='Euler',
         diff=False,
         snu=None,
         name_grd=None,
         qgiter=1):
    """ QG Shallow Water model

    Args:
        Hi (2D array): Initial SSH field.
        c (same size as Hi): Rossby first baroclinic phase speed
        lon (2D array): longitudes
        lat (2D array): latitudes
        tint (scalar): Time of propagator integration in seconds. Can be positive (forward integration) or negative (backward integration)
        dtout (scalar): Time period of outputs
        dt (scalar): Propagator time step

    Returns:
        SSH: 3D array with dimensions (timesteps, height, width), SSH forecast  
    """
    way = np.sign(tint)

    ##############
    # Setups
    ##############
    if name_grd is not None:
        name_grd += '_QGSW'
        if not os.path.isfile(name_grd):
            grd = modgrid.grid(Hi, lon, lat)
            if os.path.exists(os.path.dirname(name_grd)) is False:
                os.makedirs(os.path.dirname(name_grd))
            with open(name_grd, 'wb') as grd_file:
                pickle.dump(grd, grd_file)
                grd_file.close()
        else:
            with open(name_grd, 'rb') as grd_file:
                grd = pickle.load(grd_file)
                grd_file.close()
    else:
        grd = modgrid.grid(Hi, lon, lat)

    time_abs = 0.
    index_time = 0
    if obsspace is not None:
        hg = np.empty((np.shape(obsspace)[0]))
        hg[:] = np.NAN
        iobs = np.where((way * obsspace[:, 2] >= time_abs - dt / 2)
                        & (way * obsspace[:, 2] < time_abs + dt / 2))
        if np.size(iobs) > 0:
            hg[iobs] = griddata(
                (lon.ravel(), lat.ravel()), Hi.ravel(),
                (obsspace[iobs, 0].squeeze(), obsspace[iobs, 1].squeeze()))
    else:
        hg = None

    nindex_time = np.int(abs(tint) / dtout + 1)
    grdnx = np.int(grd.nx)
    grdny = np.int(grd.ny)
    SSH = np.empty((nindex_time, grdny, grdnx))
    SSH[index_time, :, :] = Hi
    PV = np.empty((nindex_time, grdny, grdnx))
    PV[index_time, :, :] = PVi

    nstep = int(abs(tint) / dt)
    stepout = int(dtout / dt)

    ############################
    # Active variable initializations
    ############################
    h = +Hi
    if PVi is not None:
        q = PVi
    else:
        q, = modelliptic.h2pv(h, grd, c)

    hb = +h  # just for hguess

    ############################
    # Time loop
    ############################

    for step in range(nstep):
        time_abs = (step + 1) * dt
        if (np.mod(step + 1, stepout) == 0):
            index_time += 1

        ############################
        #Initialization of previous fields
        ############################

        hguess = 2 * h - hb

        hb = +h
        qb = +q

        ########################
        # Main routines
        ########################

        # 1/
        if diff:
            u = np.zeros((grd.ny, grd.nx))
            v = np.zeros((grd.ny, grd.nx))
        else:
            u, v, = moddyn.h2uv(h, grd)

        # 2/ Advection
        rq, = moddyn.qrhs(u, v, qb, grd, way, diff, snu)

        # 3/ Time integration
        if scheme == 'Euler':
            q = qb + dt * rq

        elif scheme == 'Runge-Kutta':
            # k1
            k1 = rq * dt
            # k2
            q2 = qb + 0.5 * k1
            h2, = modelliptic.pv2h(q2, hguess, grd, qgiter)
            u2, v2, = moddyn.h2uv(h2, grd)
            rq2, = moddyn.qrhs(u2, v2, q2, grd, way)
            k2 = rq2 * dt
            # k3
            q3 = qb + 0.5 * k2
            h3, = modelliptic.pv2h(q3, hguess, grd, qgiter)
            u3, v3, = moddyn.h2uv(h3, grd)
            rq3, = moddyn.qrhs(u3, v3, q3, grd, way)
            k3 = rq3 * dt
            # k4
            q4 = qb + k2
            h4, = modelliptic.pv2h(q4, hguess, grd, qgiter)
            u4, v4, = moddyn.h2uv(h4, grd)
            rq4, = moddyn.qrhs(u4, v4, q4, grd, way)
            k4 = rq4 * dt
            # q increment
            q = qb + (k1 + 2 * k2 + 2 * k3 + k4) / 6.

        # 4/
        h, = modelliptic.pv2h(q, hguess, grd, c, qgiter)

        ############################
        #Saving outputs
        ############################

        if (np.mod(step + 1, stepout) == 0):
            SSH[index_time, :, :] = h
            PV[index_time, :, :] = q

        if obsspace is not None:
            iobs = np.where((way * obsspace[:, 2] >= time_abs - dt / 2)
                            & (way * obsspace[:, 2] < time_abs + dt / 2))
            if np.size(iobs) > 0:
                hg[iobs] = griddata(
                    (lon.ravel(), lat.ravel()), h.ravel(),
                    (obsspace[iobs, 0].squeeze(), obsspace[iobs, 1].squeeze()))

    ############################
    #Returning variables
    ############################
    if PVi is not None:
        return SSH, PV, hg
    else:
        return SSH, hg
Example #5
0
def qgsw_adj(Htraj=None,
             c=None,
             lon=None,
             lat=None,
             tint=None,
             dtout=None,
             dt=None,
             obsspace=None,
             sens=None,
             rappel=None,
             snu=None):

    adHf = Htraj[0, :, :] * 0.

    way = np.sign(tint)

    ##############
    # Setups
    ##############

    grd = modgrid.grid(adHf, c, snu, lon, lat)
    time_abs = 0.
    index_time = -1
    nindex_time = np.abs(tint) / dtout + 1

    adSSH = np.empty((nindex_time, grd.ny, grd.nx))

    nstep = int(abs(tint) / dt)
    stepout = int(dtout / dt)

    deltat = 1 * dt

    ############################
    # Active variable initializations
    ############################

    azeros = adHf * 0.
    adh = +azeros
    adhb = +azeros
    adq = +azeros
    adqb = +azeros
    adrq = +azeros
    adu = +azeros
    adv = +azeros

    adqguess = +azeros

    #adqguess=- 1./grd.g*1./grd.f0*(grd.c**2) *adh

    ############################
    # Time loop
    ############################

    Jdp = +azeros

    for step in range(nstep):
        #print step
        time_abs = step * dt
        if (np.mod(step, stepout) == 0):
            index_time += 1

        ############################
        #Tangent update on current trajectory
        ############################

        h = Htraj[-index_time - 1, :, :].squeeze()
        q, = modelliptic.h2pv(h, grd)
        u, v, = moddyn.h2uv(h, grd)

        ########################
        # Adjoint forcing
        ########################

        iobs = np.where(
            (way * obsspace[:, 2] >= np.abs(tint) - time_abs - dt / 2)
            & (way * obsspace[:, 2] < np.abs(tint) - time_abs + dt / 2))[0]
        if np.shape(iobs)[0] > 0:
            #print 'times: ', min(obsspace[iobs,2]), max(obsspace[iobs,2])
            Jd = sensongrid(obsspace, sens, iobs, grd)
            #print 'some forcing at step ', step
            #adq_forcing_guess=-1./grd.g*1./grd.f0*(grd.c**2) *Jd
            #adq_forcing, = modelliptic.pv2h(Jd,adq_forcing_guess,grd,nitr=10)
            #adq=adq+adq_forcing
            nitr = 10
        else:
            Jd = +azeros
            nitr = 10

        #adh=adh+Jd
        adh = adh + Jd
        #Jdp=Jd

        ########################
        # Main routines
        ########################

        # 4/
        fguess = -grd.c**2. / (grd.g * grd.f0) * Jd
        adqguess = adqguess + fguess
        adqguess[grd.mask == 1] = +fguess[grd.mask == 1]
        adq_tmp, = modelliptic.pv2h(adh, adqguess, grd, nitr=nitr)
        if step == 0: adq_tmpb = +adq_tmp
        adqguess = +2 * adq_tmp - adq_tmpb
        #adqguess=adq_tmp*0.
        adq_tmpb = +adq_tmp
        adq = adq + adq_tmp
        adh = +azeros
        ## Local adj test
        #if step==5:
        #  Madqguess=- grd.c**2/(grd.g*grd.f0) *adq
        #  Madq, = modelliptic.pv2h(adq,Madqguess,grd,nitr=50)
        #  #Madq, = modelliptic.h2pv(adq,grd)
        #  MtMadqguess= (grd.c**2/(grd.g*grd.f0))**2 *adq
        #  MtMadq, = modelliptic.pv2h(Madq,MtMadqguess,grd,nitr=50)
        #  #MtMadq, = modelliptic.h2pv(Madq,grd)
        #  ind=np.where((grd.mask>=1))
        #  print np.inner(Madq[ind],Madq[ind])
        #  print np.inner(adq[ind],MtMadq[ind])
        #  pdb.set_trace()

        # 3/
        if rappel is not None:
            adqb = adqb + (1 - deltat * rappel) * adq
            adrq = adrq + deltat * adq
        else:
            adqb = adqb + adq
            adrq = adrq + deltat * adq
        adq = +azeros

        # 2/
        adu_tmp, adv_tmp, adqb_tmp, = moddyn.qrhs_adj(adrq, u, v, q, grd, way)
        adu = adu + adu_tmp
        adv = adv + adv_tmp
        adqb = adqb + adqb_tmp
        adrq = +azeros
        ## local adjoint test : Exact success
        #if step==100:
        #  #adv=adv*0.
        #  #adu=adu*0.
        ##  #u=u*0.
        ##  #v=v*0.
        #  adqb=adqb*0.
        #  Mdrq, = moddyn.qrhs_tgl(adu,adv,adqb,u,v,q,grd,way)
        #  MtMadu,MtMadv,MtMadqb, = moddyn.qrhs_adj(Mdrq,u,v,q,grd,way)
        #  ind=np.where((grd.mask>=1))
        #  print np.inner(Mdrq[ind],Mdrq[ind])
        #  print np.inner(MtMadu[ind],adu[ind])
        #  print np.inner(MtMadv[ind],adv[ind])
        #  print np.inner(MtMadqb[ind],adqb[ind])
        #  pdb.set_trace()

        # 1/
        adhb_tmp = moddyn.aduv2adh(adu, adv, grd)
        adhb = adhb + adhb_tmp
        #if step==200: pdb.set_trace()
        adu = +azeros
        adv = +azeros
        ## local adjoint test : Exact success
        #Madu,Madv, =moddyn.h2uv(adhb,grd)
        #MtMadhb =moddyn.aduv2adh(Madu,Madv,grd)
        #ind=np.where((grd.mask>=1))
        #print np.inner(Madu[ind],Madu[ind])
        #print np.inner(Madv[ind],Madv[ind])
        #print np.inner(adhb[ind],MtMadhb[ind])
        #pdb.set_trace()

        ############################
        #Saving outputs
        ############################
        adhout, = modelliptic.h2pv(adqb, grd)
        adhout[grd.mask == 0] = np.nan
        if (np.mod(step, stepout) == 0):
            adSSH[index_time, :, :] = +adhout

        ############################
        # Update previous fields
        ############################

        adq = adq + adqb
        adh = adh + adhb
        adqb = +azeros
        adhb = +azeros

        #if step==100: pdb.set_trace()

    ############################
    #Saving final outputs
    ############################
    index_time += 1
    time_abs = nstep * dt
    adhout, = modelliptic.h2pv(adq, grd)
    adhout[grd.mask == 0] = np.nan
    iobs = np.where(
        (way * obsspace[:, 2] >= np.abs(tint) - time_abs - dt / 2)
        & (way * obsspace[:, 2] < np.abs(tint) - time_abs + dt / 2))[0]
    if np.shape(iobs)[0] > 0:
        #print 'times: ', min(obsspace[iobs,2]), max(obsspace[iobs,2])
        Jd = sensongrid(obsspace, sens, iobs, grd)
    else:
        Jd = +azeros
    if (np.mod(step, stepout) == 0):
        adSSH[index_time, :, :] = +adhout + Jd

# pdb.set_trace()
    return adSSH,