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
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
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
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
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,
