L96_var.py

# 2017 JA
import numpy as np
from scipy.linalg import pinv
from scipy.optimize import fmin, fsolve
from L96_model import lorenz96


###############################################################################
def var3d(x0,t,tobs,y,H,B,R,model,N):
    """Data assimilation routine for both Lorenz 1963 & 1996 using 3DVar.
    Inputs:  - x0, the real initial conditions
             - t, time array of the model (should be evenly spaced)
             - tobs, time array of the observations (should be evenly
               spaced with a timestep that is a multiple of the model
               timestep)
             - y, the observations
             - H, observation matrix
             - B, the background error covariance matrix
             - R, the observational error covariance matrix
             - model, a string indicating the name of the model: 'lor63'
               or 'lor96'
             - N, the number of variables
    Outputs: - x_b, the background
             - x_a, the analysis"""

    # General settings
    Nsteps = np.size(t)
    # For the true time
    tstep_truth = t[1]-t[0]
    # For the analysis
    tstep_obs = tobs[1]-tobs[0]
    # The ratio
    o2t = int(tstep_obs/tstep_truth+0.5)

    # Precreate the arrays for background and analysis
    x_b = np.empty((Nsteps,N)); x_b.fill(np.nan)
    x_a = np.empty((Nsteps,N)); x_a.fill(np.nan)

    invB = pinv(B)
    invR = pinv(R)

    # For the original background ensemble let's start close from the truth
    orig_bgd = 'fixed'
    #orig_bgd = 'random'

    if orig_bgd=='fixed':
        indaux = np.arange(N)
        x0_aux = x0 + (-1)**indaux
    elif orig_bgd=='random':
        x0_aux = x0 + np.random.randn(N)

    # For the first instant b and a are equal
    x_b[0,:] = x0_aux
    x_a[0,:] = x0_aux

    # The following cycle contains evolution and assimilation
    for j in range(len(tobs)-1):
        yaux = y[j+1,:]

        # First compute background; our initial condition is the
        # forecast from the analysis at the previous observational
        # time
        xb0 = x_a[j*o2t,:]
        #if model=='lor63':
        #    taux,xbaux = lorenz63(xb0,o2t*tstep_truth)
        if model=='L96':
            taux,xbaux = lorenz96(o2t*tstep_truth,xb0,N)

        x_b[j*o2t+1:(j+1)*o2t+1,:] = xbaux[1:,:]
        x_a[j*o2t+1:(j+1)*o2t+1,:] = xbaux[1:,:]

        xa_aux = one3dvar(xbaux[o2t,:],yaux,H,B,R,invB,invR)
        x_a[(j+1)*o2t,:] = xa_aux
        print('t =', t[o2t*(j+1)])

    return x_b,x_a


def one3dvar(xb,y,H,B,R,invB,invR):
    "The 3DVar algorithm for one assimilation window."
    y = np.mat(y).T      # array -> column vector
    xbvec = np.mat(xb).T # array -> column vector

    #opcmin = 'dirmin'
    opcmin = 'gradeq0'

    # The cost function in case the minimization is direct
    def costfun(xa):
        xa = np.mat(xa).T
        Jback = (xa-xbvec).T*invB*(xa-xbvec)
        Jobs = (y-H*xa).T*invR*(y-H*xa)
        J = Jback + Jobs
        return J

    # The gradient in case we prefer to find its roots
    def gradJ(xa):
        xa = np.mat(xa).T
        # The background term
        gJb = np.mat(invB)*(xa-xbvec)
        # The observational term
        gJo = -H.T*np.mat(invR)*(y-H*xa)
        gJ = gJb + gJo
        return gJ.A.flatten()

    if opcmin=='dirmin':
        xa = fmin(costfun,xb,xtol=1e-3)
    elif opcmin=='gradeq0':
        xa = fsolve(gradJ,xb,xtol=1e-4)

    return xa



