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