"""Ensemble square-root filters for the 1-d Kuramoto-Sivashinsky eqn"""

import numpy as np

import sys

from KS import KS

from enkf import serial_ensrf, bulk_ensrf, etkf, letkf, etkf_modens,\

serial_ensrf_modens, bulk_enkf, getkf, getkf_modens

np.seterr(all='raise') # raise error when overflow occurs

if len(sys.argv) < 3:

msg="""python KSensrf.py covlocal method smoothlen covinflate1 covfinflate2

all variables are observed, assimilation interval given by dtassim,

nens ensemble members, observation error standard deviation = oberrstdev,

observation operator is smooth_len pt boxcar running mean or gaussian.

time mean error and spread stats printed to standard output.

covlocal: localization distance (distance at which Gaspari-Cohn polynomial goes

to zero).

method: =0 for serial Potter method

=1 for bulk Potter method (all obs at once)

=2 for ETKF (no localization applied)

=3 for LETKF (using observation localization)

=4 for serial Potter method with localization via modulation ensemble

=5 for ETKF with modulation ensemble

=6 for ETKF with modulation ensemble and perturbed obs

=7 for serial Potter method using sqrt of localized Pb ensemble

=8 for bulk EnKF (all obs at once) with perturbed obs.

=9 for GETKF (no localization applied)

=10 for GETKF (with modulated ensemble)

smoothlen: averaging kernel width for forward operator (0 if no averaging)

covinflate1,covinflate2: (optional) inflation parameters corresponding

to a and b in Hodyss and Campbell. If not specified, a=b=1. If covinflate2

<=0, relaxation to prior spread (RTPS) inflation used with a relaxation

coefficient equal to covinflate1."""

raise SystemExit(msg)

corrl = float(sys.argv[1])

method = int(sys.argv[2])

smooth_len = int(sys.argv[3])

use_gaussian = bool(int(sys.argv[4]))

covinflate1=1.; covinflate2=1.

if len(sys.argv) > 5:

# if covinflate2 > 0, use Hodyss and Campbell inflation,

# otherwise use RTPS inflation.

covinflate1 = float(sys.argv[5])

covinflate2 = float(sys.argv[6])

ntstart = 1000 # time steps to spin up truth run

ntimes = 81000 # ob times

nens = 5 # ensemble members

oberrstdev = 1.0; oberrvar = oberrstdev**2 # ob error

verbose = False # print error stats every time if True

dtassim = 2.0 # assimilation interval

# Gaussian or running average smoothing in H.

# for running average, smooth_len is half-width of boxcar.

# for gaussian, smooth_len is standard deviation.

thresh = 0.99 # threshold for modulated ensemble eigenvalue truncation.

# model parameters...

# for truth run

dt = 0.5; npts = 128

diffusion_truth = 1.0

# for forecast model (same as above for perfect model expt)

# for simplicity, assume dt and npts stay the same.

#diffusion = 0.9

diffusion = diffusion_truth

rstruth = np.random.RandomState(42) # fixed seed for truth run

rsens = np.random.RandomState() # varying seed for ob noise and ensemble initial conditions

# model instance for truth (nature) run

model = KS(N=npts,dt=dt,diffusion=diffusion_truth,rs=rstruth)

# mode instance for forecast ensemble

ensemble = KS(N=npts,members=nens,dt=dt,diffusion=diffusion,rs=rsens)

for nt in range(ntstart): # spinup truth run

model.advance()

# sample obs from truth, compute climo stats for model.

xx = []; tt = []

for nt in range(ntimes):

model.advance()

xx.append(model.x[0]) # single member

tt.append(float(nt)*model.dt)

xtruth = np.array(xx,np.float)

timetruth = np.array(tt,np.float)

xtruth_mean = xtruth.mean()

xprime = xtruth - xtruth_mean

xvar = np.sum(xprime**2,axis=0)/(ntimes-1)

xtruth_stdev = np.sqrt(xvar.mean())

if verbose:

print('climo for truth run:')

print('x mean =',xtruth_mean)

print('x stdev =',xtruth_stdev)

# forward operator.

# identity obs.

ndim = ensemble.n

h = np.eye(ndim)

# smoothing in forward operator

# gaussian or heaviside kernel.

if smooth_len > 0:

for j in range(ndim):

for i in range(ndim):

rr = float(i-j)