##############################################################################
# 4DVar
##############################################################################
def var4d(x0,t,tobs,anawin,y,H,B,R,model,N):
    """Data assimilation routine for both Lorenz 1963 & 1996 using 4DVar.
    Inputs:  - x0, the real initial conditions
             - t, time array of the model (should be evenly spaced)
             - tobs, time array of the observations (should be evenly
               spaced with a timestep that is a multiple of the model
               timestep)
             - anawin, length of the 4D assim window, expressed as
               number of future obs included
             - y, the observations
             - H, observation matrix
             - B, the background error covariance matrix
             - R, the observational error covariance matrix
             - model, a string indicating the name of the model: 'lor63'
               or 'lor96'
             - N, the number of variables
    Outputs: - x_b, the background
             - x_a, the analysis"""

    # General settings
    Nsteps = np.size(t)
    Ntobs = np.size(tobs)
    # For the true time
    tstep_truth = t[1]-t[0]
    # For the analysis
    tstep_obs = tobs[1]-tobs[0]
    # The ratio
    o2t = int(tstep_obs/tstep_truth+0.5)

    totana = (Ntobs-1)/anawin

    # Precreate the arrays for background and analysis
    x_b = np.empty((Nsteps,N)); x_b.fill(np.nan)
    x_a = np.empty((Nsteps,N)); x_a.fill(np.nan)

    invB = pinv(B)
    invR = pinv(R)

    # For the original background ensemble let's start close from the truth
    orig_bgd = 'fixed'
    #orig_bgd = 'random'

    if orig_bgd=='fixed':
        indaux = np.arange(N)
        x0_aux = x0 + (-1)**indaux
    elif orig_bgd=='random':
        x0_aux = x0 + np.random.randn(N)

    # For the first instant b and a are equal
    x_b[0,:] = x0_aux
    x_a[0,:] = x0_aux

    # The following cycle contains evolution and assimilation
    for j in range(int(totana)):
        # Get the observations; these are distributed all over the
        # assimilation window
        yaux = y[anawin*j+1:anawin*(j+1)+1,:] # [anawin,L]

        # First compute background; our background is the forecast
        # from the analysis
        xb0 = x_a[j*anawin*o2t,:]
        #if model=='lor63':
        #    taux,xbaux = lorenz63(xb0,o2t*anawin*tstep_truth)
        if model=='L96':
            taux,xbaux = lorenz96(o2t*anawin*tstep_truth,xb0,N)

        x_b[j*o2t*anawin:(j+1)*o2t*anawin+1,:] = xbaux
        xa0 = one4dvar(tstep_truth,o2t,anawin,xb0,yaux,H,B,R,model,N,invB,invR)

        #if model=='lor63':
        #    taux,xaaux = lorenz63(xa0,o2t*anawin*tstep_truth)
        if model=='L96':
            taux,xaaux = lorenz96(o2t*anawin*tstep_truth,xa0,N)

        x_a[j*o2t*anawin:(j+1)*o2t*anawin+1,:] = xaaux
        print('t =', tobs[anawin*(j+1)])

    return x_b,x_a


