def _setup(self, **kwargs): for name, field in kwargs.items(): if name in self._fields.keys(): self._fields[name].assign(field) else: if isinstance(field, firedrake.Constant): self._fields[name] = firedrake.Constant(field) elif isinstance(field, firedrake.Function): self._fields[name] = field.copy(deepcopy=True) else: raise TypeError( "Input %s field has type %s, must be Constant or Function!" % (name, type(field)) ) # Create symbolic representations of the flux and sources of damage dt = firedrake.Constant(1.0) flux = self.model.flux(**self.fields) # Create the finite element mass matrix D = self.fields["damage"] Q = D.function_space() φ, ψ = firedrake.TrialFunction(Q), firedrake.TestFunction(Q) M = φ * ψ * dx L1 = -dt * flux D1 = firedrake.Function(Q) D2 = firedrake.Function(Q) L2 = firedrake.replace(L1, {D: D1}) L3 = firedrake.replace(L1, {D: D2}) dD = firedrake.Function(Q) parameters = { "solver_parameters": { "ksp_type": "preonly", "pc_type": "bjacobi", "sub_pc_type": "ilu", } } problem1 = LinearVariationalProblem(M, L1, dD) problem2 = LinearVariationalProblem(M, L2, dD) problem3 = LinearVariationalProblem(M, L3, dD) solver1 = LinearVariationalSolver(problem1, **parameters) solver2 = LinearVariationalSolver(problem2, **parameters) solver3 = LinearVariationalSolver(problem3, **parameters) self._solvers = [solver1, solver2, solver3] self._stages = [D1, D2] self._damage_change = dD self._timestep = dt
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 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 L = (- inner(state.perp(grad(gamma)), u))*dx 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 __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 _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 _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.solver_parameters, options_prefix='SWimplicit')
def solver(self): # setup solver using lhs and rhs defined in derived class problem = LinearVariationalProblem(self.lhs, self.rhs, self.dq) solver_name = self.field.name( ) + self.equation.__class__.__name__ + self.__class__.__name__ return LinearVariationalSolver( problem, solver_parameters=self.solver_parameters, options_prefix=solver_name)
def _setup(self, **kwargs): for name, field in kwargs.items(): if name in self.fields.keys(): self.fields[name].assign(field) else: self.fields[name] = utilities.copy(field) # Create symbolic representations of the flux and sources of damage dt = firedrake.Constant(1.) flux = self.model.flux(**self.fields) # Create the finite element mass matrix D = self.fields.get('damage', self.fields.get('D')) Q = D.function_space() φ, ψ = firedrake.TrialFunction(Q), firedrake.TestFunction(Q) M = φ * ψ * dx L1 = -dt * flux D1 = firedrake.Function(Q) D2 = firedrake.Function(Q) L2 = firedrake.replace(L1, {D: D1}) L3 = firedrake.replace(L1, {D: D2}) dD = firedrake.Function(Q) parameters = { 'solver_parameters': { 'ksp_type': 'preonly', 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu' } } problem1 = LinearVariationalProblem(M, L1, dD) problem2 = LinearVariationalProblem(M, L2, dD) problem3 = LinearVariationalProblem(M, L3, dD) solver1 = LinearVariationalSolver(problem1, **parameters) solver2 = LinearVariationalSolver(problem2, **parameters) solver3 = LinearVariationalSolver(problem3, **parameters) self._solvers = [solver1, solver2, solver3] self._stages = [D1, D2] self._damage_change = dD self._timestep = dt
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 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 _build_forcing_solvers(self): a = self.mass_term() L = self.forcing_term() if self.Vu.extruded: bcs = [ DirichletBC(self.Vu, 0.0, "bottom"), DirichletBC(self.Vu, 0.0, "top") ] else: bcs = None u_forcing_problem = LinearVariationalProblem(a, L, self.uF, bcs=bcs) solver_parameters = {} if self.state.output.log_level == DEBUG: solver_parameters["ksp_monitor_true_residual"] = True self.u_forcing_solver = LinearVariationalSolver( u_forcing_problem, solver_parameters=solver_parameters, options_prefix="UForcingSolver")
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 self.state.output.log_level == DEBUG: solver_parameters["ksp_monitor_true_residual"] = True self.divergence_solver = LinearVariationalSolver( divergence_problem, solver_parameters=solver_parameters, options_prefix="DivergenceSolver")
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 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 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 __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')
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 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 _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 _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']
T = thermodynamics.T(state.parameters, theta0, pie, water_v0) r_v_expr = thermodynamics.r_v(state.parameters, H, T, p) # make expressions to evaluate residual pi_ev = thermodynamics.pi(state.parameters, rho_averaged, theta0) p_ev = thermodynamics.p(state.parameters, pi_ev) T_ev = thermodynamics.T(state.parameters, theta0, pi_ev, water_v0) RH_ev = thermodynamics.RH(state.parameters, water_v0, T_ev, p_ev) RH = Function(Vt) # set-up rho problem to keep Pi constant gamma = TestFunction(Vr) rho_trial = TrialFunction(Vr) a = gamma * rho_trial * dxp L = gamma * (rho_b * theta_b / theta0) * dxp rho_problem = LinearVariationalProblem(a, L, rho_h) cg_ilu_params = {'ksp_type': 'cg', 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'} if args.debug: cg_ilu_params['ksp_monitor_true_residual'] = None rho_solver = LinearVariationalSolver(rho_problem, solver_parameters=cg_ilu_params) max_outer_solve_count = 20 max_inner_solve_count = 10 PETSc.Sys.Print("Starting rho solver loop...\n") for i in range(max_outer_solve_count):
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 Vu = state.spaces("HDiv") Vb = state.spaces("HDiv_v") Vp = state.spaces("DG") # Store time-stepping coefficients as UFL Constants dt = Constant(Dt) beta = Constant(beta_) # 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((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 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) bcs = None if len(self.state.bcs) == 0 else self.state.bcs # 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)
rc = 2000. Tdash = 2.0 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,
uw = lambda u, v: (s(u)('+') * v('+') + s(u)('-') * v('-')) if mesh.geometric_dimension() == 2: perp = lambda u: as_vector([-u[1], u[0]]) p_uw = lambda u, v: perp(uw(u, v)) else: 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_)
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')
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))
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 __init__(self, mesh: object, kappa: float, comm_space: MPI.Comm, mu: float = 5., *args, **kwargs): """ Constructor :param mesh: spatial domain :param kappa: diffusion coefficient :param mu: penalty weighting function """ super(Diffusion2D, self).__init__(*args, **kwargs) # Spatial domain and function space self.mesh = mesh V = FunctionSpace(self.mesh, "DG", 1) self.function_space = V self.comm_space = comm_space # Placeholder for time step - will be updated in the update method self.dt = Constant(0.) # Things we need for the form gamma = TestFunction(V) phi = TrialFunction(V) self.f = Function(V) n = FacetNormal(mesh) # Set up the rhs and bilinear form of the equation a = (inner(gamma, phi) * dx + self.dt * (inner(grad(gamma), grad(phi) * kappa) * dx - inner(2 * avg(outer(phi, n)), avg(grad(gamma) * kappa)) * dS - inner(avg(grad(phi) * kappa), 2 * avg(outer(gamma, n))) * dS + mu * inner(2 * avg(outer(phi, n)), 2 * avg(outer(gamma, n) * kappa)) * dS)) rhs = inner(gamma, self.f) * dx # Function to hold the solution self.soln = Function(V) # Setup problem and solver prob = LinearVariationalProblem(a, rhs, self.soln) self.solver = NonlinearVariationalSolver(prob) # Set the data structure for any user-defined time point self.vector_template = VectorDiffusion2D(size=len(self.function_space), comm_space=self.comm_space) # Set initial condition: # Setting up a Gaussian blob in the centre of the domain. self.vector_t_start = VectorDiffusion2D(size=len(self.function_space), comm_space=self.comm_space) x = SpatialCoordinate(self.mesh) initial_tracer = exp(-((x[0] - 5)**2 + (x[1] - 5)**2)) tmp = Function(self.function_space) tmp.interpolate(initial_tracer) self.vector_t_start.set_values(np.copy(tmp.dat.data))
def __init__(self, equation, alpha): self.field_name = equation.field_name implicit_terms = ["incompressibility", "sponge"] dt = equation.state.dt W = equation.function_space self.x0 = Function(W) self.xF = Function(W) # set up boundary conditions on the u subspace of W bcs = [ DirichletBC(W.sub(0), bc.function_arg, bc.sub_domain) for bc in equation.bcs['u'] ] # drop terms relating to transport and diffusion residual = equation.residual.label_map( lambda t: any(t.has_label(transport, diffusion, return_tuple=True) ), drop) # the lhs of both of the explicit and implicit solvers is just # the time derivative form trials = TrialFunctions(W) a = residual.label_map(lambda t: t.has_label(time_derivative), replace_subject(trials), map_if_false=drop) # the explicit forms are multiplied by (1-alpha) and moved to the rhs L_explicit = -(1 - alpha) * dt * residual.label_map( lambda t: t.has_label(time_derivative) or t.get(name) in implicit_terms or t.get(name) == "hydrostatic_form", drop, replace_subject(self.x0)) # the implicit forms are multiplied by alpha and moved to the rhs L_implicit = -alpha * dt * residual.label_map( lambda t: t.has_label(time_derivative) or t.get(name) in implicit_terms or t.get(name) == "hydrostatic_form", drop, replace_subject(self.x0)) # now add the terms that are always fully implicit if any(t.get(name) in implicit_terms for t in residual): L_implicit -= dt * residual.label_map( lambda t: t.get(name) in implicit_terms, replace_subject(self.x0), drop) # the hydrostatic equations require some additional forms: if any([t.has_label(hydrostatic) for t in residual]): L_explicit += residual.label_map( lambda t: t.get(name) == "hydrostatic_form", replace_subject(self.x0), drop) L_implicit -= residual.label_map( lambda t: t.get(name) == "hydrostatic_form", replace_subject(self.x0), drop) # now we can set up the explicit and implicit problems explicit_forcing_problem = LinearVariationalProblem(a.form, L_explicit.form, self.xF, bcs=bcs) implicit_forcing_problem = LinearVariationalProblem(a.form, L_implicit.form, self.xF, bcs=bcs) solver_parameters = {} if logger.isEnabledFor(DEBUG): solver_parameters["ksp_monitor_true_residual"] = None self.solvers = {} self.solvers["explicit"] = LinearVariationalSolver( explicit_forcing_problem, solver_parameters=solver_parameters, options_prefix="ExplicitForcingSolver") self.solvers["implicit"] = LinearVariationalSolver( implicit_forcing_problem, solver_parameters=solver_parameters, options_prefix="ImplicitForcingSolver")
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)