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)
