class Thermalwind(Param):
""" Thermal wind model
It provides the step(t,dt) function
and 'var' containing the mode state
"""
def __init__(self, param, grid):
self.list_param = ['forcing', 'noslip', 'timestepping']
param.copy(self, self.list_param)
# for variables
param.varname_list = ('vorticity', 'psi', 'u', 'v', 'buoyancy', 'V')
param.sizevar = [grid.nyl, grid.nxl]
self.var = Var(param)
# for operators
param.tracer_list = ['vorticity', 'buoyancy', 'V']
param.whosetspsi = ('vorticity')
self.ope = Operators(param, grid)
# for timescheme
self.tscheme = Timescheme(param, self.var.state)
def step(self, t, dt):
# 1/ integrate advection
self.tscheme.set(self.dynamics, self.timestepping)
self.tscheme.forward(self.var.state, t, dt)
# 2/ integrate source
if self.forcing or self.noslip:
self.tscheme.set(self.sources, 'EF')
self.tscheme.forward(self.var.state, t, dt)
def dynamics(self, x, t, dxdt):
self.ope.rhs_adv(x, t, dxdt)
self.ope.rhs_thermalwind(x, t, dxdt) # add the r.h.s. terms
self.ope.invert_vorticity(dxdt, flag='fast')
def sources(self, x, t, dxdt):
if self.forcing:
self.ope.rhs_forcing(x, t, dxdt)
self.ope.invert_vorticity(dxdt, flag='fast')
if self.noslip:
self.ope.rhs_noslip(x, t, dxdt)
self.ope.invert_vorticity(x, flag='full')
def set_psi_from_vorticity(self):
self.ope.invert_vorticity(self.var.state)
class LES(object):
def __init__(self, param):
# Create a State object for the LES model
self.state = get_state(param)
# Make U and all its components prognostic. This should
# actually be done before the state is created, but I don't want
# to mess with variables.py at the moment. This must be done
# before the timescheme is created.
self.state.U.prognostic = True
for i in "ijk":
self.state.U[i].prognostic = True
# Create and set-up a handler for the timescheme
self.timescheme = Timescheme(param, self.state)
self.timescheme.set(compute_rhs)
# TODO: also inform the timescheme about the function to compute diagnostic variables
def forward(self, t, dt):
self.timescheme.forward(self.state, t, dt)
class Thermalwind(Param):
""" Thermal wind model
It provides the step(t,dt) function
and 'var' containing the mode state
"""
def __init__(self, param, grid):
self.list_param = [
'forcing', 'noslip', 'timestepping', 'forcing_module',
'additional_tracer', 'myrank', 'gravity', 'diffusion', 'Kdiff',
'f0'
]
param.copy(self, self.list_param)
# for potential energy
self.list_param = [
'xr', 'yr', 'nh', 'Lx', 'msk', 'area', 'mpitools', 'dx'
]
grid.copy(self, self.list_param)
# for variables
param.varname_list = [
'vorticity', 'psi', 'u', 'v', 'buoyancy', 'V', 'qE'
]
param.tracer_list = ['vorticity', 'buoyancy', 'V']
param.whosetspsi = ('vorticity')
if hasattr(self, 'additional_tracer'):
for k in range(len(self.additional_tracer)):
trac = self.additional_tracer[k]
param.varname_list.append(trac)
param.tracer_list.append(trac)
param.sizevar = [grid.nyl, grid.nxl]
self.var = Var(param)
# for operators
self.ope = Operators(param, grid)
# for timescheme
self.tscheme = Timescheme(param, self.var.state)
if self.forcing:
try:
f = import_module(self.forcing_module)
except:
print('module %s for forcing cannot be found' %
self.forcing_module)
print('make sure file **%s.py** exists' % self.forcing_module)
exit(0)
self.forc = f.Forcing(param, grid)
self.diags = {}
def step(self, t, dt):
# 1/ integrate advection
self.tscheme.set(self.dynamics, self.timestepping)
self.tscheme.forward(self.var.state, t, dt)
# 2/ integrate source
if self.forcing or self.noslip:
self.tscheme.set(self.sources, 'EF')
self.tscheme.forward(self.var.state, t, dt)
self.compute_pv()
def compute_pv(self):
qE = self.var.get('qE')
v = self.var.get('V')
b = self.var.get('buoyancy')
# Ertel PV
y = qE * 0.
y[1:-1, 1:-1] = self.ope.jacobian(v + self.f0 * self.xr, b)
y *= self.msk
self.ope.fill_halo(y)
qE[:, :] = y
def dynamics(self, x, t, dxdt):
self.ope.rhs_adv(x, t, dxdt)
self.ope.rhs_thermalwind(x, t, dxdt) # add the r.h.s. terms
if (self.tscheme.kstage == self.tscheme.kforcing):
if self.forcing:
self.forc.add_forcing(x, t, dxdt)
if self.diffusion:
self.ope.rhs_diffusion(x, t, dxdt)
self.ope.invert_vorticity(dxdt, flag='fast')
def sources(self, x, t, dxdt):
if self.noslip:
self.ope.rhs_noslip(x, t, dxdt)
self.ope.invert_vorticity(x, flag='full')
def set_psi_from_vorticity(self):
self.ope.invert_vorticity(self.var.state)
def diagnostics(self, var, t):
""" should provide at least 'maxspeed' (for cfl determination) """
nh = self.nh
u = var.get('u')
v = var.get('v')
vort = var.get('vorticity')
buoy = var.get('buoyancy')
V = var.get('V')
qE = var.get('qE')
qEneg = qE.copy()
qEneg[qE > 0] = 0.
ke, maxu = fd.computekemaxu(self.msk, u, v, self.nh)
z, z2 = fd.computesumandnorm(self.msk, vort, self.nh)
b, b2 = fd.computesumandnorm(self.msk, buoy, self.nh)
vm, v2 = fd.computesumandnorm(self.msk, V, self.nh)
q, q2 = fd.computesumandnorm(self.msk, qE, self.nh)
qn, qn2 = fd.computesumandnorm(self.msk, qEneg, self.nh)
# potential energy (minus sign because buoyancy is minus density)
pe = fd.computesum(self.msk, buoy * self.yr, nh)
pe = -self.gravity * pe
cst = self.mpitools.local_to_global([(maxu, 'max'), (ke, 'sum'),
(z, 'sum'), (z2, 'sum'),
(pe, 'sum'), (b, 'sum'),
(b2, 'sum'), (q, 'sum'),
(q2, 'sum'), (qn, 'sum'),
(qn2, 'sum'), (v2, 'sum')])
a = self.area
self.diags['maxspeed'] = cst[0]
self.diags['ke'] = (cst[1]) / a
self.diags['keV'] = (0.5 * cst[11]) / a
self.diags['pe'] = cst[4] / a
self.diags[
'energy'] = self.diags['ke'] + self.diags['pe'] + self.diags['keV']
self.diags['vorticity'] = cst[2] / a
self.diags['enstrophy'] = 0.5 * cst[3] / a
bm = cst[5] / a
self.diags['buoyancy'] = bm
self.diags['brms'] = sqrt(cst[6] / a - bm**2)
pvm = cst[7] / a
pvneg_mean = cst[9] / a
self.diags['pv_mean'] = pvm
self.diags['pv_std'] = sqrt(cst[8] / a - pvm**2)
self.diags['pvneg_mean'] = pvneg_mean
self.diags['pvneg_std'] = sqrt(cst[10] / a - pvneg_mean**2)
class Euler(object):
""" Euler model
It provides the step(t,dt) function
and 'var' containing the mode state
"""
def __init__(self, param, grid):
self.list_param = [
'forcing', 'noslip', 'timestepping', 'diffusion', 'Kdiff',
'forcing_module', 'additional_tracer', 'enforce_momentum',
'var_to_save', 'customized', 'custom_module', 'spongelayer'
]
param.copy(self, self.list_param)
# for diagnostics
self.list_param = [
'yr', 'nh', 'msk', 'area', 'mpitools', 'dx', 'xr', 'yr', 'r2',
'x0', 'y0', 'x2', 'y2', 'isisland', 'Lx', 'ny'
]
grid.copy(self, self.list_param)
# for variables
param.varname_list = ['vorticity', 'psi', 'u', 'v',
'source'] # ,'tauw','wshear']
param.tracer_list = ['vorticity']
param.whosetspsi = ('vorticity')
if 'tauw' in self.var_to_save:
param.varname_list.append('tauw')
if 'wshear' in self.var_to_save:
param.varname_list.append('wshear')
if hasattr(self, 'additional_tracer'):
for k in range(len(self.additional_tracer)):
trac = self.additional_tracer[k]
param.varname_list.append(trac)
param.tracer_list.append(trac)
param.sizevar = [grid.nyl, grid.nxl]
self.var = Var(param)
self.work = np.zeros(param.sizevar)
# for timing the code
self.timers = Timers(param)
# for operators
self.ope = Operators(param, grid)
# for timescheme
self.tscheme = Timescheme(param, self.var.state)
self.tscheme.set(self.dynamics, self.timestepping)
if self.forcing:
if self.forcing_module == 'embedded':
print(
'Warning: check that you have indeed added the forcing to the model'
)
print('Right below the line : model = f2d.model')
print(
'you should have the line: model.forc = Forcing(param, grid)'
)
pass
else:
try:
f = import_module(self.forcing_module)
except ImportError:
print('module %s for forcing cannot be found' %
self.forcing_module)
print('make sure file **%s.py** exists' %
self.forcing_module)
exit(0)
self.forc = f.Forcing(param, grid)
if self.spongelayer:
# sponge layer zone [0 = full sponge, 1 = no sponge]
self.spongemsk = (1 - (1 + np.tanh(
(self.xr - self.Lx) / 0.1)) * 0.5)
self.diags = {}
if self.customized:
try:
f = import_module(self.custom_module)
self.extrastep = f.Step(param, grid)
except ImportError:
print('module %s for forcing cannot be found' %
self.custom_module)
print('make sure file **%s.py** exists' % self.custom_module)
exit(0)
def step(self, t, dt):
# 1/ integrate dynamics
self.tscheme.forward(self.var.state, t, dt)
# 2/ integrate source
if self.noslip:
self.add_noslip(self.var.state)
source = self.var.get('source')
source /= dt
if 'tauw' in self.var_to_save:
tauw = self.var.get('tauw')
tauw[:, :] = source * self.dx * self.dx
if 'wshear' in self.var_to_save:
wshear = self.var.get('wshear')
self.ope.wallshear(self.var.state, self.work)
wshear[:, :] = self.work
if self.customized:
self.extrastep.do(self.var, t, dt)
# # 3/ dye tracer
if 'dye' in self.var.varname_list:
i, jp, jm = 1, self.ny // 2 + 3, self.ny // 2 - 3
dye = self.var.get('dye')
dye[jp + self.nh, i] = 1
dye[jm + self.nh, i] = -1
if 'age' in self.var.varname_list:
age = self.var.get('age')
age += dt * self.msk
age[:, self.nh] = 0.
if self.spongelayer:
# 4/ sponge layer
w = self.var.get('vorticity')
w *= self.spongemsk
if 'dye' in self.var.varname_list:
dye *= self.spongemsk
if 'age' in self.var.varname_list:
age *= self.spongemsk
#self.set_psi_from_vorticity()
def dynamics(self, x, t, dxdt):
"""Compute the tendency terms + invert the streamfunction"""
self.timers.tic('rhs_adv')
self.ope.rhs_adv(x, t, dxdt)
self.timers.toc('rhs_adv')
if (self.tscheme.kstage == self.tscheme.kforcing):
if self.forcing:
self.forc.add_forcing(x, t, dxdt)
if self.diffusion:
self.ope.rhs_diffusion(x, t, dxdt)
# else:
self.timers.tic('invert')
self.ope.invert_vorticity(dxdt, flag='fast')
self.timers.toc('invert')
def add_noslip(self, x):
self.timers.tic('noslip')
source = self.var.get('source')
self.ope.rhs_noslip(x, source)
self.timers.toc('noslip')
# if not(self.enforce_momentum):
self.timers.tic('invert')
self.ope.invert_vorticity(x, flag='fast', island=self.isisland)
self.timers.toc('invert')
def set_psi_from_vorticity(self):
self.ope.invert_vorticity(self.var.state, island=self.isisland)
def diagnostics(self, var, t):
""" should provide at least 'maxspeed' (for cfl determination) """
self.timers.tic('diag')
u = var.get('u')
v = var.get('v')
trac = var.get('vorticity')
psi = var.get('psi')
source = self.var.get('source')
xr = self.xr
yr = self.yr
ke, maxu = fd.computekemaxu(self.msk, u, v, self.nh)
# ke = fd.computekewithpsi(self.msk, trac, psi, self.nh)
z, z2 = fd.computesumandnorm(self.msk, trac, self.nh)
px = fd.computedotprod(self.msk, trac, xr, self.nh)
py = fd.computedotprod(self.msk, trac, yr, self.nh)
angmom = fd.computesum(self.msk, psi, self.nh)
sce = fd.computedotprod(self.msk, trac, source, self.nh)
cst = self.mpitools.local_to_global([(maxu, 'max'), (ke, 'sum'),
(z, 'sum'), (z2, 'sum'),
(px, 'sum'), (py, 'sum'),
(angmom, 'sum'), (sce, 'sum')])
self.diags['maxspeed'] = cst[0]
self.diags['ke'] = cst[1] / self.area
self.diags['vorticity'] = cst[2] / self.area
self.diags['enstrophy'] = 0.5 * cst[3] / self.area
self.diags['px'] = cst[4] / self.area
self.diags['py'] = cst[5] / self.area
self.diags['angmom'] = cst[6] / self.area
self.diags['source'] = cst[7] / self.area
self.timers.toc('diag')
class Advection(object):
"""Advection model
It solves the 2D transport equation for a passive tracer with a
prescribed steady velocity field, which derives from a
streamfunction
"""
def __init__(self, param, grid):
self.list_param = ['timestepping', 'diffusion', 'Kdiff']
param.copy(self, self.list_param)
self.list_grid = ['msk', 'nh', 'area', 'mpitools']
grid.copy(self, self.list_grid)
# for variables
param.varname_list = ['tracer', 'psi', 'u', 'v', 'vorticity']
param.sizevar = [grid.nyl, grid.nxl]
self.var = Var(param)
# for operators
param.tracer_list = ['tracer']
param.whosetspsi = ('tracer')
self.ope = Operators(param, grid)
# for timescheme
self.tscheme = Timescheme(param, self.var.state)
self.tscheme.set(self.advection, self.timestepping)
self.diags = {}
def step(self, t, dt):
self.tscheme.forward(self.var.state, t, dt)
self.diagnostics(self.var, t)
def advection(self, x, t, dxdt):
self.ope.rhs_adv(x, t, dxdt)
if (self.tscheme.kstage == self.tscheme.kforcing):
if self.diffusion:
self.ope.rhs_diffusion(x, t, dxdt)
def set_psi_from_tracer(self):
self.ope.invert_vorticity(self.var.state)
def diagnostics(self, var, t):
""" should provide at least 'maxspeed' (for cfl determination) """
if t == 0.:
# since the flow is prescribed, compute it only at t=0
u = var.get('u')
v = var.get('v')
ke, maxu = fd.computekemaxu(self.msk, u, v, self.nh)
cst = self.mpitools.local_to_global([(maxu, 'max'), (ke, 'sum')])
self.diags['maxspeed'] = cst[0]
self.diags['ke'] = cst[1] / self.area
self.diags['enstrophy'] = 0.
trac = var.get('tracer')
z, z2 = fd.computesumandnorm(self.msk, trac, self.nh)
cst = self.mpitools.local_to_global([(z, 'sum'), (z2, 'sum')])
z = cst[0] / self.area
z2 = cst[1] / self.area
self.diags['mean'] = z
self.diags['rms'] = np.sqrt(z2)
class Boussinesq(Param):
""" Boussinesq model
It provides the step(t,dt) function
and 'var' containing the mode state
"""
def __init__(self,param,grid):
self.list_param=['forcing','noslip','timestepping','diffusion','Kdiff',
'forcing_module','gravity','isisland',
'customized','custom_module','additional_tracer']
param.copy(self,self.list_param)
# for potential energy
self.list_param=['xr','yr','nh','Lx','msk','area','mpitools']
grid.copy(self,self.list_param)
# for variables
param.varname_list=['vorticity','psi','u','v','buoyancy']
param.tracer_list=['vorticity','buoyancy']
param.whosetspsi=('vorticity')
if hasattr(self,'additional_tracer'):
for k in range(len(self.additional_tracer)):
trac=self.additional_tracer[k]
param.varname_list.append(trac)
param.tracer_list.append(trac)
print('Tracers are :',param.tracer_list)
param.sizevar =[grid.nyl,grid.nxl]
self.var = Var(param)
self.source = zeros(param.sizevar)
# for operators
self.ope = Operators(param,grid)
# for timescheme
self.tscheme = Timescheme(param,self.var.state)
self.tscheme.set(self.dynamics, self.timestepping)
if self.forcing:
try:
f = import_module(self.forcing_module)
except:
print('module %s for forcing cannot be found'%self.forcing_module)
print('make sure file **%s.py** exists'%self.forcing_module)
exit(0)
self.forc = f.Forcing(param,grid)
self.diags={}
if self.customized:
try:
f = import_module(self.custom_module)
print(f)
self.extrastep = f.Step(param,grid)
except:
print('module %s for forcing cannot be found'%self.custom_module)
print('make sure file **%s.py** exists'%self.custom_module)
exit(0)
def step(self,t,dt):
# 1/ integrate advection
self.tscheme.forward(self.var.state,t,dt)
# 2/ integrate source
if self.noslip:
self.add_noslip(self.var.state)
if self.customized:
self.extrastep.do(self.var,t,dt)
def dynamics(self,x,t,dxdt):
self.ope.rhs_adv(x,t,dxdt)
self.ope.rhs_torque(x,t,dxdt) # db/dx is a source term for the vorticity
if (self.tscheme.kstage==self.tscheme.kforcing):
if self.forcing:
self.forc.add_forcing(x,t,dxdt)
if self.diffusion:
self.ope.rhs_diffusion(x,t,dxdt)
self.ope.invert_vorticity(dxdt,flag='fast')
def add_noslip(self,x):
self.ope.rhs_noslip(x,self.source)
self.ope.invert_vorticity(x,flag='fast',island=self.isisland)
def set_psi_from_vorticity(self):
self.ope.invert_vorticity(self.var.state,island=self.isisland)
def diagnostics(self,var,t):
""" should provide at least 'maxspeed' (for cfl determination) """
nh = self.nh
u = var.get('u')
v = var.get('v')
vort = var.get('vorticity')
buoy = var.get('buoyancy')
ke,maxu = computekemaxu(self.msk,u,v,self.nh)
z,z2 = computesumandnorm(self.msk,vort,self.nh)
b,b2 = computesumandnorm(self.msk,buoy,self.nh)
# potential energy (minus sign because buoyancy is minus density)
pe = computesum(self.msk,buoy*self.yr,nh)
pe = - self.gravity*pe
cst = self.mpitools.local_to_global( [(maxu,'max'),(ke,'sum'),(z,'sum'),(z2,'sum'),
(pe,'sum'),(b,'sum'),(b2,'sum')] )
self.diags['maxspeed'] = cst[0]
self.diags['ke'] = cst[1] / self.area
self.diags['pe'] = cst[4] / self.area
self.diags['energy'] = (cst[1]+cst[4]) / self.area
self.diags['vorticity']= cst[2] / self.area
self.diags['enstrophy']= cst[3] / self.area
self.diags['buoyancy'] = cst[5] / self.area
self.diags['brms']= sqrt( cst[6] / self.area -(cst[5] / self.area)**2 )
class Fluxes():
"""Class to diagnostic the irreversible part of advective fluxes
The development status is
- 'euler' : ok
- 'boussinesq' : ok
- 'thermalwind' : ok
- 'qg' : not done
Note that for 'thermalwind' we should change sign of both V and f0,
but we don't, because the equations are still invariant under time
reversal if we don't do it. The rationale is that the implementation
is easier.
"""
def __init__(self, param, grid, ope):
self.list_param = [
'timestepping', 'varname_list', 'tracer_list', 'order', 'aparab',
'sizevar', 'flux_splitting_method', 'modelname'
]
p = copy.deepcopy(param)
# we add two new variables in the state vector: uc and vc
# the velocities at cell centers
# their fluxes are used to estimate the KE dissipation
newvariables = ['uc', 'vc']
p.tracer_list += newvariables
p.varname_list += newvariables
flx_list = []
for v in p.tracer_list:
for d in ['x', 'y']:
flx_list += ['flx_%s_%s' % (v, d)]
self.flx_list = flx_list
self.nvarstate = len(p.varname_list)
p.varname_list += flx_list
fullflx_list = []
for r in ['rev', 'irr']:
for v in p.tracer_list:
for d in ['x', 'y']:
fullflx_list += ['%s_%s_%s' % (r, d, v)]
self.fullflx_list = fullflx_list
p.copy(self, self.list_param)
self.list_grid = ['nh', 'dx', 'dy', 'msk']
grid.copy(self, self.list_grid)
self.ope = ope
varsize = [len(self.varname_list)] + self.sizevar
self.x = np.zeros(varsize)
# extended sate vector to host the model state vector and (uc,vc)
self.xe = np.zeros([self.nvarstate] + self.sizevar)
self.xwork = np.zeros(varsize)
tracsize = [len(self.fullflx_list)] + self.sizevar
# full stack of fluxes (including reversible and irreversible)
self.flx = np.zeros(tracsize)
# select the proper flux discretization
if self.order % 2 == 0:
self.fortran_adv = fa.adv_centered
else:
self.fortran_adv = fa.adv_upwind
self.cst = np.zeros(5, )
self.cst[0] = grid.dx
self.cst[1] = grid.dy
self.cst[2] = 0.05
self.cst[3] = 0 # umax
self.cst[4] = self.aparab
# for timescheme
#p.timestepping = 'EF'
self.tscheme = Timescheme(p, self.x)
self.tscheme.set(self.advection, p.timestepping)
# controls the flux splitting method
# 0 = min/max
# 1 = parabolic
list_fs_method = ['minmax', 'parabolic']
if self.flux_splitting_method in list_fs_method:
self.fs_method = list_fs_method.index(self.flux_splitting_method)
else:
print('Warning: %s does not exist' % self.flux_splitting_method)
print('replaced with the default: parabolic')
self.fs_method = list_fs_method.index('parabolic')
def diag_fluxes(self, x, t, dt):
"""advance x from one time step, then reverse the time and advance
another time step. Diagnose the reversible and the
irreversible part of the fluxes
"""
iu = self.varname_list.index('u')
iv = self.varname_list.index('v')
ip = self.varname_list.index('psi')
iw = self.varname_list.index('vorticity')
# set locally dt to a very small value to make the time scheme
# dissipation vanishing (hard coded so far)
#
# dt = dt*1e-6
self.xe[:self.nvarstate - 2, :, :] = x
#
# we add two more variables 'uc' and 'vc', the cell centered velocities
# we store them in the extended state self.xe
#
# i) uc
y = 0.5 * (x[iu] + np.roll(x[iu], 1, axis=1))
self.ope.fill_halo(y)
self.xe[self.nvarstate - 2, :, :] = y
# ii) vc
y = 0.5 * (x[iv] + np.roll(x[iv], 1, axis=0))
self.ope.fill_halo(y)
self.xe[self.nvarstate - 1, :, :] = y
# we copy this extended state into self.x
self.x[:self.nvarstate, :, :] = self.xe
# set fluxes entries to zero
self.x[self.nvarstate:, :, :] = 0.
self.tscheme.forward(self.x, t, dt)
self.xwork[:, :, :] = self.x
# reverse velocity
self.x[:self.nvarstate, :, :] = self.xe
self.x[iu] *= -1
self.x[iv] *= -1
self.x[ip] *= -1
self.x[iw] *= -1
self.x[self.nvarstate - 2] *= -1 # uc
self.x[self.nvarstate - 1] *= -1 # vc
# set fluxes entries to zero
self.x[self.nvarstate:, :, :] = 0.
self.tscheme.forward(self.x, t + dt, -dt)
# don't forget to divide by 'dt' to get the tendency
cff = 0.5 / (dt)
nflx = len(self.flx_list)
# this loop below should be made more readable
# k sweeps across all reversible fluxes
# (two components per advected quantity)
for k in range(nflx):
# l points to the flux index in the self.x array
l = self.nvarstate + k
# j points to the associated irreversible flux
j = nflx + k
if (k < 2) or (k >= (nflx - 4)):
# k = 0, 1 are the vorticity fluxes
# k = nflx-4, nflx-3 are the 'uc' fluxes
# k = nflx-2, nflx-1 are the 'vc' fluxes
sign = -1 #for vorticity, 'uc' and 'vc'
else:
sign = 1 # for buoyancy and tracer (if present)
self.flx[k, :, :] = cff * (self.xwork[l] + sign * self.x[l])
self.flx[j, :, :] = cff * (self.xwork[l] - sign * self.x[l])
def advection(self, x, t, dxdt):
self.rhs_adv(x, t, dxdt)
if self.modelname == 'boussinesq':
self.ope.rhs_torque(x, t, dxdt)
elif self.modelname == 'thermalwind':
self.ope.rhs_thermalwind(x, t, dxdt)
self.ope.invert_vorticity(dxdt, flag='fast')
def rhs_adv(self, x, t, dxdt):
""" compute -div(u*tracer) using finite volume flux discretization
the flux is computed at edge cells using p-th order interpolation
for p even, the flux is centered
for p odd, the flux is upwinded (more points on the upwind side) """
iu = self.varname_list.index('u')
iv = self.varname_list.index('v')
u = x[iu]
v = x[iv]
# don't forget to adjust the maxspeed
self.cst[3] = self.ope.cst[3]
for itrac, trac in enumerate(self.tracer_list):
ik = self.varname_list.index(trac)
y = dxdt[ik]
ix = self.nvarstate + itrac * 2
iy = self.nvarstate + itrac * 2 + 1
xf = dxdt[ix]
yf = dxdt[iy]
self.fortran_adv(self.msk, x[ik], y, u, v, xf, yf, self.cst,
self.nh, self.fs_method, self.order)
self.ope.fill_halo(y)
self.ope.fill_halo(xf)
self.ope.fill_halo(yf)
# for an unknown reason dxdt[ik] is
# not updated by the Fortran routine
# it should be done manually
# (this yields an excessive data movement)
dxdt[ik, :, :] = y
dxdt[ix, :, :] = xf
dxdt[iy, :, :] = yf
class Euler(Param):
""" Euler model
It provides the step(t,dt) function
and 'var' containing the mode state
"""
def __init__(self, param, grid):
self.list_param = [
'forcing', 'noslip', 'timestepping', 'diffusion', 'Kdiff',
'forcing_module', 'additional_tracer', 'enforce_momentum',
'var_to_save', 'customized', 'custom_module'
]
param.copy(self, self.list_param)
# for diagnostics
self.list_param = [
'yr', 'nh', 'msk', 'area', 'mpitools', 'dx', 'xr', 'yr', 'r2',
'x0', 'y0', 'x2', 'y2', 'isisland', 'Lx'
]
grid.copy(self, self.list_param)
# for variables
param.varname_list = ['vorticity', 'psi', 'u', 'v'] #,'tauw','wshear']
param.tracer_list = ['vorticity']
param.whosetspsi = ('vorticity')
if 'tauw' in self.var_to_save:
param.varname_list.append('tauw')
if 'wshear' in self.var_to_save:
param.varname_list.append('wshear')
if hasattr(self, 'additional_tracer'):
for k in range(len(self.additional_tracer)):
trac = self.additional_tracer[k]
param.varname_list.append(trac)
param.tracer_list.append(trac)
print('Tracers are :', param.tracer_list)
param.sizevar = [grid.nyl, grid.nxl]
self.var = Var(param)
self.source = zeros(param.sizevar)
# for operators
self.ope = Operators(param, grid)
# for timescheme
self.tscheme = Timescheme(param, self.var.state)
self.tscheme.set(self.advection, self.timestepping)
if self.forcing:
try:
f = import_module(self.forcing_module)
except:
print('module %s for forcing cannot be found' %
self.forcing_module)
print('make sure file **%s.py** exists' % self.forcing_module)
exit(0)
self.forc = f.Forcing(param, grid)
self.diags = {}
if self.customized:
try:
f = import_module(self.custom_module)
self.extrastep = f.step(param, grid)
except:
print('module %s for forcing cannot be found' %
self.custom_module)
print('make sure file **%s.py** exists' % self.custom_module)
exit(0)
def step(self, t, dt):
# 1/ integrate advection
self.tscheme.forward(self.var.state, t, dt)
# 2/ integrate source
if self.noslip:
self.add_noslip(self.var.state)
if 'tauw' in self.var_to_save:
tauw = self.var.get('tauw')
tauw[:, :] = self.source / dt * self.dx * self.dx
if 'wshear' in self.var_to_save:
wshear = self.var.get('wshear')
self.ope.wallshear(self.var.state, self.source)
wshear[:, :] = self.source
if self.customized:
self.extrastep.do(self.var, t, dt)
# # 3/ dye tracer
# i,jp,jm=1,67,61
# dye = self.var.get('dye')
# dye[jp+self.nh,i]=1
# dye[jm+self.nh,i]=-1
# age = self.var.get('age')
# age += dt*self.msk
# age[:,self.nh]=0.
# # 4/ sponge layer
# damping = ( 1- (1+tanh( (self.xr - self.Lx)/0.1))*0.5)
# w = self.var.get('vorticity')
# w *= damping
# dye *= damping
# age *= damping
def advection(self, x, t, dxdt):
self.ope.rhs_adv(x, t, dxdt)
if (self.tscheme.kstage == self.tscheme.kforcing):
if self.forcing:
self.forc.add_forcing(x, t, dxdt)
if self.diffusion:
self.ope.rhs_diffusion(x, t, dxdt)
self.ope.invert_vorticity(dxdt, flag='fast')
def add_noslip(self, x):
self.ope.rhs_noslip(x, self.source)
if not (self.enforce_momentum):
self.ope.invert_vorticity(x, flag='fast', island=self.isisland)
def set_psi_from_vorticity(self):
self.ope.invert_vorticity(self.var.state, island=self.isisland)
def diagnostics(self, var, t):
""" should provide at least 'maxspeed' (for cfl determination) """
nh = self.nh
u = var.get('u')
v = var.get('v')
trac = var.get('vorticity')
xr = self.xr
yr = self.yr
r2 = self.r2
ke, maxu = computekemaxu(self.msk, u, v, self.nh)
z, z2 = computesumandnorm(self.msk, trac, self.nh)
px = computedotprod(self.msk, trac, xr, self.nh)
py = computedotprod(self.msk, trac, yr, self.nh)
angmom = computedotprod(self.msk, trac, r2, self.nh)
cst = self.mpitools.local_to_global([(maxu, 'max'), (ke, 'sum'),
(z, 'sum'), (z2, 'sum'),
(px, 'sum'), (py, 'sum'),
(angmom, 'sum')])
self.diags['maxspeed'] = cst[0]
self.diags['ke'] = cst[1] / self.area
self.diags['vorticity'] = cst[2] / self.area
self.diags['enstrophy'] = cst[3] / self.area
self.diags['px'] = cst[4] / self.area
self.diags['py'] = cst[5] / self.area
self.diags['angmom'] = cst[6] / self.area
class QG(object):
"""Quasi-geostrophic model
The prognostic variable is 'pv'. It is the full PV, which includes
the planetary background PV 'pvback'. Full PV is materially
conserved unlike the PV anamoly that has a beta*v source term
"""
def __init__(self, param, grid):
self.list_param = ['forcing', 'diffusion', 'Kdiff', 'noslip',
'timestepping', 'beta', 'Rd', 'ageostrophic',
'bottom_torque',
'forcing_module']
param.copy(self, self.list_param)
# for diagnostics
self.list_param = ['yr', 'nh', 'msk', 'area', 'mpitools', 'isisland']
grid.copy(self, self.list_param)
# for variables
param.varname_list = ['pv', 'psi', 'u', 'v', 'pvanom', 'vorticity']
if param.bottom_torque:
param.varname_list += ['btorque']
if param.ageostrophic:
param.varname_list += ['ua', 'va']
param.sizevar = [grid.nyl, grid.nxl]
self.var = Var(param)
self.source = np.zeros(param.sizevar)
self.ipva = self.var.varname_list.index('pvanom')
self.ipv = self.var.varname_list.index('pv')
self.ivor = self.var.varname_list.index('vorticity')
self.ipsi = self.var.varname_list.index('psi')
# background pv
self.pvback = self.beta*(grid.yr-grid.Ly*.5)*grid.msk
# for operators
param.tracer_list = ['pv']
if param.bottom_torque:
param.tracer_list += ['btorque']
if param.ageostrophic:
param.tracer_list += ['ua', 'va']
param.whosetspsi = ('pvanom')
param.qgoperator = True
self.ope = Operators(param, grid)
# for timescheme
self.tscheme = Timescheme(param, self.var.state)
self.dx0 = self.tscheme.dx0
self.kt = 0
if self.forcing:
if self.forcing_module == 'embedded':
print('Warning: check that you have indeed added the forcing to the model')
print('Right below the line : model = f2d.model')
print('you should have the line: model.forc = Forcing(param, grid)')
pass
else:
try:
f = import_module(self.forcing_module)
except ImportError:
print('module %s for forcing cannot be found'
% self.forcing_module)
print('make sure file **%s.py** exists' % self.forcing_module)
sys.exit(0)
self.forc = f.Forcing(param, grid)
self.diags = {}
self.tscheme.set(self.dynamics, self.timestepping)
def step(self, t, dt):
"""Integrate the model over one time step dt"""
if self.bottom_torque:
ibt = self.var.varname_list.index('btorque')
self.var.state[ibt][:, :] = self.pvback
if self.ageostrophic:
iu = self.var.varname_list.index('u')
iv = self.var.varname_list.index('v')
iua = self.var.varname_list.index('ua')
iva = self.var.varname_list.index('va')
self.var.state[iua][:, :] = self.var.state[iu].copy()
self.var.state[iva][:, :] = self.var.state[iv].copy()
self.dt = dt
# 1/ integrate advection
self.tscheme.forward(self.var.state, t, dt)
# 2/ integrate source
if self.noslip:
self.add_noslip(self.var.state)
# self.set_psi_from_pv()
# 3/ diagnostic fields
self.var.state[self.ipva] = self.var.state[self.ipv] - self.pvback
self.var.state[self.ivor] = (self.var.state[self.ipva]
- (self.Rd**-2)*self.var.state[self.ipsi])
if self.bottom_torque:
""" diagnose the bottom torque btorque = -J(psi, htopo)
in practice it is obtained as (htopo^{n+1}-htopo^n)/dt by
using the 'btorque' entry in var.state and transport it"""
self.var.state[ibt][:, :] = (self.var.state[ibt][:, :]
- self.pvback)/dt
if self.ageostrophic:
""" diagnose the ageostrophic velocity (factor 1/f missing)
du/dt + J(psi, u) = +f va + (pvback)*vg
dv/dt + J(psi, v) = -f ua - (pvback)*ug
implementation:
- store u in 'ua'
- step ua
- compare update u (real u diagnosed from psi) with updated 'ua'
- the difference is the ageostrophic velocity
Yet to be done: implement the proper staggering, 'u' and
'v' sit on edges while the advection scheme assumes cell
centers quantities """
ug = self.var.state[iu]
vg = self.var.state[iv]
ua = -(vg - self.var.state[iva])/dt - self.pvback*ug
va = +(ug - self.var.state[iua])/dt - self.pvback*vg
self.var.state[iva][:, :] = va.copy()
self.var.state[iua][:, :] = ua.copy()
def dynamics(self, x, t, dxdt):
"""Compute the tendency terms + invert the streamfunction"""
dxdt[:] = 0.
self.ope.rhs_adv(x, t, dxdt)
if (self.tscheme.kstage == self.tscheme.kforcing):
if self.forcing:
self.forc.add_forcing(x, t, dxdt)
if self.diffusion:
self.ope.rhs_diffusion(x, t, dxdt)
else:
dxdt[self.ipva] = dxdt[self.ipv]
self.ope.invert_vorticity(dxdt, flag='fast', island=self.isisland)
def add_noslip(self, x):
"""Add the no-slip term """
self.ope.rhs_noslip(x, self.source)
self.ope.invert_vorticity(x, flag='fast', island=self.isisland)
def add_backgroundpv(self):
self.var.state[self.ipv] += self.pvback
def set_psi_from_pv(self):
"""Solve the elliptic equation for the streamfunction
The unknown of the Helmholtz equation is the PV anomaly,
not the full PV"""
state = self.var.state
state[self.ipva] = state[self.ipv] - self.pvback
self.ope.invert_vorticity(self.var.state, flag='full', island=self.isisland)
def diagnostics(self, var, t):
""" Integral diagnostics for the QG model
should provide at least 'maxspeed' (for cfl determination) """
u = var.get('u')
v = var.get('v')
trac = var.get('pv')
psi = var.get('psi')
fo.cornertocell(psi, self.ope.work)
psim, psi2 = fd.computesumandnorm(self.msk, self.ope.work, self.nh)
ape = 0.5 * psi2 / self.Rd**2
ke, maxu = fd.computekemaxu(self.msk, u, v, self.nh)
z, z2 = fd.computesumandnorm(self.msk, trac, self.nh)
cst = self.mpitools.local_to_global([(maxu, 'max'),
(ke, 'sum'),
(z, 'sum'),
(z2, 'sum'),
(ape, 'sum')])
self.diags['maxspeed'] = cst[0]
self.diags['ke'] = cst[1] / self.area
self.diags['pv'] = cst[2] / self.area
self.diags['pv2'] = 0.5*cst[3] / self.area
self.diags['ape'] = cst[4] / self.area
self.diags['energy'] = (cst[1]+cst[4]) / self.area
class TwoLayerQG(Param):
""" Quasi-geostrophic model
It provides the step(t,dt) function
and 'var' containing the mode state
"""
def __init__(self, param, grid):
self.list_param = [
'forcing', 'diffusion', 'Kdiff', 'noslip', 'timestepping', 'beta',
'Rd', 'forcing_module'
]
param.copy(self, self.list_param)
# for diagnostics
self.list_param = ['yr', 'nh', 'msk', 'area', 'mpitools']
grid.copy(self, self.list_param)
# for variables
param.varname_list = ('pv1', 'psi1', 'u1', 'v1', 'pv2', 'psi2', 'u2',
'v2')
param.sizevar = [grid.nyl, grid.nxl]
self.var = Var(param)
self.source = zeros(param.sizevar)
# self.ipva = self.var.varname_list.index('pvanom')
# self.ipv = self.var.varname_list.index('pv')
# self.ivor = self.var.varname_list.index('vorticity')
# self.ipsi = self.var.varname_list.index('psi')
# background pv
self.pvback = self.beta * (grid.yr - grid.Ly * .5) * grid.msk
# for operators
param.tracer_list = ['pv1', 'pv2']
param.whosetspsi = ('pvanom')
param.qgoperator = True
self.ope = Operators(param, grid)
# for timescheme
self.tscheme = Timescheme(param, self.var.state)
self.dx0 = self.tscheme.dx0
self.kt = 0
if self.forcing:
try:
f = import_module(self.forcing_module)
except:
print('module %s for forcing cannot be found' %
self.forcing_module)
print('make sure file **%s.py** exists' % self.forcing_module)
exit(0)
self.forc = f.Forcing(param, grid)
self.diags = {}
self.tscheme.set(self.dynamics, self.timestepping)
def step(self, t, dt):
self.dt = dt
# 1/ integrate advection
self.tscheme.forward(self.var.state, t, dt)
# 2/ integrate source
if self.noslip:
self.add_noslip(self.var.state)
# 3/ diagnostic fields
self.var.state[self.ipva] = self.var.state[self.ipv] - self.pvback
self.var.state[self.ivor] = self.var.state[
self.ipva] + self.Rd**2 * self.var.state[self.ipsi]
def dynamics(self, x, t, dxdt):
dxdt[:] = 0.
self.ope.rhs_adv(x, t, dxdt)
if (self.tscheme.kstage == self.tscheme.kforcing):
if self.forcing:
self.forc.add_forcing(x, t, dxdt)
if self.diffusion:
self.ope.rhs_diffusion(x, t, dxdt)
# substract the background pv
dxdt[self.ipva][:, :] = dxdt[self.ipv] #- self.pvback*self.dt
# self.ope.invert_vorticity(dxdt,flag='fast')
self.ope.invert_twolayers(dxdt, flag='fast')
# self.kt +=1
def add_noslip(self, x):
self.ope.rhs_noslip(x, self.source)
self.ope.invert_vorticity(x, flag='fast')
def add_backgroundpv(self):
self.var.state[
self.ipv1][:, :] = self.var.state[self.ipv1] + self.pvback
self.var.state[
self.ipv2][:, :] = self.var.state[self.ipv2] + self.pvback
def substract_backgroundpv(self):
self.var.state[
self.ipv1][:, :] = self.var.state[self.ipv1] - self.pvback
self.var.state[
self.ipv2][:, :] = self.var.state[self.ipv2] - self.pvback
def set_psi_from_pv(self):
state = self.var.state
self.substract_backgroundpv()
self.ope.invert_twolayers(self.var.state, flag='full')
self.add_backgroundpv()
def diagnostics(self, var, t):
""" should provide at least 'maxspeed' (for cfl determination) """
nh = self.nh
u = var.get('u')
v = var.get('v')
trac = var.get('pv')
ke, maxu = computekemaxu(self.msk, u, v, self.nh)
z, z2 = computesumandnorm(self.msk, trac, self.nh)
cst = self.mpitools.local_to_global([(maxu, 'max'), (ke, 'sum'),
(z, 'sum'), (z2, 'sum')])
self.diags['maxspeed'] = cst[0]
self.diags['ke'] = cst[1] / self.area
self.diags['pv'] = cst[2] / self.area
self.diags['pv2'] = cst[3] / self.area
class SQG(object):
""" Surface Quasi-Geostrophic model
The prognostic variable is 'pv'. It is the surface PV
The streamfunction is computed in the Fourier space
The geometry must be biperiodic
This model does not support multicores (because of the fft)
"""
def __init__(self, param, grid):
self.list_param = [
'forcing', 'diffusion', 'Kdiff', 'timestepping', 'ageostrophic',
'forcing_module', 'geometry'
]
param.copy(self, self.list_param)
# for diagnostics
self.list_param = ['yr', 'nh', 'msk', 'area', 'mpitools']
grid.copy(self, self.list_param)
assert grid.geometry == "perio", "SQG imposes a biperiodic domain"
assert param.ageostrophic == False, "Ageostrophic velocity not yet tested"
# for variables
param.varname_list = ['pv', 'psi', 'u', 'v', 'vorticity']
if param.ageostrophic:
# TODO: not yet tested for SQG
param.varname_list += ['ua', 'va']
param.sizevar = [grid.nyl, grid.nxl]
self.var = Var(param)
self.source = np.zeros(param.sizevar)
self.ipv = self.var.varname_list.index('pv')
self.ivor = self.var.varname_list.index('vorticity')
self.ipsi = self.var.varname_list.index('psi')
# for operators
param.tracer_list = ['pv']
if param.ageostrophic:
param.tracer_list += ['ua', 'va']
param.whosetspsi = ('pv')
param.sqgoperator = True
self.ope = Operators(param, grid)
# for timescheme
self.tscheme = Timescheme(param, self.var.state)
self.dx0 = self.tscheme.dx0
self.kt = 0
if self.forcing:
if self.forcing_module == 'embedded':
print(
'Warning: check that you have indeed added the forcing to the model'
)
print('Right below the line : model = f2d.model')
print(
'you should have the line: model.forc = Forcing(param, grid)'
)
pass
else:
try:
f = import_module(self.forcing_module)
except ImportError:
print('module %s for forcing cannot be found' %
self.forcing_module)
print('make sure file **%s.py** exists' %
self.forcing_module)
sys.exit(0)
self.forc = f.Forcing(param, grid)
self.diags = {}
self.tscheme.set(self.dynamics, self.timestepping)
def step(self, t, dt):
"""Integrate the model over one time step dt"""
if self.ageostrophic:
iu = self.var.varname_list.index('u')
iv = self.var.varname_list.index('v')
iua = self.var.varname_list.index('ua')
iva = self.var.varname_list.index('va')
self.var.state[iua][:, :] = self.var.state[iu].copy()
self.var.state[iva][:, :] = self.var.state[iv].copy()
self.dt = dt
# 1/ integrate advection
self.tscheme.forward(self.var.state, t, dt)
self.set_psi_from_pv()
if self.ageostrophic:
""" diagnose the ageostrophic velocity (factor 1/f missing)
du/dt + J(psi, u) = +f va + (pvback)*vg
dv/dt + J(psi, v) = -f ua - (pvback)*ug
implementation:
- store u in 'ua'
- step ua
- compare update u (real u diagnosed from psi) with updated 'ua'
- the difference is the ageostrophic velocity
Yet to be done: implement the proper staggering, 'u' and
'v' sit on edges while the advection scheme assumes cell
centers quantities """
ug = self.var.state[iu]
vg = self.var.state[iv]
ua = -(vg - self.var.state[iva]) / dt - self.pvback * ug
va = +(ug - self.var.state[iua]) / dt - self.pvback * vg
self.var.state[iva][:, :] = va.copy()
self.var.state[iua][:, :] = ua.copy()
def dynamics(self, x, t, dxdt):
"""Compute the tendency terms + invert the streamfunction"""
dxdt[:] = 0.
self.ope.rhs_adv(x, t, dxdt)
if (self.tscheme.kstage == self.tscheme.kforcing):
if self.forcing:
self.forc.add_forcing(x, t, dxdt)
if self.diffusion:
self.ope.rhs_diffusion(x, t, dxdt)
else:
self.ope.fourier_invert_vorticity(dxdt, flag='fast')
def set_psi_from_pv(self):
"""Solve the elliptic equation for the streamfunction
The unknown of the Helmholtz equation is the PV anomaly,
not the full PV"""
self.ope.fourier_invert_vorticity(self.var.state, flag='full')
def diagnostics(self, var, t):
"""Integral diagnostics for the SQG model
should provide at least 'maxspeed' (for cfl determination)
WARNING: the diag for KE and APE are wrong. A SQG flow is 3D
the integration should be carried over in the Fourier space
The total PV should be conserved (in the absence of external
forcing) and the enstrophy should be conserved as long as the
flow is not developping filaments (grid-scale enstrophy). The
numerical dissipation wipes out Filaments which yields
enstrophy dissipation.
"""
u = var.get('u')
v = var.get('v')
trac = var.get('pv')
psi = var.get('psi')
#fo.cornertocell(psi, self.ope.work)
#psim, psi2 = fd.computesumandnorm(self.msk, self.ope.work, self.nh)
#ape = 0.5 * psi2 / self.Rd**2
ke, maxu = fd.computekemaxu(self.msk, u, v, self.nh)
z, z2 = fd.computesumandnorm(self.msk, trac, self.nh)
cst = self.mpitools.local_to_global([(maxu, 'max'), (ke, 'sum'),
(z, 'sum'), (z2, 'sum')])
self.diags['maxspeed'] = cst[0]
self.diags['ke'] = cst[1] / self.area
self.diags['pv'] = cst[2] / self.area
self.diags['pv2'] = 0.5 * cst[3] / self.area
class BoussinesqTS(object):
""" Boussinesq model with temperature and salinity
It provides the step(t,dt) function
and 'var' containing the mode state
"""
def __init__(self, param, grid):
self.list_param = [
'forcing', 'noslip', 'timestepping', 'alphaT', 'betaS',
'diffusion', 'Kdiff', 'myrank', 'forcing_module', 'gravity',
'isisland', 'customized', 'custom_module', 'additional_tracer'
]
param.copy(self, self.list_param)
# for potential energy
self.list_param = ['xr', 'yr', 'nh', 'Lx', 'msk', 'area', 'mpitools']
grid.copy(self, self.list_param)
# for variables
param.varname_list = [
'vorticity', 'psi', 'u', 'v', 'density', 'danom', 'T', 'S'
]
param.tracer_list = ['vorticity', 'T', 'S']
param.whosetspsi = ('vorticity')
if hasattr(self, 'additional_tracer'):
for k in range(len(self.additional_tracer)):
trac = self.additional_tracer[k]
param.varname_list.append(trac)
param.tracer_list.append(trac)
self.varname_list = param.varname_list
param.sizevar = [grid.nyl, grid.nxl]
self.var = Var(param)
dref = self.var.get('density').copy()
self.dref = dref
self.source = np.zeros(param.sizevar)
# for operators
self.ope = Operators(param, grid)
# for timescheme
self.tscheme = Timescheme(param, self.var.state)
self.tscheme.set(self.dynamics, self.timestepping)
if self.forcing:
if self.forcing_module == 'embedded':
print(
'Warning: check that you have indeed added the forcing to the model'
)
print('Right below the line : model = f2d.model')
print(
'you should have the line: model.forc = Forcing(param, grid)'
)
pass
else:
try:
f = import_module(self.forcing_module)
except ImportError:
print('module %s for forcing cannot be found' %
self.forcing_module)
print('make sure file **%s.py** exists' %
self.forcing_module)
sys.exit(0)
self.forc = f.Forcing(param, grid)
self.diags = {}
if self.customized:
try:
f = import_module(self.custom_module)
print(f)
self.extrastep = f.Step(param, grid)
except ImportError:
print('module %s for forcing cannot be found' %
self.custom_module)
print('make sure file **%s.py** exists' % self.custom_module)
exit(0)
def step(self, t, dt):
# 1/ integrate advection
self.tscheme.forward(self.var.state, t, dt)
self.set_density()
# 2/ integrate source
if self.noslip:
self.add_noslip(self.var.state)
if self.customized:
self.extrastep.do(self.var, t, dt)
danom = self.var.get('danom')
danom[:, :] = self.var.get('density') - self.dref
def dynamics(self, x, t, dxdt):
self.ope.rhs_adv(x, t, dxdt)
self.eos(dxdt)
# db/dx is a source term for the vorticity
self.ope.rhs_torque_density(x, t, dxdt)
if (self.tscheme.kstage == self.tscheme.kforcing):
coef = self.tscheme.dtcoef
if self.forcing:
self.forc.add_forcing(x, t, dxdt, coef=coef)
if self.diffusion:
self.ope.rhs_diffusion(x, t, dxdt, coef=coef)
self.eos(dxdt)
self.ope.invert_vorticity(dxdt, flag='fast')
def add_noslip(self, x):
self.ope.rhs_noslip(x, self.source)
self.ope.invert_vorticity(x, flag='fast', island=self.isisland)
def eos(self, x):
""" compute density from T and S """
idens = self.varname_list.index('density')
itemp = self.varname_list.index('T')
isalt = self.varname_list.index('S')
x[idens] = -self.alphaT * x[itemp] + self.betaS * x[isalt]
def set_density(self):
self.eos(self.var.state)
def set_psi_from_vorticity(self):
self.ope.invert_vorticity(self.var.state, island=self.isisland)
def diagnostics(self, var, t):
""" should provide at least 'maxspeed' (for cfl determination) """
nh = self.nh
u = var.get('u')
v = var.get('v')
vort = var.get('vorticity')
dens = var.get('density')
ke, maxu = fd.computekemaxu(self.msk, u, v, self.nh)
z, z2 = fd.computesumandnorm(self.msk, vort, self.nh)
b, b2 = fd.computesumandnorm(self.msk, dens, self.nh)
# potential energy
pe = +self.gravity * fd.computesum(self.msk, dens * self.yr, nh)
cst = self.mpitools.local_to_global([(maxu, 'max'), (ke, 'sum'),
(z, 'sum'), (z2, 'sum'),
(pe, 'sum'), (b, 'sum'),
(b2, 'sum')])
self.diags['maxspeed'] = cst[0]
self.diags['ke'] = cst[1] / self.area
self.diags['pe'] = cst[4] / self.area
self.diags['energy'] = (cst[1] + cst[4]) / self.area
self.diags['vorticity'] = cst[2] / self.area
self.diags['enstrophy'] = 0.5 * cst[3] / self.area
self.diags['density'] = cst[5] / self.area
self.diags['drms'] = np.sqrt(cst[6] / self.area -
(cst[5] / self.area)**2)
class QG2L(object):
"""Two-Layers Quasi-geostrophic model
The prognostic variable is 'pv'. It is the full PV, which includes
the planetary background PV 'pvback'. Full PV is materially
conserved unlike the PV anamoly that has a beta*v source term
"""
def __init__(self, param, grid):
self.list_param = ['forcing', 'diffusion', 'Kdiff', 'noslip',
'timestepping', 'beta', 'Rd', 'ageostrophic',
'forcing_module']
param.copy(self, self.list_param)
# for diagnostics
self.list_param = ['yr', 'nh', 'msk', 'area', 'mpitools']
grid.copy(self, self.list_param)
# for variables
param.varname_list = ('pv', 'psi', 'u', 'v',
'pvanom', 'vorticity')
if param.ageostrophic:
param.varname_list += ('ua', 'va')
print(param.varname_list)
param.sizevar = [2, grid.nyl, grid.nxl]
self.var = Var(param)
self.source = np.zeros(param.sizevar)
self.ipva = self.var.varname_list.index('pvanom')
self.ipv = self.var.varname_list.index('pv')
self.ivor = self.var.varname_list.index('vorticity')
self.ipsi = self.var.varname_list.index('psi')
# background pv
self.pvback = self.beta*(grid.yr-grid.Ly*.5)*grid.msk
# for operators
param.tracer_list = ['pv']
if param.ageostrophic:
param.tracer_list += ['ua', 'va']
param.whosetspsi = ('pvanom')
param.qgoperator = True
self.ope = Operators(param, grid)
# for timescheme
self.tscheme = Timescheme(param, self.var.state)
self.dx0 = self.tscheme.dx0
self.kt = 0
if self.forcing:
try:
f = import_module(self.forcing_module)
except:
print('module %s for forcing cannot be found'
% self.forcing_module)
print('make sure file **%s.py** exists' % self.forcing_module)
exit(0)
self.forc = f.Forcing(param, grid)
self.diags = {}
self.tscheme.set(self.dynamics, self.timestepping)
def step(self, t, dt):
"""Integrate the model over one time step dt"""
if self.ageostrophic:
iu = self.var.varname_list.index('u')
iv = self.var.varname_list.index('v')
iua = self.var.varname_list.index('ua')
iva = self.var.varname_list.index('va')
self.var.state[iua][:, :] = self.var.state[iu].copy()
self.var.state[iva][:, :] = self.var.state[iv].copy()
self.dt = dt
# 1/ integrate advection
self.tscheme.forward(self.var.state, t, dt)
# 2/ integrate source
if self.noslip:
self.add_noslip(self.var.state)
# 3/ diagnostic fields
self.var.state[self.ipva] = self.var.state[self.ipv] - self.pvback
self.var.state[self.ivor] = (self.var.state[self.ipva]
- (self.Rd**-2)*self.var.state[self.ipsi])
if self.ageostrophic:
""" diagnose the ageostrophic velocity (factor 1/f missing)
du/dt + J(psi, u) = -f va
dv/dt + J(psi, v) = +f ua
implementation:
- store u in 'ua'
- step ua
- compare update u (real u diagnosed from psi) with updated 'ua'
- the difference is the ageostrophic velocity
Yet to be done: implement the proper staggering, 'u' and 'v' sit
on edges while the advection scheme assumes cell centers quantities
"""
ua = -(self.var.state[iv] - self.var.state[iva])/dt
va = +(self.var.state[iu] - self.var.state[iua])/dt
self.var.state[iva][:, :] = va.copy()
self.var.state[iua][:, :] = ua.copy()
def dynamics(self, x, t, dxdt):
"""Compute the tendency terms + invert the streamfunction"""
dxdt[:] = 0.
self.ope.rhs_adv_multilayers(x, t, dxdt)
if (self.tscheme.kstage == self.tscheme.kforcing):
if self.forcing:
self.forc.add_forcing(x, t, dxdt)
if self.diffusion:
self.ope.rhs_diffusion(x, t, dxdt)
# substract the background pv
dxdt[self.ipva][:, :] = dxdt[self.ipv] # -self.pvback*self.dt
# self.ope.invert_vorticity(dxdt,flag='fast')
self.ope.invert_vorticity(dxdt, flag='fast')
# self.kt +=1
def add_noslip(self, x):
"""Add the no-slip term """
self.ope.rhs_noslip(x, self.source)
self.ope.invert_vorticity(x, flag='fast', island=self.isisland)
def add_backgroundpv(self):
self.var.state[self.ipv] += self.pvback
def set_psi_from_pv(self):
"""Solve the elliptic equation for the streamfunction
The unknown of the Helmholtz equation is the PV anomaly,
not the full PV"""
state = self.var.state
state[self.ipva] = state[self.ipv] - self.pvback
self.ope.invert_vorticity(self.var.state, flag='full')
def diagnostics(self, var, t):
""" Integral diagnostics for the QG model
should provide at least 'maxspeed' (for cfl determination) """
u = var.get('u')
v = var.get('v')
trac = var.get('pv')
ke, maxu = fd.computekemaxu(self.msk, u, v, self.nh)
z, z2 = fd.computesumandnorm(self.msk, trac, self.nh)
cst = self.mpitools.local_to_global([(maxu, 'max'),
(ke, 'sum'),
(z, 'sum'),
(z2, 'sum')])
self.diags['maxspeed'] = cst[0]
self.diags['ke'] = cst[1] / self.area
self.diags['pv'] = cst[2] / self.area
self.diags['pv2'] = cst[3] / self.area
