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