def one4dvar(tstep_truth,o2t,anawin,x0b,y,H,B,R,model,N,invB,invR):
    "The 4DVar algorithm for one assimilation window."
    x0bvec = np.mat(x0b).T # array -> column vector
    y = np.mat(y).T        # array -> matrix [N, nobs in assim window]

    #opcmin = 'dirmin'
    opcmin = 'gradeq0'

    # The SciPy optimization functions don't work well with matrices,
    # so in the following the initial guesses and function inputs and
    # outputs are arrays or scalars.

    # The cost function in case the minimization is direct
    def costfun(xa):
        xavec = np.mat(xa).T
        Jback = (xavec-x0bvec).T*invB*(xavec-x0bvec)

        #if model=='lor63':
        #    taux,xaaux = lorenz63(xa,o2t*anawin*tstep_truth)
        if model=='L96':
            taux,xaaux = lorenz96(o2t*anawin*tstep_truth,xa,N)

        indobs = range(o2t,anawin*o2t+1,o2t)
        xobs = xaaux[indobs,:]
        xobs = np.mat(xobs).T
        Jobs = np.empty(len(indobs))
        Jobs.fill(np.nan)
        for iJobs in range(len(indobs)):
            Jobs[iJobs] = (y[:,iJobs]-H*(xobs[:,iJobs])).T*invR \
                          *(y[:,iJobs]-H*(xobs[:,iJobs]))
        J = Jback + np.sum(Jobs)
        return J

    # The gradient in case we prefer to find its roots
    def gradJ(xa):
        xavec = np.mat(xa).T
        # The background term
        gJb = invB*(xavec-x0bvec)

        # For the model and observation error we need the TLM and adjoint
        # Choose the TLM/adjoint depending on the model
        #if model=='lor63':
        #    xaux,M = transmat_lor63(xa,o2t*anawin*tstep_truth)
        if model=='L96':
            xaux,M = transmat_lor96(xa,o2t*anawin*tstep_truth)

        # The observation error term, evaluated at different times
        gJok = np.mat(np.empty((N,anawin)))
        gJok.fill(np.nan)
        for j in range(anawin):
            gJok[:,j] = -np.mat(M[:,:,o2t*(j+1)]).T*H.T*invR \
                        *(y[:,j]-H*np.mat(xaux[o2t*(j+1),:]).T)

        # Adding the two
        gJ = gJb + np.sum(gJok,axis=1)
        return gJ.A.flatten()

    if opcmin=='dirmin':
        xa = fmin(costfun,x0b,xtol=1e-3)
    elif opcmin=='gradeq0':
        xa = fsolve(gradJ,x0b,xtol=1e-4)

    return xa




##############################################################################
def transmat_lor96(x0,tmax):
    """The transition matrix required for the TLM and the adjoint.

    Inputs:  - x0, an array containing the initial state
             - tmax, the maximum time.  Should be a multiple of the
               timestep 0.01.
    Outputs: - x, the evolved state [len(t) x 3]
               M, the transition matrix for small perturbations from
               the initial state to the evolved state [3 x 3 x len(t)]"""

    global N
    N = len(x0)
    tstep = 0.01
    taux = np.arange(0,tmax+tstep/2,tstep)
    # Settings for M
    M0 = np.eye(N)

    xaux = np.empty((len(taux),N))
    xaux.fill(np.nan)
    xaux[0,:] = x0
    M = np.empty((N,N,len(taux)))
    M.fill(np.nan)
    M[:,:,0] = M0

    for i in range(len(taux)-1): # for each time
        xaux[i+1,:],M[:,:,i+1] = integr(xaux[i,:],M[:,:,i],tstep)

    return xaux,M

def integr(xin,Min,tstep):
    global N

    # The integration is for both the model and the TLM at the same time!
    Varsold = np.concatenate((xin,Min.flatten()))

    k1 = faux(Varsold)
    k2 = faux(Varsold+1/2.0*tstep*k1)
    k3 = faux(Varsold+1/2.0*tstep*k2)
    k4 = faux(Varsold+tstep*k3)
    Varsnew = Varsold + 1/6.0*tstep*(k1+2*k2+2*k3+k4)

    xout = Varsnew[:N]
    Mout = np.reshape(Varsnew[N:],(N,N))

    return xout,Mout

def faux(Varsin):
    "The Lorenz 1996 model and its TLM."
    global N
    dx = f(Varsin[:N])
    F = tlm(Varsin[:N],N)
    dM = F*np.mat(np.reshape(Varsin[N:],(N,N)))
    dxaux = np.concatenate((dx,dM.A.flatten()))
    return dxaux

def f(x):
    "The actual Lorenz 1996 model."
    global N
    forc = 8.0
    k=np.empty_like(x)
    k.fill(np.nan)
    # Remember it is a cyclical model, hence we need modular algebra
    for j in range(N):
        k[j]=(x[(j+1)%N]-x[(j-2)%N])*x[(j-1)%N]-x[j]+forc
    return k

def tlm(u,nx):
    F = np.zeros((nx,nx))
    for j in range(nx):
        F[j,(j-2)%nx] = -u[(j-1)%nx] 
        F[j,(j-1)%nx] = -u[(j-2)%nx] + u[(j+1)%nx]
        F[j,j%nx] = -1
        F[j,(j+1)%nx] = u[(j-1)%nx]
    return F