class Gradient(DiagnosticField): def __init__(self, name): super().__init__() self.fname = name @property def name(self): return self.fname+"_gradient" def setup(self, state): if not self._initialised: mesh_dim = state.mesh.geometric_dimension() try: field_dim = state.fields(self.fname).ufl_shape[0] except IndexError: field_dim = 1 shape = (mesh_dim, ) * field_dim space = TensorFunctionSpace(state.mesh, "CG", 1, shape=shape) super().setup(state, space=space) f = state.fields(self.fname) test = TestFunction(space) trial = TrialFunction(space) n = FacetNormal(state.mesh) a = inner(test, trial)*dx L = -inner(div(test), f)*dx if space.extruded: L += dot(dot(test, n), f)*(ds_t + ds_b) prob = LinearVariationalProblem(a, L, self.field) self.solver = LinearVariationalSolver(prob) def compute(self, state): self.solver.solve() return self.field
class HydrostaticImbalance(DiagnosticField): name = "HydrostaticImbalance" def setup(self, state): if not self._initialised: space = state.spaces("Vv") super().setup(state, space=space) rho = state.fields("rho") rhobar = state.fields("rhobar") theta = state.fields("theta") thetabar = state.fields("thetabar") pi = thermodynamics.pi(state.parameters, rho, theta) pibar = thermodynamics.pi(state.parameters, rhobar, thetabar) cp = Constant(state.parameters.cp) n = FacetNormal(state.mesh) F = TrialFunction(space) w = TestFunction(space) a = inner(w, F)*dx L = (- cp*div((theta-thetabar)*w)*pibar*dx + cp*jump((theta-thetabar)*w, n)*avg(pibar)*dS_v - cp*div(thetabar*w)*(pi-pibar)*dx + cp*jump(thetabar*w, n)*avg(pi-pibar)*dS_v) bcs = [DirichletBC(space, 0.0, "bottom"), DirichletBC(space, 0.0, "top")] imbalanceproblem = LinearVariationalProblem(a, L, self.field, bcs=bcs) self.imbalance_solver = LinearVariationalSolver(imbalanceproblem) def compute(self, state): self.imbalance_solver.solve() return self.field[1]
class Precipitation(DiagnosticField): name = "Precipitation" def setup(self, state): if not self._initialised: space = state.spaces("DG0", state.mesh, "DG", 0) super().setup(state, space=space) rain = state.fields('rain') rho = state.fields('rho') v = state.fields('rainfall_velocity') self.phi = TestFunction(space) flux = TrialFunction(space) n = FacetNormal(state.mesh) un = 0.5 * (dot(v, n) + abs(dot(v, n))) self.flux = Function(space) a = self.phi * flux * dx L = self.phi * rain * un * rho if space.extruded: L = L * (ds_b + ds_t + ds_v) else: L = L * ds # setup solver problem = LinearVariationalProblem(a, L, self.flux) self.solver = LinearVariationalSolver(problem) def compute(self, state): self.solver.solve() self.field.assign(self.field + assemble(self.flux * self.phi * dx)) return self.field
class TrueResidualV(DiagnosticField): name = "TrueResidualV" def setup(self, state): super(TrueResidualV, self).setup(state) unew, pnew, bnew = state.xn.split() uold, pold, bold = state.xb.split() ubar = 0.5 * (unew + uold) H = state.parameters.H f = state.parameters.f dbdy = state.parameters.dbdy dt = state.timestepping.dt x, y, z = SpatialCoordinate(state.mesh) V = FunctionSpace(state.mesh, "DG", 0) wv = TestFunction(V) v = TrialFunction(V) vlhs = wv * v * dx vrhs = wv * ((unew[1] - uold[1]) / dt + ubar[0] * ubar[1].dx(0) + ubar[2] * ubar[1].dx(2) + f * ubar[0] + dbdy * (z - H / 2)) * dx self.vtres = Function(V) vtresproblem = LinearVariationalProblem(vlhs, vrhs, self.vtres) self.v_residual_solver = LinearVariationalSolver( vtresproblem, solver_parameters={'ksp_type': 'cg'}) def compute(self, state): self.v_residual_solver.solve() v_residual = self.vtres return self.field.interpolate(v_residual)
class GeostrophicImbalance(DiagnosticField): name = "GeostrophicImbalance" def setup(self, state): super(GeostrophicImbalance, self).setup(state) u = state.fields("u") b = state.fields("b") p = state.fields("p") f = state.parameters.f Vu = u.function_space() v = TrialFunction(Vu) w = TestFunction(Vu) a = inner(w, v) * dx L = (div(w) * p + inner(w, as_vector([f * u[1], 0.0, b]))) * dx bcs = [DirichletBC(Vu, 0.0, "bottom"), DirichletBC(Vu, 0.0, "top")] self.imbalance = Function(Vu) imbalanceproblem = LinearVariationalProblem(a, L, self.imbalance, bcs=bcs) self.imbalance_solver = LinearVariationalSolver( imbalanceproblem, solver_parameters={'ksp_type': 'cg'}) def compute(self, state): f = state.parameters.f self.imbalance_solver.solve() geostrophic_imbalance = self.imbalance[0] / f return self.field.interpolate(geostrophic_imbalance)
class CompressibleEadyForcing(CompressibleForcing): """ Forcing class for compressible Eady equations. """ def forcing_term(self): # L = super(EadyForcing, self).forcing_term() L = Forcing.forcing_term(self) dthetady = self.state.parameters.dthetady Pi0 = self.state.parameters.Pi0 cp = self.state.parameters.cp _, rho0, theta0 = split(self.x0) Pi = thermodynamics.pi(self.state.parameters, rho0, theta0) Pi_0 = Constant(Pi0) L += self.scaling * cp * dthetady * (Pi - Pi_0) * inner( self.test, as_vector([0., 1., 0.])) * dx # Eady forcing return L def _build_forcing_solvers(self): super(CompressibleEadyForcing, self)._build_forcing_solvers() # theta_forcing dthetady = self.state.parameters.dthetady Vt = self.state.spaces("HDiv_v") F = TrialFunction(Vt) gamma = TestFunction(Vt) self.thetaF = Function(Vt) u0, _, _ = split(self.x0) a = gamma * F * dx L = -self.scaling * gamma * (dthetady * inner(u0, as_vector([0., 1., 0.]))) * dx theta_forcing_problem = LinearVariationalProblem(a, L, self.thetaF) solver_parameters = {} if logger.isEnabledFor(DEBUG): solver_parameters["ksp_monitor_true_residual"] = None self.theta_forcing_solver = LinearVariationalSolver( theta_forcing_problem, solver_parameters=solver_parameters, options_prefix="ThetaForcingSolver") def apply(self, scaling, x_in, x_nl, x_out, **kwargs): Forcing.apply(self, scaling, x_in, x_nl, x_out, **kwargs) self.theta_forcing_solver.solve() # places forcing in self.thetaF _, _, theta_out = x_out.split() theta_out += self.thetaF
class EadyForcing(IncompressibleForcing): """ Forcing class for Eady Boussinesq equations. """ def forcing_term(self): L = Forcing.forcing_term(self) dbdy = self.state.parameters.dbdy H = self.state.parameters.H Vp = self.state.spaces("DG") _, _, z = SpatialCoordinate(self.state.mesh) eady_exp = Function(Vp).interpolate(z - H / 2.) L -= self.scaling * dbdy * eady_exp * inner( self.test, as_vector([0., 1., 0.])) * dx return L def _build_forcing_solvers(self): super(EadyForcing, self)._build_forcing_solvers() # b_forcing dbdy = self.state.parameters.dbdy Vb = self.state.spaces("HDiv_v") F = TrialFunction(Vb) gamma = TestFunction(Vb) self.bF = Function(Vb) u0, _, b0 = split(self.x0) a = gamma * F * dx L = -self.scaling * gamma * (dbdy * inner(u0, as_vector([0., 1., 0.]))) * dx b_forcing_problem = LinearVariationalProblem(a, L, self.bF) solver_parameters = {} if logger.isEnabledFor(DEBUG): solver_parameters["ksp_monitor_true_residual"] = None self.b_forcing_solver = LinearVariationalSolver( b_forcing_problem, solver_parameters=solver_parameters, options_prefix="BForcingSolver") def apply(self, scaling, x_in, x_nl, x_out, **kwargs): super(EadyForcing, self).apply(scaling, x_in, x_nl, x_out, **kwargs) self.b_forcing_solver.solve() # places forcing in self.bF _, _, b_out = x_out.split() b_out += self.bF
class Vorticity(DiagnosticField): """Base diagnostic field class for vorticity.""" def setup(self, state, vorticity_type=None): """Solver for vorticity. :arg state: The state containing model. :arg vorticity_type: must be "relative", "absolute" or "potential" """ if not self._initialised: vorticity_types = ["relative", "absolute", "potential"] if vorticity_type not in vorticity_types: raise ValueError("vorticity type must be one of %s, not %s" % (vorticity_types, vorticity_type)) try: space = state.spaces("CG") except AttributeError: dgspace = state.spaces("DG") cg_degree = dgspace.ufl_element().degree() + 2 space = FunctionSpace(state.mesh, "CG", cg_degree) super().setup(state, space=space) u = state.fields("u") gamma = TestFunction(space) q = TrialFunction(space) if vorticity_type == "potential": D = state.fields("D") a = q*gamma*D*dx else: a = q*gamma*dx if state.on_sphere: cell_normals = CellNormal(state.mesh) gradperp = lambda psi: cross(cell_normals, grad(psi)) L = (- inner(gradperp(gamma), u))*dx else: raise NotImplementedError("The vorticity diagnostics have only been implemented for 2D spherical geometries.") if vorticity_type != "relative": f = state.fields("coriolis") L += gamma*f*dx problem = LinearVariationalProblem(a, L, self.field) self.solver = LinearVariationalSolver(problem, solver_parameters={"ksp_type": "cg"}) def compute(self, state): """Computes the vorticity. """ self.solver.solve() return self.field
class ShallowWaterForcing(Forcing): def __init__(self, state, linear=False): self.state = state g = state.parameters.g f = state.f Vu = state.V[0] W = state.W self.x0 = Function(W) # copy x to here u0, D0 = split(self.x0) n = FacetNormal(state.mesh) un = 0.5 * (dot(u0, n) + abs(dot(u0, n))) F = TrialFunction(Vu) w = TestFunction(Vu) self.uF = Function(Vu) outward_normals = CellNormal(state.mesh) perp = lambda u: cross(outward_normals, u) a = inner(w, F) * dx L = (-f * inner(w, perp(u0)) + g * div(w) * D0) * dx - g * inner( jump(w, n), un("+") * D0("+") - un("-") * D0("-") ) * dS if not linear: L -= 0.5 * div(w) * inner(u0, u0) * dx u_forcing_problem = LinearVariationalProblem(a, L, self.uF) self.u_forcing_solver = LinearVariationalSolver(u_forcing_problem) def apply(self, scaling, x_in, x_nl, x_out, **kwargs): self.x0.assign(x_nl) self.u_forcing_solver.solve() # places forcing in self.uF self.uF *= scaling uF, _ = x_out.split() x_out.assign(x_in) uF += self.uF
class IncompressibleForcing(Forcing): """ Forcing class for incompressible Euler Boussinesq equations. """ def pressure_gradient_term(self): _, p0, _ = split(self.x0) L = div(self.test) * p0 * dx return L def gravity_term(self): _, _, b0 = split(self.x0) L = b0 * inner(self.test, self.state.k) * dx return L def _build_forcing_solvers(self): super(IncompressibleForcing, self)._build_forcing_solvers() Vp = self.state.spaces("DG") p = TrialFunction(Vp) q = TestFunction(Vp) self.divu = Function(Vp) u0, _, _ = split(self.x0) a = p * q * dx L = q * div(u0) * dx divergence_problem = LinearVariationalProblem(a, L, self.divu) solver_parameters = {} if logger.isEnabledFor(DEBUG): solver_parameters["ksp_monitor_true_residual"] = None self.divergence_solver = LinearVariationalSolver( divergence_problem, solver_parameters=solver_parameters, options_prefix="DivergenceSolver") def apply(self, scaling, x_in, x_nl, x_out, **kwargs): super(IncompressibleForcing, self).apply(scaling, x_in, x_nl, x_out, **kwargs) if 'incompressible' in kwargs and kwargs['incompressible']: _, p_out, _ = x_out.split() self.divergence_solver.solve() p_out.assign(self.divu)
class InteriorPenalty(Diffusion): """ Interior penalty diffusion method :arg state: :class:`.State` object. :arg V: Function space of diffused field :arg direction: list containing directions in which function space :arg: mu: the penalty weighting function, which is :recommended to be proportional to 1/dx :arg: kappa: strength of diffusion :arg: bcs: (optional) a list of boundary conditions to apply """ def __init__(self, state, V, kappa, mu, bcs=None): super(InteriorPenalty, self).__init__(state) dt = state.timestepping.dt gamma = TestFunction(V) phi = TrialFunction(V) self.phi1 = Function(V) n = FacetNormal(state.mesh) a = inner(gamma, phi) * dx + dt * inner(grad(gamma), grad(phi) * kappa) * dx def get_flux_form(dS, M): fluxes = (-inner(2 * avg(outer(phi, n)), avg(grad(gamma) * M)) - inner(avg(grad(phi) * M), 2 * avg(outer(gamma, n))) + mu * inner(2 * avg(outer(phi, n)), 2 * avg(outer(gamma, n) * kappa))) * dS return fluxes a += dt * get_flux_form(dS_v, kappa) a += dt * get_flux_form(dS_h, kappa) L = inner(gamma, phi) * dx problem = LinearVariationalProblem(a, action(L, self.phi1), self.phi1, bcs=bcs) self.solver = LinearVariationalSolver(problem) def apply(self, x_in, x_out): self.phi1.assign(x_in) self.solver.solve() x_out.assign(self.phi1)
class InteriorPenalty(Diffusion): """ Interior penalty diffusion method :arg state: :class:`.State` object. :arg V: Function space of diffused field :arg direction: list containing directions in which function space is discontinuous: 1 corresponds to vertical, 2 to horizontal. :arg params: dictionary containing the interior penalty parameters :mu and kappa where mu is the penalty weighting function, which is :recommended to be proportional to 1/dx """ def __init__(self, state, V, direction=[1,2], params=None): super(InteriorPenalty, self).__init__(state) dt = state.timestepping.dt kappa = params['kappa'] mu = params['mu'] gamma = TestFunction(V) phi = TrialFunction(V) self.phi1 = Function(V) n = FacetNormal(state.mesh) a = inner(gamma,phi)*dx + dt*inner(grad(gamma), grad(phi)*kappa)*dx def get_flux_form(dS, M): fluxes = (-inner(2*avg(outer(phi, n)), avg(grad(gamma)*M)) - inner(avg(grad(phi)*M), 2*avg(outer(gamma, n))) + mu*inner(2*avg(outer(phi, n)), 2*avg(outer(gamma, n)*kappa)))*dS return fluxes if 1 in direction: a += dt*get_flux_form(dS_v, kappa) if 2 in direction: a += dt*get_flux_form(dS_h, kappa) L = inner(gamma,phi)*dx problem = LinearVariationalProblem(a, action(L,self.phi1), self.phi1) self.solver = LinearVariationalSolver(problem) def apply(self, x_in, x_out): self.phi1.assign(x_in) self.solver.solve() x_out.assign(self.phi1)
class ShallowWaterSolver(TimesteppingSolver): def _setup_solver(self): state = self.state H = state.parameters.H g = state.parameters.g beta = state.timestepping.dt*state.timestepping.alpha # Split up the rhs vector (symbolically) u_in, D_in = split(state.xrhs) W = state.W w, phi = TestFunctions(W) u, D = TrialFunctions(W) eqn = ( inner(w, u) - beta*g*div(w)*D - inner(w, u_in) + phi*D + beta*H*phi*div(u) - phi*D_in )*dx aeqn = lhs(eqn) Leqn = rhs(eqn) # Place to put result of u rho solver self.uD = Function(W) # Solver for u, D uD_problem = LinearVariationalProblem( aeqn, Leqn, self.state.dy) self.uD_solver = LinearVariationalSolver(uD_problem, solver_parameters=self.params, options_prefix='SWimplicit') def solve(self): """ Apply the solver with rhs state.xrhs and result state.dy. """ self.uD_solver.solve()
class LinearAdvection_V3(Advection): """ An advection scheme that uses the linearised background state in evaluation of the advection term for a DG space. :arg state: :class:`.State` object. :arg V:class:`.FunctionSpace` object. The DG Function space. :arg qbar: :class:`.Function` object. The reference function that we are linearising around. :arg options: a PETSc options dictionary """ def __init__(self, state, V, qbar, options=None): super(LinearAdvection_V3, self).__init__(state) p = TestFunction(V) q = TrialFunction(V) self.dq = Function(V) n = FacetNormal(state.mesh) a = p*q*dx L = (dot(grad(p), self.ubar)*qbar*dx - jump(self.ubar*p, n)*avg(qbar)*(dS_v + dS_h)) aProblem = LinearVariationalProblem(a,L,self.dq) if options is None: options = {'ksp_type':'cg', 'pc_type':'bjacobi', 'sub_pc_type':'ilu'} self.solver = LinearVariationalSolver(aProblem, solver_parameters=options, options_prefix='LinearAdvectionV3') def apply(self, x_in, x_out): dt = self.state.timestepping.dt self.solver.solve() x_out.assign(x_in + dt*self.dq)
class LinearAdvection_Vt(Advection): """ An advection scheme that uses the linearised background state in evaluation of the advection term for the Vt temperature space. :arg state: :class:`.State` object. :arg V:class:`.FunctionSpace` object. The Function space for temperature. :arg qbar: :class:`.Function` object. The reference function that we are linearising around. :arg options: a PETSc options dictionary """ def __init__(self, state, V, qbar, options=None): super(LinearAdvection_Vt, self).__init__(state) p = TestFunction(V) q = TrialFunction(V) self.dq = Function(V) a = p*q*dx k = state.k # Upward pointing unit vector L = -p*dot(self.ubar,k)*dot(k,grad(qbar))*dx aProblem = LinearVariationalProblem(a,L,self.dq) if options is None: options = {'ksp_type':'cg', 'pc_type':'bjacobi', 'sub_pc_type':'ilu'} self.solver = LinearVariationalSolver(aProblem, solver_parameters=options, options_prefix='LinearAdvectionVt') def apply(self, x_in, x_out): dt = self.state.timestepping.dt self.solver.solve() x_out.assign(x_in + dt*self.dq)
r = sqrt((x - xc)**2 + (z - zc)**2) theta_pert = Function(Vt).interpolate( conditional(r > rc, 0.0, Tdash * (cos(pi * r / (2.0 * rc)))**2)) # define initial theta theta0.assign(theta_b * (theta_pert / 300.0 + 1.0)) # find perturbed rho gamma = TestFunction(Vr) rho_trial = TrialFunction(Vr) lhs = gamma * rho_trial * dx rhs = gamma * (rho_b * theta_b / theta0) * dx rho_problem = LinearVariationalProblem(lhs, rhs, rho0) rho_solver = LinearVariationalSolver(rho_problem) rho_solver.solve() state.set_reference_profiles([('rho', rho_b), ('theta', theta_b)]) # Set up transport schemes VDG1 = state.spaces("DG1_equispaced") VCG1 = FunctionSpace(mesh, "CG", 1) Vt_brok = FunctionSpace(mesh, BrokenElement(Vt.ufl_element())) Vu_DG1 = VectorFunctionSpace(mesh, VDG1.ufl_element()) Vu_CG1 = VectorFunctionSpace(mesh, "CG", 1) Vu_brok = FunctionSpace(mesh, BrokenElement(Vu.ufl_element())) u_opts = RecoveredOptions(embedding_space=Vu_DG1, recovered_space=Vu_CG1, broken_space=Vu_brok, boundary_method=Boundary_Method.dynamics)
class ShallowWaterSolver(TimesteppingSolver): """ Timestepping linear solver object for the nonlinear shallow water equations with prognostic variables u and D. The linearized system is solved using a hybridized-mixed method. """ solver_parameters = { 'ksp_type': 'preonly', 'mat_type': 'matfree', 'pc_type': 'python', 'pc_python_type': 'firedrake.HybridizationPC', 'hybridization': {'ksp_type': 'cg', 'pc_type': 'gamg', 'ksp_rtol': 1e-8, 'mg_levels': {'ksp_type': 'chebyshev', 'ksp_max_it': 2, 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'}} } @timed_function("Gusto:SolverSetup") def _setup_solver(self): state = self.state H_ = state.parameters.H g_ = state.parameters.g beta_ = state.timestepping.dt*state.timestepping.alpha # Store time-stepping coefficients as UFL Constants beta = Constant(beta_) H = Constant(H_) g = Constant(g_) # Split up the rhs vector (symbolically) u_in, D_in = split(state.xrhs) W = state.W w, phi = TestFunctions(W) u, D = TrialFunctions(W) eqn = ( inner(w, u) - beta*g*div(w)*D - inner(w, u_in) + phi*D + beta*H*phi*div(u) - phi*D_in )*dx aeqn = lhs(eqn) Leqn = rhs(eqn) # Place to put result of u rho solver self.uD = Function(W) # Solver for u, D uD_problem = LinearVariationalProblem( aeqn, Leqn, self.state.dy) self.uD_solver = LinearVariationalSolver(uD_problem, solver_parameters=self.solver_parameters, options_prefix='SWimplicit') @timed_function("Gusto:LinearSolve") def solve(self): """ Apply the solver with rhs state.xrhs and result state.dy. """ self.uD_solver.solve()
class OldCompressibleSolver(TimesteppingSolver): """ Timestepping linear solver object for the compressible equations in theta-pi formulation with prognostic variables u,rho,theta. This solver follows the following strategy: (1) Analytically eliminate theta (introduces error near topography) (2) Solve resulting system for (u,rho) using a Schur preconditioner (3) Reconstruct theta :arg state: a :class:`.State` object containing everything else. :arg quadrature degree: tuple (q_h, q_v) where q_h is the required quadrature degree in the horizontal direction and q_v is that in the vertical direction :arg solver_parameters (optional): solver parameters :arg overwrite_solver_parameters: boolean, if True use only the solver_parameters that have been passed in, if False then update the default solver parameters with the solver_parameters passed in. :arg moisture (optional): list of names of moisture fields. """ solver_parameters = { 'pc_type': 'fieldsplit', 'pc_fieldsplit_type': 'schur', 'ksp_type': 'gcr', 'ksp_monitor_true_residual': None, 'ksp_max_it': 100, 'pc_fieldsplit_schur_fact_type': 'FULL', 'pc_fieldsplit_schur_precondition': 'selfp', 'fieldsplit_0': {'ksp_type': 'preonly', 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'}, 'fieldsplit_1': {'ksp_type': 'fgmres', 'ksp_monitor_true_residual': None, 'ksp_rtol': 1.0e-8, 'ksp_atol': 1.0e-8, 'ksp_max_it': 100, 'pc_type': 'gamg', 'pc_gamg_sym_graph': None, 'mg_levels': {'ksp_type': 'gmres', 'ksp_max_it': 5, 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'}} } def __init__(self, state, quadrature_degree=None, solver_parameters=None, overwrite_solver_parameters=False, moisture=None): self.moisture = moisture if quadrature_degree is not None: self.quadrature_degree = quadrature_degree else: dgspace = state.spaces("DG") if any(deg > 2 for deg in dgspace.ufl_element().degree()): logger.warning("default quadrature degree most likely not sufficient for this degree element") self.quadrature_degree = (5, 5) super().__init__(state, solver_parameters, overwrite_solver_parameters) @timed_function("Gusto:SolverSetup") def _setup_solver(self): state = self.state # just cutting down line length a bit Dt = state.timestepping.dt beta_ = Dt*state.timestepping.alpha cp = state.parameters.cp mu = state.mu Vu = state.spaces("HDiv") Vtheta = state.spaces("HDiv_v") Vrho = state.spaces("DG") # Store time-stepping coefficients as UFL Constants dt = Constant(Dt) beta = Constant(beta_) beta_cp = Constant(beta_ * cp) # Split up the rhs vector (symbolically) u_in, rho_in, theta_in = split(state.xrhs) # Build the reduced function space for u,rho M = MixedFunctionSpace((Vu, Vrho)) w, phi = TestFunctions(M) u, rho = TrialFunctions(M) n = FacetNormal(state.mesh) # Get background fields thetabar = state.fields("thetabar") rhobar = state.fields("rhobar") pibar = thermodynamics.pi(state.parameters, rhobar, thetabar) pibar_rho = thermodynamics.pi_rho(state.parameters, rhobar, thetabar) pibar_theta = thermodynamics.pi_theta(state.parameters, rhobar, thetabar) # Analytical (approximate) elimination of theta k = state.k # Upward pointing unit vector theta = -dot(k, u)*dot(k, grad(thetabar))*beta + theta_in # Only include theta' (rather than pi') in the vertical # component of the gradient # the pi prime term (here, bars are for mean and no bars are # for linear perturbations) pi = pibar_theta*theta + pibar_rho*rho # vertical projection def V(u): return k*inner(u, k) # specify degree for some terms as estimated degree is too large dxp = dx(degree=(self.quadrature_degree)) dS_vp = dS_v(degree=(self.quadrature_degree)) # add effect of density of water upon theta if self.moisture is not None: water_t = Function(Vtheta).assign(0.0) for water in self.moisture: water_t += self.state.fields(water) theta_w = theta / (1 + water_t) thetabar_w = thetabar / (1 + water_t) else: theta_w = theta thetabar_w = thetabar eqn = ( inner(w, (state.h_project(u) - u_in))*dx - beta_cp*div(theta_w*V(w))*pibar*dxp # following does nothing but is preserved in the comments # to remind us why (because V(w) is purely vertical). # + beta_cp*jump(theta*V(w), n)*avg(pibar)*dS_v - beta_cp*div(thetabar_w*w)*pi*dxp + beta_cp*jump(thetabar_w*w, n)*avg(pi)*dS_vp + (phi*(rho - rho_in) - beta*inner(grad(phi), u)*rhobar)*dx + beta*jump(phi*u, n)*avg(rhobar)*(dS_v + dS_h) ) if mu is not None: eqn += dt*mu*inner(w, k)*inner(u, k)*dx aeqn = lhs(eqn) Leqn = rhs(eqn) # Place to put result of u rho solver self.urho = Function(M) # Boundary conditions (assumes extruded mesh) bcs = [DirichletBC(M.sub(0), 0.0, "bottom"), DirichletBC(M.sub(0), 0.0, "top")] # Solver for u, rho urho_problem = LinearVariationalProblem( aeqn, Leqn, self.urho, bcs=bcs) self.urho_solver = LinearVariationalSolver(urho_problem, solver_parameters=self.solver_parameters, options_prefix='ImplicitSolver') # Reconstruction of theta theta = TrialFunction(Vtheta) gamma = TestFunction(Vtheta) u, rho = self.urho.split() self.theta = Function(Vtheta) theta_eqn = gamma*(theta - theta_in + dot(k, u)*dot(k, grad(thetabar))*beta)*dx theta_problem = LinearVariationalProblem(lhs(theta_eqn), rhs(theta_eqn), self.theta) self.theta_solver = LinearVariationalSolver(theta_problem, options_prefix='thetabacksubstitution') @timed_function("Gusto:SchurCompLinearSolve") def solve(self): """ Apply the solver with rhs state.xrhs and result state.dy. """ with timed_region("Gusto:VelocityDensitySolve"): self.urho_solver.solve() u1, rho1 = self.urho.split() u, rho, theta = self.state.dy.split() u.assign(u1) rho.assign(rho1) with timed_region("Gusto:ThetaRecon"): self.theta_solver.solve() theta.assign(self.theta)
perp = lambda u: cross(CellNormal(mesh), u) out_n = CellNormal(mesh) p_uw = lambda u, v: (s(u)('+') * cross(out_n('+'), v('+')) + s(u) ('-') * cross(out_n('-'), v('-'))) # Initial solve for vorticity for output eta = TestFunction(W2) q_ = TrialFunction(W2) q_eqn = eta * q_ * Dn * dx + inner(perp(grad(eta)), un) * dx - eta * f * dx q_p = LinearVariationalProblem(lhs(q_eqn), rhs(q_eqn), qn) q_solver = LinearVariationalSolver(q_p, solver_parameters={ "ksp_type": "preonly", "pc_type": "lu" }) q_solver.solve() # Build advection, forcing forms K = inner(un, un) / 3. + inner(un, unk) / 3. + inner(unk, unk) / 3. Prhs = g * (0.5 * (Dn + Dp) + b) + 0.5 * K # D advection solver v, psi = TestFunctions(M) F_, D_ = TrialFunctions(M) D_bar = 0.5 * (Dn + D_) Frhs = unk * D_ / 3. + un * D_ / 6. + unk * Dn / 6. + un * Dn / 3. D_ad = -psi * div(F_) * dx D_eqn = (D_ - Dn) * psi * dx - dt * D_ad + inner(v, F_ - Frhs) * dx Dad_p = LinearVariationalProblem(lhs(D_eqn), rhs(D_eqn), xp) D_ad_solver = LinearVariationalSolver(Dad_p,
perp = lambda u: cross(CellNormal(mesh), u) out_n = CellNormal(mesh) p_uw = lambda u, v: (s(u)('+') * cross(out_n('+'), v('+')) + s(u) ('-') * cross(out_n('-'), v('-'))) # Initial solve for vorticity for output eta = TestFunction(W2) q_ = TrialFunction(W2) q_eqn = eta * q_ * Dn * dx + inner(perp(grad(eta)), un) * dx - eta * f * dx q_p = LinearVariationalProblem(lhs(q_eqn), rhs(q_eqn), qn) q_solver = LinearVariationalSolver(q_p, solver_parameters={ "ksp_type": "preonly", "pc_type": "lu" }) q_solver.solve() q0.assign(qn) # Build advection, forcing forms M_ext = MixedFunctionSpace((W1, W0, W0)) xp_ext = Function(M_ext) u_, D_, P_ = TrialFunctions(M_ext) v, psi, chi = TestFunctions(M_ext) D_bar = 0.5 * (Dn + D_) u_bar = 0.5 * (un + u_) Frhs = unk * Dnk / 3. + un * Dnk / 6. + unk * Dn / 6. + un * Dn / 3. K = inner(un, un) / 3. + inner(un, unk) / 3. + inner(unk, unk) / 3. Prhs = g * (D_bar + b) + 0.5 * K
def compressible_hydrostatic_balance(state, theta0, rho0, pi0=None, top=False, pi_boundary=Constant(1.0), solve_for_rho=False, params=None): """ Compute a hydrostatically balanced density given a potential temperature profile. :arg state: The :class:`State` object. :arg theta0: :class:`.Function`containing the potential temperature. :arg rho0: :class:`.Function` to write the initial density into. :arg top: If True, set a boundary condition at the top. Otherwise, set it at the bottom. :arg pi_boundary: a field or expression to use as boundary data for pi on the top or bottom as specified. """ # Calculate hydrostatic Pi W = MixedFunctionSpace((state.Vv,state.V[1])) v, pi = TrialFunctions(W) dv, dpi = TestFunctions(W) n = FacetNormal(state.mesh) cp = state.parameters.cp alhs = ( (cp*inner(v,dv) - cp*div(dv*theta0)*pi)*dx + dpi*div(theta0*v)*dx ) if top: bmeasure = ds_t bstring = "bottom" else: bmeasure = ds_b bstring = "top" arhs = -cp*inner(dv,n)*theta0*pi_boundary*bmeasure if state.parameters.geopotential: Phi = state.Phi arhs += div(dv)*Phi*dx - inner(dv,n)*Phi*bmeasure else: g = state.parameters.g arhs -= g*inner(dv,state.k)*dx if(state.mesh.geometric_dimension() == 2): bcs = [DirichletBC(W.sub(0), Expression(("0.", "0.")), bstring)] elif(state.mesh.geometric_dimension() == 3): bcs = [DirichletBC(W.sub(0), Expression(("0.", "0.", "0.")), bstring)] w = Function(W) PiProblem = LinearVariationalProblem(alhs, arhs, w, bcs=bcs) if(params is None): params = {'pc_type': 'fieldsplit', 'pc_fieldsplit_type': 'schur', 'ksp_type': 'gmres', 'ksp_monitor_true_residual': True, 'ksp_max_it': 100, 'ksp_gmres_restart': 50, 'pc_fieldsplit_schur_fact_type': 'FULL', 'pc_fieldsplit_schur_precondition': 'selfp', 'fieldsplit_0_ksp_type': 'richardson', 'fieldsplit_0_ksp_max_it': 5, 'fieldsplit_0_pc_type': 'gamg', 'fieldsplit_1_pc_gamg_sym_graph': True, 'fieldsplit_1_mg_levels_ksp_type': 'chebyshev', 'fieldsplit_1_mg_levels_ksp_chebyshev_estimate_eigenvalues': True, 'fieldsplit_1_mg_levels_ksp_chebyshev_estimate_eigenvalues_random': True, 'fieldsplit_1_mg_levels_ksp_max_it': 5, 'fieldsplit_1_mg_levels_pc_type': 'bjacobi', 'fieldsplit_1_mg_levels_sub_pc_type': 'ilu'} PiSolver = LinearVariationalSolver(PiProblem, solver_parameters=params) PiSolver.solve() v, Pi = w.split() if pi0 is not None: pi0.assign(Pi) kappa = state.parameters.kappa R_d = state.parameters.R_d p_0 = state.parameters.p_0 if solve_for_rho: w1 = Function(W) v, rho = w1.split() rho.interpolate(p_0*(Pi**((1-kappa)/kappa))/R_d/theta0) v, rho = split(w1) dv, dpi = TestFunctions(W) pi = ((R_d/p_0)*rho*theta0)**(kappa/(1.-kappa)) F = ( (cp*inner(v,dv) - cp*div(dv*theta0)*pi)*dx + dpi*div(theta0*v)*dx + cp*inner(dv,n)*theta0*pi_boundary*bmeasure ) if state.parameters.geopotential: F += - div(dv)*Phi*dx + inner(dv,n)*Phi*bmeasure else: F += g*inner(dv,state.k)*dx rhoproblem = NonlinearVariationalProblem(F, w1, bcs=bcs) rhosolver = NonlinearVariationalSolver(rhoproblem, solver_parameters=params) rhosolver.solve() v, rho_ = w1.split() rho0.assign(rho_) else: rho0.interpolate(p_0*(Pi**((1-kappa)/kappa))/R_d/theta0)
class DGAdvection(Advection): """ DG 3 step SSPRK advection scheme that can be applied to a scalar or vector field :arg state: :class:`.State` object. :arg V: function space of advected field - should be DG :arg continuity: optional boolean. If ``True``, the advection equation is of the form: :math: `D_t +\nabla \cdot(uD) = 0`. If ``False``, the advection equation is of the form: :math: `D_t + (u \cdot \nabla)D = 0`. """ def __init__(self, state, V, continuity=False): super(DGAdvection, self).__init__(state) element = V.fiat_element assert element.entity_dofs() == element.entity_closure_dofs(), "Provided space is not discontinuous" dt = state.timestepping.dt if V.extruded: surface_measure = (dS_h + dS_v) else: surface_measure = dS phi = TestFunction(V) D = TrialFunction(V) self.D1 = Function(V) self.dD = Function(V) n = FacetNormal(state.mesh) # ( dot(v, n) + |dot(v, n)| )/2.0 un = 0.5*(dot(self.ubar, n) + abs(dot(self.ubar, n))) a_mass = inner(phi,D)*dx if continuity: a_int = -inner(grad(phi), outer(D, self.ubar))*dx else: a_int = -inner(div(outer(phi,self.ubar)),D)*dx a_flux = (dot(jump(phi), un('+')*D('+') - un('-')*D('-')))*surface_measure arhs = a_mass - dt*(a_int + a_flux) DGproblem = LinearVariationalProblem(a_mass, action(arhs,self.D1), self.dD) self.DGsolver = LinearVariationalSolver(DGproblem, solver_parameters={ 'ksp_type':'preonly', 'pc_type':'bjacobi', 'sub_pc_type': 'ilu'}, options_prefix='DGAdvection') def apply(self, x_in, x_out): # SSPRK Stage 1 self.D1.assign(x_in) self.DGsolver.solve() self.D1.assign(self.dD) # SSPRK Stage 2 self.DGsolver.solve() self.D1.assign(0.75*x_in + 0.25*self.dD) # SSPRK Stage 3 self.DGsolver.solve() x_out.assign((1.0/3.0)*x_in + (2.0/3.0)*self.dD)
class SawyerEliassenU(DiagnosticField): name = "SawyerEliassenU" def setup(self, state): space = state.spaces("HDiv") super(SawyerEliassenU, self).setup(state, space=space) u = state.fields("u") b = state.fields("b") v = inner(u, as_vector([0., 1., 0.])) # spaces V0 = FunctionSpace(state.mesh, "CG", 2) Vu = u.function_space() # project b to V0 self.b_v0 = Function(V0) btri = TrialFunction(V0) btes = TestFunction(V0) a = inner(btes, btri) * dx L = inner(btes, b) * dx projectbproblem = LinearVariationalProblem(a, L, self.b_v0) self.project_b_solver = LinearVariationalSolver( projectbproblem, solver_parameters={'ksp_type': 'cg'}) # project v to V0 self.v_v0 = Function(V0) vtri = TrialFunction(V0) vtes = TestFunction(V0) a = inner(vtes, vtri) * dx L = inner(vtes, v) * dx projectvproblem = LinearVariationalProblem(a, L, self.v_v0) self.project_v_solver = LinearVariationalSolver( projectvproblem, solver_parameters={'ksp_type': 'cg'}) # stm/psi is a stream function self.stm = Function(V0) psi = TrialFunction(V0) xsi = TestFunction(V0) f = state.parameters.f H = state.parameters.H L = state.parameters.L dbdy = state.parameters.dbdy x, y, z = SpatialCoordinate(state.mesh) bcs = [DirichletBC(V0, 0., "bottom"), DirichletBC(V0, 0., "top")] Mat = as_matrix([[b.dx(2), 0., -f * self.v_v0.dx(2)], [0., 0., 0.], [-self.b_v0.dx(0), 0., f**2 + f * self.v_v0.dx(0)]]) Equ = (inner(grad(xsi), dot(Mat, grad(psi))) - dbdy * inner(grad(xsi), as_vector([-v, 0., f * (z - H / 2)]))) * dx # fourth-order terms if state.parameters.fourthorder: eps = Constant(0.0001) brennersigma = Constant(10.0) n = FacetNormal(state.mesh) deltax = Constant(state.parameters.deltax) deltaz = Constant(state.parameters.deltaz) nn = as_matrix([[sqrt(brennersigma / Constant(deltax)), 0., 0.], [0., 0., 0.], [0., 0., sqrt(brennersigma / Constant(deltaz))]]) mu = as_matrix([[1., 0., 0.], [0., 0., 0.], [0., 0., H / L]]) # anisotropic form Equ += eps * ( div(dot(mu, grad(psi))) * div(dot(mu, grad(xsi))) * dx - (avg(dot(dot(grad(grad(psi)), n), n)) * jump(grad(xsi), n=n) + avg(dot(dot(grad(grad(xsi)), n), n)) * jump(grad(psi), n=n) - jump(nn * grad(psi), n=n) * jump(nn * grad(xsi), n=n)) * (dS_h + dS_v)) Au = lhs(Equ) Lu = rhs(Equ) stmproblem = LinearVariationalProblem(Au, Lu, self.stm, bcs=bcs) self.stream_function_solver = LinearVariationalSolver( stmproblem, solver_parameters={'ksp_type': 'cg'}) # solve for sawyer_eliassen u self.u = Function(Vu) utrial = TrialFunction(Vu) w = TestFunction(Vu) a = inner(w, utrial) * dx L = (w[0] * (-self.stm.dx(2)) + w[2] * (self.stm.dx(0))) * dx ugproblem = LinearVariationalProblem(a, L, self.u) self.sawyer_eliassen_u_solver = LinearVariationalSolver( ugproblem, solver_parameters={'ksp_type': 'cg'}) def compute(self, state): self.project_b_solver.solve() self.project_v_solver.solve() self.stream_function_solver.solve() self.sawyer_eliassen_u_solver.solve() sawyer_eliassen_u = self.u return self.field.project(sawyer_eliassen_u)
def compressible_hydrostatic_balance(state, theta0, rho0, exner0=None, top=False, exner_boundary=Constant(1.0), mr_t=None, solve_for_rho=False, params=None): """ Compute a hydrostatically balanced density given a potential temperature profile. By default, this uses a vertically-oriented hybridization procedure for solving the resulting discrete systems. :arg state: The :class:`State` object. :arg theta0: :class:`.Function`containing the potential temperature. :arg rho0: :class:`.Function` to write the initial density into. :arg top: If True, set a boundary condition at the top. Otherwise, set it at the bottom. :arg exner_boundary: a field or expression to use as boundary data for exner on the top or bottom as specified. :arg mr_t: the initial total water mixing ratio field. """ # Calculate hydrostatic Pi VDG = state.spaces("DG") Vu = state.spaces("HDiv") Vv = FunctionSpace(state.mesh, Vu.ufl_element()._elements[-1]) W = MixedFunctionSpace((Vv, VDG)) v, exner = TrialFunctions(W) dv, dexner = TestFunctions(W) n = FacetNormal(state.mesh) cp = state.parameters.cp # add effect of density of water upon theta theta = theta0 if mr_t is not None: theta = theta0 / (1 + mr_t) alhs = ((cp * inner(v, dv) - cp * div(dv * theta) * exner) * dx + dexner * div(theta * v) * dx) if top: bmeasure = ds_t bstring = "bottom" else: bmeasure = ds_b bstring = "top" arhs = -cp * inner(dv, n) * theta * exner_boundary * bmeasure # Possibly make g vary with spatial coordinates? g = state.parameters.g arhs -= g * inner(dv, state.k) * dx bcs = [DirichletBC(W.sub(0), zero(), bstring)] w = Function(W) exner_problem = LinearVariationalProblem(alhs, arhs, w, bcs=bcs) if params is None: params = { 'ksp_type': 'preonly', 'pc_type': 'python', 'mat_type': 'matfree', 'pc_python_type': 'gusto.VerticalHybridizationPC', # Vertical trace system is only coupled vertically in columns # block ILU is a direct solver! 'vert_hybridization': { 'ksp_type': 'preonly', 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu' } } exner_solver = LinearVariationalSolver(exner_problem, solver_parameters=params, options_prefix="exner_solver") exner_solver.solve() v, exner = w.split() if exner0 is not None: exner0.assign(exner) if solve_for_rho: w1 = Function(W) v, rho = w1.split() rho.interpolate(thermodynamics.rho(state.parameters, theta0, exner)) v, rho = split(w1) dv, dexner = TestFunctions(W) exner = thermodynamics.exner_pressure(state.parameters, rho, theta0) F = ((cp * inner(v, dv) - cp * div(dv * theta) * exner) * dx + dexner * div(theta0 * v) * dx + cp * inner(dv, n) * theta * exner_boundary * bmeasure) F += g * inner(dv, state.k) * dx rhoproblem = NonlinearVariationalProblem(F, w1, bcs=bcs) rhosolver = NonlinearVariationalSolver(rhoproblem, solver_parameters=params, options_prefix="rhosolver") rhosolver.solve() v, rho_ = w1.split() rho0.assign(rho_) else: rho0.interpolate(thermodynamics.rho(state.parameters, theta0, exner))
class CompressibleSolver(TimesteppingSolver): """ Timestepping linear solver object for the compressible equations in theta-exner formulation with prognostic variables u, rho, and theta. This solver follows the following strategy: (1) Analytically eliminate theta (introduces error near topography) (2a) Formulate the resulting mixed system for u and rho using a hybridized mixed method. This breaks continuity in the linear perturbations of u, and introduces a new unknown on the mesh interfaces approximating the average of the Exner pressure perturbations. These trace unknowns also act as Lagrange multipliers enforcing normal continuity of the "broken" u variable. (2b) Statically condense the block-sparse system into a single system for the Lagrange multipliers. This is the only globally coupled system requiring a linear solver. (2c) Using the computed trace variables, we locally recover the broken velocity and density perturbations. This is accomplished in two stages: (i): Recover rho locally using the multipliers. (ii): Recover "broken" u locally using rho and the multipliers. (2d) Project the "broken" velocity field into the HDiv-conforming space using local averaging. (3) Reconstruct theta :arg state: a :class:`.State` object containing everything else. :arg quadrature degree: tuple (q_h, q_v) where q_h is the required quadrature degree in the horizontal direction and q_v is that in the vertical direction. :arg solver_parameters (optional): solver parameters for the trace system. :arg overwrite_solver_parameters: boolean, if True use only the solver_parameters that have been passed in, if False then update. the default solver parameters with the solver_parameters passed in. :arg moisture (optional): list of names of moisture fields. """ solver_parameters = { 'mat_type': 'matfree', 'ksp_type': 'preonly', 'pc_type': 'python', 'pc_python_type': 'firedrake.SCPC', 'pc_sc_eliminate_fields': '0, 1', # The reduced operator is not symmetric 'condensed_field': { 'ksp_type': 'fgmres', 'ksp_rtol': 1.0e-8, 'ksp_atol': 1.0e-8, 'ksp_max_it': 100, 'pc_type': 'gamg', 'pc_gamg_sym_graph': None, 'mg_levels': { 'ksp_type': 'gmres', 'ksp_max_it': 5, 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu' } } } def __init__(self, state, equations, alpha=0.5, quadrature_degree=None, solver_parameters=None, overwrite_solver_parameters=False, moisture=None): self.moisture = moisture if quadrature_degree is not None: self.quadrature_degree = quadrature_degree else: dgspace = state.spaces("DG") if any(deg > 2 for deg in dgspace.ufl_element().degree()): logger.warning( "default quadrature degree most likely not sufficient for this degree element" ) self.quadrature_degree = (5, 5) if logger.isEnabledFor(DEBUG): # Set outer solver to FGMRES and turn on KSP monitor for the outer system self.solver_parameters["ksp_type"] = "fgmres" self.solver_parameters["mat_type"] = "aij" self.solver_parameters["pmat_type"] = "matfree" self.solver_parameters["ksp_monitor_true_residual"] = None # Turn monitor on for the trace system self.solver_parameters["condensed_field"][ "ksp_monitor_true_residual"] = None super().__init__(state, equations, alpha, solver_parameters, overwrite_solver_parameters) @timed_function("Gusto:SolverSetup") def _setup_solver(self): import numpy as np state = self.state dt = state.dt beta_ = dt * self.alpha cp = state.parameters.cp Vu = state.spaces("HDiv") Vu_broken = FunctionSpace(state.mesh, BrokenElement(Vu.ufl_element())) Vtheta = state.spaces("theta") Vrho = state.spaces("DG") # Store time-stepping coefficients as UFL Constants beta = Constant(beta_) beta_cp = Constant(beta_ * cp) h_deg = Vrho.ufl_element().degree()[0] v_deg = Vrho.ufl_element().degree()[1] Vtrace = FunctionSpace(state.mesh, "HDiv Trace", degree=(h_deg, v_deg)) # Split up the rhs vector (symbolically) self.xrhs = Function(self.equations.function_space) u_in, rho_in, theta_in = split(self.xrhs)[0:3] # Build the function space for "broken" u, rho, and pressure trace M = MixedFunctionSpace((Vu_broken, Vrho, Vtrace)) w, phi, dl = TestFunctions(M) u, rho, l0 = TrialFunctions(M) n = FacetNormal(state.mesh) # Get background fields thetabar = state.fields("thetabar") rhobar = state.fields("rhobar") exnerbar = thermodynamics.exner_pressure(state.parameters, rhobar, thetabar) exnerbar_rho = thermodynamics.dexner_drho(state.parameters, rhobar, thetabar) exnerbar_theta = thermodynamics.dexner_dtheta(state.parameters, rhobar, thetabar) # Analytical (approximate) elimination of theta k = state.k # Upward pointing unit vector theta = -dot(k, u) * dot(k, grad(thetabar)) * beta + theta_in # Only include theta' (rather than exner') in the vertical # component of the gradient # The exner prime term (here, bars are for mean and no bars are # for linear perturbations) exner = exnerbar_theta * theta + exnerbar_rho * rho # Vertical projection def V(u): return k * inner(u, k) # hydrostatic projection h_project = lambda u: u - k * inner(u, k) # Specify degree for some terms as estimated degree is too large dxp = dx(degree=(self.quadrature_degree)) dS_vp = dS_v(degree=(self.quadrature_degree)) dS_hp = dS_h(degree=(self.quadrature_degree)) ds_vp = ds_v(degree=(self.quadrature_degree)) ds_tbp = (ds_t(degree=(self.quadrature_degree)) + ds_b(degree=(self.quadrature_degree))) # Add effect of density of water upon theta if self.moisture is not None: water_t = Function(Vtheta).assign(0.0) for water in self.moisture: water_t += self.state.fields(water) theta_w = theta / (1 + water_t) thetabar_w = thetabar / (1 + water_t) else: theta_w = theta thetabar_w = thetabar _l0 = TrialFunction(Vtrace) _dl = TestFunction(Vtrace) a_tr = _dl('+') * _l0('+') * ( dS_vp + dS_hp) + _dl * _l0 * ds_vp + _dl * _l0 * ds_tbp def L_tr(f): return _dl('+') * avg(f) * ( dS_vp + dS_hp) + _dl * f * ds_vp + _dl * f * ds_tbp cg_ilu_parameters = { 'ksp_type': 'cg', 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu' } # Project field averages into functions on the trace space rhobar_avg = Function(Vtrace) exnerbar_avg = Function(Vtrace) rho_avg_prb = LinearVariationalProblem(a_tr, L_tr(rhobar), rhobar_avg) exner_avg_prb = LinearVariationalProblem(a_tr, L_tr(exnerbar), exnerbar_avg) rho_avg_solver = LinearVariationalSolver( rho_avg_prb, solver_parameters=cg_ilu_parameters, options_prefix='rhobar_avg_solver') exner_avg_solver = LinearVariationalSolver( exner_avg_prb, solver_parameters=cg_ilu_parameters, options_prefix='exnerbar_avg_solver') with timed_region("Gusto:HybridProjectRhobar"): rho_avg_solver.solve() with timed_region("Gusto:HybridProjectExnerbar"): exner_avg_solver.solve() # "broken" u, rho, and trace system # NOTE: no ds_v integrals since equations are defined on # a periodic (or sphere) base mesh. if any([t.has_label(hydrostatic) for t in self.equations.residual]): u_mass = inner(w, (h_project(u) - u_in)) * dx else: u_mass = inner(w, (u - u_in)) * dx eqn = ( # momentum equation u_mass - beta_cp * div(theta_w * V(w)) * exnerbar * dxp # following does nothing but is preserved in the comments # to remind us why (because V(w) is purely vertical). # + beta_cp*jump(theta_w*V(w), n=n)*exnerbar_avg('+')*dS_vp + beta_cp * jump(theta_w * V(w), n=n) * exnerbar_avg('+') * dS_hp + beta_cp * dot(theta_w * V(w), n) * exnerbar_avg * ds_tbp - beta_cp * div(thetabar_w * w) * exner * dxp # trace terms appearing after integrating momentum equation + beta_cp * jump(thetabar_w * w, n=n) * l0('+') * (dS_vp + dS_hp) + beta_cp * dot(thetabar_w * w, n) * l0 * (ds_tbp + ds_vp) # mass continuity equation + (phi * (rho - rho_in) - beta * inner(grad(phi), u) * rhobar) * dx + beta * jump(phi * u, n=n) * rhobar_avg('+') * (dS_v + dS_h) # term added because u.n=0 is enforced weakly via the traces + beta * phi * dot(u, n) * rhobar_avg * (ds_tb + ds_v) # constraint equation to enforce continuity of the velocity # through the interior facets and weakly impose the no-slip # condition + dl('+') * jump(u, n=n) * (dS_vp + dS_hp) + dl * dot(u, n) * (ds_tbp + ds_vp)) # contribution of the sponge term if hasattr(self.equations, "mu"): eqn += dt * self.equations.mu * inner(w, k) * inner(u, k) * dx aeqn = lhs(eqn) Leqn = rhs(eqn) # Function for the hybridized solutions self.urhol0 = Function(M) hybridized_prb = LinearVariationalProblem(aeqn, Leqn, self.urhol0) hybridized_solver = LinearVariationalSolver( hybridized_prb, solver_parameters=self.solver_parameters, options_prefix='ImplicitSolver') self.hybridized_solver = hybridized_solver # Project broken u into the HDiv space using facet averaging. # Weight function counting the dofs of the HDiv element: shapes = { "i": Vu.finat_element.space_dimension(), "j": np.prod(Vu.shape, dtype=int) } weight_kernel = """ for (int i=0; i<{i}; ++i) for (int j=0; j<{j}; ++j) w[i*{j} + j] += 1.0; """.format(**shapes) self._weight = Function(Vu) par_loop(weight_kernel, dx, {"w": (self._weight, INC)}) # Averaging kernel self._average_kernel = """ for (int i=0; i<{i}; ++i) for (int j=0; j<{j}; ++j) vec_out[i*{j} + j] += vec_in[i*{j} + j]/w[i*{j} + j]; """.format(**shapes) # HDiv-conforming velocity self.u_hdiv = Function(Vu) # Reconstruction of theta theta = TrialFunction(Vtheta) gamma = TestFunction(Vtheta) self.theta = Function(Vtheta) theta_eqn = gamma * (theta - theta_in + dot(k, self.u_hdiv) * dot(k, grad(thetabar)) * beta) * dx theta_problem = LinearVariationalProblem(lhs(theta_eqn), rhs(theta_eqn), self.theta) self.theta_solver = LinearVariationalSolver( theta_problem, solver_parameters=cg_ilu_parameters, options_prefix='thetabacksubstitution') # Store boundary conditions for the div-conforming velocity to apply # post-solve self.bcs = self.equations.bcs['u'] @timed_function("Gusto:LinearSolve") def solve(self, xrhs, dy): """ Apply the solver with rhs xrhs and result dy. """ self.xrhs.assign(xrhs) # Solve the hybridized system self.hybridized_solver.solve() broken_u, rho1, _ = self.urhol0.split() u1 = self.u_hdiv # Project broken_u into the HDiv space u1.assign(0.0) with timed_region("Gusto:HybridProjectHDiv"): par_loop( self._average_kernel, dx, { "w": (self._weight, READ), "vec_in": (broken_u, READ), "vec_out": (u1, INC) }) # Reapply bcs to ensure they are satisfied for bc in self.bcs: bc.apply(u1) # Copy back into u and rho cpts of dy u, rho, theta = dy.split()[0:3] u.assign(u1) rho.assign(rho1) # Reconstruct theta with timed_region("Gusto:ThetaRecon"): self.theta_solver.solve() # Copy into theta cpt of dy theta.assign(self.theta)
class IncompressibleForcing(Forcing): """ Forcing class for incompressible Euler Boussinesq equations. """ def __init__(self, state, linear=False): self.state = state self._build_forcing_solver(linear) def _build_forcing_solver(self, linear): """ Only put forcing terms into the u equation. """ state = self.state self.scaling = Constant(1.0) Vu = state.V[0] W = state.W self.x0 = Function(W) # copy x to here u0, p0, b0 = split(self.x0) F = TrialFunction(Vu) w = TestFunction(Vu) self.uF = Function(Vu) Omega = state.Omega mu = state.mu a = inner(w, F) * dx L = ( self.scaling * div(w) * p0 * dx # pressure gradient + self.scaling * b0 * inner(w, state.k) * dx # gravity term ) if not linear: L -= self.scaling * 0.5 * div(w) * inner(u0, u0) * dx if Omega is not None: L -= self.scaling * inner(w, cross(2 * Omega, u0)) * dx # Coriolis term if mu is not None: self.mu_scaling = Constant(1.0) L -= self.mu_scaling * mu * inner(w, state.k) * inner(u0, state.k) * dx bcs = [DirichletBC(Vu, 0.0, "bottom"), DirichletBC(Vu, 0.0, "top")] u_forcing_problem = LinearVariationalProblem(a, L, self.uF, bcs=bcs) self.u_forcing_solver = LinearVariationalSolver(u_forcing_problem) Vp = state.V[1] p = TrialFunction(Vp) q = TestFunction(Vp) self.divu = Function(Vp) a = p * q * dx L = q * div(u0) * dx divergence_problem = LinearVariationalProblem(a, L, self.divu) self.divergence_solver = LinearVariationalSolver(divergence_problem) def apply(self, scaling, x_in, x_nl, x_out, **kwargs): self.x0.assign(x_nl) self.scaling.assign(scaling) if "mu_alpha" in kwargs and kwargs["mu_alpha"] is not None: self.mu_scaling.assign(kwargs["mu_alpha"]) self.u_forcing_solver.solve() # places forcing in self.uF u_out, p_out, _ = x_out.split() x_out.assign(x_in) u_out += self.uF if "incompressible" in kwargs and kwargs["incompressible"]: self.divergence_solver.solve() p_out.assign(self.divu)
class CompressibleForcing(Forcing): """ Forcing class for compressible Euler equations. """ def __init__(self, state, linear=False): self.state = state self._build_forcing_solver(linear) def _build_forcing_solver(self, linear): """ Only put forcing terms into the u equation. """ state = self.state self.scaling = Constant(1.0) Vu = state.V[0] W = state.W self.x0 = Function(W) # copy x to here u0, rho0, theta0 = split(self.x0) F = TrialFunction(Vu) w = TestFunction(Vu) self.uF = Function(Vu) Omega = state.Omega cp = state.parameters.cp mu = state.mu n = FacetNormal(state.mesh) pi = exner(theta0, rho0, state) a = inner(w, F) * dx L = self.scaling * ( +cp * div(theta0 * w) * pi * dx # pressure gradient [volume] - cp * jump(w * theta0, n) * avg(pi) * dS_v # pressure gradient [surface] ) if state.parameters.geopotential: Phi = state.Phi L += self.scaling * div(w) * Phi * dx # gravity term else: g = state.parameters.g L -= self.scaling * g * inner(w, state.k) * dx # gravity term if not linear: L -= self.scaling * 0.5 * div(w) * inner(u0, u0) * dx if Omega is not None: L -= self.scaling * inner(w, cross(2 * Omega, u0)) * dx # Coriolis term if mu is not None: self.mu_scaling = Constant(1.0) L -= self.mu_scaling * mu * inner(w, state.k) * inner(u0, state.k) * dx bcs = [DirichletBC(Vu, 0.0, "bottom"), DirichletBC(Vu, 0.0, "top")] u_forcing_problem = LinearVariationalProblem(a, L, self.uF, bcs=bcs) self.u_forcing_solver = LinearVariationalSolver(u_forcing_problem) def apply(self, scaling, x_in, x_nl, x_out, **kwargs): self.x0.assign(x_nl) self.scaling.assign(scaling) if "mu_alpha" in kwargs and kwargs["mu_alpha"] is not None: self.mu_scaling.assign(kwargs["mu_alpha"]) self.u_forcing_solver.solve() # places forcing in self.uF u_out, _, _ = x_out.split() x_out.assign(x_in) u_out += self.uF
class EulerPoincareForm(Advection): def __init__(self, state, V): super(EulerPoincareForm, self).__init__(state) dt = state.timestepping.dt w = TestFunction(V) u = TrialFunction(V) self.u0 = Function(V) ustar = 0.5*(self.u0 + u) n = FacetNormal(state.mesh) Upwind = 0.5*(sign(dot(self.ubar, n))+1) if state.mesh.geometric_dimension() == 3: if V.extruded: surface_measure = (dS_h + dS_v) else: surface_measure = dS # <w,curl(u) cross ubar + grad( u.ubar)> # =<curl(u),ubar cross w> - <div(w), u.ubar> # =<u,curl(ubar cross w)> - # <<u_upwind, [[n cross(ubar cross w)cross]]>> both = lambda u: 2*avg(u) Eqn = ( inner(w, u-self.u0)*dx + dt*inner(ustar, curl(cross(self.ubar, w)))*dx - dt*inner(both(Upwind*ustar), both(cross(n, cross(self.ubar, w))))*surface_measure - dt*div(w)*inner(ustar, self.ubar)*dx ) # define surface measure and terms involving perp differently # for slice (i.e. if V.extruded is True) and shallow water # (V.extruded is False) else: if V.extruded: surface_measure = (dS_h + dS_v) perp = lambda u: as_vector([-u[1], u[0]]) perp_u_upwind = Upwind('+')*perp(ustar('+')) + Upwind('-')*perp(ustar('-')) else: surface_measure = dS outward_normals = CellNormal(state.mesh) perp = lambda u: cross(outward_normals, u) perp_u_upwind = Upwind('+')*cross(outward_normals('+'),ustar('+')) + Upwind('-')*cross(outward_normals('-'),ustar('-')) Eqn = ( (inner(w, u-self.u0) - dt*inner(w, div(perp(ustar))*perp(self.ubar)) - dt*div(w)*inner(ustar, self.ubar))*dx - dt*inner(jump(inner(w, perp(self.ubar)), n), perp_u_upwind)*surface_measure + dt*jump(inner(w, perp(self.ubar))*perp(ustar), n)*surface_measure ) a = lhs(Eqn) L = rhs(Eqn) self.u1 = Function(V) uproblem = LinearVariationalProblem(a, L, self.u1) self.usolver = LinearVariationalSolver(uproblem, options_prefix='EPAdvection') def apply(self, x_in, x_out): self.u0.assign(x_in) self.usolver.solve() x_out.assign(self.u1)
class Forcing(object, metaclass=ABCMeta): """ Base class for forcing terms for Gusto. :arg state: x :class:`.State` object. :arg euler_poincare: if True then the momentum equation is in Euler Poincare form and we need to add 0.5*grad(u^2) to the forcing term. If False then this term is not added. :arg linear: if True then we are solving a linear equation so nonlinear terms (namely the Euler Poincare term) should not be added. :arg extra_terms: extra terms to add to the u component of the forcing term - these will be multiplied by the appropriate test function. """ def __init__(self, state, euler_poincare=True, linear=False, extra_terms=None, moisture=None): self.state = state if linear: self.euler_poincare = False logger.warning( 'Setting euler_poincare to False because you have set linear=True' ) else: self.euler_poincare = euler_poincare # set up functions self.Vu = state.spaces("HDiv") # this is the function that the forcing term is applied to self.x0 = Function(state.W) self.test = TestFunction(self.Vu) self.trial = TrialFunction(self.Vu) # this is the function that contains the result of solving # <test, trial> = <test, F(x0)>, where F is the forcing term self.uF = Function(self.Vu) # find out which terms we need self.extruded = self.Vu.extruded self.coriolis = state.Omega is not None or hasattr( state.fields, "coriolis") self.sponge = state.mu is not None self.hydrostatic = state.hydrostatic self.topography = hasattr(state.fields, "topography") self.extra_terms = extra_terms self.moisture = moisture # some constants to use for scaling terms self.scaling = Constant(1.) self.impl = Constant(1.) self._build_forcing_solvers() def mass_term(self): return inner(self.test, self.trial) * dx def coriolis_term(self): u0 = split(self.x0)[0] return -inner(self.test, cross(2 * self.state.Omega, u0)) * dx def sponge_term(self): u0 = split(self.x0)[0] return self.state.mu * inner(self.test, self.state.k) * inner( u0, self.state.k) * dx def euler_poincare_term(self): u0 = split(self.x0)[0] return -0.5 * div(self.test) * inner(self.state.h_project(u0), u0) * dx def hydrostatic_term(self): u0 = split(self.x0)[0] return inner(u0, self.state.k) * inner(self.test, self.state.k) * dx @abstractmethod def pressure_gradient_term(self): pass def forcing_term(self): L = self.pressure_gradient_term() if self.extruded: L += self.gravity_term() if self.coriolis: L += self.coriolis_term() if self.euler_poincare: L += self.euler_poincare_term() if self.topography: L += self.topography_term() if self.extra_terms is not None: L += inner(self.test, self.extra_terms) * dx # scale L L = self.scaling * L # sponge term has a separate scaling factor as it is always implicit if self.sponge: L -= self.impl * self.state.timestepping.dt * self.sponge_term() # hydrostatic term has no scaling factor if self.hydrostatic: L += (2 * self.impl - 1) * self.hydrostatic_term() return L def _build_forcing_solvers(self): a = self.mass_term() L = self.forcing_term() bcs = None if len(self.state.bcs) == 0 else self.state.bcs u_forcing_problem = LinearVariationalProblem(a, L, self.uF, bcs=bcs) solver_parameters = {} if logger.isEnabledFor(DEBUG): solver_parameters["ksp_monitor_true_residual"] = None self.u_forcing_solver = LinearVariationalSolver( u_forcing_problem, solver_parameters=solver_parameters, options_prefix="UForcingSolver") def apply(self, scaling, x_in, x_nl, x_out, **kwargs): """ Function takes x as input, computes F(x_nl) and returns x_out = x + scale*F(x_nl) as output. :arg x_in: :class:`.Function` object :arg x_nl: :class:`.Function` object :arg x_out: :class:`.Function` object :arg implicit: forcing stage for sponge and hydrostatic terms, if present """ self.scaling.assign(scaling) self.x0.assign(x_nl) implicit = kwargs.get("implicit") if implicit is not None: self.impl.assign(int(implicit)) self.u_forcing_solver.solve() # places forcing in self.uF uF = x_out.split()[0] x_out.assign(x_in) uF += self.uF
def compressible_hydrostatic_balance(state, theta0, rho0, pi0=None, top=False, pi_boundary=Constant(1.0), water_t=None, solve_for_rho=False, params=None): """ Compute a hydrostatically balanced density given a potential temperature profile. :arg state: The :class:`State` object. :arg theta0: :class:`.Function`containing the potential temperature. :arg rho0: :class:`.Function` to write the initial density into. :arg top: If True, set a boundary condition at the top. Otherwise, set it at the bottom. :arg pi_boundary: a field or expression to use as boundary data for pi on the top or bottom as specified. :arg water_t: the initial total water mixing ratio field. """ # Calculate hydrostatic Pi VDG = state.spaces("DG") Vv = state.spaces("Vv") W = MixedFunctionSpace((Vv, VDG)) v, pi = TrialFunctions(W) dv, dpi = TestFunctions(W) n = FacetNormal(state.mesh) cp = state.parameters.cp # add effect of density of water upon theta theta = theta0 if water_t is not None: theta = theta0 / (1 + water_t) alhs = ( (cp*inner(v, dv) - cp*div(dv*theta)*pi)*dx + dpi*div(theta*v)*dx ) if top: bmeasure = ds_t bstring = "bottom" else: bmeasure = ds_b bstring = "top" arhs = -cp*inner(dv, n)*theta*pi_boundary*bmeasure g = state.parameters.g arhs -= g*inner(dv, state.k)*dx bcs = [DirichletBC(W.sub(0), 0.0, bstring)] w = Function(W) PiProblem = LinearVariationalProblem(alhs, arhs, w, bcs=bcs) if params is None: params = {'pc_type': 'fieldsplit', 'pc_fieldsplit_type': 'schur', 'ksp_type': 'gmres', 'ksp_monitor_true_residual': True, 'ksp_max_it': 100, 'ksp_gmres_restart': 50, 'pc_fieldsplit_schur_fact_type': 'FULL', 'pc_fieldsplit_schur_precondition': 'selfp', 'fieldsplit_0_ksp_type': 'richardson', 'fieldsplit_0_ksp_max_it': 5, 'fieldsplit_0_pc_type': 'gamg', 'fieldsplit_1_pc_gamg_sym_graph': True, 'fieldsplit_1_mg_levels_ksp_type': 'chebyshev', 'fieldsplit_1_mg_levels_ksp_chebyshev_esteig': True, 'fieldsplit_1_mg_levels_ksp_max_it': 5, 'fieldsplit_1_mg_levels_pc_type': 'bjacobi', 'fieldsplit_1_mg_levels_sub_pc_type': 'ilu'} PiSolver = LinearVariationalSolver(PiProblem, solver_parameters=params) PiSolver.solve() v, Pi = w.split() if pi0 is not None: pi0.assign(Pi) if solve_for_rho: w1 = Function(W) v, rho = w1.split() rho.interpolate(thermodynamics.rho(state.parameters, theta0, Pi)) v, rho = split(w1) dv, dpi = TestFunctions(W) pi = thermodynamics.pi(state.parameters, rho, theta0) F = ( (cp*inner(v, dv) - cp*div(dv*theta)*pi)*dx + dpi*div(theta0*v)*dx + cp*inner(dv, n)*theta*pi_boundary*bmeasure ) F += g*inner(dv, state.k)*dx rhoproblem = NonlinearVariationalProblem(F, w1, bcs=bcs) rhosolver = NonlinearVariationalSolver(rhoproblem, solver_parameters=params) rhosolver.solve() v, rho_ = w1.split() rho0.assign(rho_) else: rho0.interpolate(thermodynamics.rho(state.parameters, theta0, Pi))
class CompressibleForcing(Forcing): """ Forcing class for compressible Euler equations. """ def pressure_gradient_term(self): u0, rho0, theta0 = split(self.x0) cp = self.state.parameters.cp n = FacetNormal(self.state.mesh) Vtheta = self.state.spaces("HDiv_v") # introduce new theta so it can be changed by moisture theta = theta0 # add effect of density of water upon theta if self.moisture is not None: water_t = Function(Vtheta).assign(0.0) for water in self.moisture: water_t += self.state.fields(water) theta = theta / (1 + water_t) pi = thermodynamics.pi(self.state.parameters, rho0, theta0) L = (+cp * div(theta * self.test) * pi * dx - cp * jump(self.test * theta, n) * avg(pi) * dS_v) return L def gravity_term(self): g = self.state.parameters.g L = -g * inner(self.test, self.state.k) * dx return L def theta_forcing(self): cv = self.state.parameters.cv cp = self.state.parameters.cp c_vv = self.state.parameters.c_vv c_pv = self.state.parameters.c_pv c_pl = self.state.parameters.c_pl R_d = self.state.parameters.R_d R_v = self.state.parameters.R_v u0, _, theta0 = split(self.x0) water_v = self.state.fields('water_v') water_c = self.state.fields('water_c') c_vml = cv + water_v * c_vv + water_c * c_pl c_pml = cp + water_v * c_pv + water_c * c_pl R_m = R_d + water_v * R_v L = -theta0 * (R_m / c_vml - (R_d * c_pml) / (cp * c_vml)) * div(u0) return self.scaling * L def _build_forcing_solvers(self): super(CompressibleForcing, self)._build_forcing_solvers() # build forcing for theta equation if self.moisture is not None: _, _, theta0 = split(self.x0) Vt = self.state.spaces("HDiv_v") p = TrialFunction(Vt) q = TestFunction(Vt) self.thetaF = Function(Vt) a = p * q * dx L = self.theta_forcing() L = q * L * dx theta_problem = LinearVariationalProblem(a, L, self.thetaF) solver_parameters = {} if logger.isEnabledFor(DEBUG): solver_parameters["ksp_monitor_true_residual"] = None self.theta_solver = LinearVariationalSolver( theta_problem, solver_parameters=solver_parameters, option_prefix="ThetaForcingSolver") def apply(self, scaling, x_in, x_nl, x_out, **kwargs): super(CompressibleForcing, self).apply(scaling, x_in, x_nl, x_out, **kwargs) if self.moisture is not None: self.theta_solver.solve() _, _, theta_out = x_out.split() theta_out += self.thetaF
class SUPGAdvection(Advection): """ An SUPG advection scheme that can apply DG upwinding (in the direction specified by the direction arg) if the function space is only partially continuous. :arg state: :class:`.State` object. :arg V:class:`.FunctionSpace` object. The advected field function space. :arg direction: list containing the directions in which the function space is discontinuous. 1 corresponds to the vertical direction, 2 to the horizontal direction :arg supg_params: dictionary containing SUPG parameters tau for each direction. If not supplied tau is set to dt/sqrt(15.) """ def __init__(self, state, V, direction=[], supg_params=None): super(SUPGAdvection, self).__init__(state) dt = state.timestepping.dt params = supg_params.copy() if supg_params else {} params.setdefault('a0', dt/sqrt(15.)) params.setdefault('a1', dt/sqrt(15.)) gamma = TestFunction(V) theta = TrialFunction(V) self.theta0 = Function(V) # make SUPG test function taus = [params["a0"], params["a1"]] for i in direction: taus[i] = 0.0 tau = Constant(((taus[0], 0.), (0., taus[1]))) dgamma = dot(dot(self.ubar, tau), grad(gamma)) gammaSU = gamma + dgamma n = FacetNormal(state.mesh) un = 0.5*(dot(self.ubar, n) + abs(dot(self.ubar, n))) a_mass = gammaSU*theta*dx arhs = a_mass - dt*gammaSU*dot(self.ubar, grad(theta))*dx if 1 in direction: arhs -= ( dt*dot(jump(gammaSU), (un('+')*theta('+') - un('-')*theta('-')))*dS_v - dt*(gammaSU('+')*dot(self.ubar('+'), n('+'))*theta('+') + gammaSU('-')*dot(self.ubar('-'), n('-'))*theta('-'))*dS_v ) if 2 in direction: arhs -= ( dt*dot(jump(gammaSU), (un('+')*theta('+') - un('-')*theta('-')))*dS_h - dt*(gammaSU('+')*dot(self.ubar('+'), n('+'))*theta('+') + gammaSU('-')*dot(self.ubar('-'), n('-'))*theta('-'))*dS_h ) self.theta1 = Function(V) self.dtheta = Function(V) problem = LinearVariationalProblem(a_mass, action(arhs,self.theta1), self.dtheta) self.solver = LinearVariationalSolver(problem, options_prefix='SUPGAdvection') def apply(self, x_in, x_out): # SSPRK Stage 1 self.theta1.assign(x_in) self.solver.solve() self.theta1.assign(self.dtheta) # SSPRK Stage 2 self.solver.solve() self.theta1.assign(0.75*x_in + 0.25*self.dtheta) # SSPRK Stage 3 self.solver.solve() x_out.assign((1.0/3.0)*x_in + (2.0/3.0)*self.dtheta)
def _setup_solver(self): import numpy as np state = self.state dt = state.dt beta_ = dt * self.alpha cp = state.parameters.cp Vu = state.spaces("HDiv") Vu_broken = FunctionSpace(state.mesh, BrokenElement(Vu.ufl_element())) Vtheta = state.spaces("theta") Vrho = state.spaces("DG") # Store time-stepping coefficients as UFL Constants beta = Constant(beta_) beta_cp = Constant(beta_ * cp) h_deg = Vrho.ufl_element().degree()[0] v_deg = Vrho.ufl_element().degree()[1] Vtrace = FunctionSpace(state.mesh, "HDiv Trace", degree=(h_deg, v_deg)) # Split up the rhs vector (symbolically) self.xrhs = Function(self.equations.function_space) u_in, rho_in, theta_in = split(self.xrhs)[0:3] # Build the function space for "broken" u, rho, and pressure trace M = MixedFunctionSpace((Vu_broken, Vrho, Vtrace)) w, phi, dl = TestFunctions(M) u, rho, l0 = TrialFunctions(M) n = FacetNormal(state.mesh) # Get background fields thetabar = state.fields("thetabar") rhobar = state.fields("rhobar") exnerbar = thermodynamics.exner_pressure(state.parameters, rhobar, thetabar) exnerbar_rho = thermodynamics.dexner_drho(state.parameters, rhobar, thetabar) exnerbar_theta = thermodynamics.dexner_dtheta(state.parameters, rhobar, thetabar) # Analytical (approximate) elimination of theta k = state.k # Upward pointing unit vector theta = -dot(k, u) * dot(k, grad(thetabar)) * beta + theta_in # Only include theta' (rather than exner') in the vertical # component of the gradient # The exner prime term (here, bars are for mean and no bars are # for linear perturbations) exner = exnerbar_theta * theta + exnerbar_rho * rho # Vertical projection def V(u): return k * inner(u, k) # hydrostatic projection h_project = lambda u: u - k * inner(u, k) # Specify degree for some terms as estimated degree is too large dxp = dx(degree=(self.quadrature_degree)) dS_vp = dS_v(degree=(self.quadrature_degree)) dS_hp = dS_h(degree=(self.quadrature_degree)) ds_vp = ds_v(degree=(self.quadrature_degree)) ds_tbp = (ds_t(degree=(self.quadrature_degree)) + ds_b(degree=(self.quadrature_degree))) # Add effect of density of water upon theta if self.moisture is not None: water_t = Function(Vtheta).assign(0.0) for water in self.moisture: water_t += self.state.fields(water) theta_w = theta / (1 + water_t) thetabar_w = thetabar / (1 + water_t) else: theta_w = theta thetabar_w = thetabar _l0 = TrialFunction(Vtrace) _dl = TestFunction(Vtrace) a_tr = _dl('+') * _l0('+') * ( dS_vp + dS_hp) + _dl * _l0 * ds_vp + _dl * _l0 * ds_tbp def L_tr(f): return _dl('+') * avg(f) * ( dS_vp + dS_hp) + _dl * f * ds_vp + _dl * f * ds_tbp cg_ilu_parameters = { 'ksp_type': 'cg', 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu' } # Project field averages into functions on the trace space rhobar_avg = Function(Vtrace) exnerbar_avg = Function(Vtrace) rho_avg_prb = LinearVariationalProblem(a_tr, L_tr(rhobar), rhobar_avg) exner_avg_prb = LinearVariationalProblem(a_tr, L_tr(exnerbar), exnerbar_avg) rho_avg_solver = LinearVariationalSolver( rho_avg_prb, solver_parameters=cg_ilu_parameters, options_prefix='rhobar_avg_solver') exner_avg_solver = LinearVariationalSolver( exner_avg_prb, solver_parameters=cg_ilu_parameters, options_prefix='exnerbar_avg_solver') with timed_region("Gusto:HybridProjectRhobar"): rho_avg_solver.solve() with timed_region("Gusto:HybridProjectExnerbar"): exner_avg_solver.solve() # "broken" u, rho, and trace system # NOTE: no ds_v integrals since equations are defined on # a periodic (or sphere) base mesh. if any([t.has_label(hydrostatic) for t in self.equations.residual]): u_mass = inner(w, (h_project(u) - u_in)) * dx else: u_mass = inner(w, (u - u_in)) * dx eqn = ( # momentum equation u_mass - beta_cp * div(theta_w * V(w)) * exnerbar * dxp # following does nothing but is preserved in the comments # to remind us why (because V(w) is purely vertical). # + beta_cp*jump(theta_w*V(w), n=n)*exnerbar_avg('+')*dS_vp + beta_cp * jump(theta_w * V(w), n=n) * exnerbar_avg('+') * dS_hp + beta_cp * dot(theta_w * V(w), n) * exnerbar_avg * ds_tbp - beta_cp * div(thetabar_w * w) * exner * dxp # trace terms appearing after integrating momentum equation + beta_cp * jump(thetabar_w * w, n=n) * l0('+') * (dS_vp + dS_hp) + beta_cp * dot(thetabar_w * w, n) * l0 * (ds_tbp + ds_vp) # mass continuity equation + (phi * (rho - rho_in) - beta * inner(grad(phi), u) * rhobar) * dx + beta * jump(phi * u, n=n) * rhobar_avg('+') * (dS_v + dS_h) # term added because u.n=0 is enforced weakly via the traces + beta * phi * dot(u, n) * rhobar_avg * (ds_tb + ds_v) # constraint equation to enforce continuity of the velocity # through the interior facets and weakly impose the no-slip # condition + dl('+') * jump(u, n=n) * (dS_vp + dS_hp) + dl * dot(u, n) * (ds_tbp + ds_vp)) # contribution of the sponge term if hasattr(self.equations, "mu"): eqn += dt * self.equations.mu * inner(w, k) * inner(u, k) * dx aeqn = lhs(eqn) Leqn = rhs(eqn) # Function for the hybridized solutions self.urhol0 = Function(M) hybridized_prb = LinearVariationalProblem(aeqn, Leqn, self.urhol0) hybridized_solver = LinearVariationalSolver( hybridized_prb, solver_parameters=self.solver_parameters, options_prefix='ImplicitSolver') self.hybridized_solver = hybridized_solver # Project broken u into the HDiv space using facet averaging. # Weight function counting the dofs of the HDiv element: shapes = { "i": Vu.finat_element.space_dimension(), "j": np.prod(Vu.shape, dtype=int) } weight_kernel = """ for (int i=0; i<{i}; ++i) for (int j=0; j<{j}; ++j) w[i*{j} + j] += 1.0; """.format(**shapes) self._weight = Function(Vu) par_loop(weight_kernel, dx, {"w": (self._weight, INC)}) # Averaging kernel self._average_kernel = """ for (int i=0; i<{i}; ++i) for (int j=0; j<{j}; ++j) vec_out[i*{j} + j] += vec_in[i*{j} + j]/w[i*{j} + j]; """.format(**shapes) # HDiv-conforming velocity self.u_hdiv = Function(Vu) # Reconstruction of theta theta = TrialFunction(Vtheta) gamma = TestFunction(Vtheta) self.theta = Function(Vtheta) theta_eqn = gamma * (theta - theta_in + dot(k, self.u_hdiv) * dot(k, grad(thetabar)) * beta) * dx theta_problem = LinearVariationalProblem(lhs(theta_eqn), rhs(theta_eqn), self.theta) self.theta_solver = LinearVariationalSolver( theta_problem, solver_parameters=cg_ilu_parameters, options_prefix='thetabacksubstitution') # Store boundary conditions for the div-conforming velocity to apply # post-solve self.bcs = self.equations.bcs['u']
class LinearTimesteppingSolver(object): """ Timestepping linear solver object for the nonlinear shallow water equations with prognostic variables u and D. The linearized system is solved using a hybridized-mixed method. """ solver_parameters = { 'ksp_type': 'preonly', 'mat_type': 'matfree', 'pc_type': 'python', 'pc_python_type': 'firedrake.HybridizationPC', 'hybridization': { 'ksp_type': 'cg', 'pc_type': 'gamg', 'ksp_rtol': 1e-8, 'mg_levels': { 'ksp_type': 'chebyshev', 'ksp_max_it': 2, 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu' } } } def __init__(self, equation, alpha): residual = equation.residual.label_map( lambda t: t.has_label(linearisation), lambda t: Term(t.get(linearisation).form, t.labels), drop) dt = equation.state.dt W = equation.function_space beta = dt * alpha # Split up the rhs vector (symbolically) self.xrhs = Function(W) aeqn = residual.label_map( lambda t: (t.has_label(time_derivative) and t.has_label(linearisation)), map_if_false=lambda t: beta * t) Leqn = residual.label_map( lambda t: (t.has_label(time_derivative) and t.has_label(linearisation)), map_if_false=drop) # Place to put result of solver self.dy = Function(W) # Solver bcs = equation.bcs['u'] problem = LinearVariationalProblem(aeqn.form, action(Leqn.form, self.xrhs), self.dy, bcs=bcs) self.solver = LinearVariationalSolver( problem, solver_parameters=self.solver_parameters, options_prefix='linear_solver') @timed_function("Gusto:LinearSolve") def solve(self, xrhs, dy): """ Apply the solver with rhs xrhs and result dy. """ self.xrhs.assign(xrhs) self.solver.solve() dy.assign(self.dy)
class IncompressibleSolver(TimesteppingSolver): """Timestepping linear solver object for the incompressible Boussinesq equations with prognostic variables u, p, b. This solver follows the following strategy: (1) Analytically eliminate b (introduces error near topography) (2) Solve resulting system for (u,p) using a block Hdiv preconditioner (3) Reconstruct b This currently requires a (parallel) direct solver so is probably a bit memory-hungry, we'll improve this with a hybridised solver soon. :arg state: a :class:`.State` object containing everything else. :arg L: the width of the domain, used in the preconditioner. :arg params: Solver parameters. """ def __init__(self, state, L, params=None): self.state = state if params is None: self.params = {'ksp_type':'gmres', 'pc_type':'fieldsplit', 'pc_fieldsplit_type':'additive', 'fieldsplit_0_pc_type':'lu', 'fieldsplit_1_pc_type':'lu', 'fieldsplit_0_pc_factor_mat_solver_package': 'mumps', 'fieldsplit_0_pc_factor_mat_solver_package': 'mumps', 'fieldsplit_0_ksp_type':'preonly', 'fieldsplit_1_ksp_type':'preonly'} else: self.params = params self.L = L # setup the solver self._setup_solver() def _setup_solver(self): state = self.state # just cutting down line length a bit dt = state.timestepping.dt beta = dt*state.timestepping.alpha mu = state.mu # Split up the rhs vector (symbolically) u_in, p_in, b_in = split(state.xrhs) # Build the reduced function space for u,p M = MixedFunctionSpace((state.V[0], state.V[1])) w, phi = TestFunctions(M) u, p = TrialFunctions(M) # Get background fields bbar = state.bbar # Analytical (approximate) elimination of theta k = state.k # Upward pointing unit vector b = -dot(k,u)*dot(k,grad(bbar))*beta + b_in # vertical projection def V(u): return k*inner(u,k) eqn = ( inner(w, (u - u_in))*dx - beta*div(w)*p*dx - beta*inner(w,k)*b*dx + phi*div(u)*dx ) if mu is not None: eqn += dt*mu*inner(w,k)*inner(u,k)*dx aeqn = lhs(eqn) Leqn = rhs(eqn) # Place to put result of u p solver self.up = Function(M) # Boundary conditions (assumes extruded mesh) dim = M.sub(0).ufl_element().value_shape()[0] bc = ("0.0",)*dim bcs = [DirichletBC(M.sub(0), Expression(bc), "bottom"), DirichletBC(M.sub(0), Expression(bc), "top")] # preconditioner equation L = self.L Ap = ( inner(w,u) + L*L*div(w)*div(u) + phi*p/L/L )*dx # Solver for u, p up_problem = LinearVariationalProblem( aeqn, Leqn, self.up, bcs=bcs, aP=Ap) nullspace = MixedVectorSpaceBasis(M, [M.sub(0), VectorSpaceBasis(constant=True)]) self.up_solver = LinearVariationalSolver(up_problem, solver_parameters=self.params, nullspace=nullspace) # Reconstruction of b b = TrialFunction(state.V[2]) gamma = TestFunction(state.V[2]) u, p = self.up.split() self.b = Function(state.V[2]) b_eqn = gamma*(b - b_in + dot(k,u)*dot(k,grad(bbar))*beta)*dx b_problem = LinearVariationalProblem(lhs(b_eqn), rhs(b_eqn), self.b) self.b_solver = LinearVariationalSolver(b_problem) def solve(self): """ Apply the solver with rhs state.xrhs and result state.dy. """ self.up_solver.solve() u1, p1 = self.up.split() u, p, b = self.state.dy.split() u.assign(u1) p.assign(p1) self.b_solver.solve() b.assign(self.b)
class IncompressibleSolver(TimesteppingSolver): """Timestepping linear solver object for the incompressible Boussinesq equations with prognostic variables u, p, b. This solver follows the following strategy: (1) Analytically eliminate b (introduces error near topography) (2) Solve resulting system for (u,p) using a hybrid-mixed method (3) Reconstruct b :arg state: a :class:`.State` object containing everything else. :arg solver_parameters: (optional) Solver parameters. :arg overwrite_solver_parameters: boolean, if True use only the solver_parameters that have been passed in, if False then update the default solver parameters with the solver_parameters passed in. """ solver_parameters = { 'ksp_type': 'preonly', 'mat_type': 'matfree', 'pc_type': 'python', 'pc_python_type': 'firedrake.HybridizationPC', 'hybridization': { 'ksp_type': 'cg', 'pc_type': 'gamg', 'ksp_rtol': 1e-8, 'mg_levels': { 'ksp_type': 'chebyshev', 'ksp_max_it': 2, 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu' } } } @timed_function("Gusto:SolverSetup") def _setup_solver(self): state = self.state # just cutting down line length a bit dt = state.dt beta_ = dt * self.alpha Vu = state.spaces("HDiv") Vb = state.spaces("theta") Vp = state.spaces("DG") # Store time-stepping coefficients as UFL Constants beta = Constant(beta_) # Split up the rhs vector (symbolically) self.xrhs = Function(self.equations.function_space) u_in, p_in, b_in = split(self.xrhs) # Build the reduced function space for u,p M = MixedFunctionSpace((Vu, Vp)) w, phi = TestFunctions(M) u, p = TrialFunctions(M) # Get background fields bbar = state.fields("bbar") # Analytical (approximate) elimination of theta k = state.k # Upward pointing unit vector b = -dot(k, u) * dot(k, grad(bbar)) * beta + b_in # vertical projection def V(u): return k * inner(u, k) eqn = (inner(w, (u - u_in)) * dx - beta * div(w) * p * dx - beta * inner(w, k) * b * dx + phi * div(u) * dx) if hasattr(self.equations, "mu"): eqn += dt * self.equations.mu * inner(w, k) * inner(u, k) * dx aeqn = lhs(eqn) Leqn = rhs(eqn) # Place to put result of u p solver self.up = Function(M) # Boundary conditions (assumes extruded mesh) # BCs are declared for the plain velocity space. As we need them in # a mixed problem, we replicate the BCs but for subspace of M bcs = [ DirichletBC(M.sub(0), bc.function_arg, bc.sub_domain) for bc in self.equations.bcs['u'] ] # Solver for u, p up_problem = LinearVariationalProblem(aeqn, Leqn, self.up, bcs=bcs) # Provide callback for the nullspace of the trace system def trace_nullsp(T): return VectorSpaceBasis(constant=True) appctx = {"trace_nullspace": trace_nullsp} self.up_solver = LinearVariationalSolver( up_problem, solver_parameters=self.solver_parameters, appctx=appctx) # Reconstruction of b b = TrialFunction(Vb) gamma = TestFunction(Vb) u, p = self.up.split() self.b = Function(Vb) b_eqn = gamma * (b - b_in + dot(k, u) * dot(k, grad(bbar)) * beta) * dx b_problem = LinearVariationalProblem(lhs(b_eqn), rhs(b_eqn), self.b) self.b_solver = LinearVariationalSolver(b_problem) @timed_function("Gusto:LinearSolve") def solve(self, xrhs, dy): """ Apply the solver with rhs xrhs and result dy. """ self.xrhs.assign(xrhs) with timed_region("Gusto:VelocityPressureSolve"): self.up_solver.solve() u1, p1 = self.up.split() u, p, b = dy.split() u.assign(u1) p.assign(p1) with timed_region("Gusto:BuoyancyRecon"): self.b_solver.solve() b.assign(self.b)
class CompressibleSolver(TimesteppingSolver): """ Timestepping linear solver object for the compressible equations in theta-pi formulation with prognostic variables u,rho,theta. This solver follows the following strategy: (1) Analytically eliminate theta (introduces error near topography) (2) Solve resulting system for (u,rho) using a Schur preconditioner (3) Reconstruct theta :arg state: a :class:`.State` object containing everything else. """ def __init__(self, state, params=None): self.state = state if params is None: self.params = {'pc_type': 'fieldsplit', 'pc_fieldsplit_type': 'schur', 'ksp_type': 'gmres', 'ksp_max_it': 100, 'ksp_gmres_restart': 50, 'pc_fieldsplit_schur_fact_type': 'FULL', 'pc_fieldsplit_schur_precondition': 'selfp', 'fieldsplit_0_ksp_type': 'preonly', 'fieldsplit_0_pc_type': 'bjacobi', 'fieldsplit_0_sub_pc_type': 'ilu', 'fieldsplit_1_ksp_type': 'preonly', 'fieldsplit_1_pc_type': 'gamg', 'fieldsplit_1_mg_levels_ksp_type': 'chebyshev', 'fieldsplit_1_mg_levels_ksp_chebyshev_estimate_eigenvalues': True, 'fieldsplit_1_mg_levels_ksp_chebyshev_estimate_eigenvalues_random': True, 'fieldsplit_1_mg_levels_ksp_max_it': 1, 'fieldsplit_1_mg_levels_pc_type': 'bjacobi', 'fieldsplit_1_mg_levels_sub_pc_type': 'ilu'} else: self.params = params # setup the solver self._setup_solver() def _setup_solver(self): state = self.state # just cutting down line length a bit dt = state.timestepping.dt beta = dt*state.timestepping.alpha cp = state.parameters.cp mu = state.mu # Split up the rhs vector (symbolically) u_in, rho_in, theta_in = split(state.xrhs) # Build the reduced function space for u,rho M = MixedFunctionSpace((state.V[0], state.V[1])) w, phi = TestFunctions(M) u, rho = TrialFunctions(M) n = FacetNormal(state.mesh) # Get background fields thetabar = state.thetabar rhobar = state.rhobar pibar = exner(thetabar, rhobar, state) pibar_rho = exner_rho(thetabar, rhobar, state) pibar_theta = exner_theta(thetabar, rhobar, state) # Analytical (approximate) elimination of theta k = state.k # Upward pointing unit vector theta = -dot(k,u)*dot(k,grad(thetabar))*beta + theta_in # Only include theta' (rather than pi') in the vertical # component of the gradient # the pi prime term (here, bars are for mean and no bars are # for linear perturbations) pi = pibar_theta*theta + pibar_rho*rho # vertical projection def V(u): return k*inner(u,k) eqn = ( inner(w, (u - u_in))*dx - beta*cp*div(theta*V(w))*pibar*dx # following does nothing but is preserved in the comments # to remind us why (because V(w) is purely vertical. # + beta*cp*jump(theta*V(w),n)*avg(pibar)*dS_v - beta*cp*div(thetabar*w)*pi*dx + beta*cp*jump(thetabar*w,n)*avg(pi)*dS_v + (phi*(rho - rho_in) - beta*inner(grad(phi), u)*rhobar)*dx + beta*jump(phi*u, n)*avg(rhobar)*(dS_v + dS_h) ) if mu is not None: eqn += dt*mu*inner(w,k)*inner(u,k)*dx aeqn = lhs(eqn) Leqn = rhs(eqn) # Place to put result of u rho solver self.urho = Function(M) # Boundary conditions (assumes extruded mesh) dim = M.sub(0).ufl_element().value_shape()[0] bc = ("0.0",)*dim bcs = [DirichletBC(M.sub(0), Expression(bc), "bottom"), DirichletBC(M.sub(0), Expression(bc), "top")] # Solver for u, rho urho_problem = LinearVariationalProblem( aeqn, Leqn, self.urho, bcs=bcs) self.urho_solver = LinearVariationalSolver(urho_problem, solver_parameters=self.params, options_prefix='ImplicitSolver') # Reconstruction of theta theta = TrialFunction(state.V[2]) gamma = TestFunction(state.V[2]) u, rho = self.urho.split() self.theta = Function(state.V[2]) theta_eqn = gamma*(theta - theta_in + dot(k,u)*dot(k,grad(thetabar))*beta)*dx theta_problem = LinearVariationalProblem(lhs(theta_eqn), rhs(theta_eqn), self.theta) self.theta_solver = LinearVariationalSolver(theta_problem, options_prefix='thetabacksubstitution') def solve(self): """ Apply the solver with rhs state.xrhs and result state.dy. """ self.urho_solver.solve() u1, rho1 = self.urho.split() u, rho, theta = self.state.dy.split() u.assign(u1) rho.assign(rho1) self.theta_solver.solve() theta.assign(self.theta)
class HybridizedCompressibleSolver(TimesteppingSolver): """ Timestepping linear solver object for the compressible equations in theta-pi formulation with prognostic variables u, rho, and theta. This solver follows the following strategy: (1) Analytically eliminate theta (introduces error near topography) (2a) Formulate the resulting mixed system for u and rho using a hybridized mixed method. This breaks continuity in the linear perturbations of u, and introduces a new unknown on the mesh interfaces approximating the average of the Exner pressure perturbations. These trace unknowns also act as Lagrange multipliers enforcing normal continuity of the "broken" u variable. (2b) Statically condense the block-sparse system into a single system for the Lagrange multipliers. This is the only globally coupled system requiring a linear solver. (2c) Using the computed trace variables, we locally recover the broken velocity and density perturbations. This is accomplished in two stages: (i): Recover rho locally using the multipliers. (ii): Recover "broken" u locally using rho and the multipliers. (2d) Project the "broken" velocity field into the HDiv-conforming space using local averaging. (3) Reconstruct theta :arg state: a :class:`.State` object containing everything else. :arg quadrature degree: tuple (q_h, q_v) where q_h is the required quadrature degree in the horizontal direction and q_v is that in the vertical direction. :arg solver_parameters (optional): solver parameters for the trace system. :arg overwrite_solver_parameters: boolean, if True use only the solver_parameters that have been passed in, if False then update. the default solver parameters with the solver_parameters passed in. :arg moisture (optional): list of names of moisture fields. """ # Solver parameters for the Lagrange multiplier system # NOTE: The reduced operator is not symmetric solver_parameters = {'ksp_type': 'gmres', 'pc_type': 'gamg', 'ksp_rtol': 1.0e-8, 'mg_levels': {'ksp_type': 'richardson', 'ksp_max_it': 2, 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'}} def __init__(self, state, quadrature_degree=None, solver_parameters=None, overwrite_solver_parameters=False, moisture=None): self.moisture = moisture self.state = state if quadrature_degree is not None: self.quadrature_degree = quadrature_degree else: dgspace = state.spaces("DG") if any(deg > 2 for deg in dgspace.ufl_element().degree()): state.logger.warning("default quadrature degree most likely not sufficient for this degree element") self.quadrature_degree = (5, 5) super().__init__(state, solver_parameters, overwrite_solver_parameters) @timed_function("Gusto:SolverSetup") def _setup_solver(self): from firedrake.assemble import create_assembly_callable import numpy as np state = self.state dt = state.timestepping.dt beta = dt*state.timestepping.alpha cp = state.parameters.cp mu = state.mu Vu = state.spaces("HDiv") Vu_broken = FunctionSpace(state.mesh, BrokenElement(Vu.ufl_element())) Vtheta = state.spaces("HDiv_v") Vrho = state.spaces("DG") h_deg = state.horizontal_degree v_deg = state.vertical_degree Vtrace = FunctionSpace(state.mesh, "HDiv Trace", degree=(h_deg, v_deg)) # Split up the rhs vector (symbolically) u_in, rho_in, theta_in = split(state.xrhs) # Build the function space for "broken" u and rho # and add the trace variable M = MixedFunctionSpace((Vu_broken, Vrho)) w, phi = TestFunctions(M) u, rho = TrialFunctions(M) l0 = TrialFunction(Vtrace) dl = TestFunction(Vtrace) n = FacetNormal(state.mesh) # Get background fields thetabar = state.fields("thetabar") rhobar = state.fields("rhobar") pibar = thermodynamics.pi(state.parameters, rhobar, thetabar) pibar_rho = thermodynamics.pi_rho(state.parameters, rhobar, thetabar) pibar_theta = thermodynamics.pi_theta(state.parameters, rhobar, thetabar) # Analytical (approximate) elimination of theta k = state.k # Upward pointing unit vector theta = -dot(k, u)*dot(k, grad(thetabar))*beta + theta_in # Only include theta' (rather than pi') in the vertical # component of the gradient # The pi prime term (here, bars are for mean and no bars are # for linear perturbations) pi = pibar_theta*theta + pibar_rho*rho # Vertical projection def V(u): return k*inner(u, k) # Specify degree for some terms as estimated degree is too large dxp = dx(degree=(self.quadrature_degree)) dS_vp = dS_v(degree=(self.quadrature_degree)) dS_hp = dS_h(degree=(self.quadrature_degree)) ds_vp = ds_v(degree=(self.quadrature_degree)) ds_tbp = ds_t(degree=(self.quadrature_degree)) + ds_b(degree=(self.quadrature_degree)) # Mass matrix for the trace space tM = assemble(dl('+')*l0('+')*(dS_v + dS_h) + dl*l0*ds_v + dl*l0*(ds_t + ds_b)) Lrhobar = Function(Vtrace) Lpibar = Function(Vtrace) rhopi_solver = LinearSolver(tM, solver_parameters={'ksp_type': 'cg', 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'}, options_prefix='rhobarpibar_solver') rhobar_avg = Function(Vtrace) pibar_avg = Function(Vtrace) def _traceRHS(f): return (dl('+')*avg(f)*(dS_v + dS_h) + dl*f*ds_v + dl*f*(ds_t + ds_b)) assemble(_traceRHS(rhobar), tensor=Lrhobar) assemble(_traceRHS(pibar), tensor=Lpibar) # Project averages of coefficients into the trace space with timed_region("Gusto:HybridProjectRhobar"): rhopi_solver.solve(rhobar_avg, Lrhobar) with timed_region("Gusto:HybridProjectPibar"): rhopi_solver.solve(pibar_avg, Lpibar) # Add effect of density of water upon theta if self.moisture is not None: water_t = Function(Vtheta).assign(0.0) for water in self.moisture: water_t += self.state.fields(water) theta_w = theta / (1 + water_t) thetabar_w = thetabar / (1 + water_t) else: theta_w = theta thetabar_w = thetabar # "broken" u and rho system Aeqn = (inner(w, (state.h_project(u) - u_in))*dx - beta*cp*div(theta_w*V(w))*pibar*dxp # following does nothing but is preserved in the comments # to remind us why (because V(w) is purely vertical). # + beta*cp*dot(theta_w*V(w), n)*pibar_avg('+')*dS_vp + beta*cp*dot(theta_w*V(w), n)*pibar_avg('+')*dS_hp + beta*cp*dot(theta_w*V(w), n)*pibar_avg*ds_tbp - beta*cp*div(thetabar_w*w)*pi*dxp + (phi*(rho - rho_in) - beta*inner(grad(phi), u)*rhobar)*dx + beta*dot(phi*u, n)*rhobar_avg('+')*(dS_v + dS_h)) if mu is not None: Aeqn += dt*mu*inner(w, k)*inner(u, k)*dx # Form the mixed operators using Slate # (A K)(X) = (X_r) # (K.T 0)(l) (0 ) # where X = ("broken" u, rho) A = Tensor(lhs(Aeqn)) X_r = Tensor(rhs(Aeqn)) # Off-diagonal block matrices containing the contributions # of the Lagrange multipliers (surface terms in the momentum equation) K = Tensor(beta*cp*dot(thetabar_w*w, n)*l0('+')*(dS_vp + dS_hp) + beta*cp*dot(thetabar_w*w, n)*l0*ds_vp + beta*cp*dot(thetabar_w*w, n)*l0*ds_tbp) # X = A.inv * (X_r - K * l), # 0 = K.T * X = -(K.T * A.inv * K) * l + K.T * A.inv * X_r, # so (K.T * A.inv * K) * l = K.T * A.inv * X_r # is the reduced equation for the Lagrange multipliers. # Right-hand side expression: (Forward substitution) Rexp = K.T * A.inv * X_r self.R = Function(Vtrace) # We need to rebuild R everytime data changes self._assemble_Rexp = create_assembly_callable(Rexp, tensor=self.R) # Schur complement operator: Smatexp = K.T * A.inv * K with timed_region("Gusto:HybridAssembleTraceOp"): S = assemble(Smatexp) S.force_evaluation() # Set up the Linear solver for the system of Lagrange multipliers self.lSolver = LinearSolver(S, solver_parameters=self.solver_parameters, options_prefix='lambda_solve') # Result function for the multiplier solution self.lambdar = Function(Vtrace) # Place to put result of u rho reconstruction self.urho = Function(M) # Reconstruction of broken u and rho u_, rho_ = self.urho.split() # Split operators for two-stage reconstruction _A = A.blocks _K = K.blocks _Xr = X_r.blocks A00 = _A[0, 0] A01 = _A[0, 1] A10 = _A[1, 0] A11 = _A[1, 1] K0 = _K[0, 0] Ru = _Xr[0] Rrho = _Xr[1] lambda_vec = AssembledVector(self.lambdar) # rho reconstruction Srho = A11 - A10 * A00.inv * A01 rho_expr = Srho.solve(Rrho - A10 * A00.inv * (Ru - K0 * lambda_vec), decomposition="PartialPivLU") self._assemble_rho = create_assembly_callable(rho_expr, tensor=rho_) # "broken" u reconstruction rho_vec = AssembledVector(rho_) u_expr = A00.solve(Ru - A01 * rho_vec - K0 * lambda_vec, decomposition="PartialPivLU") self._assemble_u = create_assembly_callable(u_expr, tensor=u_) # Project broken u into the HDiv space using facet averaging. # Weight function counting the dofs of the HDiv element: shapes = (Vu.finat_element.space_dimension(), np.prod(Vu.shape)) weight_kernel = """ for (int i=0; i<%d; ++i) { for (int j=0; j<%d; ++j) { w[i][j] += 1.0; }}""" % shapes self._weight = Function(Vu) par_loop(weight_kernel, dx, {"w": (self._weight, INC)}) # Averaging kernel self._average_kernel = """ for (int i=0; i<%d; ++i) { for (int j=0; j<%d; ++j) { vec_out[i][j] += vec_in[i][j]/w[i][j]; }}""" % shapes # HDiv-conforming velocity self.u_hdiv = Function(Vu) # Reconstruction of theta theta = TrialFunction(Vtheta) gamma = TestFunction(Vtheta) self.theta = Function(Vtheta) theta_eqn = gamma*(theta - theta_in + dot(k, self.u_hdiv)*dot(k, grad(thetabar))*beta)*dx theta_problem = LinearVariationalProblem(lhs(theta_eqn), rhs(theta_eqn), self.theta) self.theta_solver = LinearVariationalSolver(theta_problem, solver_parameters={'ksp_type': 'cg', 'pc_type': 'bjacobi', 'pc_sub_type': 'ilu'}, options_prefix='thetabacksubstitution') self.bcs = [DirichletBC(Vu, 0.0, "bottom"), DirichletBC(Vu, 0.0, "top")] @timed_function("Gusto:LinearSolve") def solve(self): """ Apply the solver with rhs state.xrhs and result state.dy. """ # Solve the velocity-density system with timed_region("Gusto:VelocityDensitySolve"): # Assemble the RHS for lambda into self.R with timed_region("Gusto:HybridRHS"): self._assemble_Rexp() # Solve for lambda with timed_region("Gusto:HybridTraceSolve"): self.lSolver.solve(self.lambdar, self.R) # Reconstruct broken u and rho with timed_region("Gusto:HybridRecon"): self._assemble_rho() self._assemble_u() broken_u, rho1 = self.urho.split() u1 = self.u_hdiv # Project broken_u into the HDiv space u1.assign(0.0) with timed_region("Gusto:HybridProjectHDiv"): par_loop(self._average_kernel, dx, {"w": (self._weight, READ), "vec_in": (broken_u, READ), "vec_out": (u1, INC)}) # Reapply bcs to ensure they are satisfied for bc in self.bcs: bc.apply(u1) # Copy back into u and rho cpts of dy u, rho, theta = self.state.dy.split() u.assign(u1) rho.assign(rho1) # Reconstruct theta with timed_region("Gusto:ThetaRecon"): self.theta_solver.solve() # Copy into theta cpt of dy theta.assign(self.theta)