Skip to content
Snippets Groups Projects
Commit 31baa4b2 authored by Stefano Serafin's avatar Stefano Serafin
Browse files

first commit

parents
No related branches found
No related tags found
No related merge requests found
Expand all Hide whitespace changes
Inline Side-by-side
.gitignore 0 → 100644
+ 4
− 0
View file @ 31baa4b2
__pycache__/
figures/
*.png
\ No newline at end of file
ENDA.py 0 → 100644
+ 0
− 0
View file @ 31baa4b2
This diff is collapsed.
PE_CBL.py 0 → 100644
+ 207
− 0
View file @ 31baa4b2
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import numpy as np
from matplotlib import pyplot as p
# Own modules
from models import CBL
from ENDA import experiment, parameter_transform_log, parameter_transform_none
from graphics import *
if __name__ == '__main__':
cbl_settings ={
# Physical parameters
'g' : 9.806,
'f' : 1e-4,
'kvonk' : 0.4,
'z0' : 0.1,
'ug' : 0,
'vg' : 10,
'theta_0' : 290,
'gamma' : 3e-3,
'Hmax' : 0.1,
'exct' : 0.3,
'pfac' : 1.5,
'Kmin' : 0.1,
'Kmax' : 200,
# Model formulation
'is_bwind' : False,
'is_cgrad' : True,
'is_cycle' : False,
# Domain parameters (Kmax determines dt)
'dz' : 50,
'ztop' : 4000,
'maxtime' : 10800,
# Ensemble generation parameters (only effective for ensemble runs)
# Part 1: perturbation of initial state
# perturbation_type can be "smooth", "random", "uniform"
# the adjective refers to vertical variability
'perturb_ensemble_state' : True,
'perturbations_type' : "uniform",
'perturbations_theta_amplitude' : .1,
'perturbations_uv_amplitude' : 0.1,
'perturbations_smooth_number' : 5,
'perturbations_symmetric' : True,
# Part 2: perturbation of parameters
'perturb_ensemble_parameters' : True,
'parameter_transform' : parameter_transform_none,
'parameter_number' : 3,
'parameter_ensemble_min' : np.array([1,1,0]),
'parameter_ensemble_max' : np.array([2,2,0.1]),
'parameter_true ' : np.array([1.5,0,0.1]),
# In case of parameter estimation
'do_parameter_estimation' : False,
'parameter_inflation_rtps_alpha' : np.array([1.,1.,1.0]),
}
CASE = 1 # deterministic run
CASE = 2 # ensemble run, no data assimilation
CASE = 3 # single assimilation
CASE = 4 # assimilation cycle
CASE = 5 # parameter identifiability
CASE = 6 # parameter estimation
#CASE = 0 # eddy diffusivity plots
if CASE == 1: # deterministic run
cbl = CBL(cbl_settings)
cbl.initialize(1)
cbl.run()
fig = p.figure(141)
fig.set_size_inches(3,4)
ax1 = fig.add_subplot(111)
ax1.plot(cbl.x0[:cbl.nz],cbl.zt)
ax1.plot(cbl.x[:cbl.nz],cbl.zt)
ax1.set_title(r'$\theta$ profile')
fig.savefig('run_deterministic.png',format='png',dpi=300)
p.close(fig)
elif CASE == 2: # ensemble run, no data assimilation
cbl = CBL(cbl_settings)
cbl.initialize(20)
cbl.run()
init_color = 'royalblue'
fin_color = 'orange'
nens = cbl.x.shape[1]
fig = p.figure(141)
fig.set_size_inches(9,4)
#
ax1 = fig.add_subplot(131)
ax1.plot(cbl.x0[:cbl.nz].mean(axis=1),cbl.zt,color=init_color)
ax1.plot(cbl.x[:cbl.nz].mean(axis=1),cbl.zt,color=fin_color)
ax1.set_title(r'$\theta$, ensemble mean')
#
ax2 = fig.add_subplot(132,sharey=ax1)
ax2.plot(cbl.x0[:cbl.nz].std(axis=1),cbl.zt,color=init_color)
ax2.plot(cbl.x[:cbl.nz].std(axis=1),cbl.zt,color=fin_color)
ax2.set_title(r'$\theta$, ensemble standard deviation')
#
ax3 = fig.add_subplot(133,sharey=ax1)
for i in range(nens):
ax3.plot(cbl.x0[:cbl.nz,i],cbl.zt,color=init_color,alpha=0.2)
ax3.plot(cbl.x[:cbl.nz,i],cbl.zt,color=fin_color,alpha=0.2)
ax3.set_title(r'$\theta$, full ensemble')
#
fig.savefig('run_ensemble.png',format='png',dpi=300)
p.close(fig)
elif CASE == 3: # single assimilation
da_settings = {'cbl_settings' : cbl_settings,
'tspinup' : 10800,
'trun' : 3600,
'assimilation_interval' : 3600,
'obs_coordinates' : np.array([433.]),
'obs_kinds' : np.array(['theta']),
'obs_error_sdev_generate' : np.array([.2]),
'obs_error_sdev_assimilate' : np.array([1.]),
'nens' : 50,
'FILTER' : 'EAKF',
'inflation_rtps_alpha' : 0.1,
}
exp = experiment(da_settings)
plot_CBL_assimilation(exp,'single_assimilation.png')
elif CASE == 4: # assimilation cycle
nobs = 20
da_settings = {'cbl_settings' : cbl_settings,
'tspinup' : 3600,
'trun' : 10800,
'assimilation_interval' : 3600,
'obs_coordinates' : np.linspace(0,2000,nobs),
'obs_kinds' : np.array(['theta']*nobs),
'obs_error_sdev_generate' : np.ones(nobs)*.1,
'obs_error_sdev_assimilate' : np.ones(nobs)*.1,
'nens' : 20,
'FILTER' : 'EAKF',
'inflation_rtps_alpha' : 0.9,
}
exp = experiment(da_settings)
plot_CBL_assimilation(exp,'profile_assimilation.png',which_cycle=exp.da.ncycles)
elif CASE == 5: # parameter identifiability
cbl_settings['perturbations_theta_amplitude'] = .1
cbl_settings['perturbations_type'] = "uniform"
cbl_settings['perturb_ensemble_state'] = True
cbl_settings['perturb_ensemble_parameters'] = True
cbl_settings['perturbations_symmetric'] = True
cbl_settings['parameter_number'] = 1
cbl_settings['parameter_transform'] = [ parameter_transform_log ]
cbl_settings['parameter_ensemble_min'] = np.array([1])
cbl_settings['parameter_ensemble_max'] = np.array([4])
cbl_settings['parameter_true'] = np.array([1.5])
cbl_settings['do_parameter_estimation'] = False
cbl = CBL(cbl_settings)
cbl.initialize(20)
cbl.run(output_full_history=True)
plot_CBL_identifiability(cbl,'pfac_identifiability.png')
elif CASE == 6: # parameter estimation
cbl_settings['perturbations_theta_amplitude'] = .1
cbl_settings['perturbations_type'] = "uniform"
cbl_settings['perturb_ensemble_state'] = True
cbl_settings['perturb_ensemble_parameters'] = True
cbl_settings['perturbations_symmetric'] = True
cbl_settings['parameter_number'] = 1
cbl_settings['parameter_transform'] = [ parameter_transform_log ]
cbl_settings['parameter_ensemble_min'] = np.array([1])
cbl_settings['parameter_ensemble_max'] = np.array([4])
cbl_settings['parameter_true'] = np.array([1.5])
cbl_settings['do_parameter_estimation'] = True
cbl_settings['parameter_inflation_rtps_alpha'] = np.array([1.0])
nobs = 10
da_settings = {'cbl_settings' : cbl_settings,
'tspinup' : 3600,
'trun' : 7200,
'assimilation_interval' : 3.125*48,
'obs_coordinates' : np.linspace(0,1500,nobs),
'obs_kinds' : np.array(['theta']*nobs),
'obs_error_sdev_generate' : np.ones(nobs)*.1,
'obs_error_sdev_assimilate' : np.ones(nobs)*.1,
'localization_cutoff' : 800,
'nens' : 30,
'FILTER' : 'EAKF',
'inflation_rtps_alpha' : 0.8,
}
exp = experiment(da_settings)
plot_CBL_assimilation(exp,'profile_assimilation.png',which_cycle=exp.da.ncycles)
plot_CBL_PE(exp,'parameter_estimation.png')
elif CASE == 0: # plot K profile examples
p_factors = [0.2,0.5,1,2,4,6]
theta_profiles = []
cbl_settings['dz'] = 25
cbl_settings['ztop'] = 2000
for pfac in p_factors:
cbl_settings['pfac'] = pfac
cbl = CBL(cbl_settings)
cbl.initialize(1)
cbl.run()
theta_profiles.append(cbl.x[:cbl.nz])
plot_p(p_factors,theta_profiles,cbl.zt,'Kprofiles.png')
README.md 0 → 100644
+ 14
− 0
View file @ 31baa4b2
# Tools for idealized parameter estimation experiments
* `PE_CBL.py`
Main driver program for parameter estimation with a convective boundary layer model. Includes a few examples.
* `models.py`
Model code. At the moment, it only includes a CBL model.
* `ENDA.py`
Data assimilation code. Includes implementations of the EAKF and the LETKF. Also includes classes for data assimilation cycles, diagnostics and experiments (nature run + cycle + diagnostics).
* `graphics.py`
Includes all the code for generating figures.
\ No newline at end of file
graphics.py 0 → 100644
+ 256
− 0
View file @ 31baa4b2
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import numpy as np
from matplotlib import pyplot as p
def plot_CBL_assimilation(exp,figtitle,which_cycle=1):
naturerun_color = 'black'
ensmembers_color = 'lightskyblue'
ensmean_color = 'royalblue'
def ensplot(ax,x,nr,x1,x2,z):
nens = x.shape[1]
ax.plot(nr[x1:x2],zt,color=naturerun_color,zorder=2)
ax.plot(x[x1:x2].mean(axis=1),z,color=ensmean_color,zorder=5)
for i in range(nens): ax.plot(x[x1:x2,i],z,color=ensmembers_color,alpha=0.2,zorder=1)
# Shorthand
nr = exp.nr
o = exp.observations[which_cycle-1]
b = exp.da.backgrounds[which_cycle-1,:,:]
a = exp.da.analyses[which_cycle-1,:,:]
ocoords = exp.obs_coordinates
nz = exp.nr.nz
zt = exp.nr.zt
if exp.nr.is_bwind:
ii=3
else:
ii=1
zmax = 2000
fig = p.figure(151)
fig.set_size_inches(9,3*ii)
#
# Potential temperature
#
ax1 = fig.add_subplot(ii,4,1)
ax1.plot(nr.x[:nz],zt,color=naturerun_color)
ax1.scatter(o,ocoords,marker="o",s=20,color=naturerun_color,zorder=10)
ax1.set_title(r'$\theta$ (K), nature run')
ax1.set_ylabel(r'$z$ (m)')
ax1.set_ylim([0,zmax])
#
ax2 = fig.add_subplot(ii,4,2,sharey=ax1)
ensplot(ax2,b,nr.x,0,nz,zt)
ax2.scatter(o,ocoords,marker="o",s=20,color=naturerun_color,zorder=10)
ax2.set_title(r'$\theta$ (K), prior')
#
ax3 = fig.add_subplot(ii,4,3,sharey=ax1)
ensplot(ax3,a,nr.x,0,nz,zt)
ax3.scatter(o,ocoords,marker="o",s=20,color=naturerun_color,zorder=10)
ax3.set_title(r'$\theta$ (K), posterior')
#
ax4 = fig.add_subplot(ii,4,4,sharey=ax1)
ax4.plot(a[:nz].mean(axis=1)-b[:nz].mean(axis=1),zt,color=ensmean_color)
ax4.set_title(r'$\theta$ (K), mean increment',zorder=10)
#
# Fix axis labels
#
for ax in [ax2,ax3,ax4]:
p.setp(ax.get_yticklabels(), visible=False)
#
if exp.nr.is_bwind:
#
# u wind component
#
ax5 = fig.add_subplot(ii,4,5)
ax5.plot(nr.x[nz:nz*2],zt,color=naturerun_color)
ax5.set_title(r'$u$ (m/s), nature run')
ax5.set_ylabel(r'$z$ (m)')
ax5.set_ylim([0,zmax])
#
ax6 = fig.add_subplot(ii,4,6,sharey=ax1)
ensplot(ax6,b,nr.x,nz,nz*2,zt)
ax6.set_title(r'$u$ (m/s), prior')
#
ax7 = fig.add_subplot(ii,4,7,sharey=ax1)
ensplot(ax7,a,nr.x,nz,nz*2,zt)
ax7.set_title(r'$u$ (m/s), posterior')
#
ax8 = fig.add_subplot(ii,4,8,sharey=ax1)
ax8.plot(a[nz:nz*2].mean(axis=1)-b[nz:nz*2].mean(axis=1),zt,color=ensmean_color)
ax8.set_title(r'$u$ (m/s), mean increment',zorder=10)
#
# v wind component
#
ax9 = fig.add_subplot(ii,4,9)
ax9.plot(nr.x[nz*2:nz*3],zt,color=naturerun_color)
ax9.set_title(r'$v$ (m/s), nature run')
ax9.set_ylabel(r'$z$ (m)')
ax9.set_ylim([0,zmax])
#
axa = fig.add_subplot(ii,4,10,sharey=ax1)
ensplot(axa,b,nr.x,nz*2,nz*3,zt)
axa.set_title(r'$v$ (m/s), prior')
#
axb = fig.add_subplot(ii,4,11,sharey=ax1)
ensplot(axb,a,nr.x,nz*2,nz*3,zt)
axb.set_title(r'$v$ (m/s), posterior')
#
axc = fig.add_subplot(ii,4,12,sharey=ax1)
axc.plot(a[nz*2:nz*3].mean(axis=1)-b[nz*2:nz*3].mean(axis=1),zt,color=ensmean_color)
axc.set_title(r'$v$ (m/s), mean increment',zorder=10)
#
# Fix axis labels
#
for ax in [ax6,ax7,ax8,axa,axb,axc]:
p.setp(ax.get_yticklabels(), visible=False)
#
fig.savefig(figtitle,format='png',dpi=300)
p.close(fig)
def plot_CBL_PE(exp,figtitle,parameter_id=0,plot_spread=False):
if exp.nr.do_parameter_estimation:
spread_color = 'lightskyblue'
mean_color = 'royalblue'
mean_color_2 = 'orange'
if parameter_id==0:
true_value = exp.nr.pfac
pt = exp.nr.parameter_transform
parameter_number = exp.nr.parameter_number
b = exp.da.backgrounds[:,-parameter_number:,:][:,parameter_id,:]
par_b = pt[parameter_id](exp.da.backgrounds[:,-parameter_number:,:][:,parameter_id,:],kind='inv')
t = exp.da.time
if plot_spread:
ii=2
else:
ii=1
fig = p.figure(151)
fig.set_size_inches(4,3*ii)
#
ax1 = fig.add_subplot(ii,1,1)
ax1.step(t,par_b.mean(axis=1),color=mean_color)
for i in range(t.size-1):
bmin=(par_b.mean(axis=1)-par_b.std(axis=1))[i+1]
bmax=(par_b.mean(axis=1)+par_b.std(axis=1))[i+1]
ax1.fill_between(t[i:i+2],\
y1=[bmin,bmin],\
y2=[bmax,bmax],\
color=spread_color,
edgecolor=None,
alpha=0.3)
ax1.axhline(y=true_value,linestyle='--',color='black')
ax1.set_title('Parameter evolution')
ax1.set_xlabel('Time (s)')
#
if plot_spread:
ax2 = fig.add_subplot(ii,1,2)
ax2.plot(t,par_b.std(axis=1),color=mean_color,label='physical')
ax2.plot(t,b.std(axis=1),color=mean_color_2,label='transformed', linestyle='--')
ax2.legend()
ax2.set_title('Parameter spread')
#
fig.tight_layout()
fig.savefig(figtitle,format='png',dpi=300)
p.close(fig)
def plot_CBL_identifiability(cbl,figtitle):
# Read relevant dimensions
times = cbl.history['0000']['time']
zt = cbl.zt
nz = cbl.nz
ntimes = times.size
nens = cbl.nens
# Reconstruct_history
theta = np.zeros((nens,nz,ntimes))+np.nan
for k in range(nens):
theta[k,:,:]= cbl.history['%04u'%k]['theta']
# Read parameter values
if hasattr(cbl,'initial_perturbed_parameters'):
pfac = cbl.parameter_transform[0](cbl.initial_perturbed_parameters[0],kind='inv')
else:
pfac = np.ones(nens)*cbl.pfac
# Compute correlation coefficient
# (a little noise is added to the denominator to avoid singularities)
p_x_correlation = np.zeros((nz,ntimes))+np.nan
for i in range(nz):
for j in range(ntimes):
p_x_correlation[i,j] = np.cov(pfac,theta[:,i,j])[0,1] / (1e-9+np.sqrt(np.var(pfac)*np.var(theta[:,i,j])))
# Make plot
fig = p.figure(151)
fig.set_size_inches(9,3)
zmax = 2000
ncont = 13
#
ax1 = fig.add_subplot(1,3,1)
c1 = ax1.pcolormesh(times,zt,theta.mean(axis=0),vmin=cbl.theta_0,vmax=cbl.theta_0+cbl.gamma*zmax)
ax1.contour(times,zt,theta.mean(axis=0),
np.linspace(cbl.theta_0,cbl.theta_0+cbl.gamma*zmax,ncont),
colors='white',
linestyles='--',
linewidths=0.75)
ax1.set_ylim([0,zmax])
ax1.set_xlabel(r'time')
ax1.set_ylabel(r'$z$ (m)')
ax1.set_title(r'ensemble mean $\theta$ (K)')
p.colorbar(c1)
#
ax2 = fig.add_subplot(1,3,2)
c2 = ax2.pcolormesh(times,zt,theta.std(axis=0))
ax2.set_ylim([0,zmax])
ax2.set_xlabel(r'time')
ax2.set_ylabel(r'$z$ (m)')
ax2.set_title(r'ensemble sdev $\theta$ (K)')
p.colorbar(c2)
#
ax3 = fig.add_subplot(1,3,3,sharey=ax1)
c3 = ax3.pcolormesh(times,zt,p_x_correlation,vmin=-1,vmax=1,cmap='RdBu_r')
ax3.set_xlabel(r'time')
ax3.set_title(r'$p$-$\theta$ correlation')
p.colorbar(c3)
#
p.setp(ax2.get_yticklabels(), visible=False)
p.setp(ax3.get_yticklabels(), visible=False)
fig.tight_layout()
fig.savefig(figtitle,format='png',dpi=300)
p.close(fig)
def plot_p(p_factors,theta_profiles,zt,figtitle):
zoverh = np.linspace(0,1,101)
fig = p.figure(151)
fig.set_size_inches(6,3)
#
ax1 = fig.add_subplot(1,2,1)
for pfac in p_factors:
Koverkws = zoverh*(1-zoverh)**pfac
ax1.plot(Koverkws,zoverh,label='$p=%4.1f$'%pfac)
ax1.set_xlabel('$K/(\kappa w_s z_i)$')
ax1.set_ylabel('$z/z_i$')
#
ax2 = fig.add_subplot(1,2,2)
for i in range(len(p_factors)):
ax2.plot(theta_profiles[i],zt,label='$p=%4.1f$'%p_factors[i])
ax2.set_xlabel(r'$\theta$ (K)')
ax2.set_ylabel(r'$z$ (m)')
ax2.set_ylim([0,1500])
ax2.set_xlim([291,297])
ax2.legend(loc=4)
#
fig.tight_layout()
fig.savefig(figtitle,format='png',dpi=300)
p.close(fig)
models.py 0 → 100644
+ 398
− 0
View file @ 31baa4b2
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import numpy as np
from scipy.interpolate import CubicSpline
verbose = True
class CBL:
def __init__(self,settings):
for k, v in settings.items():
setattr(self, k, v)
def initialize(self,argin,coordinates=None):
"""
if input is scalar, interpret it as an ensemble size
if input is a 1D or 2D array, interpret it as an initial state
"""
self.zf = np.linspace(0,self.ztop,int(self.ztop/self.dz)+1)
self.zt = 0.5*(self.zf[:-1]+self.zf[1:])
self.nz = self.zt.size
self.dt = 0.25*self.dz*self.dz/self.Kmax
# scalar input is nens
if np.isscalar(argin):
self.nens = argin
assert self.nens > 0, f"ensemble size must be greater than zero"
# Define base state profile
theta_init = self.theta_0+self.gamma*self.zt
# For forced convection
if self.is_bwind:
u_init = np.zeros(self.nz)+self.ug
v_init = np.zeros(self.nz)+self.vg
x_init = np.zeros(self.nz*3)+np.nan
x_init[:self.nz] = theta_init
x_init[self.nz:self.nz*2] = u_init
x_init[self.nz*2:self.nz*3] = v_init
coordinates = np.zeros(self.nz*3)+np.nan
coordinates[:self.nz] = self.zt
coordinates[self.nz:self.nz*2] = self.zt
coordinates[self.nz*2:self.nz*3] = self.zt
else:
x_init = theta_init
coordinates = self.zt
# If desired, append parameters to state vector;
# initialize with the value in the middle of the admitted range
# Coordinates of parameters are set to NaN
if self.do_parameter_estimation:
if self.nens > 1:
for k in range(-self.parameter_number,0):
pmean = 0.5*(self.parameter_ensemble_min[k]+self.parameter_ensemble_max[k])
x_init = np.append(x_init,self.parameter_transform[k](pmean,kind='dir'))
else:
for k in range(-self.parameter_number,0):
pmean = self.parameter_true[k]
x_init = np.append(x_init,self.parameter_transform[k](pmean,kind='dir'))
coordinates = np.append(coordinates,self.parameter_ensemble_min*np.nan)
# Replicate base state profile nens times
x0 = np.tile(x_init,(self.nens,1)).T
# 1D-array input is initial state of a deterministic run
elif argin.ndim==1:
self.nens = 1
state_vector_length = self.nz
if self.is_bwind:
state_vector_length += 2*self.nz
if self.do_parameter_estimation:
state_vector_length += self.parameter_number
assert argin.size == state_vector_length, f"1D input array should have size = nz*nvars+parameter_number, got {argin.shape[0]}"
x0 = argin
assert coordinates is not None, f"you must provide a valid coordinate array"
assert x0.size == coordinates.size, f"coordinate vector must have same size as state vector"
# 2D-array input is initial ensemble state
elif argin.ndim==2:
assert coordinates is not None, f"you must provide a valid coordinate array"
state_vector_length = self.nz
if self.is_bwind:
state_vector_length += 2*self.nz
if self.do_parameter_estimation:
state_vector_length += self.parameter_number
assert argin.shape[0] == state_vector_length, f"First dimension of 2D input array should be nz*nvars+parameter_number, got {argin.shape[0]}"
self.nens = argin.shape[1]
x0 = argin
assert x0.shape[0] == coordinates.size, f"coordinate vector must have same size as state vector"
# If desired, add ensemble perturbations to model state
if self.nens>1 and self.perturb_ensemble_state:
# Random or uniform perturbations perturbations
# differ only by magnitude of random sample
if self.perturbations_type == "random":
randomsize = self.nz
elif self.perturbations_type == "uniform":
randomsize = 1
if self.perturbations_type == "random" or self.perturbations_type == "uniform":
ppt = np.random.randn(randomsize,self.nens)
if self.is_bwind:
ppu = np.random.randn(self.nens,randomsize)
ppv = np.random.randn(self.nens,randomsize)
# Smooth perturbations are slightly more complicated
if self.perturbations_type == "smooth":
# Define positions of perturbations
# enforce two at boundaries to avoid extrapolation
ipert = np.arange(0,self.perturbations_smooth_number-1)*(self.nz//(self.perturbations_smooth_number))
ipert = np.append(ipert,self.nz-1)
# Draw random perturbations, then interpolate
randomsize = self.perturbations_smooth_number
ppt = np.zeros((self.nz,self.nens))+np.nan
pert_t = np.random.randn(randomsize,self.nens)
for n in range(self.nens):
f = CubicSpline(ipert,pert_t[:,n])
ppt[:,n] = f(np.arange(self.nz))
if self.is_bwind:
ppu = np.zeros((self.nz,self.nens))+np.nan
ppv = np.zeros((self.nz,self.nens))+np.nan
pert_u = np.random.randn(randomsize,self.nens)
pert_v = np.random.randn(randomsize,self.nens)
for n in range(self.nens):
f = CubicSpline(ipert,pert_u[:,n])
ppu[:,n] = f(np.arange(self.nz))
f = CubicSpline(ipert,pert_v[:,n])
ppv[:,n] = f(np.arange(self.nz))
if self.perturbations_symmetric:
# if n is odd, the last member remains unperturbed
for n in range(self.nens//2):
x0[:self.nz,n] += self.perturbations_theta_amplitude*ppt[:,n]
x0[:self.nz,self.nens//2+n] -= self.perturbations_theta_amplitude*ppt[:,n]
if self.is_bwind:
x0[self.nz:self.nz*2,n] += self.perturbations_uv_amplitude*ppu[:,n]
x0[self.nz:self.nz*2,self.nens//2+n] -= self.perturbations_uv_amplitude*ppu[:,n]
x0[self.nz*2:self.nz*3,n] += self.perturbations_uv_amplitude*ppv[:,n]
x0[self.nz*2:self.nz*3,self.nens//2+n] -= self.perturbations_uv_amplitude*ppv[:,n]
else:
for n in range(self.nens):
x0[:self.nz,n] += self.perturbations_theta_amplitude*ppt[:,n]
if self.is_bwind:
x0[self.nz:self.nz*2,n] += self.perturbations_uv_amplitude*ppu[:,n]
x0[self.nz*2:self.nz*3,n] += self.perturbations_uv_amplitude*ppv[:,n]
# If desired, perturb parameters too
# (last parameter_number elements of state vector)
# Overwrite the pre-existing unperturbed values
# This is needed for safety, because parameter transformations may be nonlinear
if self.nens>1:
if self.perturb_ensemble_parameters or self.do_parameter_estimation:
perturbed_parameters = self.create_parameter_perturbations()
self.initial_perturbed_parameters = perturbed_parameters
if self.do_parameter_estimation:
assert (self.parameter_number==self.parameter_ensemble_min.size==self.parameter_ensemble_max.size),\
f'parameter min and max values must match parameters number,\
got {self.parameter_ensemble_min.size,self.parameter_ensemble_max.size,self.parameter_number}'
for k in range(-self.parameter_number,0):
assert (self.parameter_ensemble_min[k]<=self.parameter_ensemble_max[k]),\
f'parameter min must be less or equal to max,\
got {self.parameter_ensemble_min[k],self.parameter_ensemble_max[k]}'
x0[-self.parameter_number:,:] = perturbed_parameters
# Squeeze ensemble dimension if needed
if self.nens==1:
x0 = x0.squeeze()
# Store initial state
self.x0 = x0
self.state_coordinates = coordinates
def update(self,x0):
self.x0 = x0
def run(self,output_full_history=False):
self.nt = int(self.maxtime//self.dt)+1
if self.nens == 1:
if verbose: print('#',end='')
if output_full_history:
self.x,self.history = self.cbl_model(self.x0,output_full_history=True)
else:
self.x = self.cbl_model(self.x0)
if verbose: print(' ...done')
else:
self.x = np.zeros(self.x0.shape) + np.nan
if self.perturb_ensemble_parameters and not self.do_parameter_estimation:
pp = self.initial_perturbed_parameters
else:
pp = np.ones((1,self.nens))*np.nan
if output_full_history:
self.history = {}
for n in range(self.nens):
if verbose: print('#',end='')
self.x[:,n],self.history['%04u'%n] = self.cbl_model(self.x0[:,n],perturbed_parameters=pp[:,n],output_full_history=True)
else:
for n in range(self.nens):
if verbose: print('#',end='')
self.x[:,n] = self.cbl_model(self.x0[:,n],perturbed_parameters=pp[:,n])
if verbose: print(' ...done')
def create_parameter_perturbations(self):
# Initialize array to store perturbations
pp = np.zeros((self.parameter_number,self.nens))
if self.perturbations_symmetric:
for k in range(-self.parameter_number,0):
mean = 0.5*(self.parameter_ensemble_min[k]+self.parameter_ensemble_max[k])
dum1 = np.random.uniform(self.parameter_ensemble_min[k], self.parameter_ensemble_max[k], size=self.nens//2)
dum2 = 2*mean-dum1
dum = np.ones(self.nens)*mean
dum[:self.nens//2] = dum1
dum[self.nens//2:] = dum2
pp[k,:] = self.parameter_transform[k](dum , kind='dir')
else:
for k in range(-self.parameter_number,0):
dum = np.random.uniform(self.parameter_ensemble_min[k], self.parameter_ensemble_max[k], size=self.nens)
pp[k,:] = self.parameter_transform[k](dum , kind='dir')
return pp
def cbl_model(self,x0,output_full_history=False,perturbed_parameters=np.nan):
# Read settings in
nt = self.nt
dt = self.dt
nz = self.nz
dz = self.dz
zf = self.zf
zt = self.zt
Kmax = self.Kmax
Kmin = self.Kmin
g = self.g
f = self.f
kvonk = self.kvonk
z0 = self.z0
theta_0 = self.theta_0
Hmax = self.Hmax
exct = self.exct
ug = self.ug
vg = self.vg
pfac = self.pfac
is_bwind = self.is_bwind
is_cgrad = self.is_cgrad
is_cycle = self.is_cycle
parameter_transform = self.parameter_transform
if hasattr(self,'do_parameter_estimation'):
do_parameter_estimation = self.do_parameter_estimation
parameter_number = self.parameter_number
else:
do_parameter_estimation = False
parameter_number = 0
# Initialize values and preallocate empty arrays
xb = np.zeros(x0.shape)+np.nan
theta = np.zeros((nz,nt))+np.nan
H = np.zeros((nz+1,nt))+np.nan
# For forced convection
if is_bwind:
u = np.zeros((nz,nt))+np.nan
v = np.zeros((nz,nt))+np.nan
M = np.zeros((nz+1,nt))+np.nan
N = np.zeros((nz+1,nt))+np.nan
# Sensible first guesses for PBL depth and eddy diffusivity
h = dz*np.ones(nt)
K = np.ones((nz+1,nt))*Kmin
# Optionally, read perturbed parameters
# This is not effective in the case of parameter estimation
if not do_parameter_estimation and ~np.isnan(perturbed_parameters):
pfac = parameter_transform[0](perturbed_parameters[0],kind='inv')
# Unpack state vector
theta[:,0] = x0[:nz]
if is_bwind:
u[:,0] = x0[nz:2*nz]
v[:,0] = x0[2*nz:3*nz]
if do_parameter_estimation:
parameters = x0[-parameter_number:]
pfac = parameter_transform[0](parameters[0],kind='inv')
# Store initial state in temporary variable
thetap = theta[:,0]
if is_bwind:
up = u[:,0]
vp = v[:,0]
# Integrate forward in time (ok for diffusion equation!)
for j in range(1,nt):
# Find model level nearest to h (parcel method)
thr = thetap[0]+exct
theta_profile = thetap[:]
kh = theta_profile[theta_profile<=thr].size-1
# Interpolate h
h[j] = np.maximum(np.interp(thr,thetap[kh:kh+2],zt[kh:kh+2]),0.)
# Set the surface sensible heat flux (H0)
if is_cycle:
H0 = Hmax*np.sin(2.*np.pi/86400.*dt*(j-1))
else:
H0 = Hmax
# Set the surface momentum flux (ustar)
ustar = 0
# For forced convection:
# - you need surface momentum fluxes to integrate the momentum equations
# - therefore ustar is not zero
if is_bwind:
Cd = (kvonk/np.log(zt[1]/z0))**2
spd = (up**2+vp**2)**0.5
upwp = -Cd*up[0]*spd[0]
vpwp = -Cd*vp[0]*spd[0]
ustar = (upwp**2+vpwp**2)**0.25
# Iteration
Cd_iterative = False
if Cd_iterative:
# First guess
Cd = (kvonk/np.log(zt[1]/z0))**2
delta = 1e6
Cdmin = 1e-5
while delta>1e-6 and Cd>Cdmin:
upwp = -Cd*up[0]*spd[0]
vpwp = -Cd*vp[0]*spd[0]
ustar = (upwp**2+vpwp**2)**0.25
zeta = -kvonk*zf[1]*g*H0/(theta_0*ustar**3)
xx = (1-15*zeta)**0.25
psim = np.log((1+xx**2)/2)+2*np.log((1+xx)/2)-2*np.arctan(xx)+np.pi/2.
psim = np.minimum(500,psim)
dummy = (kvonk/(np.log(zf[1]/z0)-psim))**2
delta = dummy-Cd
Cd = np.maximum(Cdmin,dummy)
# Define the eddy viscosity profile
wstar = (g/theta_0*h[j]*np.maximum(H0,0.001))**0.3333
ws = (ustar**3+0.7*kvonk*wstar**3)**0.3333
zz = (zf<h[j-1])*zf
K[:,j] = np.minimum(kvonk*ws*zz*(1-zz/h[j-1])**pfac,Kmax)
# Determine the countergradient correction
if is_cgrad:
gammac = 6.5*H0/ws/h[j]
else:
gammac = 0
# Inverse grid spacing
rdz = 1./self.dz
# Compute sensible heat flux and integrate T equation
H[0,j] = H0
H[1:-1,j] = -K[1:-1,j]*( (thetap[1:]-thetap[:-1])*rdz - gammac)
H[nz,j] = 2*H[nz-1,j]-H[nz-2,j]
theta[:,j] = thetap[:]-dt*rdz*(H[1:,j]-H[:-1,j])
thetap = theta[:,j]
# For forced convection
if is_bwind:
# Compute U momentum flux and integrate U equation
M[0,j] = upwp
M[1:-1,j] = -K[1:-1,j]*(up[1:]-up[:-1])*rdz
M[nz,j] = 2*M[nz-1,j]-M[nz-2,j]
u[:,j] = up[:]-dt*rdz*(M[1:,j]-M[:-1,j])+dt*f*(vp[:]-vg)
up = u[:,j]
# Compute V momentum flux and integrate V equation
N[0,j] = vpwp
N[1:-1,j] = -K[1:-1,j]*(vp[1:]-vp[:-1])*rdz
N[nz,j] = 2*N[nz-1,j]-N[nz-2,j]
v[:,j] = vp[:]-dt*rdz*(N[1:,j]-N[:-1,j])-dt*f*(up[:]-ug)
vp = v[:,j]
# Form state vector
xb[:nz] = theta[:,-1]
if is_bwind:
xb[nz:2*nz] = u[:,-1]
xb[2*nz:3*nz] = v[:,-1]
if do_parameter_estimation:
parameters = parameter_transform[0](pfac,kind='dir')
xb[-parameter_number:] = parameters
# Return state vector and, optionally, history of variables
if is_bwind:
history={'time' : dt*np.arange(0,nt), 'theta' : theta, 'u' : u, 'v': v}
else:
history={'time' : dt*np.arange(0,nt), 'theta' : theta}
if output_full_history:
return xb,history
else:
return xb
\ No newline at end of file
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment