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