if i-j < -(ndim/2): rr = float(ndim-j+i)

if i-j > (ndim/2): rr = float(i-ndim-j)

r = np.fabs(rr)/smooth_len

if use_gaussian:

h[j,i] = np.exp(-r**2) # Gaussian

else: # running average (heaviside kernel)

if r <= 1:

h[j,i] = 1.

else:

h[j,i] = 0.

# normalize H so sum of weight is 1

h[j,:] = h[j,:]/h[j,:].sum()

obs = np.empty(xtruth.shape, xtruth.dtype)

for nt in range(xtruth.shape[0]):

obs[nt] = np.dot(h,xtruth[nt])

obs = obs + oberrstdev*rsens.standard_normal(size=obs.shape)

# spinup ensemble

ntot = xtruth.shape[0]

nspinup = ntstart

for n in range(ntstart):

ensemble.advance()

nsteps = np.int(dtassim/model.dt) # time steps in assimilation interval

if verbose:

print('ntstart, nspinup, ntot, nsteps =',ntstart,nspinup,ntot,nsteps)

if nsteps % 1 != 0:

raise(ValueError, 'assimilation interval must be an integer number of model time steps')

else:

nsteps = int(nsteps)

def ensrf(ensemble,xmean,xprime,h,obs,oberrvar,covlocal,method=1,z=None):

raise ValueError('illegal value for enkf method flag')

# define localization matrix.

covlocal = np.eye(ndim)

if corrl < 2*ndim:

for j in range(ndim):

for i in range(ndim):

rr = float(i-j)

if i-j < -(ndim/2): rr = float(ndim-j+i)

if i-j > (ndim/2): rr = float(i-ndim-j)

r = np.fabs(rr)/corrl; taper = 0.0

# Bohman taper (Gneiting 2002, doi:10.1006/jmva.2001.2056,

# equation, eq 21)

#if r < 1.:

# taper = (1.-r)*np.cos(np.pi*r) + np.sin(np.pi*r)/np.pi

# Gaspari-Cohn polynomial (Gneiting 2002, eq 23).

# Figure 3 of that paper compares Bohman and GC tapers.

rr = 2.*r

if r <= 0.5:

taper = ((( -0.25*rr +0.5 )*rr +0.625 )*rr -5.0/3.0 )*rr**2+1.

elif r > 0.5 and r < 1.:

taper = (((( rr/12.0 -0.5 )*rr +0.625 )*rr +5.0/3.0 )*rr -5.0 )*rr \

+ 4.0 - 2.0 / (3.0 * rr)

covlocal[j,i]=taper

# compute square root of covlocal

if method in [4,5,6,10]:

evals, eigs = np.linalg.eigh(covlocal)

evals = np.where(evals > 1.e-10, evals, 1.e-10)

evalsum = evals.sum(); neig = 0

frac = 0.0

while frac < thresh:

frac = evals[ndim-neig-1:ndim].sum()/evalsum

neig += 1

#print 'neig = ',neig

zz = (eigs*np.sqrt(evals/frac)).T

z = zz[ndim-neig:ndim,:]

else:

neig = 0

z = None

# run assimilation.

fcsterr = []

fcsterr1 = []

fcstsprd = []

analerr = []

analsprd = []

diverged = False

fsprdmean = np.zeros(ndim,np.float)

fsprdobmean = np.zeros(ndim,np.float)

asprdmean = np.zeros(ndim,np.float)

ferrmean = np.zeros(ndim,np.float)

aerrmean = np.zeros(ndim,np.float)

corrmean = np.zeros(ndim,np.float)

corrhmean = np.zeros(ndim,np.float)

for nassim in range(0,ntot,nsteps):

# assimilate obs

xmean = ensemble.x.mean(axis=0)

xmean_b = xmean.copy()

xprime = ensemble.x - xmean

# calculate background error, sprd stats.

ferr = (xmean - xtruth[nassim])**2

if np.isnan(ferr.mean()):

diverged = True

break

fsprd = (xprime**2).sum(axis=0)/(ensemble.members-1)

