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