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 _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 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. 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 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 # Possibly make g vary with spatial coordinates? g = state.parameters.g arhs -= g*inner(dv, state.k)*dx bcs = [DirichletBC(W.sub(0), Constant(0.0), bstring)] w = Function(W) PiProblem = 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'}} PiSolver = LinearVariationalSolver(PiProblem, solver_parameters=params, options_prefix="pisolver") 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, options_prefix="rhosolver") rhosolver.solve() v, rho_ = w1.split() rho0.assign(rho_) else: rho0.interpolate(thermodynamics.rho(state.parameters, theta0, Pi))
def initialize(self, pc): """ Set up the problem context. Takes the original mixed problem and transforms it into the equivalent hybrid-mixed system. A KSP object is created for the Lagrange multipliers on the top/bottom faces of the mesh cells. """ from firedrake import (FunctionSpace, Function, Constant, FiniteElement, TensorProductElement, TrialFunction, TrialFunctions, TestFunction, DirichletBC, interval, MixedElement, BrokenElement) from firedrake.assemble import (allocate_matrix, create_assembly_callable) from firedrake.formmanipulation import split_form from ufl.algorithms.replace import replace from ufl.cell import TensorProductCell # Extract PC context prefix = pc.getOptionsPrefix() + "vert_hybridization_" _, P = pc.getOperators() self.ctx = P.getPythonContext() if not isinstance(self.ctx, ImplicitMatrixContext): raise ValueError( "The python context must be an ImplicitMatrixContext") test, trial = self.ctx.a.arguments() V = test.function_space() mesh = V.mesh() # Magically determine which spaces are vector and scalar valued for i, Vi in enumerate(V): # Vector-valued spaces will have a non-empty value_shape if Vi.ufl_element().value_shape(): self.vidx = i else: self.pidx = i Vv = V[self.vidx] Vp = V[self.pidx] # Create the space of approximate traces in the vertical. # NOTE: Technically a hack since the resulting space is technically # defined in cell interiors, however the degrees of freedom will only # be geometrically defined on edges. Arguments will only be used in # surface integrals deg, _ = Vv.ufl_element().degree() # Assumes a tensor product cell (quads, triangular-prisms, cubes) if not isinstance(Vp.ufl_element().cell(), TensorProductCell): raise NotImplementedError( "Currently only implemented for tensor product discretizations" ) # Only want the horizontal cell cell, _ = Vp.ufl_element().cell()._cells DG = FiniteElement("DG", cell, deg) CG = FiniteElement("CG", interval, 1) Vv_tr_element = TensorProductElement(DG, CG) Vv_tr = FunctionSpace(mesh, Vv_tr_element) # Break the spaces broken_elements = MixedElement( [BrokenElement(Vi.ufl_element()) for Vi in V]) V_d = FunctionSpace(mesh, broken_elements) # Set up relevant functions self.broken_solution = Function(V_d) self.broken_residual = Function(V_d) self.trace_solution = Function(Vv_tr) self.unbroken_solution = Function(V) self.unbroken_residual = Function(V) # Set up transfer kernels to and from the broken velocity space # NOTE: Since this snippet of code is used in a couple places in # in Gusto, might be worth creating a utility function that is # is importable and just called where needed. shape = Vv.finat_element.space_dimension() * np.prod(Vv.shape) weight_kernel = """ for (int i=0; i<%d; ++i) { w[i] += 1.0; }""" % shape self.weight = Function(Vv) par_loop(weight_kernel, dx, {"w": (self.weight, INC)}) self.average_kernel = """ for (int i=0; i<%d; ++i) { vec_out[i] += vec_in[i]/w[i]; }""" % shape # Original mixed operator replaced with "broken" arguments arg_map = {test: TestFunction(V_d), trial: TrialFunction(V_d)} Atilde = Tensor(replace(self.ctx.a, arg_map)) gammar = TestFunction(Vv_tr) n = FacetNormal(mesh) sigma = TrialFunctions(V_d)[self.vidx] # Again, assumes tensor product structure. Why use this if you # don't have some form of vertical extrusion? Kform = gammar('+') * jump(sigma, n=n) * dS_h # Here we deal with boundary conditions if self.ctx.row_bcs: # Find all the subdomains with neumann BCS # These are Dirichlet BCs on the vidx space neumann_subdomains = set() for bc in self.ctx.row_bcs: if bc.function_space().index == self.pidx: raise NotImplementedError( "Dirichlet conditions for scalar variable not supported. Use a weak bc." ) if bc.function_space().index != self.vidx: raise NotImplementedError( "Dirichlet bc set on unsupported space.") # append the set of sub domains subdom = bc.sub_domain if isinstance(subdom, str): neumann_subdomains |= set([subdom]) else: neumann_subdomains |= set(as_tuple(subdom, int)) # separate out the top and bottom bcs extruded_neumann_subdomains = neumann_subdomains & { "top", "bottom" } neumann_subdomains = neumann_subdomains - extruded_neumann_subdomains integrand = gammar * dot(sigma, n) measures = [] trace_subdomains = [] for subdomain in sorted(extruded_neumann_subdomains): measures.append({"top": ds_t, "bottom": ds_b}[subdomain]) trace_subdomains.extend( sorted({"top", "bottom"} - extruded_neumann_subdomains)) measures.extend((ds(sd) for sd in sorted(neumann_subdomains))) markers = [int(x) for x in mesh.exterior_facets.unique_markers] dirichlet_subdomains = set(markers) - neumann_subdomains trace_subdomains.extend(sorted(dirichlet_subdomains)) for measure in measures: Kform += integrand * measure else: trace_subdomains = ["top", "bottom"] trace_bcs = [ DirichletBC(Vv_tr, Constant(0.0), subdomain) for subdomain in trace_subdomains ] # Make a SLATE tensor from Kform K = Tensor(Kform) # Assemble the Schur complement operator and right-hand side self.schur_rhs = Function(Vv_tr) self._assemble_Srhs = create_assembly_callable( K * Atilde.inv * AssembledVector(self.broken_residual), tensor=self.schur_rhs, form_compiler_parameters=self.ctx.fc_params) mat_type = PETSc.Options().getString(prefix + "mat_type", "aij") schur_comp = K * Atilde.inv * K.T self.S = allocate_matrix(schur_comp, bcs=trace_bcs, form_compiler_parameters=self.ctx.fc_params, mat_type=mat_type, options_prefix=prefix) self._assemble_S = create_assembly_callable( schur_comp, tensor=self.S, bcs=trace_bcs, form_compiler_parameters=self.ctx.fc_params, mat_type=mat_type) self._assemble_S() Smat = self.S.petscmat nullspace = self.ctx.appctx.get("vert_trace_nullspace", None) if nullspace is not None: nsp = nullspace(Vv_tr) Smat.setNullSpace(nsp.nullspace(comm=pc.comm)) # Set up the KSP for the system of Lagrange multipliers trace_ksp = PETSc.KSP().create(comm=pc.comm) trace_ksp.setOptionsPrefix(prefix) trace_ksp.setOperators(Smat) trace_ksp.setUp() trace_ksp.setFromOptions() self.trace_ksp = trace_ksp split_mixed_op = dict(split_form(Atilde.form)) split_trace_op = dict(split_form(K.form)) # Generate reconstruction calls self._reconstruction_calls(split_mixed_op, split_trace_op)
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)
B_x = B_0 + B_2 * X**2 + B_4 * X**4 + B_6 * X**6 f_c = 4e3 d_c = 500 w_c = 24e3 B_y = d_c * (1 / (1 + exp(-2 * (y - Ly / 2 - w_c) / f_c)) + 1 / (1 + exp(+2 * (y - Ly / 2 + w_c) / f_c))) z_deep = -720 z_b = interpolate(max_value(B_x + B_y, z_deep), Q) # We use units of megapascals-meters-years, whereas in the paper the parameters # are reported in units of pascals, so these values include some conversion # factors. A = Constant(20) C = Constant(1e-2) # Create the bed friction dissipation functional from icepack.constants import (ice_density as ρ_I, water_density as ρ_W, gravity as g, weertman_sliding_law as m) def friction(**kwargs): keys = ('velocity', 'thickness', 'surface', 'friction') u, h, s, C = map(kwargs.get, keys) p_W = ρ_W * g * max_value(0, -(s - h)) p_I = ρ_I * g * h N = p_I - p_W τ_c = N / 2
def initialize(self, pc): from firedrake import TrialFunction, TestFunction, dx, \ assemble, inner, grad, split, Constant, parameters from firedrake.assemble import allocate_matrix if pc.getType() != "python": raise ValueError("Expecting PC type python") prefix = pc.getOptionsPrefix() + "pcd_" # we assume P has things stuffed inside of it _, P = pc.getOperators() context = P.getPythonContext() test, trial = context.a.arguments() if test.function_space() != trial.function_space(): raise ValueError("Pressure space test and trial space differ") Q = test.function_space() p = TrialFunction(Q) q = TestFunction(Q) mass = p*q*dx # Regularisation to avoid having to think about nullspaces. stiffness = inner(grad(p), grad(q))*dx + Constant(1e-6)*p*q*dx opts = PETSc.Options() # we're inverting Mp and Kp, so default them to assembled. # Fp only needs its action, so default it to mat-free. # These can of course be overridden. # only Fp is referred to in update, so that's the only # one we stash. default = parameters["default_matrix_type"] Mp_mat_type = opts.getString(prefix+"Mp_mat_type", default) Kp_mat_type = opts.getString(prefix+"Kp_mat_type", default) self.Fp_mat_type = opts.getString(prefix+"Fp_mat_type", "matfree") Mp = assemble(mass, form_compiler_parameters=context.fc_params, mat_type=Mp_mat_type, options_prefix=prefix + "Mp_") Kp = assemble(stiffness, form_compiler_parameters=context.fc_params, mat_type=Kp_mat_type, options_prefix=prefix + "Kp_") # FIXME: Should we transfer nullspaces over. I think not. Mksp = PETSc.KSP().create(comm=pc.comm) Mksp.incrementTabLevel(1, parent=pc) Mksp.setOptionsPrefix(prefix + "Mp_") Mksp.setOperators(Mp.petscmat) Mksp.setUp() Mksp.setFromOptions() self.Mksp = Mksp Kksp = PETSc.KSP().create(comm=pc.comm) Kksp.incrementTabLevel(1, parent=pc) Kksp.setOptionsPrefix(prefix + "Kp_") Kksp.setOperators(Kp.petscmat) Kksp.setUp() Kksp.setFromOptions() self.Kksp = Kksp state = context.appctx["state"] Re = context.appctx.get("Re", 1.0) velid = context.appctx["velocity_space"] u0 = split(state)[velid] fp = 1.0/Re * inner(grad(p), grad(q))*dx + inner(u0, grad(p))*q*dx self.Re = Re self.Fp = allocate_matrix(fp, form_compiler_parameters=context.fc_params, mat_type=self.Fp_mat_type, options_prefix=prefix + "Fp_") self._assemble_Fp = functools.partial(assemble, fp, tensor=self.Fp, form_compiler_parameters=context.fc_params, mat_type=self.Fp_mat_type, assembly_type="residual") self._assemble_Fp() Fpmat = self.Fp.petscmat self.workspace = [Fpmat.createVecLeft() for i in (0, 1)]
def initialize(self, pc): """Set up the problem context. Take the original mixed problem and reformulate the problem as a hybridized mixed system. A KSP is created for the Lagrange multiplier system. """ from firedrake import (FunctionSpace, Function, Constant, TrialFunction, TrialFunctions, TestFunction, DirichletBC, assemble) from firedrake.assemble import (allocate_matrix, create_assembly_callable) from firedrake.formmanipulation import split_form from ufl.algorithms.replace import replace # Extract the problem context prefix = pc.getOptionsPrefix() + "hybridization_" _, P = pc.getOperators() self.cxt = P.getPythonContext() if not isinstance(self.cxt, ImplicitMatrixContext): raise ValueError( "The python context must be an ImplicitMatrixContext") test, trial = self.cxt.a.arguments() V = test.function_space() mesh = V.mesh() if len(V) != 2: raise ValueError("Expecting two function spaces.") if all(Vi.ufl_element().value_shape() for Vi in V): raise ValueError("Expecting an H(div) x L2 pair of spaces.") # Automagically determine which spaces are vector and scalar for i, Vi in enumerate(V): if Vi.ufl_element().sobolev_space().name == "HDiv": self.vidx = i else: assert Vi.ufl_element().sobolev_space().name == "L2" self.pidx = i # Create the space of approximate traces. W = V[self.vidx] if W.ufl_element().family() == "Brezzi-Douglas-Marini": tdegree = W.ufl_element().degree() else: try: # If we have a tensor product element h_deg, v_deg = W.ufl_element().degree() tdegree = (h_deg - 1, v_deg - 1) except TypeError: tdegree = W.ufl_element().degree() - 1 TraceSpace = FunctionSpace(mesh, "HDiv Trace", tdegree) # Break the function spaces and define fully discontinuous spaces broken_elements = ufl.MixedElement( [ufl.BrokenElement(Vi.ufl_element()) for Vi in V]) V_d = FunctionSpace(mesh, broken_elements) # Set up the functions for the original, hybridized # and schur complement systems self.broken_solution = Function(V_d) self.broken_residual = Function(V_d) self.trace_solution = Function(TraceSpace) self.unbroken_solution = Function(V) self.unbroken_residual = Function(V) # Set up the KSP for the hdiv residual projection hdiv_mass_ksp = PETSc.KSP().create(comm=pc.comm) hdiv_mass_ksp.setOptionsPrefix(prefix + "hdiv_residual_") # HDiv mass operator p = TrialFunction(V[self.vidx]) q = TestFunction(V[self.vidx]) mass = ufl.dot(p, q) * ufl.dx # TODO: Bcs? M = assemble(mass, bcs=None, form_compiler_parameters=self.cxt.fc_params) M.force_evaluation() Mmat = M.petscmat hdiv_mass_ksp.setOperators(Mmat) hdiv_mass_ksp.setUp() hdiv_mass_ksp.setFromOptions() self.hdiv_mass_ksp = hdiv_mass_ksp # Storing the result of A.inv * r, where A is the HDiv # mass matrix and r is the HDiv residual self._primal_r = Function(V[self.vidx]) tau = TestFunction(V_d[self.vidx]) self._assemble_broken_r = create_assembly_callable( ufl.dot(self._primal_r, tau) * ufl.dx, tensor=self.broken_residual.split()[self.vidx], form_compiler_parameters=self.cxt.fc_params) # Create the symbolic Schur-reduction: # Original mixed operator replaced with "broken" # arguments arg_map = {test: TestFunction(V_d), trial: TrialFunction(V_d)} Atilde = Tensor(replace(self.cxt.a, arg_map)) gammar = TestFunction(TraceSpace) n = ufl.FacetNormal(mesh) sigma = TrialFunctions(V_d)[self.vidx] # We zero out the contribution of the trace variables on the exterior # boundary. Extruded cells will have both horizontal and vertical # facets if mesh.cell_set._extruded: trace_bcs = [ DirichletBC(TraceSpace, Constant(0.0), "on_boundary"), DirichletBC(TraceSpace, Constant(0.0), "bottom"), DirichletBC(TraceSpace, Constant(0.0), "top") ] K = Tensor( gammar('+') * ufl.dot(sigma, n) * ufl.dS_h + gammar('+') * ufl.dot(sigma, n) * ufl.dS_v) else: trace_bcs = [DirichletBC(TraceSpace, Constant(0.0), "on_boundary")] K = Tensor(gammar('+') * ufl.dot(sigma, n) * ufl.dS) # If boundary conditions are contained in the ImplicitMatrixContext: if self.cxt.row_bcs: raise NotImplementedError( "Strong BCs not currently handled. Try imposing them weakly.") # Assemble the Schur complement operator and right-hand side self.schur_rhs = Function(TraceSpace) self._assemble_Srhs = create_assembly_callable( K * Atilde.inv * self.broken_residual, tensor=self.schur_rhs, form_compiler_parameters=self.cxt.fc_params) schur_comp = K * Atilde.inv * K.T self.S = allocate_matrix(schur_comp, bcs=trace_bcs, form_compiler_parameters=self.cxt.fc_params) self._assemble_S = create_assembly_callable( schur_comp, tensor=self.S, bcs=trace_bcs, form_compiler_parameters=self.cxt.fc_params) self._assemble_S() self.S.force_evaluation() Smat = self.S.petscmat # Nullspace for the multiplier problem nullspace = create_schur_nullspace(P, -K * Atilde, V, V_d, TraceSpace, pc.comm) if nullspace: Smat.setNullSpace(nullspace) # Set up the KSP for the system of Lagrange multipliers trace_ksp = PETSc.KSP().create(comm=pc.comm) trace_ksp.setOptionsPrefix(prefix) trace_ksp.setOperators(Smat) trace_ksp.setUp() trace_ksp.setFromOptions() self.trace_ksp = trace_ksp split_mixed_op = dict(split_form(Atilde.form)) split_trace_op = dict(split_form(K.form)) # Generate reconstruction calls self._reconstruction_calls(split_mixed_op, split_trace_op) # NOTE: The projection stage *might* be replaced by a Fortin # operator. We may want to allow the user to specify if they # wish to use a Fortin operator over a projection, or vice-versa. # In a future add-on, we can add a switch which chooses either # the Fortin reconstruction or the usual KSP projection. # Set up the projection KSP hdiv_projection_ksp = PETSc.KSP().create(comm=pc.comm) hdiv_projection_ksp.setOptionsPrefix(prefix + 'hdiv_projection_') # Reuse the mass operator from the hdiv_mass_ksp hdiv_projection_ksp.setOperators(Mmat) # Construct the RHS for the projection stage self._projection_rhs = Function(V[self.vidx]) self._assemble_projection_rhs = create_assembly_callable( ufl.dot(self.broken_solution.split()[self.vidx], q) * ufl.dx, tensor=self._projection_rhs, form_compiler_parameters=self.cxt.fc_params) # Finalize ksp setup hdiv_projection_ksp.setUp() hdiv_projection_ksp.setFromOptions() self.hdiv_projection_ksp = hdiv_projection_ksp
stabilisation = args.stabilisation or 'lax_friedrichs' if stabilisation == 'none' or family == 'cg-cg' or not nonlinear: stabilisation = None approach = args.approach or 'hessian' kwargs = { 'approach': approach, 'inversion_level': 2, # TODO: Avoid this hard-code # Space-time domain 'level': int(args.level or 0), 'end_time': float(args.end_time or 1440.0), 'num_meshes': int(args.num_meshes or 24), # Physics 'bathymetry_cap': 30.0, 'base_viscosity': Constant(args.base_viscosity or 1.0e-03), # Solver 'family': family, 'stabilisation': stabilisation, 'use_wetting_and_drying': False, # Mesh adaptation 'adapt_field': args.adapt_field or 'all_avg', 'hessian_time_combination': args.time_combine or 'integrate', 'hessian_timestep_lag': int(args.hessian_lag or 1), # NOTE: Not used in weighted Hessian 'normalisation': args.normalisation or 'complexity', 'norm_order': 1 if p is None else None if p == 'inf' else float(p), 'target': float(args.target or 5.0e+03), 'h_min': float(args.h_min or 1.0e+02),
diagnostics=diagnostics, fieldlist=fieldlist, diagnostic_fields=diagnostic_fields) # Initial conditions u0 = state.fields("u") rho0 = state.fields("rho") theta0 = state.fields("theta") # spaces Vu = u0.function_space() Vt = theta0.function_space() Vr = rho0.function_space() # Isentropic background state Tsurf = Constant(300.) theta_b = Function(Vt).interpolate(Tsurf) rho_b = Function(Vr) # Calculate hydrostatic Pi compressible_hydrostatic_balance(state, theta_b, rho_b, solve_for_rho=True) x = SpatialCoordinate(mesh) a = 5.0e3 deltaTheta = 1.0e-2 xc = 0.5 * L xr = 4000. zc = 3000. zr = 2000. r = sqrt(((x[0] - xc) / xr)**2 + ((x[1] - zc) / zr)**2)
"| Diagnostics:", tuple(nc_diag.keys())) print("Starting Initial condition, and function setup at", ctime()) R, Omega, = 6371220., 7.292e-5 mesh = IcosahedralSphereMesh(radius=R, refinement_level=ref_level, degree=2) mesh.init_cell_orientations(SpatialCoordinate(mesh)) x = SpatialCoordinate(mesh) f = Function(FunctionSpace(mesh, "CG", 1)) f.interpolate(2 * Omega * x[2] / R) g, H = 9.810616, 5960. u_0 = 2 * pi * R / (12 * 24 * 60 * 60.) uexpr = u_0 * as_vector([-x[1], x[0], 0.0]) / R Dexpr = H - (R * Omega * u_0 + u_0**2 / 2.) * x[2]**2 / (g * R**2) bexpr = Constant(0.) # Build function spaces degree = 2 family = ("DG", "BDM", "CG") W0 = FunctionSpace(mesh, family[0], degree - 1, family[0]) W1_elt = FiniteElement(family[1], triangle, degree) W1 = FunctionSpace(mesh, W1_elt, name="HDiv") W2 = FunctionSpace(mesh, family[2], degree + 1) M = MixedFunctionSpace((W1, W0)) # Set up functions xn = Function(M) un, Dn = xn.split() un.rename('u') Dn.rename('D')
# Define constant theta_e and water_t Tsurf = 283.0 psurf = 85000. pi_surf = (psurf / state.parameters.p_0)**state.parameters.kappa humidity = 0.2 S = 1.3e-5 theta_surf = thermodynamics.theta(state.parameters, Tsurf, psurf) theta_d = Function(Vt).interpolate(theta_surf * exp(S * z)) H = Function(Vt).assign(humidity) # Calculate hydrostatic fields unsaturated_hydrostatic_balance(state, theta_d, H, pi_boundary=Constant(pi_surf)) # make mean fields theta_b = Function(Vt).assign(theta0) rho_b = Function(Vr).assign(rho0) water_vb = Function(Vt).assign(water_v0) # define perturbation to RH xc = L / 2 zc = 800. r1 = 300. r2 = 200. r = sqrt((x - xc)**2 + (z - zc)**2) H_expr = conditional( r > r1, 0.0,
def initialize(self, pc): """Set up the problem context. Take the original mixed problem and reformulate the problem as a hybridized mixed system. A KSP is created for the Lagrange multiplier system. """ from firedrake import (FunctionSpace, Function, Constant, TrialFunction, TrialFunctions, TestFunction, DirichletBC) from firedrake.assemble import (allocate_matrix, create_assembly_callable) from firedrake.formmanipulation import split_form from ufl.algorithms.replace import replace # Extract the problem context prefix = pc.getOptionsPrefix() + "hybridization_" _, P = pc.getOperators() self.ctx = P.getPythonContext() if not isinstance(self.ctx, ImplicitMatrixContext): raise ValueError("The python context must be an ImplicitMatrixContext") test, trial = self.ctx.a.arguments() V = test.function_space() mesh = V.mesh() if len(V) != 2: raise ValueError("Expecting two function spaces.") if all(Vi.ufl_element().value_shape() for Vi in V): raise ValueError("Expecting an H(div) x L2 pair of spaces.") # Automagically determine which spaces are vector and scalar for i, Vi in enumerate(V): if Vi.ufl_element().sobolev_space().name == "HDiv": self.vidx = i else: assert Vi.ufl_element().sobolev_space().name == "L2" self.pidx = i # Create the space of approximate traces. W = V[self.vidx] if W.ufl_element().family() == "Brezzi-Douglas-Marini": tdegree = W.ufl_element().degree() else: try: # If we have a tensor product element h_deg, v_deg = W.ufl_element().degree() tdegree = (h_deg - 1, v_deg - 1) except TypeError: tdegree = W.ufl_element().degree() - 1 TraceSpace = FunctionSpace(mesh, "HDiv Trace", tdegree) # Break the function spaces and define fully discontinuous spaces broken_elements = ufl.MixedElement([ufl.BrokenElement(Vi.ufl_element()) for Vi in V]) V_d = FunctionSpace(mesh, broken_elements) # Set up the functions for the original, hybridized # and schur complement systems self.broken_solution = Function(V_d) self.broken_residual = Function(V_d) self.trace_solution = Function(TraceSpace) self.unbroken_solution = Function(V) self.unbroken_residual = Function(V) shapes = (V[self.vidx].finat_element.space_dimension(), np.prod(V[self.vidx].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(V[self.vidx]) par_loop(weight_kernel, ufl.dx, {"w": (self.weight, INC)}) 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 # Create the symbolic Schur-reduction: # Original mixed operator replaced with "broken" # arguments arg_map = {test: TestFunction(V_d), trial: TrialFunction(V_d)} Atilde = Tensor(replace(self.ctx.a, arg_map)) gammar = TestFunction(TraceSpace) n = ufl.FacetNormal(mesh) sigma = TrialFunctions(V_d)[self.vidx] if mesh.cell_set._extruded: Kform = (gammar('+') * ufl.jump(sigma, n=n) * ufl.dS_h + gammar('+') * ufl.jump(sigma, n=n) * ufl.dS_v) else: Kform = (gammar('+') * ufl.jump(sigma, n=n) * ufl.dS) # Here we deal with boundaries. If there are Neumann # conditions (which should be enforced strongly for # H(div)xL^2) then we need to add jump terms on the exterior # facets. If there are Dirichlet conditions (which should be # enforced weakly) then we need to zero out the trace # variables there as they are not active (otherwise the hybrid # problem is not well-posed). # If boundary conditions are contained in the ImplicitMatrixContext: if self.ctx.row_bcs: # Find all the subdomains with neumann BCS # These are Dirichlet BCs on the vidx space neumann_subdomains = set() for bc in self.ctx.row_bcs: if bc.function_space().index == self.pidx: raise NotImplementedError("Dirichlet conditions for scalar variable not supported. Use a weak bc") if bc.function_space().index != self.vidx: raise NotImplementedError("Dirichlet bc set on unsupported space.") # append the set of sub domains subdom = bc.sub_domain if isinstance(subdom, str): neumann_subdomains |= set([subdom]) else: neumann_subdomains |= set(as_tuple(subdom, int)) # separate out the top and bottom bcs extruded_neumann_subdomains = neumann_subdomains & {"top", "bottom"} neumann_subdomains = neumann_subdomains - extruded_neumann_subdomains integrand = gammar * ufl.dot(sigma, n) measures = [] trace_subdomains = [] if mesh.cell_set._extruded: ds = ufl.ds_v for subdomain in sorted(extruded_neumann_subdomains): measures.append({"top": ufl.ds_t, "bottom": ufl.ds_b}[subdomain]) trace_subdomains.extend(sorted({"top", "bottom"} - extruded_neumann_subdomains)) else: ds = ufl.ds if "on_boundary" in neumann_subdomains: measures.append(ds) else: measures.extend((ds(sd) for sd in sorted(neumann_subdomains))) markers = [int(x) for x in mesh.exterior_facets.unique_markers] dirichlet_subdomains = set(markers) - neumann_subdomains trace_subdomains.extend(sorted(dirichlet_subdomains)) for measure in measures: Kform += integrand*measure trace_bcs = [DirichletBC(TraceSpace, Constant(0.0), subdomain) for subdomain in trace_subdomains] else: # No bcs were provided, we assume weak Dirichlet conditions. # We zero out the contribution of the trace variables on # the exterior boundary. Extruded cells will have both # horizontal and vertical facets trace_subdomains = ["on_boundary"] if mesh.cell_set._extruded: trace_subdomains.extend(["bottom", "top"]) trace_bcs = [DirichletBC(TraceSpace, Constant(0.0), subdomain) for subdomain in trace_subdomains] # Make a SLATE tensor from Kform K = Tensor(Kform) # Assemble the Schur complement operator and right-hand side self.schur_rhs = Function(TraceSpace) self._assemble_Srhs = create_assembly_callable( K * Atilde.inv * AssembledVector(self.broken_residual), tensor=self.schur_rhs, form_compiler_parameters=self.ctx.fc_params) mat_type = PETSc.Options().getString(prefix + "mat_type", "aij") schur_comp = K * Atilde.inv * K.T self.S = allocate_matrix(schur_comp, bcs=trace_bcs, form_compiler_parameters=self.ctx.fc_params, mat_type=mat_type, options_prefix=prefix) self._assemble_S = create_assembly_callable(schur_comp, tensor=self.S, bcs=trace_bcs, form_compiler_parameters=self.ctx.fc_params, mat_type=mat_type) self._assemble_S() self.S.force_evaluation() Smat = self.S.petscmat nullspace = self.ctx.appctx.get("trace_nullspace", None) if nullspace is not None: nsp = nullspace(TraceSpace) Smat.setNullSpace(nsp.nullspace(comm=pc.comm)) # Set up the KSP for the system of Lagrange multipliers trace_ksp = PETSc.KSP().create(comm=pc.comm) trace_ksp.setOptionsPrefix(prefix) trace_ksp.setOperators(Smat) trace_ksp.setUp() trace_ksp.setFromOptions() self.trace_ksp = trace_ksp split_mixed_op = dict(split_form(Atilde.form)) split_trace_op = dict(split_form(K.form)) # Generate reconstruction calls self._reconstruction_calls(split_mixed_op, split_trace_op)
def test_ice_shelf_prognostic_solver(solver_type): ρ = ρ_I * (1 - ρ_I / ρ_W) Lx, Ly = 20.0e3, 20.0e3 h0 = 500.0 u0 = 100.0 T = 254.15 model = icepack.models.IceShelf() opts = { "dirichlet_ids": [1], "side_wall_ids": [3, 4], "prognostic_solver_type": solver_type, } delta_x, error = [], [] for N in range(16, 65, 4): delta_x.append(Lx / N) mesh = firedrake.RectangleMesh(N, N, Lx, Ly) x, y = firedrake.SpatialCoordinate(mesh) degree = 2 V = firedrake.VectorFunctionSpace(mesh, "CG", degree) Q = firedrake.FunctionSpace(mesh, "CG", degree) q = (n + 1) * (ρ * g * h0 * u0 / 4)**n * icepack.rate_factor(T) ux = (u0**(n + 1) + q * x)**(1 / (n + 1)) h = interpolate(h0 * u0 / ux, Q) h_initial = h.copy(deepcopy=True) A = Constant(icepack.rate_factor(T)) a = Constant(0) solver = icepack.solvers.FlowSolver(model, **opts) u_guess = interpolate(firedrake.as_vector((ux, 0)), V) u = solver.diagnostic_solve(velocity=u_guess, thickness=h, fluidity=A, strain_rate_min=Constant(0.0)) final_time, dt = 1.0, 1.0 / 12 num_timesteps = int(final_time / dt) for k in range(num_timesteps): h = solver.prognostic_solve(dt, thickness=h, velocity=u, accumulation=a, thickness_inflow=h_initial) u = solver.diagnostic_solve(velocity=u, thickness=h, fluidity=A, strain_rate_min=Constant(0.0)) error.append(norm(h - h_initial) / norm(h_initial)) print(delta_x[-1], error[-1]) log_delta_x = np.log2(np.array(delta_x)) log_error = np.log2(np.array(error)) slope, intercept = np.polyfit(log_delta_x, log_error, 1) print(f"log(error) ~= {slope:g} * log(dx) + {intercept:g}") assert slope > degree - 0.05
def test_hybrid_prognostic_solve(solver_type): Lx, Ly = 20e3, 20e3 h0, dh = 500.0, 100.0 T = 254.15 u_in = 100.0 model = icepack.models.HybridModel() opts = { "dirichlet_ids": [1], "side_wall_ids": [3, 4], "prognostic_solver_type": solver_type, } Nx, Ny = 32, 32 mesh2d = firedrake.RectangleMesh(Nx, Ny, Lx, Ly) mesh = firedrake.ExtrudedMesh(mesh2d, layers=1) V = firedrake.VectorFunctionSpace(mesh, "CG", 2, vfamily="GL", vdegree=1, dim=2) Q = firedrake.FunctionSpace(mesh, "CG", 2, vfamily="DG", vdegree=0) x, y, ζ = firedrake.SpatialCoordinate(mesh) height_above_flotation = 10.0 d = -ρ_I / ρ_W * (h0 - dh) + height_above_flotation ρ = ρ_I - ρ_W * d**2 / (h0 - dh)**2 Z = icepack.rate_factor(T) * (ρ * g * h0 / 4)**n q = 1 - (1 - (dh / h0) * (x / Lx))**(n + 1) ux = u_in + Z * q * Lx * (h0 / dh) / (n + 1) u0 = interpolate(firedrake.as_vector((ux, 0)), V) thickness = h0 - dh * x / Lx β = 1 / 2 α = β * ρ / ρ_I * dh / Lx h = interpolate(h0 - dh * x / Lx, Q) h_inflow = h.copy(deepcopy=True) ds = (1 + β) * ρ / ρ_I * dh s = interpolate(d + h0 - dh + ds * (1 - x / Lx), Q) b = interpolate(s - h, Q) C = interpolate(α * (ρ_I * g * thickness) * ux**(-1 / m), Q) A = Constant(icepack.rate_factor(T)) final_time, dt = 1.0, 1.0 / 12 num_timesteps = int(final_time / dt) solver = icepack.solvers.FlowSolver(model, **opts) u = solver.diagnostic_solve( velocity=u0, thickness=h, surface=s, fluidity=A, friction=C, strain_rate_min=Constant(0.0), ) a = firedrake.Function(Q) h_n = solver.prognostic_solve(dt, thickness=h, velocity=u, accumulation=a, thickness_inflow=h_inflow) a.interpolate((h_n - h) / dt) for k in range(num_timesteps): h = solver.prognostic_solve(dt, thickness=h, velocity=u, accumulation=a, thickness_inflow=h_inflow) s = icepack.compute_surface(thickness=h, bed=b) u = solver.diagnostic_solve( velocity=u, thickness=h, surface=s, fluidity=A, friction=C, strain_rate_min=Constant(0.0), ) assert icepack.norm(h, norm_type="Linfty") < np.inf
} discrete_turbines = True # discrete_turbines = False op = TurbineArrayOptions(**kwargs) op.update({ 'spun': False, 'di': create_directory(os.path.join(op.di, 'unsteady')), # Extend to time-dependent case 'timestepper': 'CrankNicolson', 'dt': 5.0, 'dt_per_export': 1, 'end_time': 600.0, # Crank down viscosity and plot vorticity 'base_viscosity': Constant(1.0e-05), 'characteristic_velocity': Constant(op.inflow_velocity), 'grad_depth_viscosity': True, 'max_reynolds_number': 5.0e+05, 'recover_vorticity': True, # Only consider the first turbine 'region_of_interest': [op.region_of_interest[0]], 'array_ids': np.array([2]), 'farm_ids': (2, ), }) # Solve forward problem tp = AdaptiveTurbineProblem(op, discrete_turbines=discrete_turbines, ramp_dir=op.di)
def setup_tracer(dirname): # declare grid shape, with length L and height H L = 1000. H = 1000. nlayers = int(H / 100.) ncolumns = int(L / 100.) # make mesh m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers)) fieldlist = ['u', 'rho', 'theta'] timestepping = TimesteppingParameters(dt=10.0, maxk=4, maxi=1) output = OutputParameters(dirname=dirname + "/tracer", dumpfreq=1, dumplist=['u'], perturbation_fields=['theta', 'rho']) parameters = CompressibleParameters() state = State(mesh, vertical_degree=1, horizontal_degree=1, family="CG", timestepping=timestepping, output=output, parameters=parameters, fieldlist=fieldlist, diagnostic_fields=[Difference('theta', 'tracer')]) # declare initial fields u0 = state.fields("u") rho0 = state.fields("rho") theta0 = state.fields("theta") # spaces Vu = u0.function_space() Vt = theta0.function_space() Vr = rho0.function_space() # declare tracer field and a background field tracer0 = state.fields("tracer", Vt) # Isentropic background state Tsurf = Constant(300.) theta_b = Function(Vt).interpolate(Tsurf) rho_b = Function(Vr) # Calculate initial rho compressible_hydrostatic_balance(state, theta_b, rho_b, solve_for_rho=True) # set up perturbation to theta xc = 500. zc = 350. rc = 250. x = SpatialCoordinate(mesh) r = sqrt((x[0] - xc)**2 + (x[1] - zc)**2) theta_pert = conditional(r > rc, 0., 0.25 * (1. + cos((pi / rc) * r))) theta0.interpolate(theta_b + theta_pert) rho0.interpolate(rho_b) tracer0.interpolate(theta0) state.initialise([('u', u0), ('rho', rho0), ('theta', theta0), ('tracer', tracer0)]) state.set_reference_profiles([('rho', rho_b), ('theta', theta_b)]) # set up advection schemes ueqn = EulerPoincare(state, Vu) rhoeqn = AdvectionEquation(state, Vr, equation_form="continuity") thetaeqn = SUPGAdvection(state, Vt, equation_form="advective") # build advection dictionary advected_fields = [] advected_fields.append(("u", ThetaMethod(state, u0, ueqn))) advected_fields.append(("rho", SSPRK3(state, rho0, rhoeqn))) advected_fields.append(("theta", SSPRK3(state, theta0, thetaeqn))) advected_fields.append(("tracer", SSPRK3(state, tracer0, thetaeqn))) # Set up linear solver linear_solver = CompressibleSolver(state) compressible_forcing = CompressibleForcing(state) # build time stepper stepper = CrankNicolson(state, advected_fields, linear_solver, compressible_forcing) return stepper, 100.0
def __init__(self, mesh, vertical_degree=None, horizontal_degree=1, family="RT", Coriolis=None, sponge_function=None, hydrostatic=None, timestepping=None, output=None, parameters=None, diagnostics=None, fieldlist=None, diagnostic_fields=None, u_bc_ids=None): self.family = family self.vertical_degree = vertical_degree self.horizontal_degree = horizontal_degree self.Omega = Coriolis self.mu = sponge_function self.hydrostatic = hydrostatic self.timestepping = timestepping if output is None: raise RuntimeError("You must provide a directory name for dumping results") else: self.output = output self.parameters = parameters if fieldlist is None: raise RuntimeError("You must provide a fieldlist containing the names of the prognostic fields") else: self.fieldlist = fieldlist if diagnostics is not None: self.diagnostics = diagnostics else: self.diagnostics = Diagnostics(*fieldlist) if diagnostic_fields is not None: self.diagnostic_fields = diagnostic_fields else: self.diagnostic_fields = [] if u_bc_ids is not None: self.u_bc_ids = u_bc_ids else: self.u_bc_ids = [] # The mesh self.mesh = mesh # Build the spaces self._build_spaces(mesh, vertical_degree, horizontal_degree, family) # Allocate state self._allocate_state() if self.output.dumplist is None: self.output.dumplist = fieldlist self.fields = FieldCreator(fieldlist, self.xn, self.output.dumplist) # set up bcs V = self.fields('u').function_space() self.bcs = [] if V.extruded: self.bcs.append(DirichletBC(V, 0.0, "bottom")) self.bcs.append(DirichletBC(V, 0.0, "top")) for id in self.u_bc_ids: self.bcs.append(DirichletBC(V, 0.0, id)) self.dumpfile = None # figure out if we're on a sphere try: self.on_sphere = (mesh._base_mesh.geometric_dimension() == 3 and mesh._base_mesh.topological_dimension() == 2) except AttributeError: self.on_sphere = (mesh.geometric_dimension() == 3 and mesh.topological_dimension() == 2) # build the vertical normal and define perp for 2d geometries dim = mesh.topological_dimension() if self.on_sphere: x = SpatialCoordinate(mesh) R = sqrt(inner(x, x)) self.k = interpolate(x/R, mesh.coordinates.function_space()) if dim == 2: outward_normals = CellNormal(mesh) self.perp = lambda u: cross(outward_normals, u) else: kvec = [0.0]*dim kvec[dim-1] = 1.0 self.k = Constant(kvec) if dim == 2: self.perp = lambda u: as_vector([-u[1], u[0]]) # project test function for hydrostatic case if self.hydrostatic: self.h_project = lambda u: u - self.k*inner(u, self.k) else: self.h_project = lambda u: u # Constant to hold current time self.t = Constant(0.0) # setup logger logger.setLevel(output.log_level) set_log_handler(mesh.comm) logger.info("Timestepping parameters that take non-default values:") logger.info(", ".join("%s: %s" % item for item in vars(timestepping).items())) if parameters is not None: logger.info("Physical parameters that take non-default values:") logger.info(", ".join("%s: %s" % item for item in vars(parameters).items()))
def streamplot(function, resolution=None, min_length=None, max_time=None, start_width=0.5, end_width=1.5, tolerance=3e-3, loc_tolerance=1e-10, seed=None, **kwargs): r"""Create a streamline plot of a vector field Similar to matplotlib :func:`streamplot <matplotlib.pyplot.streamplot>` :arg function: the Firedrake :class:`~.Function` to plot :arg resolution: minimum spacing between streamlines (defaults to domain size / 20) :arg min_length: minimum length of a streamline (defaults to 4x resolution) :arg max_time: maximum time to integrate a streamline :arg start_width: line width at beginning of streamline :arg end_width: line width at end of streamline, to convey direction :arg tolerance: dimensionless tolerance for adaptive ODE integration :arg loc_tolerance: point location tolerance for :meth:`~firedrake.functions.Function.at` :kwarg kwargs: same as for matplotlib :class:`~matplotlib.collections.LineCollection` """ import randomgen if function.ufl_shape != (2, ): raise ValueError("Streamplot only defined for 2D vector fields!") axes = kwargs.pop("axes", None) if axes is None: figure = plt.figure() axes = figure.add_subplot(111) mesh = function.ufl_domain() if resolution is None: coords = mesh.coordinates.dat.data_ro resolution = (coords.max(axis=0) - coords.min(axis=0)).max() / 20 if min_length is None: min_length = 4 * resolution if max_time is None: area = assemble(Constant(1) * dx(mesh)) average_speed = np.sqrt( assemble(inner(function, function) * dx) / area) max_time = 50 * min_length / average_speed streamplotter = Streamplotter(function, resolution, min_length, max_time, tolerance, loc_tolerance) # TODO: better way of seeding start points shape = streamplotter._grid.shape xmin = streamplotter._grid_point((0, 0)) xmax = streamplotter._grid_point((shape[0] - 2, shape[1] - 2)) X, Y = np.meshgrid(np.linspace(xmin[0], xmax[0], shape[0] - 2), np.linspace(xmin[1], xmax[1], shape[1] - 2)) start_points = np.vstack((X.ravel(), Y.ravel())).T # Randomly shuffle the start points generator = randomgen.MT19937(seed).generator for x in generator.permutation(np.array(start_points)): streamplotter.add_streamline(x) # Colors are determined by the speed, thicknesses by arc length speeds = [] widths = [] for streamline in streamplotter.streamlines: velocity = np.array(function.at(streamline, tolerance=loc_tolerance)) speed = np.sqrt(np.sum(velocity**2, axis=1)) speeds.extend(speed[:-1]) delta = np.sqrt(np.sum(np.diff(streamline, axis=0)**2, axis=1)) arc_length = np.cumsum(delta) length = arc_length[-1] s = arc_length / length linewidth = (1 - s) * start_width + s * end_width widths.extend(linewidth) points = [] for streamline in streamplotter.streamlines: pts = streamline.reshape(-1, 1, 2) points.extend(np.hstack((pts[:-1], pts[1:]))) speeds = np.array(speeds) widths = np.array(widths) vmin = kwargs.pop("vmin", speeds.min()) vmax = kwargs.pop("vmax", speeds.max()) norm = kwargs.pop("norm", matplotlib.colors.Normalize(vmin=vmin, vmax=vmax)) cmap = plt.get_cmap(kwargs.pop("cmap", None)) collection = LineCollection(points, cmap=cmap, norm=norm, linewidth=widths) collection.set_array(speeds) axes.add_collection(collection) _autoscale_view(axes, function.ufl_domain().coordinates.dat.data_ro) return collection
('rho', SSPRK3(state, rho0, rhoeqn)), ('theta', SSPRK3(state, theta0, thetaeqn, limiter=limiter)), ('water_v', SSPRK3(state, water_v0, thetaeqn, limiter=limiter)), ('water_c', SSPRK3(state, water_c0, thetaeqn, limiter=limiter))] # Set up linear solver linear_solver = CompressibleSolver(state, moisture=moisture) # Set up forcing compressible_forcing = CompressibleForcing(state, moisture=moisture, euler_poincare=euler_poincare) # diffusion bcs = [DirichletBC(Vu, 0.0, "bottom"), DirichletBC(Vu, 0.0, "top")] diffused_fields = [] if diffusion: diffused_fields.append(('u', InteriorPenalty(state, Vu, kappa=Constant(60.), mu=Constant(10./deltax), bcs=bcs))) # define condensation physics_list = [Condensation(state)] # build time stepper stepper = CrankNicolson(state, advected_fields, linear_solver, compressible_forcing, physics_list=physics_list, diffused_fields=diffused_fields) stepper.run(t=0, tmax=tmax)
else: # Fixed quadrature with 64 points gives absolute errors below 1e-13 # for a quantity of order 1e-3. v, _ = integrate.fixed_quad(D_integrand, theta_0, angles[ii], n=64) val[ii] = -(R / g) * v return val # Get coordinates to pass to Dval function D0 = Function(W0) W = VectorFunctionSpace(mesh, D0.ufl_element()) X = interpolate(mesh.coordinates, W) D0.dat.data[:] = Dval(X.dat.data_ro) # Adjust mean value of initial D C = Function(D0.function_space()).assign(Constant(1.0)) area = assemble(C * dx) Dmean = assemble(D0 * dx) / area D0 -= Dmean D0 += Constant(D_mean) # Dpert D_p = Function(W0) Dpert = D_bump * cos(theta) * exp(-(lamda / a)**2) * exp(-( (theta_2 - theta) / b)**2) D_p.interpolate(Dpert) # Dexpr = Dbar + Dpert Dexpr = D0 + D_p bexpr = Constant(0) # Set up functions
# Set up linear solver linear_solver = CompressibleSolver(state, moisture=moisture) # Set up forcing compressible_forcing = CompressibleForcing(state, moisture=moisture) # diffusion bcs = [DirichletBC(Vu, 0.0, "bottom"), DirichletBC(Vu, 0.0, "top")] diffused_fields = [] if diffusion: diffused_fields.append(('u', InteriorPenalty(state, Vu, kappa=Constant(60.), mu=Constant(10. / deltax), bcs=bcs))) # define condensation physics_list = [Condensation(state)] # build time stepper stepper = CrankNicolson(state, advected_fields, linear_solver, compressible_forcing, physics_list=physics_list, diffused_fields=diffused_fields) stepper.run(t=0, tmax=tmax)
def getForm(F, butch, t, dt, u0, bcs=None): """Given a time-dependent variational form and a :class:`ButcherTableau`, produce UFL for the s-stage RK method. :arg F: UFL form for the semidiscrete ODE/DAE :arg butch: the :class:`ButcherTableau` for the RK method being used to advance in time. :arg t: a :class:`Constant` referring to the current time level. Any explicit time-dependence in F is included :arg dt: a :class:`Constant` referring to the size of the current time step. :arg u0: a :class:`Function` referring to the state of the PDE system at time `t` :arg bcs: optionally, a :class:`DirichletBC` object (or iterable thereof) containing (possible time-dependent) boundary conditions imposed on the system. On output, we return a tuple consisting of four parts: - Fnew, the :class:`Form` - k, the :class:`firedrake.Function` holding all the stages. It lives in a :class:`firedrake.FunctionSpace` corresponding to the s-way tensor product of the space on which the semidiscrete form lives. - `bcnew`, a list of :class:`firedrake.DirichletBC` objects to be posed on the stages, - `gblah`, a list of pairs of the form (f, expr), where f is a :class:`firedrake.Function` and expr is a :class:`ufl.Expr`. at each time step, each expr needs to be re-interpolated/projected onto the corresponding f in order for Firedrake to pick up that time-dependent boundary conditions need to be re-applied. """ v = F.arguments()[0] V = v.function_space() assert V == u0.function_space() A = numpy.array([[Constant(aa) for aa in arow] for arow in butch.A]) c = numpy.array([Constant(ci) for ci in butch.c]) num_stages = len(c) num_fields = len(V) Vbig = numpy.prod([V for i in range(num_stages)]) # Silence a warning about transfer managers when we # coarsen coefficients in V push_parent(V.dm, Vbig.dm) vnew = TestFunction(Vbig) k = Function(Vbig) if len(V) == 1: u0bits = [u0] vbits = [v] if num_stages == 1: vbigbits = [vnew] kbits = [k] else: vbigbits = split(vnew) kbits = split(k) else: u0bits = split(u0) vbits = split(v) vbigbits = split(vnew) kbits = split(k) kbits_np = numpy.zeros((num_stages, num_fields), dtype="object") for i in range(num_stages): for j in range(num_fields): kbits_np[i, j] = kbits[i * num_fields + j] Ak = A @ kbits_np Fnew = Zero() for i in range(num_stages): repl = {t: t + c[i] * dt} for j, (ubit, vbit, kbit) in enumerate(zip(u0bits, vbits, kbits)): repl[ubit] = ubit + dt * Ak[i, j] repl[vbit] = vbigbits[num_fields * i + j] repl[TimeDerivative(ubit)] = kbits_np[i, j] if (len(ubit.ufl_shape) == 1): for kk, kbitbit in enumerate(kbits_np[i, j]): repl[TimeDerivative(ubit[kk])] = kbitbit repl[ubit[kk]] = repl[ubit][kk] repl[vbit[kk]] = repl[vbit][kk] Fnew += replace(F, repl) bcnew = [] gblah = [] if bcs is None: bcs = [] for bc in bcs: if isinstance(bc.domain_args[0], str): boundary = bc.domain_args[0] else: boundary = () try: for j in bc.sub_domain: boundary += j except TypeError: boundary = (bc.sub_domain, ) gfoo = expand_derivatives(diff(bc._original_arg, t)) if len(V) == 1: for i in range(num_stages): gcur = replace(gfoo, {t: t + c[i] * dt}) try: gdat = interpolate(gcur, V) except NotImplementedError: gdat = project(gcur, V) gblah.append((gdat, gcur)) bcnew.append(DirichletBC(Vbig[i], gdat, boundary)) else: sub = bc.function_space_index() for i in range(num_stages): gcur = replace(gfoo, {t: t + butch.c[i] * dt}) try: gdat = interpolate(gcur, V.sub(sub)) except NotImplementedError: gdat = project(gcur, V) gblah.append((gdat, gcur)) bcnew.append( DirichletBC(Vbig[sub + num_fields * i], gdat, boundary)) return Fnew, k, bcnew, gblah
# Define constant theta_e and water_t Tsurf = 283.0 psurf = 85000. exner_surf = (psurf / state.parameters.p_0)**state.parameters.kappa humidity = 0.2 S = 1.3e-5 theta_surf = thermodynamics.theta(state.parameters, Tsurf, psurf) theta_d = Function(Vt).interpolate(theta_surf * exp(S * z)) H = Function(Vt).assign(humidity) # Calculate hydrostatic fields unsaturated_hydrostatic_balance(state, theta_d, H, exner_boundary=Constant(exner_surf)) # make mean fields theta_b = Function(Vt).assign(theta0) rho_b = Function(Vr).assign(rho0) water_vb = Function(Vt).assign(water_v0) # define perturbation to RH xc = L / 2 zc = 800. r1 = 300. r2 = 200. r = sqrt((x - xc)**2 + (z - zc)**2) H_expr = conditional( r > r1, 0.0,
diagnostics=diagnostics, fieldlist=fieldlist, diagnostic_fields=diagnostic_fields) # Initial conditions u0 = state.fields("u") rho0 = state.fields("rho") theta0 = state.fields("theta") # spaces Vu = u0.function_space() Vt = theta0.function_space() Vr = rho0.function_space() # Isentropic background state Tsurf = Constant(300.) theta_b = Function(Vt).interpolate(Tsurf) rho_b = Function(Vr) # Calculate hydrostatic Pi compressible_hydrostatic_balance(state, theta_b, rho_b, solve_for_rho=True) x = SpatialCoordinate(mesh) xc = 500. zc = 350. rc = 250. r = sqrt((x[0]-xc)**2 + (x[1]-zc)**2) theta_pert = conditional(r > rc, 0., 0.25*(1. + cos((pi/rc)*r))) theta0.interpolate(theta_b + theta_pert)
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 _ad_assign_numpy(dst, src, offset): dst.assign(Constant(numpy.reshape(src[offset:offset + dst.value_size()], dst.ufl_shape))) offset += dst.value_size() return dst, offset
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
u_bc_ids=[1, 2]) # Initial conditions u0 = state.fields("u") rho0 = state.fields("rho") theta0 = state.fields("theta") # spaces Vu = u0.function_space() Vt = theta0.function_space() Vr = rho0.function_space() x, z = SpatialCoordinate(mesh) # Define constant theta_e and water_t Tsurf = 300.0 theta_b = Function(Vt).interpolate(Constant(Tsurf)) # Calculate hydrostatic fields compressible_hydrostatic_balance(state, theta_b, rho0, solve_for_rho=True) # make mean fields rho_b = Function(Vr).assign(rho0) # define perturbation xc = L / 2 zc = 2000. rc = 2000. Tdash = 2.0 r = sqrt((x - xc)**2 + (z - zc)**2) theta_pert = Function(Vt).interpolate( conditional(r > rc, 0.0,
def saturated_hydrostatic_balance(state, theta_e, water_t, pi0=None, top=False, pi_boundary=Constant(1.0), max_outer_solve_count=40, max_theta_solve_count=5, max_inner_solve_count=3): """ Given a wet equivalent potential temperature, theta_e, and the total moisture content, water_t, compute a hydrostatically balance virtual potential temperature, dry density and water vapour profile. The general strategy is to split up the solving into two steps: 1) finding rho to balance the theta profile 2) finding theta_v and r_v to get back theta_e and saturation We iteratively solve these steps until we (hopefully) converge to a solution. :arg state: The :class:`State` object. :arg theta_e: The initial wet equivalent potential temperature profile. :arg water_t: The total water pseudo-mixing ratio profile. :arg pi0: Optional function to put exner pressure into. :arg top: If True, set a boundary condition at the top, otherwise it will be at the bottom. :arg pi_boundary: The value of pi on the specified boundary. :arg max_outer_solve_count: Max number of outer iterations for balance solver. :arg max_theta_solve_count: Max number of iterations for theta solver (middle part of solve). :arg max_inner_solve_count: Max number of iterations on the inner most loop for the water vapour solver. """ theta0 = state.fields('theta') rho0 = state.fields('rho') water_v0 = state.fields('water_v') # Calculate hydrostatic Pi Vt = theta0.function_space() Vr = rho0.function_space() VDG = state.spaces("DG") if any(deg > 2 for deg in VDG.ufl_element().degree()): logger.warning("default quadrature degree most likely not sufficient for this degree element") theta0.interpolate(theta_e) water_v0.interpolate(water_t) if state.vertical_degree == 0: boundary_method = Boundary_Method.physics else: boundary_method = None rho_h = Function(Vr) Vt_broken = FunctionSpace(state.mesh, BrokenElement(Vt.ufl_element())) rho_averaged = Function(Vt) rho_recoverer = Recoverer(rho0, rho_averaged, VDG=Vt_broken, boundary_method=boundary_method) w_h = Function(Vt) theta_h = Function(Vt) theta_e_test = Function(Vt) delta = 0.8 # expressions for finding theta0 and water_v0 from theta_e and water_t pie = thermodynamics.pi(state.parameters, rho_averaged, theta0) p = thermodynamics.p(state.parameters, pie) T = thermodynamics.T(state.parameters, theta0, pie, water_v0) r_v_expr = thermodynamics.r_sat(state.parameters, T, p) theta_e_expr = thermodynamics.theta_e(state.parameters, T, p, water_v0, water_t) for i in range(max_outer_solve_count): # solve for rho with theta_vd and w_v guesses compressible_hydrostatic_balance(state, theta0, rho_h, top=top, pi_boundary=pi_boundary, water_t=water_t, solve_for_rho=True) # damp solution rho0.assign(rho0 * (1 - delta) + delta * rho_h) theta_e_test.assign(theta_e_expr) if errornorm(theta_e_test, theta_e) < 1e-8: break # calculate averaged rho rho_recoverer.project() # now solve for r_v for j in range(max_theta_solve_count): theta_h.interpolate(theta_e / theta_e_expr * theta0) theta0.assign(theta0 * (1 - delta) + delta * theta_h) # break when close enough if errornorm(theta_e_test, theta_e) < 1e-6: break for k in range(max_inner_solve_count): w_h.interpolate(r_v_expr) water_v0.assign(water_v0 * (1 - delta) + delta * w_h) # break when close enough theta_e_test.assign(theta_e_expr) if errornorm(theta_e_test, theta_e) < 1e-6: break if i == max_outer_solve_count: raise RuntimeError('Hydrostatic balance solve has not converged within %i' % i, 'iterations') if pi0 is not None: pie = thermodynamics.pi(state.parameters, rho0, theta0) pi0.interpolate(pie) # do one extra solve for rho compressible_hydrostatic_balance(state, theta0, rho0, top=top, pi_boundary=pi_boundary, water_t=water_t, solve_for_rho=True)
def initialize(self, pc): """Set up the problem context. Take the original mixed problem and reformulate the problem as a hybridized mixed system. A KSP is created for the Lagrange multiplier system. """ from firedrake import (FunctionSpace, Function, Constant, TrialFunction, TrialFunctions, TestFunction, DirichletBC, assemble) from firedrake.assemble import (allocate_matrix, create_assembly_callable) from firedrake.formmanipulation import split_form from ufl.algorithms.replace import replace # Extract the problem context prefix = pc.getOptionsPrefix() + "hybridization_" _, P = pc.getOperators() self.cxt = P.getPythonContext() if not isinstance(self.cxt, ImplicitMatrixContext): raise ValueError("The python context must be an ImplicitMatrixContext") test, trial = self.cxt.a.arguments() V = test.function_space() mesh = V.mesh() if len(V) != 2: raise ValueError("Expecting two function spaces.") if all(Vi.ufl_element().value_shape() for Vi in V): raise ValueError("Expecting an H(div) x L2 pair of spaces.") # Automagically determine which spaces are vector and scalar for i, Vi in enumerate(V): if Vi.ufl_element().sobolev_space().name == "HDiv": self.vidx = i else: assert Vi.ufl_element().sobolev_space().name == "L2" self.pidx = i # Create the space of approximate traces. W = V[self.vidx] if W.ufl_element().family() == "Brezzi-Douglas-Marini": tdegree = W.ufl_element().degree() else: try: # If we have a tensor product element h_deg, v_deg = W.ufl_element().degree() tdegree = (h_deg - 1, v_deg - 1) except TypeError: tdegree = W.ufl_element().degree() - 1 TraceSpace = FunctionSpace(mesh, "HDiv Trace", tdegree) # Break the function spaces and define fully discontinuous spaces broken_elements = ufl.MixedElement([ufl.BrokenElement(Vi.ufl_element()) for Vi in V]) V_d = FunctionSpace(mesh, broken_elements) # Set up the functions for the original, hybridized # and schur complement systems self.broken_solution = Function(V_d) self.broken_residual = Function(V_d) self.trace_solution = Function(TraceSpace) self.unbroken_solution = Function(V) self.unbroken_residual = Function(V) # Set up the KSP for the hdiv residual projection hdiv_mass_ksp = PETSc.KSP().create(comm=pc.comm) hdiv_mass_ksp.setOptionsPrefix(prefix + "hdiv_residual_") # HDiv mass operator p = TrialFunction(V[self.vidx]) q = TestFunction(V[self.vidx]) mass = ufl.dot(p, q)*ufl.dx # TODO: Bcs? M = assemble(mass, bcs=None, form_compiler_parameters=self.cxt.fc_params) M.force_evaluation() Mmat = M.petscmat hdiv_mass_ksp.setOperators(Mmat) hdiv_mass_ksp.setUp() hdiv_mass_ksp.setFromOptions() self.hdiv_mass_ksp = hdiv_mass_ksp # Storing the result of A.inv * r, where A is the HDiv # mass matrix and r is the HDiv residual self._primal_r = Function(V[self.vidx]) tau = TestFunction(V_d[self.vidx]) self._assemble_broken_r = create_assembly_callable( ufl.dot(self._primal_r, tau)*ufl.dx, tensor=self.broken_residual.split()[self.vidx], form_compiler_parameters=self.cxt.fc_params) # Create the symbolic Schur-reduction: # Original mixed operator replaced with "broken" # arguments arg_map = {test: TestFunction(V_d), trial: TrialFunction(V_d)} Atilde = Tensor(replace(self.cxt.a, arg_map)) gammar = TestFunction(TraceSpace) n = ufl.FacetNormal(mesh) sigma = TrialFunctions(V_d)[self.vidx] if mesh.cell_set._extruded: Kform = (gammar('+') * ufl.dot(sigma, n) * ufl.dS_h + gammar('+') * ufl.dot(sigma, n) * ufl.dS_v) else: Kform = (gammar('+') * ufl.dot(sigma, n) * ufl.dS) # Here we deal with boundaries. If there are Neumann # conditions (which should be enforced strongly for # H(div)xL^2) then we need to add jump terms on the exterior # facets. If there are Dirichlet conditions (which should be # enforced weakly) then we need to zero out the trace # variables there as they are not active (otherwise the hybrid # problem is not well-posed). # If boundary conditions are contained in the ImplicitMatrixContext: if self.cxt.row_bcs: # Find all the subdomains with neumann BCS # These are Dirichlet BCs on the vidx space neumann_subdomains = set() for bc in self.cxt.row_bcs: if bc.function_space().index == self.pidx: raise NotImplementedError("Dirichlet conditions for scalar variable not supported. Use a weak bc") if bc.function_space().index != self.vidx: raise NotImplementedError("Dirichlet bc set on unsupported space.") # append the set of sub domains subdom = bc.sub_domain if isinstance(subdom, str): neumann_subdomains |= set([subdom]) else: neumann_subdomains |= set(as_tuple(subdom, int)) # separate out the top and bottom bcs extruded_neumann_subdomains = neumann_subdomains & {"top", "bottom"} neumann_subdomains = neumann_subdomains.difference(extruded_neumann_subdomains) integrand = gammar * ufl.dot(sigma, n) measures = [] trace_subdomains = [] if mesh.cell_set._extruded: ds = ufl.ds_v for subdomain in extruded_neumann_subdomains: measures.append({"top": ufl.ds_t, "bottom": ufl.ds_b}[subdomain]) trace_subdomains.extend(sorted({"top", "bottom"} - extruded_neumann_subdomains)) else: ds = ufl.ds if "on_boundary" in neumann_subdomains: measures.append(ds) else: measures.append(ds(tuple(neumann_subdomains))) dirichlet_subdomains = set(mesh.exterior_facets.unique_markers) - neumann_subdomains trace_subdomains.append(sorted(dirichlet_subdomains)) for measure in measures: Kform += integrand*measure trace_bcs = [DirichletBC(TraceSpace, Constant(0.0), subdomain) for subdomain in trace_subdomains] else: # No bcs were provided, we assume weak Dirichlet conditions. # We zero out the contribution of the trace variables on # the exterior boundary. Extruded cells will have both # horizontal and vertical facets trace_subdomains = ["on_boundary"] if mesh.cell_set._extruded: trace_subdomains.extend(["bottom", "top"]) trace_bcs = [DirichletBC(TraceSpace, Constant(0.0), subdomain) for subdomain in trace_subdomains] # Make a SLATE tensor from Kform K = Tensor(Kform) # Assemble the Schur complement operator and right-hand side self.schur_rhs = Function(TraceSpace) self._assemble_Srhs = create_assembly_callable( K * Atilde.inv * AssembledVector(self.broken_residual), tensor=self.schur_rhs, form_compiler_parameters=self.cxt.fc_params) schur_comp = K * Atilde.inv * K.T self.S = allocate_matrix(schur_comp, bcs=trace_bcs, form_compiler_parameters=self.cxt.fc_params) self._assemble_S = create_assembly_callable(schur_comp, tensor=self.S, bcs=trace_bcs, form_compiler_parameters=self.cxt.fc_params) self._assemble_S() self.S.force_evaluation() Smat = self.S.petscmat # Nullspace for the multiplier problem nullspace = create_schur_nullspace(P, -K * Atilde, V, V_d, TraceSpace, pc.comm) if nullspace: Smat.setNullSpace(nullspace) # Set up the KSP for the system of Lagrange multipliers trace_ksp = PETSc.KSP().create(comm=pc.comm) trace_ksp.setOptionsPrefix(prefix) trace_ksp.setOperators(Smat) trace_ksp.setUp() trace_ksp.setFromOptions() self.trace_ksp = trace_ksp split_mixed_op = dict(split_form(Atilde.form)) split_trace_op = dict(split_form(K.form)) # Generate reconstruction calls self._reconstruction_calls(split_mixed_op, split_trace_op) # NOTE: The projection stage *might* be replaced by a Fortin # operator. We may want to allow the user to specify if they # wish to use a Fortin operator over a projection, or vice-versa. # In a future add-on, we can add a switch which chooses either # the Fortin reconstruction or the usual KSP projection. # Set up the projection KSP hdiv_projection_ksp = PETSc.KSP().create(comm=pc.comm) hdiv_projection_ksp.setOptionsPrefix(prefix + 'hdiv_projection_') # Reuse the mass operator from the hdiv_mass_ksp hdiv_projection_ksp.setOperators(Mmat) # Construct the RHS for the projection stage self._projection_rhs = Function(V[self.vidx]) self._assemble_projection_rhs = create_assembly_callable( ufl.dot(self.broken_solution.split()[self.vidx], q)*ufl.dx, tensor=self._projection_rhs, form_compiler_parameters=self.cxt.fc_params) # Finalize ksp setup hdiv_projection_ksp.setUp() hdiv_projection_ksp.setFromOptions() self.hdiv_projection_ksp = hdiv_projection_ksp
def unsaturated_hydrostatic_balance(state, theta_d, H, pi0=None, top=False, pi_boundary=Constant(1.0), max_outer_solve_count=40, max_inner_solve_count=20): """ Given vertical profiles for dry potential temperature and relative humidity compute hydrostatically balanced virtual potential temperature, dry density and water vapour profiles. The general strategy is to split up the solving into two steps: 1) finding rho to balance the theta profile 2) finding theta_v and r_v to get back theta_d and H We iteratively solve these steps until we (hopefully) converge to a solution. :arg state: The :class:`State` object. :arg theta_d: The initial dry potential temperature profile. :arg H: The relative humidity profile. :arg pi0: Optional function to put exner pressure into. :arg top: If True, set a boundary condition at the top, otherwise it will be at the bottom. :arg pi_boundary: The value of pi on the specified boundary. :arg max_outer_solve_count: Max number of iterations for outer loop of balance solver. :arg max_inner_solve_count: Max number of iterations for inner loop of balanace solver. """ theta0 = state.fields('theta') rho0 = state.fields('rho') water_v0 = state.fields('water_v') # Calculate hydrostatic Pi Vt = theta0.function_space() Vr = rho0.function_space() R_d = state.parameters.R_d R_v = state.parameters.R_v epsilon = R_d / R_v VDG = state.spaces("DG") if any(deg > 2 for deg in VDG.ufl_element().degree()): logger.warning("default quadrature degree most likely not sufficient for this degree element") # apply first guesses theta0.assign(theta_d * 1.01) water_v0.assign(0.01) if state.vertical_degree == 0: method = Boundary_Method.physics else: method = None rho_h = Function(Vr) rho_averaged = Function(Vt) Vt_broken = FunctionSpace(state.mesh, BrokenElement(Vt.ufl_element())) rho_recoverer = Recoverer(rho0, rho_averaged, VDG=Vt_broken, boundary_method=method) w_h = Function(Vt) delta = 1.0 # make expressions for determining water_v0 pie = thermodynamics.pi(state.parameters, rho_averaged, theta0) p = thermodynamics.p(state.parameters, pie) 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) for i in range(max_outer_solve_count): # solve for rho with theta_vd and w_v guesses compressible_hydrostatic_balance(state, theta0, rho_h, top=top, pi_boundary=pi_boundary, water_t=water_v0, solve_for_rho=True) # damp solution rho0.assign(rho0 * (1 - delta) + delta * rho_h) # calculate averaged rho rho_recoverer.project() RH.assign(RH_ev) if errornorm(RH, H) < 1e-10: break # now solve for r_v for j in range(max_inner_solve_count): w_h.interpolate(r_v_expr) water_v0.assign(water_v0 * (1 - delta) + delta * w_h) # compute theta_vd theta0.assign(theta_d * (1 + water_v0 / epsilon)) # test quality of solution by re-evaluating expression RH.assign(RH_ev) if errornorm(RH, H) < 1e-10: break if i == max_outer_solve_count: raise RuntimeError('Hydrostatic balance solve has not converged within %i' % i, 'iterations') if pi0 is not None: pie = thermodynamics.pi(state.parameters, rho0, theta0) pi0.interpolate(pie) # do one extra solve for rho compressible_hydrostatic_balance(state, theta0, rho0, top=top, pi_boundary=pi_boundary, water_t=water_v0, solve_for_rho=True)
def transmission( mesh, scatterer_bdy_id, outer_bdy_id, wave_number, options_prefix=None, solver_parameters=None, fspace=None, true_sol_grad_expr=None, ): r""" preconditioner_gamma and preconditioner_lambda are used to precondition with the following equation: \Delta u - \kappa^2 \gamma u = 0 (\partial_n - i\kappa\beta) u |_\Sigma = 0 """ # need as tuple so can use integral measure if isinstance(outer_bdy_id, int): outer_bdy_id = [outer_bdy_id] outer_bdy_id = tuple(outer_bdy_id) u = TrialFunction(fspace) v = TestFunction(fspace) a = inner(grad(u), grad(v)) * dx - Constant(wave_number**2) * inner(u, v) * dx \ - Constant(1j * wave_number) * inner(u, v) * ds(outer_bdy_id) n = FacetNormal(mesh) metadata = {'quadrature_degree': 2 * fspace.ufl_element().degree()} L = inner(inner(true_sol_grad_expr, n), v) * ds(scatterer_bdy_id, metadata=metadata) solution = Function(fspace) # {{{ Used for preconditioning if 'gamma' in solver_parameters or 'beta' in solver_parameters: solver_params = dict(solver_parameters) gamma = complex(solver_parameters.pop('gamma', 1.0)) import cmath beta = complex(solver_parameters.pop('beta', cmath.sqrt(gamma))) aP = inner(grad(u), grad(v)) * dx \ - Constant(wave_number**2 * gamma) * inner(u, v) * dx \ - Constant(1j * wave_number * beta) * inner(u, v) * ds(outer_bdy_id) else: aP = None solver_params = solver_parameters # }}} # prepare to set up pyamg preconditioner if using it using_pyamg = solver_params['pc_type'] == 'pyamg' if using_pyamg: pyamg_tol = solver_parameters.get('pyamg_tol', None) if pyamg_tol is not None: pyamg_tol = float(pyamg_tol) pyamg_maxiter = solver_params.get('pyamg_maxiter', None) if pyamg_maxiter is not None: pyamg_maxiter = int(pyamg_maxiter) del solver_params['pc_type'] # Create a solver and return the KSP object with the solution so that can get # PETSc information # Create problem problem = vs.LinearVariationalProblem(a, L, solution, aP=aP) # Create solver and call solve solver = vs.LinearVariationalSolver(problem, solver_parameters=solver_params, options_prefix=options_prefix) # prepare to set up pyamg preconditioner if using it if using_pyamg: A = assemble(a).M.handle pc = solver.snes.getKSP().pc pc.setType(pc.Type.PYTHON) pc.setPythonContext( AMGTransmissionPreconditioner(wave_number, fspace, A, tol=pyamg_tol, maxiter=pyamg_maxiter, use_plane_waves=True)) # If using pyamg as preconditioner, use it! try: solver.solve() except ConvergenceError: pass return solver.snes, solution
def simulation(simulation_parameters, diagnostic_values=None, fields_to_output=None, expected_u=False): """ This function sets up and runs a basic simulation of the stochastic Camassa-Holm equation. :arg simulation_parameters: a dictionary containing the parameters for the simulation. Each value should be a tuple, which will normally have only one element. Two-element tuples will also contain an index for outputting -- these are variables being looped over by the experiment. :arg diagnostic_values: a list of diagnostic numbers to output. :arg fields_to_output: a list of diagnostic fields to be dumped. """ # ensure simulation parameters use ufl Constant for key in ['alphasq', 'c0', 'gamma']: if not isinstance(simulation_parameters[key][-1], Constant): if len(simulation_parameters[key]) == 2: simulation_parameters[key] = ( simulation_parameters[key][0], Constant(simulation_parameters[key][1])) elif len(simulation_parameters[key]) == 1: simulation_parameters[key] = (Constant( simulation_parameters[key][-1]), ) else: raise ValueError( 'The value of simulation_parameters %s seems to be of length %d. It should be either 1 or 2.' % (key, len(simulation_parameters[key]))) mesh = simulation_parameters['mesh'][-1] dt = simulation_parameters['dt'][-1] tmax = simulation_parameters['tmax'][-1] ndump = simulation_parameters['ndump'][-1] field_ndump = simulation_parameters['field_ndump'][-1] file_name = simulation_parameters['file_name'][-1] scheme = simulation_parameters['scheme'][-1] allow_fail = simulation_parameters['allow_fail'][-1] only_peakons = simulation_parameters['only_peakons'][-1] peakon_speed = simulation_parameters['peakon_speed'][-1] peakon_method = simulation_parameters['peakon_method'][-1] num_Xis = simulation_parameters['num_Xis'][-1] allow_fail = True if only_peakons else allow_fail prognostic_variables = PrognosticVariables(scheme, mesh, num_Xis) diagnostic_variables = DiagnosticVariables( prognostic_variables, fields_to_output, diagnostic_values=diagnostic_values) build_initial_conditions(prognostic_variables, simulation_parameters) Xis = StochasticFunctions(prognostic_variables, simulation_parameters) equations = Equations(prognostic_variables, simulation_parameters) # set up peakon equations if simulation_parameters['peakon_equations'][-1] == True: peakon_equations = PeakonEquations(prognostic_variables, simulation_parameters, peakon_speed, peakon_method) elif simulation_parameters['peakon_equations'][-1] == False: peakon_equations = None else: raise ValueError('peakon_equations must be True or False, not %s' % simulation_parameters['peakon_equations'][-1]) outputting = Outputting(prognostic_variables, diagnostic_variables, simulation_parameters, peakon_equations=peakon_equations, diagnostic_values=diagnostic_values) diagnostic_equations = DiagnosticEquations(diagnostic_variables, prognostic_variables, outputting, simulation_parameters) diagnostic_equations.solve() dumpn = 0 field_dumpn = 0 # want the update to always happen on first time step (so results are consistent with earlier results) t = 0 failed = True if only_peakons else False failed_time = np.nan outputting.dump_diagnostics(t) outputting.dump_fields(t) # run simulation while (t < tmax - 0.5 * dt): t += dt # update stochastic basis functions Xis.update(t) # solve problems if allow_fail: if not failed: try: equations.solve() except ConvergenceError: failed = True failed_time = t print("Solver failed at t = %f" % t) else: equations.solve() if peakon_equations is not None: peakon_equations.update() # output if necessary dumpn += 1 field_dumpn += 1 # do interpolations/projections for any outputing if dumpn == ndump or field_dumpn == field_ndump: if not failed: diagnostic_equations.solve() # output diagnostic values if dumpn == ndump: outputting.dump_diagnostics(t, failed=failed) dumpn -= ndump # output fields if field_dumpn == field_ndump: if not failed: outputting.dump_fields(t) field_dumpn -= field_ndump outputting.store_times(failed_time) outputting.data_file.close() if expected_u: return prognostic_variables.u