corr = (xprime.T*xprime[:,ndim//2]).sum(axis=1)/float(ensemble.members-1)

hxprime = np.dot(xprime,h)

fsprdob = (hxprime**2).sum(axis=0)/(ensemble.members-1)

corrh = (xprime.T*hxprime[:,ndim//2]).sum(axis=1)/float(ensemble.members-1)

if nassim >= nspinup:

fsprdmean = fsprdmean + fsprd

fsprdobmean = fsprdobmean + fsprdob

corrmean = corrmean + corr

corrhmean = corrhmean + corrh

ferrmean = ferrmean + ferr

fcsterr.append(ferr.mean()); fcstsprd.append(fsprd.mean())

fcsterr1.append(xmean - xtruth[nassim])

# update state estimate.

xmean,xprime =\

ensrf(ensemble,xmean,xprime,h,obs[nassim,:],oberrvar,covlocal,method=method,z=z)

# calculate analysis error, sprd stats.

aerr = (xmean - xtruth[nassim])**2

asprd = (xprime**2).sum(axis=0)/(ensemble.members-1)

if nassim >= nspinup:

asprdmean = asprdmean + asprd

aerrmean = aerrmean + aerr

analerr.append(aerr.mean()); analsprd.append(asprd.mean())

if verbose:

print(nassim,timetruth[nassim],np.sqrt(ferr.mean()),np.sqrt(fsprd.mean()),np.sqrt(aerr.mean()),np.sqrt(asprd.mean()))

if covinflate2 > 0:

# Hodyss and Campbell inflation.

inc = xmean - xmean_b

inf_fact = np.sqrt(covinflate1 + \

(asprd/fsprd**2)*((fsprd/ensemble.members) + covinflate2*(2.*inc**2/(ensemble.members-1))))

else:

# relaxation to prior spread inflation

asprd = np.sqrt(asprd); fsprd = np.sqrt(fsprd)

inf_fact = 1.+covinflate1*(fsprd-asprd)/asprd

xprime *= inf_fact

# run forecast model.

ensemble.x = xmean + xprime

for n in range(nsteps):

ensemble.advance()

# print out time mean stats.

# error and spread are normalized by observation error.

if diverged:

print(method,len(fcsterr),corrl,covinflate1,covinflate2,oberrstdev,np.nan,np.nan,np.nan,np.nan,neig)

else:

ncount = len(fcstsprd)

fcsterr = np.array(fcsterr)

fcsterr1 = np.array(fcsterr1)

fcstsprd = np.array(fcstsprd)

analerr = np.array(analerr)

analsprd = np.array(analsprd)

fstdev = np.sqrt(fcstsprd.mean())

astdev = np.sqrt(analsprd.mean())

asprdmean = asprdmean/ncount

aerrmean = aerrmean/ncount

fsprdmean = fsprdmean/ncount

fsprdobmean = fsprdobmean/ncount

corrmean = corrmean/ncount

corrhmean = corrhmean/ncount

fstd = np.sqrt(fsprdmean)

fstdob = np.sqrt(fsprdobmean)

covmean = corrmean; corrmean = corrmean/(fstd*fstd[ndim//2])

covhmean = corrhmean; corrhmean = corrhmean/(fstd*fstdob[ndim//2])

fcsterrcorr =\

(fcsterr1.T*fcsterr1[:,ndim//2]).sum(axis=1)/float(fcsterr1.shape[0]-1)

ferrstd = np.sqrt((fcsterr1**2).sum(axis=0)/float(fcsterr1.shape[0]-1))

errcovmean = fcsterrcorr

fcsterrcorr = fcsterrcorr/(ferrstd*ferrstd[ndim//2])

#import matplotlib.pyplot as plt

#plt.plot(np.arange(ndim),corrmean,color='k',label='r')

#plt.plot(np.arange(ndim),corrhmean,color='b',label='r (x vs hx)')

#plt.plot(np.arange(ndim),h[:,ndim//2]/h.max(),color='r',label='H')

#plt.plot(np.arange(ndim),covlocal[:,ndim//2],'k:',label='L')

#plt.xlim(0,ndim)

#plt.legend()

#plt.figure()

#plt.plot(np.arange(ndim),covmean,color='b',label='mean ens cov')

#plt.plot(np.arange(ndim),errcovmean,color='r',label='ens mean errcov')

#plt.xlim(0,ndim)

#plt.legend()

#plt.show()

print(method,ncount,corrl,covinflate1,covinflate2,oberrstdev,np.sqrt(fcsterr.mean()),fstdev,\

np.sqrt(analerr.mean()),astdev,neig)