class IncompressibleState(BaroclinicState): def __init__(self, mesh, vertical_degree=1, horizontal_degree=1, family="RT", z=None, k=None, Omega=None, mu=None, timestepping=None, output=None, parameters=None, diagnostics=None, fieldlist=None, diagnostic_fields=[], on_sphere=False): super(IncompressibleState, self).__init__(mesh=mesh, vertical_degree=vertical_degree, horizontal_degree=horizontal_degree, family=family, z=z, k=k, Omega=Omega, mu=mu, timestepping=timestepping, output=output, parameters=parameters, diagnostics=diagnostics, fieldlist=fieldlist, diagnostic_fields=diagnostic_fields) if parameters.geopotential: raise RuntimeError("geopotential formulation is not compatible with incompressible Boussinesq") def set_reference_profiles(self, b_ref): """ Initialise reference profiles :arg b_ref: :class:`.Function` object, reference bouyancy """ self.bbar = Function(self.V[2]) self.bbar.project(b_ref)
class EmbeddedDGAdvection(DGAdvection): def __init__(self, state, V, Vdg=None, continuity=False): if Vdg is None: Vdg_elt = BrokenElement(V.ufl_element()) Vdg = FunctionSpace(state.mesh, Vdg_elt) super(EmbeddedDGAdvection, self).__init__(state, Vdg, continuity) self.xdg_in = Function(Vdg) self.xdg_out = Function(Vdg) self.x_projected = Function(V) pparameters = {'ksp_type':'cg', 'pc_type':'bjacobi', 'sub_pc_type':'ilu'} self.Projector = Projector(self.xdg_out, self.x_projected, solver_parameters=pparameters) def apply(self, x_in, x_out): self.xdg_in.interpolate(x_in) super(EmbeddedDGAdvection, self).apply(self.xdg_in, self.xdg_out) self.Projector.project() x_out.assign(self.x_projected)
def set_reference_profiles(self, rho_ref, theta_ref): """ Initialise reference profiles :arg rho_ref: :class:`.Function` object, initial rho :arg theta_ref: :class:`.Function` object, initial theta """ self.rhobar = Function(self.V[1]) self.thetabar = Function(self.V[2]) self.rhobar.project(rho_ref) self.thetabar.project(theta_ref)
def __init__(self, state, linear=False): self.state = state g = state.parameters.g f = state.f Vu = state.V[0] W = state.W self.x0 = Function(W) # copy x to here u0, D0 = split(self.x0) n = FacetNormal(state.mesh) un = 0.5 * (dot(u0, n) + abs(dot(u0, n))) F = TrialFunction(Vu) w = TestFunction(Vu) self.uF = Function(Vu) outward_normals = CellNormal(state.mesh) perp = lambda u: cross(outward_normals, u) a = inner(w, F) * dx L = (-f * inner(w, perp(u0)) + g * div(w) * D0) * dx - g * inner( jump(w, n), un("+") * D0("+") - un("-") * D0("-") ) * dS if not linear: L -= 0.5 * div(w) * inner(u0, u0) * dx u_forcing_problem = LinearVariationalProblem(a, L, self.uF) self.u_forcing_solver = LinearVariationalSolver(u_forcing_problem)
def __init__(self, state, V, direction=[1,2], params=None): super(InteriorPenalty, self).__init__(state) dt = state.timestepping.dt kappa = params['kappa'] mu = params['mu'] gamma = TestFunction(V) phi = TrialFunction(V) self.phi1 = Function(V) n = FacetNormal(state.mesh) a = inner(gamma,phi)*dx + dt*inner(grad(gamma), grad(phi)*kappa)*dx def get_flux_form(dS, M): fluxes = (-inner(2*avg(outer(phi, n)), avg(grad(gamma)*M)) - inner(avg(grad(phi)*M), 2*avg(outer(gamma, n))) + mu*inner(2*avg(outer(phi, n)), 2*avg(outer(gamma, n)*kappa)))*dS return fluxes if 1 in direction: a += dt*get_flux_form(dS_v, kappa) if 2 in direction: a += dt*get_flux_form(dS_h, kappa) L = inner(gamma,phi)*dx problem = LinearVariationalProblem(a, action(L,self.phi1), self.phi1) self.solver = LinearVariationalSolver(problem)
def get_latlon_mesh(mesh): coords_orig = mesh.coordinates mesh_dg_fs = VectorFunctionSpace(mesh, "DG", 1) coords_dg = Function(mesh_dg_fs) coords_latlon = Function(mesh_dg_fs) par_loop(""" for (int i=0; i<3; i++) { for (int j=0; j<3; j++) { dg[i][j] = cg[i][j]; } } """, dx, {'dg': (coords_dg, WRITE), 'cg': (coords_orig, READ)}) # lat-lon 'x' = atan2(y, x) coords_latlon.dat.data[:,0] = np.arctan2(coords_dg.dat.data[:,1], coords_dg.dat.data[:,0]) # lat-lon 'y' = asin(z/sqrt(x^2 + y^2 + z^2)) coords_latlon.dat.data[:,1] = np.arcsin(coords_dg.dat.data[:,2]/np.sqrt(coords_dg.dat.data[:,0]**2 + coords_dg.dat.data[:,1]**2 + coords_dg.dat.data[:,2]**2)) coords_latlon.dat.data[:,2] = 0.0 kernel = op2.Kernel(""" #define PI 3.141592653589793 #define TWO_PI 6.283185307179586 void splat_coords(double **coords) { double diff0 = (coords[0][0] - coords[1][0]); double diff1 = (coords[0][0] - coords[2][0]); double diff2 = (coords[1][0] - coords[2][0]); if (fabs(diff0) > PI || fabs(diff1) > PI || fabs(diff2) > PI) { const int sign0 = coords[0][0] < 0 ? -1 : 1; const int sign1 = coords[1][0] < 0 ? -1 : 1; const int sign2 = coords[2][0] < 0 ? -1 : 1; if (sign0 < 0) { coords[0][0] += TWO_PI; } if (sign1 < 0) { coords[1][0] += TWO_PI; } if (sign2 < 0) { coords[2][0] += TWO_PI; } } }""", "splat_coords") op2.par_loop(kernel, coords_latlon.cell_set, coords_latlon.dat(op2.RW, coords_latlon.cell_node_map())) return Mesh(coords_latlon)
class ShallowWaterForcing(Forcing): def __init__(self, state, linear=False): self.state = state g = state.parameters.g f = state.f Vu = state.V[0] W = state.W self.x0 = Function(W) # copy x to here u0, D0 = split(self.x0) n = FacetNormal(state.mesh) un = 0.5 * (dot(u0, n) + abs(dot(u0, n))) F = TrialFunction(Vu) w = TestFunction(Vu) self.uF = Function(Vu) outward_normals = CellNormal(state.mesh) perp = lambda u: cross(outward_normals, u) a = inner(w, F) * dx L = (-f * inner(w, perp(u0)) + g * div(w) * D0) * dx - g * inner( jump(w, n), un("+") * D0("+") - un("-") * D0("-") ) * dS if not linear: L -= 0.5 * div(w) * inner(u0, u0) * dx u_forcing_problem = LinearVariationalProblem(a, L, self.uF) self.u_forcing_solver = LinearVariationalSolver(u_forcing_problem) def apply(self, scaling, x_in, x_nl, x_out, **kwargs): self.x0.assign(x_nl) self.u_forcing_solver.solve() # places forcing in self.uF self.uF *= scaling uF, _ = x_out.split() x_out.assign(x_in) uF += self.uF
class InteriorPenalty(Diffusion): """ Interior penalty diffusion method :arg state: :class:`.State` object. :arg V: Function space of diffused field :arg direction: list containing directions in which function space is discontinuous: 1 corresponds to vertical, 2 to horizontal. :arg params: dictionary containing the interior penalty parameters :mu and kappa where mu is the penalty weighting function, which is :recommended to be proportional to 1/dx """ def __init__(self, state, V, direction=[1,2], params=None): super(InteriorPenalty, self).__init__(state) dt = state.timestepping.dt kappa = params['kappa'] mu = params['mu'] gamma = TestFunction(V) phi = TrialFunction(V) self.phi1 = Function(V) n = FacetNormal(state.mesh) a = inner(gamma,phi)*dx + dt*inner(grad(gamma), grad(phi)*kappa)*dx def get_flux_form(dS, M): fluxes = (-inner(2*avg(outer(phi, n)), avg(grad(gamma)*M)) - inner(avg(grad(phi)*M), 2*avg(outer(gamma, n))) + mu*inner(2*avg(outer(phi, n)), 2*avg(outer(gamma, n)*kappa)))*dS return fluxes if 1 in direction: a += dt*get_flux_form(dS_v, kappa) if 2 in direction: a += dt*get_flux_form(dS_h, kappa) L = inner(gamma,phi)*dx problem = LinearVariationalProblem(a, action(L,self.phi1), self.phi1) self.solver = LinearVariationalSolver(problem) def apply(self, x_in, x_out): self.phi1.assign(x_in) self.solver.solve() x_out.assign(self.phi1)
def set_reference_profiles(self, b_ref): """ Initialise reference profiles :arg b_ref: :class:`.Function` object, reference bouyancy """ self.bbar = Function(self.V[2]) self.bbar.project(b_ref)
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 _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 bezier_plot(function, axes=None, **kwargs): """Plot a 1D function on a function space with order no more than 4 using Bezier curve within each cell, return a matplotlib axes :arg function: 1D function for plotting :arg axes: Axes for plotting, if None, a new one will be created :arg kwargs: additional key work arguments to plot """ try: import matplotlib.pyplot as plt from matplotlib.path import Path import matplotlib.patches as patches except ImportError: raise RuntimeError("Matplotlib not importable, is it installed?") deg = function.function_space().ufl_element().degree() mesh = function.function_space().mesh() if deg == 0: V = FunctionSpace(mesh, "DG", 1) func = Function(V).interpolate(function) return bezier_plot(func, axes, **kwargs) y_vals = _bezier_calculate_points(function) x = SpatialCoordinate(mesh) coords = Function(FunctionSpace(mesh, 'DG', deg)) coords.interpolate(x[0]) x_vals = _bezier_calculate_points(coords) vals = np.dstack((x_vals, y_vals)) if axes is None: figure = plt.figure() axes = figure.add_subplot(111) codes = {1: [Path.MOVETO, Path.LINETO], 2: [Path.MOVETO, Path.CURVE3, Path.CURVE3], 3: [Path.MOVETO, Path.CURVE4, Path.CURVE4, Path.CURVE4]} vertices = vals.reshape(-1, 2) path = Path(vertices, np.tile(codes[deg], function.function_space().cell_node_list.shape[0])) patch = patches.PathPatch(path, facecolor='none', lw=2) axes.add_patch(patch) axes.plot(**kwargs) return axes
def getranks(mesh): element_rank = Function(FunctionSpace(mesh,'DG',0)) element_rank.dat.data[:] = mesh.comm.rank element_rank.rename('element_rank') vertex_rank = Function(FunctionSpace(mesh,'CG',1)) vertex_rank.dat.data[:] = mesh.comm.rank vertex_rank.rename('vertex_rank') return (element_rank,vertex_rank)
def __init__(self, state, V, direction=[], supg_params=None): super(SUPGAdvection, self).__init__(state) dt = state.timestepping.dt params = supg_params.copy() if supg_params else {} params.setdefault('a0', dt/sqrt(15.)) params.setdefault('a1', dt/sqrt(15.)) gamma = TestFunction(V) theta = TrialFunction(V) self.theta0 = Function(V) # make SUPG test function taus = [params["a0"], params["a1"]] for i in direction: taus[i] = 0.0 tau = Constant(((taus[0], 0.), (0., taus[1]))) dgamma = dot(dot(self.ubar, tau), grad(gamma)) gammaSU = gamma + dgamma n = FacetNormal(state.mesh) un = 0.5*(dot(self.ubar, n) + abs(dot(self.ubar, n))) a_mass = gammaSU*theta*dx arhs = a_mass - dt*gammaSU*dot(self.ubar, grad(theta))*dx if 1 in direction: arhs -= ( dt*dot(jump(gammaSU), (un('+')*theta('+') - un('-')*theta('-')))*dS_v - dt*(gammaSU('+')*dot(self.ubar('+'), n('+'))*theta('+') + gammaSU('-')*dot(self.ubar('-'), n('-'))*theta('-'))*dS_v ) if 2 in direction: arhs -= ( dt*dot(jump(gammaSU), (un('+')*theta('+') - un('-')*theta('-')))*dS_h - dt*(gammaSU('+')*dot(self.ubar('+'), n('+'))*theta('+') + gammaSU('-')*dot(self.ubar('-'), n('-'))*theta('-'))*dS_h ) self.theta1 = Function(V) self.dtheta = Function(V) problem = LinearVariationalProblem(a_mass, action(arhs,self.theta1), self.dtheta) self.solver = LinearVariationalSolver(problem, options_prefix='SUPGAdvection')
def __init__(self, state, V, continuity=False): super(DGAdvection, self).__init__(state) element = V.fiat_element assert element.entity_dofs() == element.entity_closure_dofs(), "Provided space is not discontinuous" dt = state.timestepping.dt if V.extruded: surface_measure = (dS_h + dS_v) else: surface_measure = dS phi = TestFunction(V) D = TrialFunction(V) self.D1 = Function(V) self.dD = Function(V) n = FacetNormal(state.mesh) # ( dot(v, n) + |dot(v, n)| )/2.0 un = 0.5*(dot(self.ubar, n) + abs(dot(self.ubar, n))) a_mass = inner(phi,D)*dx if continuity: a_int = -inner(grad(phi), outer(D, self.ubar))*dx else: a_int = -inner(div(outer(phi,self.ubar)),D)*dx a_flux = (dot(jump(phi), un('+')*D('+') - un('-')*D('-')))*surface_measure arhs = a_mass - dt*(a_int + a_flux) DGproblem = LinearVariationalProblem(a_mass, action(arhs,self.D1), self.dD) self.DGsolver = LinearVariationalSolver(DGproblem, solver_parameters={ 'ksp_type':'preonly', 'pc_type':'bjacobi', 'sub_pc_type': 'ilu'}, options_prefix='DGAdvection')
def PeriodicRectangleMesh(nx, ny, Lx, Ly, direction="both", quadrilateral=False, reorder=None, comm=COMM_WORLD): """Generate a periodic rectangular mesh :arg nx: The number of cells in the x direction :arg ny: The number of cells in the y direction :arg Lx: The extent in the x direction :arg Ly: The extent in the y direction :arg direction: The direction of the periodicity, one of ``"both"``, ``"x"`` or ``"y"``. :kwarg quadrilateral: (optional), creates quadrilateral mesh, defaults to False :kwarg reorder: (optional), should the mesh be reordered :kwarg comm: Optional communicator to build the mesh on (defaults to COMM_WORLD). If direction == "x" the boundary edges in this mesh are numbered as follows: * 1: plane y == 0 * 2: plane y == Ly If direction == "y" the boundary edges are: * 1: plane x == 0 * 2: plane x == Lx """ if direction == "both" and ny == 1 and quadrilateral: return OneElementThickMesh(nx, Lx, Ly) if direction not in ("both", "x", "y"): raise ValueError("Cannot have a periodic mesh with periodicity '%s'" % direction) if direction != "both": return PartiallyPeriodicRectangleMesh(nx, ny, Lx, Ly, direction=direction, quadrilateral=quadrilateral, reorder=reorder, comm=comm) if nx < 3 or ny < 3: raise ValueError("2D periodic meshes with fewer than 3 \ cells in each direction are not currently supported") m = TorusMesh(nx, ny, 1.0, 0.5, quadrilateral=quadrilateral, reorder=reorder, comm=comm) coord_fs = VectorFunctionSpace(m, 'DG', 1, dim=2) old_coordinates = m.coordinates new_coordinates = Function(coord_fs) periodic_kernel = """ double pi = 3.141592653589793; double eps = 1e-12; double bigeps = 1e-1; double phi, theta, Y, Z; Y = 0.0; Z = 0.0; for(int i=0; i<old_coords.dofs; i++) { Y += old_coords[i][1]; Z += old_coords[i][2]; } for(int i=0; i<new_coords.dofs; i++) { phi = atan2(old_coords[i][1], old_coords[i][0]); if (fabs(sin(phi)) > bigeps) theta = atan2(old_coords[i][2], old_coords[i][1]/sin(phi) - 1.0); else theta = atan2(old_coords[i][2], old_coords[i][0]/cos(phi) - 1.0); new_coords[i][0] = phi/(2.0*pi); if(new_coords[i][0] < -eps) { new_coords[i][0] += 1.0; } if(fabs(new_coords[i][0]) < eps && Y < 0.0) { new_coords[i][0] = 1.0; } new_coords[i][1] = theta/(2.0*pi); if(new_coords[i][1] < -eps) { new_coords[i][1] += 1.0; } if(fabs(new_coords[i][1]) < eps && Z < 0.0) { new_coords[i][1] = 1.0; } new_coords[i][0] *= Lx[0]; new_coords[i][1] *= Ly[0]; } """ cLx = Constant(Lx) cLy = Constant(Ly) par_loop(periodic_kernel, dx, {"new_coords": (new_coordinates, WRITE), "old_coords": (old_coordinates, READ), "Lx": (cLx, READ), "Ly": (cLy, READ)}) return mesh.Mesh(new_coordinates)
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)) domain = "{[i,j]: 0 <= i < %d and 0 <= j < %d}" % shapes instructions = """ for i, j w[i,j] = w[i,j] + 1 end """ self.weight = Function(V[self.vidx]) par_loop((domain, instructions), ufl.dx, {"w": (self.weight, INC)}, is_loopy_kernel=True) instructions = """ for i, j vec_out[i,j] = vec_out[i,j] + vec_in[i,j]/w[i,j] end """ self.average_kernel = (domain, instructions) # 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 OneElementThickMesh(ncells, Lx, Ly, comm=COMM_WORLD): """ Generate a rectangular mesh in the domain with corners [0,0] and [Lx, Ly] with ncells, that is periodic in the x-direction. :arg ncells: The number of cells in the mesh. :arg Lx: The width of the domain in the x-direction. :arg Ly: The width of the domain in the y-direction. :kwarg comm: Optional communicator to build the mesh on (defaults to COMM_WORLD). """ left = np.arange(ncells, dtype=np.int32) right = np.roll(left, -1) cells = np.array([left, left, right, right]).T dx = Lx/ncells X = np.arange(1.0*ncells, dtype=np.double)*dx Y = 0.*X coords = np.array([X, Y]).T # a line of coordinates, with a looped topology plex = mesh._from_cell_list(2, cells, coords, comm) mesh1 = mesh.Mesh(plex) mesh1.topology.init() cell_numbering = mesh1._cell_numbering cell_range = plex.getHeightStratum(0) cell_closure = np.zeros((cell_range[1], 9), dtype=IntType) # Get the coordinates for this process coords = plex.getCoordinatesLocal().array_r # get the PETSc section coords_sec = plex.getCoordinateSection() for e in range(*cell_range): closure, orient = plex.getTransitiveClosure(e) # get the row for this cell row = cell_numbering.getOffset(e) # run some checks assert(closure[0] == e) assert len(closure) == 7, closure edge_range = plex.getHeightStratum(1) assert(all(closure[1:5] >= edge_range[0])) assert(all(closure[1:5] < edge_range[1])) vertex_range = plex.getHeightStratum(2) assert(all(closure[5:] >= vertex_range[0])) assert(all(closure[5:] < vertex_range[1])) # enter the cell number cell_closure[row][8] = e # Get a list of unique edges edge_set = list(set(closure[1:5])) # there are two vertices in the cell cell_vertices = closure[5:] cell_X = np.array([0., 0.]) for i, v in enumerate(cell_vertices): cell_X[i] = coords[coords_sec.getOffset(v)] # Add in the edges for i in range(3): # count up how many times each edge is repeated repeats = list(closure[1:5]).count(edge_set[i]) if repeats == 2: # we have a y-periodic edge cell_closure[row][6] = edge_set[i] cell_closure[row][7] = edge_set[i] elif repeats == 1: # in this code we check if it is a right edge, or a left edge # by inspecting the x coordinates of the edge vertex (1) # and comparing with the x coordinates of the cell vertices (2) # there is only one vertex on the edge in this case edge_vertex = plex.getCone(edge_set[i])[0] # get X coordinate for this edge edge_X = coords[coords_sec.getOffset(edge_vertex)] # get X coordinates for this cell if(cell_X.min() < dx/2): if cell_X.max() < 3*dx/2: # We are in the first cell if(edge_X.min() < dx/2): # we are on left hand edge cell_closure[row][4] = edge_set[i] else: # we are on right hand edge cell_closure[row][5] = edge_set[i] else: # We are in the last cell if(edge_X.min() < dx/2): # we are on right hand edge cell_closure[row][5] = edge_set[i] else: # we are on left hand edge cell_closure[row][4] = edge_set[i] else: if(abs(cell_X.min()-edge_X.min()) < dx/2): # we are on left hand edge cell_closure[row][4] = edge_set[i] else: # we are on right hand edge cell_closure[row][5] = edge_set[i] # Add in the vertices vertices = closure[5:] v1 = vertices[0] v2 = vertices[1] x1 = coords[coords_sec.getOffset(v1)] x2 = coords[coords_sec.getOffset(v2)] # Fix orientations if(x1 > x2): if(x1 - x2 < dx*1.5): # we are not on the rightmost cell and need to swap v1, v2 = v2, v1 elif(x2 - x1 > dx*1.5): # we are on the rightmost cell and need to swap v1, v2 = v2, v1 cell_closure[row][0:4] = [v1, v1, v2, v2] mesh1.topology.cell_closure = np.array(cell_closure, dtype=IntType) mesh1.init() Vc = VectorFunctionSpace(mesh1, 'DQ', 1) fc = Function(Vc).interpolate(mesh1.coordinates) mash = mesh.Mesh(fc) topverts = Vc.cell_node_list[:, 1::2].flatten() mash.coordinates.dat.data_with_halos[topverts, 1] = Ly # search for the last cell mcoords_ro = mash.coordinates.dat.data_ro_with_halos mcoords = mash.coordinates.dat.data_with_halos for e in range(*cell_range): cell = cell_numbering.getOffset(e) cell_nodes = Vc.cell_node_list[cell, :] Xvals = mcoords_ro[cell_nodes, 0] if(Xvals.max() - Xvals.min() > Lx/2): mcoords[cell_nodes[2:], 0] = Lx else: mcoords local_facet_dat = mash.topology.interior_facets.local_facet_dat local_facet_number = mash.topology.interior_facets.local_facet_number lfd_ro = local_facet_dat.data_ro for i in range(lfd_ro.shape[0]): if all(lfd_ro[i, :] == np.array([3, 3])): local_facet_dat.data[i, :] = [2, 3] local_facet_number[i, :] = [2, 3] return mash
def setup_sk(dirname): nlayers = 10 # horizontal layers columns = 30 # number of columns L = 1.e5 m = PeriodicIntervalMesh(columns, L) dt = 6.0 # build volume mesh H = 1.0e4 # Height position of the model top mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers) # Set up points for output at the centre of the domain, edges and corners. # The point at x=L*(13.0/30) is in the halo region for a two-way MPI decomposition points_x = [0.0, L * (13.0 / 30), L / 2.0, L] points_z = [0.0, H / 2.0, H] points = np.array([p for p in itertools.product(points_x, points_z)]) fieldlist = ['u', 'rho', 'theta'] timestepping = TimesteppingParameters(dt=dt) output = OutputParameters(dirname=dirname + "/sk_nonlinear", dumplist=['u'], dumpfreq=5, log_level=INFO, point_data=[('rho', points), ('u', points)]) parameters = CompressibleParameters() diagnostic_fields = [CourantNumber()] state = State(mesh, vertical_degree=1, horizontal_degree=1, family="CG", timestepping=timestepping, output=output, parameters=parameters, 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() # Thermodynamic constants required for setting initial conditions # and reference profiles g = parameters.g N = parameters.N # N^2 = (g/theta)dtheta/dz => dtheta/dz = theta N^2g => theta=theta_0exp(N^2gz) x, z = SpatialCoordinate(mesh) Tsurf = 300. thetab = Tsurf * exp(N**2 * z / g) theta_b = Function(Vt).interpolate(thetab) rho_b = Function(Vr) # Calculate hydrostatic Pi compressible_hydrostatic_balance(state, theta_b, rho_b) a = 5.0e3 deltaTheta = 1.0e-2 theta_pert = deltaTheta * sin(np.pi * z / H) / (1 + (x - L / 2)**2 / a**2) theta0.interpolate(theta_b + theta_pert) rho0.assign(rho_b) u0.project(as_vector([20.0, 0.0])) state.initialise([('u', u0), ('rho', rho0), ('theta', theta0)]) 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) 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))) # Set up linear solver linear_solver = CompressibleSolver(state) # Set up forcing compressible_forcing = CompressibleForcing(state) # build time stepper stepper = CrankNicolson(state, advected_fields, linear_solver, compressible_forcing) return stepper, 2 * dt
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)) domain = "{[i,j]: 0 <= i < %d and 0 <= j < %d}" % shapes instructions = """ for i, j w[i,j] = w[i,j] + 1 end """ self.weight = Function(V[self.vidx]) par_loop((domain, instructions), ufl.dx, {"w": (self.weight, INC)}, is_loopy_kernel=True) instructions = """ for i, j vec_out[i,j] = vec_out[i,j] + vec_in[i,j]/w[i,j] end """ self.average_kernel = (domain, instructions) # 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, numbers.Integral)) # 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) with timed_region("HybridOperatorAssembly"): 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)
p0 = state.fields("p") # spaces Vu = u0.function_space() Vb = b0.function_space() x, z = SpatialCoordinate(mesh) # first setup the background buoyancy profile # z.grad(bref) = N**2 # the following is symbolic algebra, using the default buoyancy frequency # from the parameters class. N = parameters.N bref = z * (N**2) # interpolate the expression to the function b_b = Function(Vb).interpolate(bref) # setup constants a = 5.0e3 deltab = 1.0e-2 b_pert = deltab * sin(np.pi * z / H) / (1 + (x - L / 2)**2 / a**2) # interpolate the expression to the function b0.interpolate(b_b + b_pert) incompressible_hydrostatic_balance(state, b_b, p0) # interpolate velocity to vector valued function space W_VectorCG1 = VectorFunctionSpace(mesh, "CG", 1) uinit = Function(W_VectorCG1).interpolate(as_vector([20.0, 0.0])) # project to the function space we actually want to use # this step is purely because it is not yet possible to interpolate to the
def get_latlon_mesh(mesh): coords_orig = mesh.coordinates coords_fs = coords_orig.function_space() if coords_fs.extruded: cell = mesh._base_mesh.ufl_cell().cellname() DG1_hori_elt = FiniteElement("DG", cell, 1, variant="equispaced") DG1_vert_elt = FiniteElement("DG", interval, 1, variant="equispaced") DG1_elt = TensorProductElement(DG1_hori_elt, DG1_vert_elt) else: cell = mesh.ufl_cell().cellname() DG1_elt = FiniteElement("DG", cell, 1, variant="equispaced") vec_DG1 = VectorFunctionSpace(mesh, DG1_elt) coords_dg = Function(vec_DG1).interpolate(coords_orig) coords_latlon = Function(vec_DG1) shapes = {"nDOFs": vec_DG1.finat_element.space_dimension(), 'dim': 3} radius = np.min( np.sqrt(coords_dg.dat.data[:, 0]**2 + coords_dg.dat.data[:, 1]**2 + coords_dg.dat.data[:, 2]**2)) # lat-lon 'x' = atan2(y, x) coords_latlon.dat.data[:, 0] = np.arctan2(coords_dg.dat.data[:, 1], coords_dg.dat.data[:, 0]) # lat-lon 'y' = asin(z/sqrt(x^2 + y^2 + z^2)) coords_latlon.dat.data[:, 1] = np.arcsin( coords_dg.dat.data[:, 2] / np.sqrt(coords_dg.dat.data[:, 0]**2 + coords_dg.dat.data[:, 1]**2 + coords_dg.dat.data[:, 2]**2)) # our vertical coordinate is radius - the minimum radius coords_latlon.dat.data[:, 2] = np.sqrt(coords_dg.dat.data[:, 0]**2 + coords_dg.dat.data[:, 1]**2 + coords_dg.dat.data[:, 2]**2) - radius # We need to ensure that all points in a cell are on the same side of the branch cut in longitude coords # This kernel amends the longitude coords so that all longitudes in one cell are close together kernel = op2.Kernel( """ #define PI 3.141592653589793 #define TWO_PI 6.283185307179586 void splat_coords(double *coords) {{ double max_diff = 0.0; double diff = 0.0; for (int i=0; i<{nDOFs}; i++) {{ for (int j=0; j<{nDOFs}; j++) {{ diff = coords[i*{dim}] - coords[j*{dim}]; if (fabs(diff) > max_diff) {{ max_diff = diff; }} }} }} if (max_diff > PI) {{ for (int i=0; i<{nDOFs}; i++) {{ if (coords[i*{dim}] < 0) {{ coords[i*{dim}] += TWO_PI; }} }} }} }} """.format(**shapes), "splat_coords") op2.par_loop(kernel, coords_latlon.cell_set, coords_latlon.dat(op2.RW, coords_latlon.cell_node_map())) return Mesh(coords_latlon)
def setup(self, equation, uadv=None, apply_bcs=True, *active_labels): self.residual = equation.residual if self.field_name is not None: self.idx = equation.field_names.index(self.field_name) self.fs = self.state.fields(self.field_name).function_space() self.residual = self.residual.label_map( lambda t: t.get(prognostic) == self.field_name, lambda t: Term( split_form(t.form)[self.idx].form, t.labels), drop) bcs = equation.bcs[self.field_name] else: self.field_name = equation.field_name self.fs = equation.function_space self.idx = None if type(self.fs.ufl_element()) is MixedElement: bcs = [bc for _, bcs in equation.bcs.items() for bc in bcs] else: bcs = equation.bcs[self.field_name] if len(active_labels) > 0: self.residual = self.residual.label_map( lambda t: any(t.has_label(time_derivative, *active_labels)), map_if_false=drop) options = self.options # -------------------------------------------------------------------- # # Routines relating to transport # -------------------------------------------------------------------- # if hasattr(self.options, 'ibp'): self.replace_transport_term() self.replace_transporting_velocity(uadv) # -------------------------------------------------------------------- # # Wrappers for embedded / recovery methods # -------------------------------------------------------------------- # if self.discretisation_option in ["embedded_dg", "recovered"]: # construct the embedding space if not specified if options.embedding_space is None: V_elt = BrokenElement(self.fs.ufl_element()) self.fs = FunctionSpace(self.state.mesh, V_elt) else: self.fs = options.embedding_space self.xdg_in = Function(self.fs) self.xdg_out = Function(self.fs) if self.idx is None: self.x_projected = Function(equation.function_space) else: self.x_projected = Function(self.state.fields(self.field_name).function_space()) new_test = TestFunction(self.fs) parameters = {'ksp_type': 'cg', 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'} # -------------------------------------------------------------------- # # Make boundary conditions # -------------------------------------------------------------------- # if not apply_bcs: self.bcs = None elif self.discretisation_option in ["embedded_dg", "recovered"]: # Transfer boundary conditions onto test function space self.bcs = [DirichletBC(self.fs, bc.function_arg, bc.sub_domain) for bc in bcs] else: self.bcs = bcs # -------------------------------------------------------------------- # # Modify test function for SUPG methods # -------------------------------------------------------------------- # if self.discretisation_option == "supg": # construct tau, if it is not specified dim = self.state.mesh.topological_dimension() if options.tau is not None: # if tau is provided, check that is has the right size tau = options.tau assert as_ufl(tau).ufl_shape == (dim, dim), "Provided tau has incorrect shape!" else: # create tuple of default values of size dim default_vals = [options.default*self.dt]*dim # check for directions is which the space is discontinuous # so that we don't apply supg in that direction if is_cg(self.fs): vals = default_vals else: space = self.fs.ufl_element().sobolev_space() if space.name in ["HDiv", "DirectionalH"]: vals = [default_vals[i] if space[i].name == "H1" else 0. for i in range(dim)] else: raise ValueError("I don't know what to do with space %s" % space) tau = Constant(tuple([ tuple( [vals[j] if i == j else 0. for i, v in enumerate(vals)] ) for j in range(dim)]) ) self.solver_parameters = {'ksp_type': 'gmres', 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'} test = TestFunction(self.fs) new_test = test + dot(dot(uadv, tau), grad(test)) if self.discretisation_option is not None: # replace the original test function with one defined on # the embedding space, as this is the space where the # the problem will be solved self.residual = self.residual.label_map( all_terms, map_if_true=replace_test_function(new_test)) if self.discretisation_option == "embedded_dg": if self.limiter is None: self.x_out_projector = Projector(self.xdg_out, self.x_projected, solver_parameters=parameters) else: self.x_out_projector = Recoverer(self.xdg_out, self.x_projected) if self.discretisation_option == "recovered": # set up the necessary functions self.x_in = Function(self.state.fields(self.field_name).function_space()) x_rec = Function(options.recovered_space) x_brok = Function(options.broken_space) # set up interpolators and projectors self.x_rec_projector = Recoverer(self.x_in, x_rec, VDG=self.fs, boundary_method=options.boundary_method) # recovered function self.x_brok_projector = Projector(x_rec, x_brok) # function projected back self.xdg_interpolator = Interpolator(self.x_in + x_rec - x_brok, self.xdg_in) if self.limiter is not None: self.x_brok_interpolator = Interpolator(self.xdg_out, x_brok) self.x_out_projector = Recoverer(x_brok, self.x_projected) else: self.x_out_projector = Projector(self.xdg_out, self.x_projected) # setup required functions self.dq = Function(self.fs) self.q1 = Function(self.fs)
def __init__(self, state, V): super(EulerPoincareForm, self).__init__(state) dt = state.timestepping.dt w = TestFunction(V) u = TrialFunction(V) self.u0 = Function(V) ustar = 0.5*(self.u0 + u) n = FacetNormal(state.mesh) Upwind = 0.5*(sign(dot(self.ubar, n))+1) if state.mesh.geometric_dimension() == 3: if V.extruded: surface_measure = (dS_h + dS_v) else: surface_measure = dS # <w,curl(u) cross ubar + grad( u.ubar)> # =<curl(u),ubar cross w> - <div(w), u.ubar> # =<u,curl(ubar cross w)> - # <<u_upwind, [[n cross(ubar cross w)cross]]>> both = lambda u: 2*avg(u) Eqn = ( inner(w, u-self.u0)*dx + dt*inner(ustar, curl(cross(self.ubar, w)))*dx - dt*inner(both(Upwind*ustar), both(cross(n, cross(self.ubar, w))))*surface_measure - dt*div(w)*inner(ustar, self.ubar)*dx ) # define surface measure and terms involving perp differently # for slice (i.e. if V.extruded is True) and shallow water # (V.extruded is False) else: if V.extruded: surface_measure = (dS_h + dS_v) perp = lambda u: as_vector([-u[1], u[0]]) perp_u_upwind = Upwind('+')*perp(ustar('+')) + Upwind('-')*perp(ustar('-')) else: surface_measure = dS outward_normals = CellNormal(state.mesh) perp = lambda u: cross(outward_normals, u) perp_u_upwind = Upwind('+')*cross(outward_normals('+'),ustar('+')) + Upwind('-')*cross(outward_normals('-'),ustar('-')) Eqn = ( (inner(w, u-self.u0) - dt*inner(w, div(perp(ustar))*perp(self.ubar)) - dt*div(w)*inner(ustar, self.ubar))*dx - dt*inner(jump(inner(w, perp(self.ubar)), n), perp_u_upwind)*surface_measure + dt*jump(inner(w, perp(self.ubar))*perp(ustar), n)*surface_measure ) a = lhs(Eqn) L = rhs(Eqn) self.u1 = Function(V) uproblem = LinearVariationalProblem(a, L, self.u1) self.usolver = LinearVariationalSolver(uproblem, options_prefix='EPAdvection')
class RK4(ExplicitTimeDiscretisation): """ Class to implement the 4-stage Runge-Kutta timestepping method: k1 = f(y_n) k2 = f(y_n + 1/2*dt*k1) k3 = f(y_n + 1/2*dt*k2) k4 = f(y_n + dt*k3) y_(n+1) = y_n + (1/6) * dt * (k1 + 2*k2 + 2*k3 + k4) where subscripts indicate the timelevel, superscripts indicate the stage number and f is the RHS. """ def setup(self, equation, uadv, *active_labels): super().setup(equation, uadv, *active_labels) self.k1 = Function(self.fs) self.k2 = Function(self.fs) self.k3 = Function(self.fs) self.k4 = Function(self.fs) @cached_property def lhs(self): l = self.residual.label_map( lambda t: t.has_label(time_derivative), map_if_true=replace_subject(self.dq, self.idx), map_if_false=drop) return l.form @cached_property def rhs(self): r = self.residual.label_map( all_terms, map_if_true=replace_subject(self.q1, self.idx)) r = r.label_map( lambda t: t.has_label(time_derivative), map_if_true=drop, map_if_false=lambda t: -1*t) return r.form def solve_stage(self, x_in, stage): if stage == 0: self.solver.solve() self.k1.assign(self.dq) self.q1.assign(x_in + 0.5 * self.dt * self.k1) elif stage == 1: self.solver.solve() self.k2.assign(self.dq) self.q1.assign(x_in + 0.5 * self.dt * self.k2) elif stage == 2: self.solver.solve() self.k3.assign(self.dq) self.q1.assign(x_in + self.dt * self.k3) elif stage == 3: self.solver.solve() self.k4.assign(self.dq) self.q1.assign(x_in + 1/6 * self.dt * (self.k1 + 2*self.k2 + 2*self.k3 + self.k4)) def apply_cycle(self, x_in, x_out): if self.limiter is not None: self.limiter.apply(x_in) self.q1.assign(x_in) for i in range(4): self.solve_stage(x_in, i) x_out.assign(self.q1)
class TimeDiscretisation(object, metaclass=ABCMeta): """ Base class for time discretisation schemes. :arg state: :class:`.State` object. :arg field: field to be evolved :arg equation: :class:`.Equation` object, specifying the equation that field satisfies :arg solver_parameters: solver_parameters :arg limiter: :class:`.Limiter` object. :arg options: :class:`.DiscretisationOptions` object """ def __init__(self, state, field_name=None, solver_parameters=None, limiter=None, options=None): self.state = state self.field_name = field_name self.dt = self.state.dt self.limiter = limiter self.options = options if options is not None: self.discretisation_option = options.name else: self.discretisation_option = None # get default solver options if none passed in if solver_parameters is None: self.solver_parameters = {'ksp_type': 'cg', 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'} else: self.solver_parameters = solver_parameters if logger.isEnabledFor(DEBUG): self.solver_parameters["ksp_monitor_true_residual"] = None def setup(self, equation, uadv=None, apply_bcs=True, *active_labels): self.residual = equation.residual if self.field_name is not None: self.idx = equation.field_names.index(self.field_name) self.fs = self.state.fields(self.field_name).function_space() self.residual = self.residual.label_map( lambda t: t.get(prognostic) == self.field_name, lambda t: Term( split_form(t.form)[self.idx].form, t.labels), drop) bcs = equation.bcs[self.field_name] else: self.field_name = equation.field_name self.fs = equation.function_space self.idx = None if type(self.fs.ufl_element()) is MixedElement: bcs = [bc for _, bcs in equation.bcs.items() for bc in bcs] else: bcs = equation.bcs[self.field_name] if len(active_labels) > 0: self.residual = self.residual.label_map( lambda t: any(t.has_label(time_derivative, *active_labels)), map_if_false=drop) options = self.options # -------------------------------------------------------------------- # # Routines relating to transport # -------------------------------------------------------------------- # if hasattr(self.options, 'ibp'): self.replace_transport_term() self.replace_transporting_velocity(uadv) # -------------------------------------------------------------------- # # Wrappers for embedded / recovery methods # -------------------------------------------------------------------- # if self.discretisation_option in ["embedded_dg", "recovered"]: # construct the embedding space if not specified if options.embedding_space is None: V_elt = BrokenElement(self.fs.ufl_element()) self.fs = FunctionSpace(self.state.mesh, V_elt) else: self.fs = options.embedding_space self.xdg_in = Function(self.fs) self.xdg_out = Function(self.fs) if self.idx is None: self.x_projected = Function(equation.function_space) else: self.x_projected = Function(self.state.fields(self.field_name).function_space()) new_test = TestFunction(self.fs) parameters = {'ksp_type': 'cg', 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'} # -------------------------------------------------------------------- # # Make boundary conditions # -------------------------------------------------------------------- # if not apply_bcs: self.bcs = None elif self.discretisation_option in ["embedded_dg", "recovered"]: # Transfer boundary conditions onto test function space self.bcs = [DirichletBC(self.fs, bc.function_arg, bc.sub_domain) for bc in bcs] else: self.bcs = bcs # -------------------------------------------------------------------- # # Modify test function for SUPG methods # -------------------------------------------------------------------- # if self.discretisation_option == "supg": # construct tau, if it is not specified dim = self.state.mesh.topological_dimension() if options.tau is not None: # if tau is provided, check that is has the right size tau = options.tau assert as_ufl(tau).ufl_shape == (dim, dim), "Provided tau has incorrect shape!" else: # create tuple of default values of size dim default_vals = [options.default*self.dt]*dim # check for directions is which the space is discontinuous # so that we don't apply supg in that direction if is_cg(self.fs): vals = default_vals else: space = self.fs.ufl_element().sobolev_space() if space.name in ["HDiv", "DirectionalH"]: vals = [default_vals[i] if space[i].name == "H1" else 0. for i in range(dim)] else: raise ValueError("I don't know what to do with space %s" % space) tau = Constant(tuple([ tuple( [vals[j] if i == j else 0. for i, v in enumerate(vals)] ) for j in range(dim)]) ) self.solver_parameters = {'ksp_type': 'gmres', 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'} test = TestFunction(self.fs) new_test = test + dot(dot(uadv, tau), grad(test)) if self.discretisation_option is not None: # replace the original test function with one defined on # the embedding space, as this is the space where the # the problem will be solved self.residual = self.residual.label_map( all_terms, map_if_true=replace_test_function(new_test)) if self.discretisation_option == "embedded_dg": if self.limiter is None: self.x_out_projector = Projector(self.xdg_out, self.x_projected, solver_parameters=parameters) else: self.x_out_projector = Recoverer(self.xdg_out, self.x_projected) if self.discretisation_option == "recovered": # set up the necessary functions self.x_in = Function(self.state.fields(self.field_name).function_space()) x_rec = Function(options.recovered_space) x_brok = Function(options.broken_space) # set up interpolators and projectors self.x_rec_projector = Recoverer(self.x_in, x_rec, VDG=self.fs, boundary_method=options.boundary_method) # recovered function self.x_brok_projector = Projector(x_rec, x_brok) # function projected back self.xdg_interpolator = Interpolator(self.x_in + x_rec - x_brok, self.xdg_in) if self.limiter is not None: self.x_brok_interpolator = Interpolator(self.xdg_out, x_brok) self.x_out_projector = Recoverer(x_brok, self.x_projected) else: self.x_out_projector = Projector(self.xdg_out, self.x_projected) # setup required functions self.dq = Function(self.fs) self.q1 = Function(self.fs) def pre_apply(self, x_in, discretisation_option): """ Extra steps to discretisation if using an embedded method, which might be either the plain embedded method or the recovered space scheme. :arg x_in: the input set of prognostic fields. :arg discretisation option: string specifying which scheme to use. """ if discretisation_option == "embedded_dg": try: self.xdg_in.interpolate(x_in) except NotImplementedError: self.xdg_in.project(x_in) elif discretisation_option == "recovered": self.x_in.assign(x_in) self.x_rec_projector.project() self.x_brok_projector.project() self.xdg_interpolator.interpolate() def post_apply(self, x_out, discretisation_option): """ The projection steps, returning a field to its original space for an embedded DG scheme. For the case of the recovered scheme, there are two options dependent on whether the scheme is limited or not. :arg x_out: the outgoing field. :arg discretisation_option: string specifying which option to use. """ if discretisation_option == "recovered" and self.limiter is not None: self.x_brok_interpolator.interpolate() self.x_out_projector.project() x_out.assign(self.x_projected) @abstractproperty def lhs(self): l = self.residual.label_map( lambda t: t.has_label(time_derivative), map_if_true=replace_subject(self.dq, self.idx), map_if_false=drop) return l.form @abstractproperty def rhs(self): r = self.residual.label_map( all_terms, map_if_true=replace_subject(self.q1, self.idx)) r = r.label_map( lambda t: t.has_label(time_derivative), map_if_false=lambda t: -self.dt*t) return r.form def replace_transport_term(self): """ This routine allows the default transport term to be replaced with a different one, specified through the transport options. This is necessary because when the prognostic equations are declared, the whole transport """ # Extract transport term of equation old_transport_term_list = self.residual.label_map( lambda t: t.has_label(transport), map_if_false=drop) # If there are more transport terms, extract only the one for this variable if len(old_transport_term_list.terms) > 1: raise NotImplementedError('Cannot replace transport terms when there are more than one') # Then we should only have one transport term old_transport_term = old_transport_term_list.terms[0] # If the transport term has an ibp label, then it could be replaced if old_transport_term.has_label(ibp_label) and hasattr(self.options, 'ibp'): # Do the options specify a different ibp to the old transport term? if old_transport_term.labels['ibp'] != self.options.ibp: # Set up a new transport term field = self.state.fields(self.field_name) test = TestFunction(self.fs) # Set up new transport term (depending on the type of transport equation) if old_transport_term.labels['transport'] == TransportEquationType.advective: new_transport_term = advection_form(self.state, test, field, ibp=self.options.ibp) elif old_transport_term.labels['transport'] == TransportEquationType.conservative: new_transport_term = continuity_form(self.state, test, field, ibp=self.options.ibp) else: raise NotImplementedError(f'Replacement of transport term not implemented yet for {old_transport_term.labels["transport"]}') # Finally, drop the old transport term and add the new one self.residual = self.residual.label_map( lambda t: t.has_label(transport), map_if_true=drop) self.residual += subject(new_transport_term, field) def replace_transporting_velocity(self, uadv): # replace the transporting velocity in any terms that contain it if any([t.has_label(transporting_velocity) for t in self.residual]): assert uadv is not None if uadv == "prognostic": self.residual = self.residual.label_map( lambda t: t.has_label(transporting_velocity), map_if_true=lambda t: Term(ufl.replace( t.form, {t.get(transporting_velocity): split(t.get(subject))[0]}), t.labels) ) else: self.residual = self.residual.label_map( lambda t: t.has_label(transporting_velocity), map_if_true=lambda t: Term(ufl.replace( t.form, {t.get(transporting_velocity): uadv}), t.labels) ) self.residual = transporting_velocity.update_value(self.residual, uadv) @cached_property def solver(self): # setup solver using lhs and rhs defined in derived class problem = NonlinearVariationalProblem(self.lhs-self.rhs, self.dq, bcs=self.bcs) solver_name = self.field_name+self.__class__.__name__ return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name) @abstractmethod def apply(self, x_in, x_out): """ Function takes x as input, computes L(x) as defined by the equation, and returns x_out as output. :arg x: :class:`.Function` object, the input Function. :arg x_out: :class:`.Function` object, the output Function. """ pass
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. shapes = { "i": Vv.finat_element.space_dimension(), "j": np.prod(Vv.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(Vv) 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) # 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() self.S.force_evaluation() 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 VerticalHybridizationPC(PCBase): """ A Slate-based python preconditioner for solving the hydrostatic pressure equation (after rewriting as a mixed vertical HDiv x L2 system). This preconditioner hybridizes a mixed system in the vertical direction. This means that velocities are rendered discontinuous in the vertical and Lagrange multipliers are introduced on the top/bottom facets to weakly enforce continuity through the top/bottom faces of each cell. This PC assembles a statically condensed problem for the multipliers and inverts the resulting system using the provided solver options. The original unknowns are recovered element-wise by solving local linear systems. All elimination and recovery kernels are generated using the Slate DSL in Firedrake. """ @timed_function("VertHybridInit") 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. shapes = { "i": Vv.finat_element.space_dimension(), "j": np.prod(Vv.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(Vv) 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) # 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() self.S.force_evaluation() 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) def _reconstruction_calls(self, split_mixed_op, split_trace_op): """This generates the reconstruction calls for the unknowns using the Lagrange multipliers. :arg split_mixed_op: a ``dict`` of split forms that make up the broken mixed operator from the original problem. :arg split_trace_op: a ``dict`` of split forms that make up the trace contribution in the hybridized mixed system. """ from firedrake.assemble import create_assembly_callable # We always eliminate the velocity block first id0, id1 = (self.vidx, self.pidx) # TODO: When PyOP2 is able to write into mixed dats, # the reconstruction expressions can simplify into # one clean expression. A = Tensor(split_mixed_op[(id0, id0)]) B = Tensor(split_mixed_op[(id0, id1)]) C = Tensor(split_mixed_op[(id1, id0)]) D = Tensor(split_mixed_op[(id1, id1)]) K_0 = Tensor(split_trace_op[(0, id0)]) K_1 = Tensor(split_trace_op[(0, id1)]) # Split functions and reconstruct each bit separately split_residual = self.broken_residual.split() split_sol = self.broken_solution.split() g = AssembledVector(split_residual[id0]) f = AssembledVector(split_residual[id1]) sigma = split_sol[id0] u = split_sol[id1] lambdar = AssembledVector(self.trace_solution) M = D - C * A.inv * B R = K_1.T - C * A.inv * K_0.T u_rec = M.solve(f - C * A.inv * g - R * lambdar, decomposition="PartialPivLU") self._sub_unknown = create_assembly_callable( u_rec, tensor=u, form_compiler_parameters=self.ctx.fc_params) sigma_rec = A.solve(g - B * AssembledVector(u) - K_0.T * lambdar, decomposition="PartialPivLU") self._elim_unknown = create_assembly_callable( sigma_rec, tensor=sigma, form_compiler_parameters=self.ctx.fc_params) @timed_function("VertHybridRecon") def _reconstruct(self): """Reconstructs the system unknowns using the multipliers. Note that the reconstruction calls are assumed to be initialized at this point. """ # We assemble the unknown which is an expression # of the first eliminated variable. self._sub_unknown() # Recover the eliminated unknown self._elim_unknown() @timed_function("VertHybridUpdate") def update(self, pc): """Update by assembling into the operator. No need to reconstruct symbolic objects. """ self._assemble_S() self.S.force_evaluation() def apply(self, pc, x, y): """We solve the forward eliminated problem for the approximate traces of the scalar solution (the multipliers) and reconstruct the "broken flux and scalar variable." Lastly, we project the broken solutions into the mimetic non-broken finite element space. """ with timed_region("VertHybridBreak"): with self.unbroken_residual.dat.vec_wo as v: x.copy(v) # Transfer unbroken_rhs into broken_rhs # NOTE: Scalar space is already "broken" so no need for # any projections unbroken_scalar_data = self.unbroken_residual.split()[self.pidx] broken_scalar_data = self.broken_residual.split()[self.pidx] unbroken_scalar_data.dat.copy(broken_scalar_data.dat) with timed_region("VertHybridRHS"): # Assemble the new "broken" hdiv residual # We need a residual R' in the broken space that # gives R'[w] = R[w] when w is in the unbroken space. # We do this by splitting the residual equally between # basis functions that add together to give unbroken # basis functions. unbroken_res_hdiv = self.unbroken_residual.split()[self.vidx] broken_res_hdiv = self.broken_residual.split()[self.vidx] broken_res_hdiv.assign(0) par_loop( self.average_kernel, dx, { "w": (self.weight, READ), "vec_in": (unbroken_res_hdiv, READ), "vec_out": (broken_res_hdiv, INC) }) # Compute the rhs for the multiplier system self._assemble_Srhs() with timed_region("VertHybridSolve"): # Solve the system for the Lagrange multipliers with self.schur_rhs.dat.vec_ro as b: if self.trace_ksp.getInitialGuessNonzero(): acc = self.trace_solution.dat.vec else: acc = self.trace_solution.dat.vec_wo with acc as x_trace: self.trace_ksp.solve(b, x_trace) # Reconstruct the unknowns self._reconstruct() with timed_region("VertHybridRecover"): # Project the broken solution into non-broken spaces broken_pressure = self.broken_solution.split()[self.pidx] unbroken_pressure = self.unbroken_solution.split()[self.pidx] broken_pressure.dat.copy(unbroken_pressure.dat) # Compute the hdiv projection of the broken hdiv solution broken_hdiv = self.broken_solution.split()[self.vidx] unbroken_hdiv = self.unbroken_solution.split()[self.vidx] unbroken_hdiv.assign(0) par_loop( self.average_kernel, dx, { "w": (self.weight, READ), "vec_in": (broken_hdiv, READ), "vec_out": (unbroken_hdiv, INC) }) with self.unbroken_solution.dat.vec_ro as v: v.copy(y) def applyTranspose(self, pc, x, y): """Apply the transpose of the preconditioner.""" raise NotImplementedError( "The transpose application of the PC is not implemented.") def view(self, pc, viewer=None): """Viewer calls for the various configurable objects in this PC.""" super(VerticalHybridizationPC, self).view(pc, viewer) viewer.pushASCIITab() viewer.printfASCII("Solves K * P^-1 * K.T using local eliminations.\n") viewer.printfASCII("KSP solver for the multipliers:\n") viewer.pushASCIITab() self.trace_ksp.view(viewer) viewer.popASCIITab() viewer.printfASCII( "Locally reconstructing the broken solutions from the multipliers.\n" ) viewer.pushASCIITab() viewer.printfASCII( "Project the broken hdiv solution into the HDiv space.\n") viewer.popASCIITab()
Vu = u0.function_space() Vt = theta0.function_space() Vr = rho0.function_space() # first setup the background buoyancy profile # z.grad(bref) = N**2 # the following is symbolic algebra, using the default buoyancy frequency # from the parameters class. x, y, z = SpatialCoordinate(mesh) g = parameters.g Nsq = parameters.Nsq theta_surf = parameters.theta_surf # N^2 = (g/theta)dtheta/dz => dtheta/dz = theta N^2g => theta=theta_0exp(N^2gz) theta_ref = theta_surf * exp(Nsq * (z - H / 2) / g) theta_b = Function(Vt).interpolate(theta_ref) # set theta_pert def coth(x): return cosh(x) / sinh(x) def Z(z): return Bu * ((z / H) - 0.5) def n(): return Bu**(-1) * sqrt( (Bu * 0.5 - tanh(Bu * 0.5)) * (coth(Bu * 0.5) - Bu * 0.5))
Test run: %s.\n """ % (dt, deltax, deltaz, bool(args.test))) PETSc.Sys.Print("Initializing problem with dt: %s and tmax: %s.\n" % (dt, tmax)) PETSc.Sys.Print("Creating mesh with %s columns and %s layers...\n" % (columns, nlayers)) m = PeriodicIntervalMesh(columns, L) # build volume mesh ext_mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers) Vc = VectorFunctionSpace(ext_mesh, "DG", 2) coord = SpatialCoordinate(ext_mesh) x = Function(Vc).interpolate(as_vector([coord[0], coord[1]])) a = 10000. xc = L / 2. x, z = SpatialCoordinate(ext_mesh) hm = 1. zs = hm * a**2 / ((x - xc)**2 + a**2) smooth_z = True dirname = 'h_mountain' if smooth_z: dirname += '_smootherz' zh = 5000. xexpr = as_vector( [x, conditional(z < zh, z + cos(0.5 * pi * z / zh)**6 * zs, z)]) else: xexpr = as_vector([x, z + ((H - z) / H) * zs])
def __init__(self, state, iterations=1): super().__init__(state) self.iterations = iterations # obtain our fields self.theta = state.fields('theta') self.water_v = state.fields('water_v') self.water_c = state.fields('water_c') rho = state.fields('rho') try: rain = state.fields('rain') water_l = self.water_c + rain except NotImplementedError: water_l = self.water_c # declare function space Vt = self.theta.function_space() # make rho variables # we recover rho into theta space if state.vertical_degree == 0 and state.horizontal_degree == 0: boundary_method = Boundary_Method.physics else: boundary_method = None Vt_broken = FunctionSpace(state.mesh, BrokenElement(Vt.ufl_element())) rho_averaged = Function(Vt) self.rho_recoverer = Recoverer(rho, rho_averaged, VDG=Vt_broken, boundary_method=boundary_method) # define some parameters as attributes dt = state.timestepping.dt R_d = state.parameters.R_d cp = state.parameters.cp cv = state.parameters.cv c_pv = state.parameters.c_pv c_pl = state.parameters.c_pl c_vv = state.parameters.c_vv R_v = state.parameters.R_v # make useful fields Pi = thermodynamics.pi(state.parameters, rho_averaged, self.theta) T = thermodynamics.T(state.parameters, self.theta, Pi, r_v=self.water_v) p = thermodynamics.p(state.parameters, Pi) L_v = thermodynamics.Lv(state.parameters, T) R_m = R_d + R_v * self.water_v c_pml = cp + c_pv * self.water_v + c_pl * water_l c_vml = cv + c_vv * self.water_v + c_pl * water_l # use Teten's formula to calculate w_sat w_sat = thermodynamics.r_sat(state.parameters, T, p) # make appropriate condensation rate dot_r_cond = ((self.water_v - w_sat) / (dt * (1.0 + ((L_v**2.0 * w_sat) / (cp * R_v * T**2.0))))) # make cond_rate function, that needs to be the same for all updates in one time step cond_rate = Function(Vt) # adjust cond rate so negative concentrations don't occur self.lim_cond_rate = Interpolator( conditional(dot_r_cond < 0, max_value(dot_r_cond, -self.water_c / dt), min_value(dot_r_cond, self.water_v / dt)), cond_rate) # tell the prognostic fields what to update to self.water_v_new = Interpolator(self.water_v - dt * cond_rate, Vt) self.water_c_new = Interpolator(self.water_c + dt * cond_rate, Vt) self.theta_new = Interpolator( self.theta * (1.0 + dt * cond_rate * (cv * L_v / (c_vml * cp * T) - R_v * cv * c_pml / (R_m * cp * c_vml))), Vt)
def nonlocal_integral_eq( mesh, scatterer_bdy_id, outer_bdy_id, wave_number, options_prefix=None, solver_parameters=None, fspace=None, vfspace=None, true_sol_grad=None, queue=None, fspace_analog=None, qbx_kwargs=None, ): r""" see run_method for descriptions of unlisted args args: :arg queue: A command queue for the computing context gamma and beta are used to precondition with the following equation: \Delta u - \kappa^2 \gamma u = 0 (\partial_n - i\kappa\beta) u |_\Sigma = 0 """ with_refinement = True # away from the excluded region, but firedrake and meshmode point # into pyt_inner_normal_sign = -1 ambient_dim = mesh.geometric_dimension() # {{{ Create operator from pytential import sym r""" ..math: x \in \Sigma grad_op(x) = \nabla( \int_\Gamma( u(y) \partial_n H_0^{(1)}(\kappa |x - y|) )d\gamma(y) ) """ grad_op = pyt_inner_normal_sign * sym.grad( ambient_dim, sym.D(HelmholtzKernel(ambient_dim), sym.var("u"), k=sym.var("k"), qbx_forced_limit=None)) r""" ..math: x \in \Sigma op(x) = i \kappa \cdot \int_\Gamma( u(y) \partial_n H_0^{(1)}(\kappa |x - y|) )d\gamma(y) """ op = pyt_inner_normal_sign * 1j * sym.var("k") * (sym.D( HelmholtzKernel(ambient_dim), sym.var("u"), k=sym.var("k"), qbx_forced_limit=None)) pyt_grad_op = fd2mm.fd_bind( queue.context, fspace_analog, grad_op, source=(fspace, scatterer_bdy_id), target=(vfspace, outer_bdy_id), with_refinement=with_refinement, qbx_kwargs=qbx_kwargs, ) pyt_op = fd2mm.fd_bind( queue.context, fspace_analog, op, source=(fspace, scatterer_bdy_id), target=(fspace, outer_bdy_id), with_refinement=with_refinement, qbx_kwargs=qbx_kwargs, ) # }}} class MatrixFreeB(object): def __init__(self, A, pyt_grad_op, pyt_op, queue, kappa): """ :arg kappa: The wave number """ self.queue = queue self.k = kappa self.pyt_op = pyt_op self.pyt_grad_op = pyt_grad_op self.A = A # {{{ Create some functions needed for multing self.x_fntn = Function(fspace) self.potential_int = Function(fspace) self.potential_int.dat.data[:] = 0.0 self.grad_potential_int = Function(vfspace) self.grad_potential_int.dat.data[:] = 0.0 self.pyt_result = Function(fspace) self.n = FacetNormal(mesh) self.v = TestFunction(fspace) # }}} def mult(self, mat, x, y): # Perform pytential operation self.x_fntn.dat.data[:] = x[:] self.pyt_op(self.queue, self.potential_int, u=self.x_fntn, k=self.k) self.pyt_grad_op(self.queue, self.grad_potential_int, u=self.x_fntn, k=self.k) # Integrate the potential r""" Compute the inner products using firedrake. Note this will be subtracted later, hence appears off by a sign. .. math:: \langle n(x) \cdot \nabla( \int_\Gamma( u(y) \partial_n H_0^{(1)}(\kappa |x - y|) )d\gamma(y) ), v \rangle_\Sigma - \langle i \kappa \cdot \int_\Gamma( u(y) \partial_n H_0^{(1)}(\kappa |x - y|) )d\gamma(y), v \rangle_\Sigma """ self.pyt_result = assemble( inner(inner(self.grad_potential_int, self.n), self.v) * ds(outer_bdy_id) - inner(self.potential_int, self.v) * ds(outer_bdy_id)) # y <- Ax - evaluated potential self.A.mult(x, y) with self.pyt_result.dat.vec_ro as ep: y.axpy(-1, ep) # {{{ Compute normal helmholtz operator u = TrialFunction(fspace) v = TestFunction(fspace) r""" .. math:: \langle \nabla u, \nabla v \rangle - \kappa^2 \cdot \langle u, v \rangle - i \kappa \langle u, v \rangle_\Sigma """ 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) # get the concrete matrix from a general bilinear form A = assemble(a).M.handle # }}} # {{{ Setup Python matrix B = PETSc.Mat().create() # build matrix context Bctx = MatrixFreeB(A, pyt_grad_op, pyt_op, queue, wave_number) # set up B as same size as A B.setSizes(*A.getSizes()) B.setType(B.Type.PYTHON) B.setPythonContext(Bctx) B.setUp() # }}} # {{{ Create rhs # Remember f is \partial_n(true_sol)|_\Gamma # so we just need to compute \int_\Gamma\partial_n(true_sol) H(x-y) from pytential import sym sigma = sym.make_sym_vector("sigma", ambient_dim) r""" ..math: x \in \Sigma grad_op(x) = \nabla( \int_\Gamma( f(y) H_0^{(1)}(\kappa |x - y|) )d\gamma(y) ) """ grad_op = pyt_inner_normal_sign * \ sym.grad(ambient_dim, sym.S(HelmholtzKernel(ambient_dim), sym.n_dot(sigma), k=sym.var("k"), qbx_forced_limit=None)) r""" ..math: x \in \Sigma op(x) = i \kappa \cdot \int_\Gamma( f(y) H_0^{(1)}(\kappa |x - y|) )d\gamma(y) ) """ op = 1j * sym.var("k") * pyt_inner_normal_sign * \ sym.S(HelmholtzKernel(ambient_dim), sym.n_dot(sigma), k=sym.var("k"), qbx_forced_limit=None) rhs_grad_op = fd2mm.fd_bind( queue.context, fspace_analog, grad_op, source=(vfspace, scatterer_bdy_id), target=(vfspace, outer_bdy_id), with_refinement=with_refinement, qbx_kwargs=qbx_kwargs, ) rhs_op = fd2mm.fd_bind( queue.context, fspace_analog, op, source=(vfspace, scatterer_bdy_id), target=(fspace, outer_bdy_id), with_refinement=with_refinement, qbx_kwargs=qbx_kwargs, ) f_grad_convoluted = Function(vfspace) f_convoluted = Function(fspace) rhs_grad_op(queue, f_grad_convoluted, sigma=true_sol_grad, k=wave_number) rhs_op(queue, f_convoluted, sigma=true_sol_grad, k=wave_number) r""" \langle f, v \rangle_\Gamma + \langle i \kappa \cdot \int_\Gamma( f(y) H_0^{(1)}(\kappa |x - y|) )d\gamma(y), v \rangle_\Sigma - \langle n(x) \cdot \nabla( \int_\Gamma( f(y) H_0^{(1)}(\kappa |x - y|) )d\gamma(y) ), v \rangle_\Sigma """ rhs_form = inner(inner(true_sol_grad, FacetNormal(mesh)), v) * ds(scatterer_bdy_id) \ + inner(f_convoluted, v) * ds(outer_bdy_id) \ - inner(inner(f_grad_convoluted, FacetNormal(mesh)), v) * ds(outer_bdy_id) rhs = assemble(rhs_form) # {{{ set up a solver: solution = Function(fspace, name="Computed Solution") # {{{ Used for preconditioning if 'gamma' in solver_parameters or 'beta' in solver_parameters: gamma = complex(solver_parameters.pop('gamma', 1.0)) import cmath beta = complex(solver_parameters.pop('beta', cmath.sqrt(gamma))) p = 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) P = assemble(p).M.handle else: P = A # }}} # Set up options to contain solver parameters: ksp = PETSc.KSP().create() if solver_parameters['pc_type'] == 'pyamg': del solver_parameters['pc_type'] # We are using the AMG preconditioner pyamg_tol = solver_parameters.get('pyamg_tol', None) if pyamg_tol is not None: pyamg_tol = float(pyamg_tol) pyamg_maxiter = solver_parameters.get('pyamg_maxiter', None) if pyamg_maxiter is not None: pyamg_maxiter = int(pyamg_maxiter) ksp.setOperators(B) ksp.setUp() pc = ksp.pc pc.setType(pc.Type.PYTHON) pc.setPythonContext( AMGTransmissionPreconditioner(wave_number, fspace, A, tol=pyamg_tol, maxiter=pyamg_maxiter, use_plane_waves=True)) # Otherwise use regular preconditioner else: ksp.setOperators(B, P) options_manager = OptionsManager(solver_parameters, options_prefix) options_manager.set_from_options(ksp) with rhs.dat.vec_ro as b: with solution.dat.vec as x: ksp.solve(b, x) # }}} return ksp, solution
def setup_3d_recovery(dirname): L = 3. H = 3. W = 3. deltax = L / 3. deltay = W / 3. deltaz = H / 3. nlayers = int(H / deltaz) ncolumnsx = int(L / deltax) ncolumnsy = int(W / deltay) m = PeriodicRectangleMesh(ncolumnsx, ncolumnsy, L, W, direction='y', quadrilateral=True) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers) x, y, z = SpatialCoordinate(mesh) # horizontal base spaces cell = mesh._base_mesh.ufl_cell().cellname() u_hori = FiniteElement("RTCF", cell, 1) w_hori = FiniteElement("DG", cell, 0) # vertical base spaces u_vert = FiniteElement("DG", interval, 0) w_vert = FiniteElement("CG", interval, 1) # build elements u_element = HDiv(TensorProductElement(u_hori, u_vert)) w_element = HDiv(TensorProductElement(w_hori, w_vert)) theta_element = TensorProductElement(w_hori, w_vert) v_element = u_element + w_element # spaces VDG0 = FunctionSpace(mesh, "DG", 0) VCG1 = FunctionSpace(mesh, "CG", 1) VDG1 = FunctionSpace(mesh, "DG", 1) Vt = FunctionSpace(mesh, theta_element) Vt_brok = FunctionSpace(mesh, BrokenElement(theta_element)) Vu = FunctionSpace(mesh, v_element) VuCG1 = VectorFunctionSpace(mesh, "CG", 1) VuDG1 = VectorFunctionSpace(mesh, "DG", 1) # set up initial conditions np.random.seed(0) expr = np.random.randn() + np.random.randn() * x + np.random.randn( ) * z + np.random.randn() * x * z # our actual theta and rho and v rho_CG1_true = Function(VCG1).interpolate(expr) theta_CG1_true = Function(VCG1).interpolate(expr) v_CG1_true = Function(VuCG1).interpolate(as_vector([expr, expr, expr])) rho_Vt_true = Function(Vt).interpolate(expr) # make the initial fields by projecting expressions into the lowest order spaces rho_DG0 = Function(VDG0).interpolate(expr) rho_CG1 = Function(VCG1) theta_Vt = Function(Vt).interpolate(expr) theta_CG1 = Function(VCG1) v_Vu = Function(Vu).project(as_vector([expr, expr, expr])) v_CG1 = Function(VuCG1) rho_Vt = Function(Vt) # make the recoverers and do the recovery rho_recoverer = Recoverer(rho_DG0, rho_CG1, VDG=VDG1, boundary_method=Boundary_Method.dynamics) theta_recoverer = Recoverer(theta_Vt, theta_CG1, VDG=VDG1, boundary_method=Boundary_Method.dynamics) v_recoverer = Recoverer(v_Vu, v_CG1, VDG=VuDG1, boundary_method=Boundary_Method.dynamics) rho_Vt_recoverer = Recoverer(rho_DG0, rho_Vt, VDG=Vt_brok, boundary_method=Boundary_Method.physics) rho_recoverer.project() theta_recoverer.project() v_recoverer.project() rho_Vt_recoverer.project() rho_diff = errornorm(rho_CG1, rho_CG1_true) / norm(rho_CG1_true) theta_diff = errornorm(theta_CG1, theta_CG1_true) / norm(theta_CG1_true) v_diff = errornorm(v_CG1, v_CG1_true) / norm(v_CG1_true) rho_Vt_diff = errornorm(rho_Vt, rho_Vt_true) / norm(rho_Vt_true) return (rho_diff, theta_diff, v_diff, rho_Vt_diff)
def setup_dump(self, t, tmax, pickup=False): """ Setup dump files Check for existence of directory so as not to overwrite output files Setup checkpoint file :arg tmax: model stop time :arg pickup: recover state from the checkpointing file if true, otherwise dump and checkpoint to disk. (default is False). """ if any([ self.output.dump_vtus, self.output.dumplist_latlon, self.output.dump_diagnostics, self.output.point_data, self.output.checkpoint and not pickup ]): # setup output directory and check that it does not already exist self.dumpdir = path.join("results", self.output.dirname) running_tests = '--running-tests' in sys.argv or "pytest" in self.output.dirname if self.mesh.comm.rank == 0: if not running_tests and path.exists( self.dumpdir) and not pickup: raise IOError("results directory '%s' already exists" % self.dumpdir) else: if not running_tests: makedirs(self.dumpdir) if self.output.dump_vtus: # setup pvd output file outfile = path.join(self.dumpdir, "field_output.pvd") self.dumpfile = File(outfile, project_output=self.output.project_fields, comm=self.mesh.comm) # make list of fields to dump self.to_dump = [field for field in self.fields if field.dump] # make dump counter self.dumpcount = itertools.count() # if there are fields to be dumped in latlon coordinates, # setup the latlon coordinate mesh and make output file if len(self.output.dumplist_latlon) > 0: mesh_ll = get_latlon_mesh(self.mesh) outfile_ll = path.join(self.dumpdir, "field_output_latlon.pvd") self.dumpfile_ll = File(outfile_ll, project_output=self.output.project_fields, comm=self.mesh.comm) # make functions on latlon mesh, as specified by dumplist_latlon self.to_dump_latlon = [] for name in self.output.dumplist_latlon: f = self.fields(name) field = Function(functionspaceimpl.WithGeometry( f.function_space(), mesh_ll), val=f.topological, name=name + '_ll') self.to_dump_latlon.append(field) # we create new netcdf files to write to, unless pickup=True, in # which case we just need the filenames if self.output.dump_diagnostics: diagnostics_filename = self.dumpdir + "/diagnostics.nc" self.diagnostic_output = DiagnosticsOutput(diagnostics_filename, self.diagnostics, self.output.dirname, self.mesh.comm, create=not pickup) if len(self.output.point_data) > 0: pointdata_filename = self.dumpdir + "/point_data.nc" ndt = int(tmax / self.timestepping.dt) self.pointdata_output = PointDataOutput(pointdata_filename, ndt, self.output.point_data, self.output.dirname, self.fields, self.mesh.comm, create=not pickup) # if we want to checkpoint and are not picking up from a previous # checkpoint file, setup the dumb checkpointing if self.output.checkpoint and not pickup: self.chkpt = DumbCheckpoint(path.join(self.dumpdir, "chkpt"), mode=FILE_CREATE) # make list of fields to pickup (this doesn't include # diagnostic fields) self.to_pickup = [field for field in self.fields if field.pickup] # if we want to checkpoint then make a checkpoint counter if self.output.checkpoint: self.chkptcount = itertools.count() # dump initial fields self.dump(t)
class Advection(object, metaclass=ABCMeta): """ Base class for advection schemes. :arg state: :class:`.State` object. :arg field: field to be advected :arg equation: :class:`.Equation` object, specifying the equation that field satisfies :arg solver_parameters: solver_parameters :arg limiter: :class:`.Limiter` object. :arg options: :class:`.AdvectionOptions` object """ def __init__(self, state, field, equation=None, *, solver_parameters=None, limiter=None): if equation is not None: self.state = state self.field = field self.equation = equation # get ubar from the equation class self.ubar = self.equation.ubar self.dt = self.state.timestepping.dt # get default solver options if none passed in if solver_parameters is None: self.solver_parameters = equation.solver_parameters else: self.solver_parameters = solver_parameters if logger.isEnabledFor(DEBUG): self.solver_parameters["ksp_monitor_true_residual"] = True self.limiter = limiter if hasattr(equation, "options"): self.discretisation_option = equation.options.name self._setup(state, field, equation.options) else: self.discretisation_option = None self.fs = field.function_space() # setup required functions self.dq = Function(self.fs) self.q1 = Function(self.fs) def _setup(self, state, field, options): if options.name in ["embedded_dg", "recovered"]: self.fs = options.embedding_space self.xdg_in = Function(self.fs) self.xdg_out = Function(self.fs) self.x_projected = Function(field.function_space()) parameters = {'ksp_type': 'cg', 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'} self.Projector = Projector(self.xdg_out, self.x_projected, solver_parameters=parameters) if options.name == "recovered": # set up the necessary functions self.x_in = Function(field.function_space()) x_rec = Function(options.recovered_space) x_brok = Function(options.broken_space) # set up interpolators and projectors self.x_rec_projector = Recoverer(self.x_in, x_rec, VDG=self.fs, boundary_method=options.boundary_method) # recovered function self.x_brok_projector = Projector(x_rec, x_brok) # function projected back self.xdg_interpolator = Interpolator(self.x_in + x_rec - x_brok, self.xdg_in) if self.limiter is not None: self.x_brok_interpolator = Interpolator(self.xdg_out, x_brok) self.x_out_projector = Recoverer(x_brok, self.x_projected) def pre_apply(self, x_in, discretisation_option): """ Extra steps to advection if using an embedded method, which might be either the plain embedded method or the recovered space advection scheme. :arg x_in: the input set of prognostic fields. :arg discretisation option: string specifying which scheme to use. """ if discretisation_option == "embedded_dg": try: self.xdg_in.interpolate(x_in) except NotImplementedError: self.xdg_in.project(x_in) elif discretisation_option == "recovered": self.x_in.assign(x_in) self.x_rec_projector.project() self.x_brok_projector.project() self.xdg_interpolator.interpolate() def post_apply(self, x_out, discretisation_option): """ The projection steps, returning a field to its original space for an embedded DG advection scheme. For the case of the recovered scheme, there are two options dependent on whether the scheme is limited or not. :arg x_out: the outgoing field. :arg discretisation_option: string specifying which option to use. """ if discretisation_option == "embedded_dg": self.Projector.project() elif discretisation_option == "recovered": if self.limiter is not None: self.x_brok_interpolator.interpolate() self.x_out_projector.project() else: self.Projector.project() x_out.assign(self.x_projected) @abstractproperty def lhs(self): return self.equation.mass_term(self.equation.trial) @abstractproperty def rhs(self): return self.equation.mass_term(self.q1) - self.dt*self.equation.advection_term(self.q1) def update_ubar(self, xn, xnp1, alpha): un = xn.split()[0] unp1 = xnp1.split()[0] self.ubar.assign(un + alpha*(unp1-un)) @cached_property 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) @abstractmethod def apply(self, x_in, x_out): """ Function takes x as input, computes L(x) as defined by the equation, and returns x_out as output. :arg x: :class:`.Function` object, the input Function. :arg x_out: :class:`.Function` object, the output Function. """ pass
def __init__(self, state): super().__init__(state) # obtain our fields self.theta = state.fields('theta') self.water_v = state.fields('water_v') self.rain = state.fields('rain') rho = state.fields('rho') try: water_c = state.fields('water_c') water_l = self.rain + water_c except NotImplementedError: water_l = self.rain # declare function space Vt = self.theta.function_space() # make rho variables # we recover rho into theta space if state.vertical_degree == 0 and state.horizontal_degree == 0: boundary_method = Boundary_Method.physics else: boundary_method = None Vt_broken = FunctionSpace(state.mesh, BrokenElement(Vt.ufl_element())) rho_averaged = Function(Vt) self.rho_recoverer = Recoverer(rho, rho_averaged, VDG=Vt_broken, boundary_method=boundary_method) # define some parameters as attributes dt = state.timestepping.dt R_d = state.parameters.R_d cp = state.parameters.cp cv = state.parameters.cv c_pv = state.parameters.c_pv c_pl = state.parameters.c_pl c_vv = state.parameters.c_vv R_v = state.parameters.R_v # make useful fields Pi = thermodynamics.pi(state.parameters, rho_averaged, self.theta) T = thermodynamics.T(state.parameters, self.theta, Pi, r_v=self.water_v) p = thermodynamics.p(state.parameters, Pi) L_v = thermodynamics.Lv(state.parameters, T) R_m = R_d + R_v * self.water_v c_pml = cp + c_pv * self.water_v + c_pl * water_l c_vml = cv + c_vv * self.water_v + c_pl * water_l # use Teten's formula to calculate w_sat w_sat = thermodynamics.r_sat(state.parameters, T, p) # expression for ventilation factor a = Constant(1.6) b = Constant(124.9) c = Constant(0.2046) C = a + b * (rho_averaged * self.rain)**c # make appropriate condensation rate f = Constant(5.4e5) g = Constant(2.55e6) h = Constant(0.525) dot_r_evap = (((1 - self.water_v / w_sat) * C * (rho_averaged * self.rain)**h) / (rho_averaged * (f + g / (p * w_sat)))) # make evap_rate function, needs to be the same for all updates in one time step evap_rate = Function(Vt) # adjust evap rate so negative rain doesn't occur self.lim_evap_rate = Interpolator( conditional( dot_r_evap < 0, 0.0, conditional(self.rain < 0.0, 0.0, min_value(dot_r_evap, self.rain / dt))), evap_rate) # tell the prognostic fields what to update to self.water_v_new = Interpolator(self.water_v + dt * evap_rate, Vt) self.rain_new = Interpolator(self.rain - dt * evap_rate, Vt) self.theta_new = Interpolator( self.theta * (1.0 - dt * evap_rate * (cv * L_v / (c_vml * cp * T) - R_v * cv * c_pml / (R_m * cp * c_vml))), Vt)
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 setup_vert_limiters(dirname): # declare grid shape L = 200. H = 800. ncolumns = int(L / 20.) nlayers = int(H / 20.) # make mesh m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers)) x, z = SpatialCoordinate(mesh) fieldlist = ['u'] timestepping = TimesteppingParameters(dt=1.0, maxk=4, maxi=1) output = OutputParameters(dirname=dirname + "/limiting_vert", dumpfreq=5, dumplist=['u'], perturbation_fields=['theta0', 'theta1']) parameters = CompressibleParameters() state = State(mesh, vertical_degree=1, horizontal_degree=1, family="CG", timestepping=timestepping, output=output, parameters=parameters, fieldlist=fieldlist) # make elements # v is continuous in vertical, h is horizontal cell = mesh._base_mesh.ufl_cell().cellname() DG0_element = FiniteElement("DG", cell, 0) CG1_element = FiniteElement("CG", cell, 1) DG1_element = FiniteElement("DG", cell, 1) CG2_element = FiniteElement("CG", cell, 2) V0_element = TensorProductElement(DG0_element, CG1_element) V1_element = TensorProductElement(DG1_element, CG2_element) # spaces VDG1 = FunctionSpace(mesh, "DG", 1) VCG1 = FunctionSpace(mesh, "CG", 1) V0 = FunctionSpace(mesh, V0_element) V1 = FunctionSpace(mesh, V1_element) V0_brok = FunctionSpace(mesh, BrokenElement(V0.ufl_element())) V0_spaces = (VDG1, VCG1, V0_brok) # declare initial fields u0 = state.fields("u") theta0 = state.fields("theta0", V0) theta1 = state.fields("theta1", V1) # Isentropic background state Tsurf = 0. thetab = Constant(Tsurf) theta_b1 = Function(V1).interpolate(thetab) theta_b0 = Function(V0).interpolate(thetab) # set up bubble xc = L / 2 zc = 700. rc = 80. theta_expr = conditional( sqrt((x - xc)**2.0) < rc, conditional(sqrt((z - zc)**2.0) < rc, Constant(2.0), Constant(0.0)), Constant(0.0)) theta_pert1 = Function(V1).interpolate(theta_expr) theta_pert0 = Function(V0).interpolate(theta_expr) # set up velocity field u_max = Constant(5.0) u0.project(as_vector([0, -u_max])) theta0.interpolate(theta_b0 + theta_pert0) theta1.interpolate(theta_b1 + theta_pert1) state.initialise([('u', u0), ('theta1', theta1), ('theta0', theta0)]) state.set_reference_profiles([('theta1', theta_b1), ('theta0', theta_b0)]) # set up advection schemes thetaeqn1 = EmbeddedDGAdvection(state, V1, equation_form="advective") thetaeqn0 = EmbeddedDGAdvection(state, V0, equation_form="advective", recovered_spaces=V0_spaces) # build advection dictionary advected_fields = [] advected_fields.append(('u', NoAdvection(state, u0, None))) advected_fields.append(('theta1', SSPRK3(state, theta1, thetaeqn1, limiter=ThetaLimiter(thetaeqn1)))) advected_fields.append(('theta0', SSPRK3(state, theta0, thetaeqn0, limiter=VertexBasedLimiter(VDG1)))) # build time stepper stepper = AdvectionDiffusion(state, advected_fields) return stepper, 40.0
def _diagonal(self): from firedrake import Function assert self.on_diag return Function(self._x.function_space())
eqns = CompressibleEulerEquations(state, "CG", 1) # Initial conditions u0 = state.fields("u") rho0 = state.fields("rho") theta0 = state.fields("theta") # spaces Vu = state.spaces("HDiv") Vt = state.spaces("theta") Vr = state.spaces("DG") # Isentropic background state Tsurf = Constant(300.) theta_b = Function(Vt).interpolate(Tsurf) rho_b = Function(Vr) # Calculate hydrostatic exner 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) rho0.interpolate(rho_b)
class DGAdvection(Advection): """ DG 3 step SSPRK advection scheme that can be applied to a scalar or vector field :arg state: :class:`.State` object. :arg V: function space of advected field - should be DG :arg continuity: optional boolean. If ``True``, the advection equation is of the form: :math: `D_t +\nabla \cdot(uD) = 0`. If ``False``, the advection equation is of the form: :math: `D_t + (u \cdot \nabla)D = 0`. """ def __init__(self, state, V, continuity=False): super(DGAdvection, self).__init__(state) element = V.fiat_element assert element.entity_dofs() == element.entity_closure_dofs(), "Provided space is not discontinuous" dt = state.timestepping.dt if V.extruded: surface_measure = (dS_h + dS_v) else: surface_measure = dS phi = TestFunction(V) D = TrialFunction(V) self.D1 = Function(V) self.dD = Function(V) n = FacetNormal(state.mesh) # ( dot(v, n) + |dot(v, n)| )/2.0 un = 0.5*(dot(self.ubar, n) + abs(dot(self.ubar, n))) a_mass = inner(phi,D)*dx if continuity: a_int = -inner(grad(phi), outer(D, self.ubar))*dx else: a_int = -inner(div(outer(phi,self.ubar)),D)*dx a_flux = (dot(jump(phi), un('+')*D('+') - un('-')*D('-')))*surface_measure arhs = a_mass - dt*(a_int + a_flux) DGproblem = LinearVariationalProblem(a_mass, action(arhs,self.D1), self.dD) self.DGsolver = LinearVariationalSolver(DGproblem, solver_parameters={ 'ksp_type':'preonly', 'pc_type':'bjacobi', 'sub_pc_type': 'ilu'}, options_prefix='DGAdvection') def apply(self, x_in, x_out): # SSPRK Stage 1 self.D1.assign(x_in) self.DGsolver.solve() self.D1.assign(self.dD) # SSPRK Stage 2 self.DGsolver.solve() self.D1.assign(0.75*x_in + 0.25*self.dD) # SSPRK Stage 3 self.DGsolver.solve() x_out.assign((1.0/3.0)*x_in + (2.0/3.0)*self.dD)
def __init__(self, state, field, equation=None, *, solver_parameters=None, limiter=None): if equation is not None: self.state = state self.field = field self.equation = equation # get ubar from the equation class self.ubar = self.equation.ubar self.dt = self.state.timestepping.dt # get default solver options if none passed in if solver_parameters is None: self.solver_parameters = equation.solver_parameters else: self.solver_parameters = solver_parameters if state.output.log_level == DEBUG: self.solver_parameters["ksp_monitor_true_residual"] = True self.limiter = limiter # check to see if we are using an embedded DG method - if we are then # the projector and output function will have been set up in the # equation class and we can get the correct function space from # the output function. if isinstance(equation, EmbeddedDGAdvection): # check that the field and the equation are compatible if equation.V0 != field.function_space(): raise ValueError('The field to be advected is not compatible with the equation used.') self.embedded_dg = True fs = equation.space self.xdg_in = Function(fs) self.xdg_out = Function(fs) self.x_projected = Function(field.function_space()) parameters = {'ksp_type': 'cg', 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'} self.Projector = Projector(self.xdg_out, self.x_projected, solver_parameters=parameters) self.recovered = equation.recovered if self.recovered: # set up the necessary functions self.x_in = Function(field.function_space()) x_adv = Function(fs) x_rec = Function(equation.V_rec) x_brok = Function(equation.V_brok) # set up interpolators and projectors self.x_adv_interpolator = Interpolator(self.x_in, x_adv) # interpolate before recovery self.x_rec_projector = Recoverer(x_adv, x_rec) # recovered function # when the "average" method comes into firedrake master, this will be # self.x_rec_projector = Projector(self.x_in, equation.Vrec, method="average") self.x_brok_projector = Projector(x_rec, x_brok) # function projected back self.xdg_interpolator = Interpolator(self.x_in + x_rec - x_brok, self.xdg_in) if self.limiter is not None: self.x_brok_interpolator = Interpolator(self.xdg_out, x_brok) self.x_out_projector = Recoverer(x_brok, self.x_projected) # when the "average" method comes into firedrake master, this will be # self.x_out_projector = Projector(x_brok, self.x_projected, method="average") else: self.embedded_dg = False fs = field.function_space() # setup required functions self.fs = fs self.dq = Function(fs) self.q1 = Function(fs)
class SUPGAdvection(Advection): """ An SUPG advection scheme that can apply DG upwinding (in the direction specified by the direction arg) if the function space is only partially continuous. :arg state: :class:`.State` object. :arg V:class:`.FunctionSpace` object. The advected field function space. :arg direction: list containing the directions in which the function space is discontinuous. 1 corresponds to the vertical direction, 2 to the horizontal direction :arg supg_params: dictionary containing SUPG parameters tau for each direction. If not supplied tau is set to dt/sqrt(15.) """ def __init__(self, state, V, direction=[], supg_params=None): super(SUPGAdvection, self).__init__(state) dt = state.timestepping.dt params = supg_params.copy() if supg_params else {} params.setdefault('a0', dt/sqrt(15.)) params.setdefault('a1', dt/sqrt(15.)) gamma = TestFunction(V) theta = TrialFunction(V) self.theta0 = Function(V) # make SUPG test function taus = [params["a0"], params["a1"]] for i in direction: taus[i] = 0.0 tau = Constant(((taus[0], 0.), (0., taus[1]))) dgamma = dot(dot(self.ubar, tau), grad(gamma)) gammaSU = gamma + dgamma n = FacetNormal(state.mesh) un = 0.5*(dot(self.ubar, n) + abs(dot(self.ubar, n))) a_mass = gammaSU*theta*dx arhs = a_mass - dt*gammaSU*dot(self.ubar, grad(theta))*dx if 1 in direction: arhs -= ( dt*dot(jump(gammaSU), (un('+')*theta('+') - un('-')*theta('-')))*dS_v - dt*(gammaSU('+')*dot(self.ubar('+'), n('+'))*theta('+') + gammaSU('-')*dot(self.ubar('-'), n('-'))*theta('-'))*dS_v ) if 2 in direction: arhs -= ( dt*dot(jump(gammaSU), (un('+')*theta('+') - un('-')*theta('-')))*dS_h - dt*(gammaSU('+')*dot(self.ubar('+'), n('+'))*theta('+') + gammaSU('-')*dot(self.ubar('-'), n('-'))*theta('-'))*dS_h ) self.theta1 = Function(V) self.dtheta = Function(V) problem = LinearVariationalProblem(a_mass, action(arhs,self.theta1), self.dtheta) self.solver = LinearVariationalSolver(problem, options_prefix='SUPGAdvection') def apply(self, x_in, x_out): # SSPRK Stage 1 self.theta1.assign(x_in) self.solver.solve() self.theta1.assign(self.dtheta) # SSPRK Stage 2 self.solver.solve() self.theta1.assign(0.75*x_in + 0.25*self.dtheta) # SSPRK Stage 3 self.solver.solve() x_out.assign((1.0/3.0)*x_in + (2.0/3.0)*self.dtheta)
class IncompressibleForcing(Forcing): """ Forcing class for incompressible Euler Boussinesq equations. """ def __init__(self, state, linear=False): self.state = state self._build_forcing_solver(linear) def _build_forcing_solver(self, linear): """ Only put forcing terms into the u equation. """ state = self.state self.scaling = Constant(1.0) Vu = state.V[0] W = state.W self.x0 = Function(W) # copy x to here u0, p0, b0 = split(self.x0) F = TrialFunction(Vu) w = TestFunction(Vu) self.uF = Function(Vu) Omega = state.Omega mu = state.mu a = inner(w, F) * dx L = ( self.scaling * div(w) * p0 * dx # pressure gradient + self.scaling * b0 * inner(w, state.k) * dx # gravity term ) if not linear: L -= self.scaling * 0.5 * div(w) * inner(u0, u0) * dx if Omega is not None: L -= self.scaling * inner(w, cross(2 * Omega, u0)) * dx # Coriolis term if mu is not None: self.mu_scaling = Constant(1.0) L -= self.mu_scaling * mu * inner(w, state.k) * inner(u0, state.k) * dx bcs = [DirichletBC(Vu, 0.0, "bottom"), DirichletBC(Vu, 0.0, "top")] u_forcing_problem = LinearVariationalProblem(a, L, self.uF, bcs=bcs) self.u_forcing_solver = LinearVariationalSolver(u_forcing_problem) Vp = state.V[1] p = TrialFunction(Vp) q = TestFunction(Vp) self.divu = Function(Vp) a = p * q * dx L = q * div(u0) * dx divergence_problem = LinearVariationalProblem(a, L, self.divu) self.divergence_solver = LinearVariationalSolver(divergence_problem) def apply(self, scaling, x_in, x_nl, x_out, **kwargs): self.x0.assign(x_nl) self.scaling.assign(scaling) if "mu_alpha" in kwargs and kwargs["mu_alpha"] is not None: self.mu_scaling.assign(kwargs["mu_alpha"]) self.u_forcing_solver.solve() # places forcing in self.uF u_out, p_out, _ = x_out.split() x_out.assign(x_in) u_out += self.uF if "incompressible" in kwargs and kwargs["incompressible"]: self.divergence_solver.solve() p_out.assign(self.divu)
class IncompressibleSolver(TimesteppingSolver): """Timestepping linear solver object for the incompressible Boussinesq equations with prognostic variables u, p, b. This solver follows the following strategy: (1) Analytically eliminate b (introduces error near topography) (2) Solve resulting system for (u,p) using a block Hdiv preconditioner (3) Reconstruct b This currently requires a (parallel) direct solver so is probably a bit memory-hungry, we'll improve this with a hybridised solver soon. :arg state: a :class:`.State` object containing everything else. :arg L: the width of the domain, used in the preconditioner. :arg params: Solver parameters. """ def __init__(self, state, L, params=None): self.state = state if params is None: self.params = {'ksp_type':'gmres', 'pc_type':'fieldsplit', 'pc_fieldsplit_type':'additive', 'fieldsplit_0_pc_type':'lu', 'fieldsplit_1_pc_type':'lu', 'fieldsplit_0_pc_factor_mat_solver_package': 'mumps', 'fieldsplit_0_pc_factor_mat_solver_package': 'mumps', 'fieldsplit_0_ksp_type':'preonly', 'fieldsplit_1_ksp_type':'preonly'} else: self.params = params self.L = L # setup the solver self._setup_solver() def _setup_solver(self): state = self.state # just cutting down line length a bit dt = state.timestepping.dt beta = dt*state.timestepping.alpha mu = state.mu # Split up the rhs vector (symbolically) u_in, p_in, b_in = split(state.xrhs) # Build the reduced function space for u,p M = MixedFunctionSpace((state.V[0], state.V[1])) w, phi = TestFunctions(M) u, p = TrialFunctions(M) # Get background fields bbar = state.bbar # Analytical (approximate) elimination of theta k = state.k # Upward pointing unit vector b = -dot(k,u)*dot(k,grad(bbar))*beta + b_in # vertical projection def V(u): return k*inner(u,k) eqn = ( inner(w, (u - u_in))*dx - beta*div(w)*p*dx - beta*inner(w,k)*b*dx + phi*div(u)*dx ) if mu is not None: eqn += dt*mu*inner(w,k)*inner(u,k)*dx aeqn = lhs(eqn) Leqn = rhs(eqn) # Place to put result of u p solver self.up = Function(M) # Boundary conditions (assumes extruded mesh) dim = M.sub(0).ufl_element().value_shape()[0] bc = ("0.0",)*dim bcs = [DirichletBC(M.sub(0), Expression(bc), "bottom"), DirichletBC(M.sub(0), Expression(bc), "top")] # preconditioner equation L = self.L Ap = ( inner(w,u) + L*L*div(w)*div(u) + phi*p/L/L )*dx # Solver for u, p up_problem = LinearVariationalProblem( aeqn, Leqn, self.up, bcs=bcs, aP=Ap) nullspace = MixedVectorSpaceBasis(M, [M.sub(0), VectorSpaceBasis(constant=True)]) self.up_solver = LinearVariationalSolver(up_problem, solver_parameters=self.params, nullspace=nullspace) # Reconstruction of b b = TrialFunction(state.V[2]) gamma = TestFunction(state.V[2]) u, p = self.up.split() self.b = Function(state.V[2]) b_eqn = gamma*(b - b_in + dot(k,u)*dot(k,grad(bbar))*beta)*dx b_problem = LinearVariationalProblem(lhs(b_eqn), rhs(b_eqn), self.b) self.b_solver = LinearVariationalSolver(b_problem) def solve(self): """ Apply the solver with rhs state.xrhs and result state.dy. """ self.up_solver.solve() u1, p1 = self.up.split() u, p, b = self.state.dy.split() u.assign(u1) p.assign(p1) self.b_solver.solve() b.assign(self.b)
def _setup_solver(self): import numpy as np state = self.state Dt = state.timestepping.dt beta_ = Dt*state.timestepping.alpha cp = state.parameters.cp mu = state.mu Vu = state.spaces("HDiv") Vu_broken = FunctionSpace(state.mesh, BrokenElement(Vu.ufl_element())) Vtheta = state.spaces("HDiv_v") Vrho = state.spaces("DG") # Store time-stepping coefficients as UFL Constants dt = Constant(Dt) beta = Constant(beta_) beta_cp = Constant(beta_ * cp) h_deg = state.horizontal_degree v_deg = state.vertical_degree Vtrace = FunctionSpace(state.mesh, "HDiv Trace", degree=(h_deg, v_deg)) # Split up the rhs vector (symbolically) u_in, rho_in, theta_in = split(state.xrhs) # Build the function space for "broken" u, 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") pibar = thermodynamics.pi(state.parameters, rhobar, thetabar) pibar_rho = thermodynamics.pi_rho(state.parameters, rhobar, thetabar) pibar_theta = thermodynamics.pi_theta(state.parameters, rhobar, thetabar) # Analytical (approximate) elimination of theta k = state.k # Upward pointing unit vector theta = -dot(k, u)*dot(k, grad(thetabar))*beta + theta_in # Only include theta' (rather than pi') in the vertical # component of the gradient # The pi prime term (here, bars are for mean and no bars are # for linear perturbations) pi = pibar_theta*theta + pibar_rho*rho # Vertical projection def V(u): return k*inner(u, k) # Specify degree for some terms as estimated degree is too large dxp = dx(degree=(self.quadrature_degree)) dS_vp = dS_v(degree=(self.quadrature_degree)) dS_hp = dS_h(degree=(self.quadrature_degree)) ds_vp = ds_v(degree=(self.quadrature_degree)) ds_tbp = (ds_t(degree=(self.quadrature_degree)) + ds_b(degree=(self.quadrature_degree))) # 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) pibar_avg = Function(Vtrace) rho_avg_prb = LinearVariationalProblem(a_tr, L_tr(rhobar), rhobar_avg) pi_avg_prb = LinearVariationalProblem(a_tr, L_tr(pibar), pibar_avg) rho_avg_solver = LinearVariationalSolver(rho_avg_prb, solver_parameters=cg_ilu_parameters, options_prefix='rhobar_avg_solver') pi_avg_solver = LinearVariationalSolver(pi_avg_prb, solver_parameters=cg_ilu_parameters, options_prefix='pibar_avg_solver') with timed_region("Gusto:HybridProjectRhobar"): rho_avg_solver.solve() with timed_region("Gusto:HybridProjectPibar"): pi_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. eqn = ( # momentum equation 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_w*V(w), n=n)*pibar_avg('+')*dS_vp + beta_cp*jump(theta_w*V(w), n=n)*pibar_avg('+')*dS_hp + beta_cp*dot(theta_w*V(w), n)*pibar_avg*ds_tbp - beta_cp*div(thetabar_w*w)*pi*dxp # 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 mu is not None: eqn += dt*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.state.bcs
class CompressibleSolver(TimesteppingSolver): """ Timestepping linear solver object for the compressible equations in theta-pi formulation with prognostic variables u,rho,theta. This solver follows the following strategy: (1) Analytically eliminate theta (introduces error near topography) (2) Solve resulting system for (u,rho) using a Schur preconditioner (3) Reconstruct theta :arg state: a :class:`.State` object containing everything else. """ def __init__(self, state, params=None): self.state = state if params is None: self.params = {'pc_type': 'fieldsplit', 'pc_fieldsplit_type': 'schur', 'ksp_type': 'gmres', 'ksp_max_it': 100, 'ksp_gmres_restart': 50, 'pc_fieldsplit_schur_fact_type': 'FULL', 'pc_fieldsplit_schur_precondition': 'selfp', 'fieldsplit_0_ksp_type': 'preonly', 'fieldsplit_0_pc_type': 'bjacobi', 'fieldsplit_0_sub_pc_type': 'ilu', 'fieldsplit_1_ksp_type': 'preonly', 'fieldsplit_1_pc_type': 'gamg', 'fieldsplit_1_mg_levels_ksp_type': 'chebyshev', 'fieldsplit_1_mg_levels_ksp_chebyshev_estimate_eigenvalues': True, 'fieldsplit_1_mg_levels_ksp_chebyshev_estimate_eigenvalues_random': True, 'fieldsplit_1_mg_levels_ksp_max_it': 1, 'fieldsplit_1_mg_levels_pc_type': 'bjacobi', 'fieldsplit_1_mg_levels_sub_pc_type': 'ilu'} else: self.params = params # setup the solver self._setup_solver() def _setup_solver(self): state = self.state # just cutting down line length a bit dt = state.timestepping.dt beta = dt*state.timestepping.alpha cp = state.parameters.cp mu = state.mu # Split up the rhs vector (symbolically) u_in, rho_in, theta_in = split(state.xrhs) # Build the reduced function space for u,rho M = MixedFunctionSpace((state.V[0], state.V[1])) w, phi = TestFunctions(M) u, rho = TrialFunctions(M) n = FacetNormal(state.mesh) # Get background fields thetabar = state.thetabar rhobar = state.rhobar pibar = exner(thetabar, rhobar, state) pibar_rho = exner_rho(thetabar, rhobar, state) pibar_theta = exner_theta(thetabar, rhobar, state) # Analytical (approximate) elimination of theta k = state.k # Upward pointing unit vector theta = -dot(k,u)*dot(k,grad(thetabar))*beta + theta_in # Only include theta' (rather than pi') in the vertical # component of the gradient # the pi prime term (here, bars are for mean and no bars are # for linear perturbations) pi = pibar_theta*theta + pibar_rho*rho # vertical projection def V(u): return k*inner(u,k) eqn = ( inner(w, (u - u_in))*dx - beta*cp*div(theta*V(w))*pibar*dx # following does nothing but is preserved in the comments # to remind us why (because V(w) is purely vertical. # + beta*cp*jump(theta*V(w),n)*avg(pibar)*dS_v - beta*cp*div(thetabar*w)*pi*dx + beta*cp*jump(thetabar*w,n)*avg(pi)*dS_v + (phi*(rho - rho_in) - beta*inner(grad(phi), u)*rhobar)*dx + beta*jump(phi*u, n)*avg(rhobar)*(dS_v + dS_h) ) if mu is not None: eqn += dt*mu*inner(w,k)*inner(u,k)*dx aeqn = lhs(eqn) Leqn = rhs(eqn) # Place to put result of u rho solver self.urho = Function(M) # Boundary conditions (assumes extruded mesh) dim = M.sub(0).ufl_element().value_shape()[0] bc = ("0.0",)*dim bcs = [DirichletBC(M.sub(0), Expression(bc), "bottom"), DirichletBC(M.sub(0), Expression(bc), "top")] # Solver for u, rho urho_problem = LinearVariationalProblem( aeqn, Leqn, self.urho, bcs=bcs) self.urho_solver = LinearVariationalSolver(urho_problem, solver_parameters=self.params, options_prefix='ImplicitSolver') # Reconstruction of theta theta = TrialFunction(state.V[2]) gamma = TestFunction(state.V[2]) u, rho = self.urho.split() self.theta = Function(state.V[2]) theta_eqn = gamma*(theta - theta_in + dot(k,u)*dot(k,grad(thetabar))*beta)*dx theta_problem = LinearVariationalProblem(lhs(theta_eqn), rhs(theta_eqn), self.theta) self.theta_solver = LinearVariationalSolver(theta_problem, options_prefix='thetabacksubstitution') def solve(self): """ Apply the solver with rhs state.xrhs and result state.dy. """ self.urho_solver.solve() u1, rho1 = self.urho.split() u, rho, theta = self.state.dy.split() u.assign(u1) rho.assign(rho1) self.theta_solver.solve() theta.assign(self.theta)
class CompressibleForcing(Forcing): """ Forcing class for compressible Euler equations. """ def __init__(self, state, linear=False): self.state = state self._build_forcing_solver(linear) def _build_forcing_solver(self, linear): """ Only put forcing terms into the u equation. """ state = self.state self.scaling = Constant(1.0) Vu = state.V[0] W = state.W self.x0 = Function(W) # copy x to here u0, rho0, theta0 = split(self.x0) F = TrialFunction(Vu) w = TestFunction(Vu) self.uF = Function(Vu) Omega = state.Omega cp = state.parameters.cp mu = state.mu n = FacetNormal(state.mesh) pi = exner(theta0, rho0, state) a = inner(w, F) * dx L = self.scaling * ( +cp * div(theta0 * w) * pi * dx # pressure gradient [volume] - cp * jump(w * theta0, n) * avg(pi) * dS_v # pressure gradient [surface] ) if state.parameters.geopotential: Phi = state.Phi L += self.scaling * div(w) * Phi * dx # gravity term else: g = state.parameters.g L -= self.scaling * g * inner(w, state.k) * dx # gravity term if not linear: L -= self.scaling * 0.5 * div(w) * inner(u0, u0) * dx if Omega is not None: L -= self.scaling * inner(w, cross(2 * Omega, u0)) * dx # Coriolis term if mu is not None: self.mu_scaling = Constant(1.0) L -= self.mu_scaling * mu * inner(w, state.k) * inner(u0, state.k) * dx bcs = [DirichletBC(Vu, 0.0, "bottom"), DirichletBC(Vu, 0.0, "top")] u_forcing_problem = LinearVariationalProblem(a, L, self.uF, bcs=bcs) self.u_forcing_solver = LinearVariationalSolver(u_forcing_problem) def apply(self, scaling, x_in, x_nl, x_out, **kwargs): self.x0.assign(x_nl) self.scaling.assign(scaling) if "mu_alpha" in kwargs and kwargs["mu_alpha"] is not None: self.mu_scaling.assign(kwargs["mu_alpha"]) self.u_forcing_solver.solve() # places forcing in self.uF u_out, _, _ = x_out.split() x_out.assign(x_in) u_out += self.uF
class IncompressibleSolver(TimesteppingSolver): """Timestepping linear solver object for the incompressible Boussinesq equations with prognostic variables u, p, b. This solver follows the following strategy: (1) Analytically eliminate b (introduces error near topography) (2) Solve resulting system for (u,p) using a hybrid-mixed method (3) Reconstruct b :arg state: a :class:`.State` object containing everything else. :arg solver_parameters: (optional) Solver parameters. :arg overwrite_solver_parameters: boolean, if True use only the solver_parameters that have been passed in, if False then update the default solver parameters with the solver_parameters passed in. """ solver_parameters = { 'ksp_type': 'preonly', 'mat_type': 'matfree', 'pc_type': 'python', 'pc_python_type': 'firedrake.HybridizationPC', 'hybridization': {'ksp_type': 'cg', 'pc_type': 'gamg', 'ksp_rtol': 1e-8, 'mg_levels': {'ksp_type': 'chebyshev', 'ksp_max_it': 2, 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'}} } def __init__(self, state, solver_parameters=None, overwrite_solver_parameters=False): super().__init__(state, solver_parameters, overwrite_solver_parameters) @timed_function("Gusto:SolverSetup") def _setup_solver(self): state = self.state # just cutting down line length a bit Dt = state.timestepping.dt beta_ = Dt*state.timestepping.alpha 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) @timed_function("Gusto:LinearSolve") def solve(self): """ Apply the solver with rhs state.xrhs and result state.dy. """ with timed_region("Gusto:VelocityPressureSolve"): self.up_solver.solve() u1, p1 = self.up.split() u, p, b = self.state.dy.split() u.assign(u1) p.assign(p1) with timed_region("Gusto:BuoyancyRecon"): self.b_solver.solve() b.assign(self.b)
def compressible_hydrostatic_balance(state, theta0, rho0, pi0=None, top=False, pi_boundary=Constant(1.0), solve_for_rho=False, params=None): """ Compute a hydrostatically balanced density given a potential temperature profile. :arg state: The :class:`State` object. :arg theta0: :class:`.Function`containing the potential temperature. :arg rho0: :class:`.Function` to write the initial density into. :arg top: If True, set a boundary condition at the top. Otherwise, set it at the bottom. :arg pi_boundary: a field or expression to use as boundary data for pi on the top or bottom as specified. """ # Calculate hydrostatic Pi W = MixedFunctionSpace((state.Vv,state.V[1])) v, pi = TrialFunctions(W) dv, dpi = TestFunctions(W) n = FacetNormal(state.mesh) cp = state.parameters.cp alhs = ( (cp*inner(v,dv) - cp*div(dv*theta0)*pi)*dx + dpi*div(theta0*v)*dx ) if top: bmeasure = ds_t bstring = "bottom" else: bmeasure = ds_b bstring = "top" arhs = -cp*inner(dv,n)*theta0*pi_boundary*bmeasure if state.parameters.geopotential: Phi = state.Phi arhs += div(dv)*Phi*dx - inner(dv,n)*Phi*bmeasure else: g = state.parameters.g arhs -= g*inner(dv,state.k)*dx if(state.mesh.geometric_dimension() == 2): bcs = [DirichletBC(W.sub(0), Expression(("0.", "0.")), bstring)] elif(state.mesh.geometric_dimension() == 3): bcs = [DirichletBC(W.sub(0), Expression(("0.", "0.", "0.")), bstring)] w = Function(W) PiProblem = LinearVariationalProblem(alhs, arhs, w, bcs=bcs) if(params is None): params = {'pc_type': 'fieldsplit', 'pc_fieldsplit_type': 'schur', 'ksp_type': 'gmres', 'ksp_monitor_true_residual': True, 'ksp_max_it': 100, 'ksp_gmres_restart': 50, 'pc_fieldsplit_schur_fact_type': 'FULL', 'pc_fieldsplit_schur_precondition': 'selfp', 'fieldsplit_0_ksp_type': 'richardson', 'fieldsplit_0_ksp_max_it': 5, 'fieldsplit_0_pc_type': 'gamg', 'fieldsplit_1_pc_gamg_sym_graph': True, 'fieldsplit_1_mg_levels_ksp_type': 'chebyshev', 'fieldsplit_1_mg_levels_ksp_chebyshev_estimate_eigenvalues': True, 'fieldsplit_1_mg_levels_ksp_chebyshev_estimate_eigenvalues_random': True, 'fieldsplit_1_mg_levels_ksp_max_it': 5, 'fieldsplit_1_mg_levels_pc_type': 'bjacobi', 'fieldsplit_1_mg_levels_sub_pc_type': 'ilu'} PiSolver = LinearVariationalSolver(PiProblem, solver_parameters=params) PiSolver.solve() v, Pi = w.split() if pi0 is not None: pi0.assign(Pi) kappa = state.parameters.kappa R_d = state.parameters.R_d p_0 = state.parameters.p_0 if solve_for_rho: w1 = Function(W) v, rho = w1.split() rho.interpolate(p_0*(Pi**((1-kappa)/kappa))/R_d/theta0) v, rho = split(w1) dv, dpi = TestFunctions(W) pi = ((R_d/p_0)*rho*theta0)**(kappa/(1.-kappa)) F = ( (cp*inner(v,dv) - cp*div(dv*theta0)*pi)*dx + dpi*div(theta0*v)*dx + cp*inner(dv,n)*theta0*pi_boundary*bmeasure ) if state.parameters.geopotential: F += - div(dv)*Phi*dx + inner(dv,n)*Phi*bmeasure else: F += g*inner(dv,state.k)*dx rhoproblem = NonlinearVariationalProblem(F, w1, bcs=bcs) rhosolver = NonlinearVariationalSolver(rhoproblem, solver_parameters=params) rhosolver.solve() v, rho_ = w1.split() rho0.assign(rho_) else: rho0.interpolate(p_0*(Pi**((1-kappa)/kappa))/R_d/theta0)
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)
class CompressibleSolver(TimesteppingSolver): """ Timestepping linear solver object for the compressible equations in theta-pi formulation with prognostic variables u, rho, and theta. This solver follows the following strategy: (1) Analytically eliminate theta (introduces error near topography) (2a) Formulate the resulting mixed system for u and rho using a hybridized mixed method. This breaks continuity in the linear perturbations of u, and introduces a new unknown on the mesh interfaces approximating the average of the Exner pressure perturbations. These trace unknowns also act as Lagrange multipliers enforcing normal continuity of the "broken" u variable. (2b) Statically condense the block-sparse system into a single system for the Lagrange multipliers. This is the only globally coupled system requiring a linear solver. (2c) Using the computed trace variables, we locally recover the broken velocity and density perturbations. This is accomplished in two stages: (i): Recover rho locally using the multipliers. (ii): Recover "broken" u locally using rho and the multipliers. (2d) Project the "broken" velocity field into the HDiv-conforming space using local averaging. (3) Reconstruct theta :arg state: a :class:`.State` object containing everything else. :arg quadrature degree: tuple (q_h, q_v) where q_h is the required quadrature degree in the horizontal direction and q_v is that in the vertical direction. :arg solver_parameters (optional): solver parameters for the trace system. :arg overwrite_solver_parameters: boolean, if True use only the solver_parameters that have been passed in, if False then update. the default solver parameters with the solver_parameters passed in. :arg moisture (optional): list of names of moisture fields. """ solver_parameters = {'mat_type': 'matfree', 'ksp_type': 'preonly', 'pc_type': 'python', 'pc_python_type': 'firedrake.SCPC', 'pc_sc_eliminate_fields': '0, 1', # The reduced operator is not symmetric 'condensed_field': {'ksp_type': 'fgmres', 'ksp_rtol': 1.0e-8, 'ksp_atol': 1.0e-8, 'ksp_max_it': 100, 'pc_type': 'gamg', 'pc_gamg_sym_graph': None, 'mg_levels': {'ksp_type': 'gmres', 'ksp_max_it': 5, 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'}}} def __init__(self, state, quadrature_degree=None, solver_parameters=None, overwrite_solver_parameters=False, moisture=None): self.moisture = moisture self.state = state if quadrature_degree is not None: self.quadrature_degree = quadrature_degree else: dgspace = state.spaces("DG") if any(deg > 2 for deg in dgspace.ufl_element().degree()): logger.warning("default quadrature degree most likely not sufficient for this degree element") self.quadrature_degree = (5, 5) if logger.isEnabledFor(DEBUG): # Set outer solver to FGMRES and turn on KSP monitor for the outer system self.solver_parameters["ksp_type"] = "fgmres" self.solver_parameters["mat_type"] = "aij" self.solver_parameters["pmat_type"] = "matfree" self.solver_parameters["ksp_monitor_true_residual"] = None # Turn monitor on for the trace system self.solver_parameters["condensed_field"]["ksp_monitor_true_residual"] = None super().__init__(state, solver_parameters, overwrite_solver_parameters) @timed_function("Gusto:SolverSetup") def _setup_solver(self): import numpy as np state = self.state Dt = state.timestepping.dt beta_ = Dt*state.timestepping.alpha cp = state.parameters.cp mu = state.mu Vu = state.spaces("HDiv") Vu_broken = FunctionSpace(state.mesh, BrokenElement(Vu.ufl_element())) Vtheta = state.spaces("HDiv_v") Vrho = state.spaces("DG") # Store time-stepping coefficients as UFL Constants dt = Constant(Dt) beta = Constant(beta_) beta_cp = Constant(beta_ * cp) h_deg = state.horizontal_degree v_deg = state.vertical_degree Vtrace = FunctionSpace(state.mesh, "HDiv Trace", degree=(h_deg, v_deg)) # Split up the rhs vector (symbolically) u_in, rho_in, theta_in = split(state.xrhs) # Build the function space for "broken" u, 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") pibar = thermodynamics.pi(state.parameters, rhobar, thetabar) pibar_rho = thermodynamics.pi_rho(state.parameters, rhobar, thetabar) pibar_theta = thermodynamics.pi_theta(state.parameters, rhobar, thetabar) # Analytical (approximate) elimination of theta k = state.k # Upward pointing unit vector theta = -dot(k, u)*dot(k, grad(thetabar))*beta + theta_in # Only include theta' (rather than pi') in the vertical # component of the gradient # The pi prime term (here, bars are for mean and no bars are # for linear perturbations) pi = pibar_theta*theta + pibar_rho*rho # Vertical projection def V(u): return k*inner(u, k) # Specify degree for some terms as estimated degree is too large dxp = dx(degree=(self.quadrature_degree)) dS_vp = dS_v(degree=(self.quadrature_degree)) dS_hp = dS_h(degree=(self.quadrature_degree)) ds_vp = ds_v(degree=(self.quadrature_degree)) ds_tbp = (ds_t(degree=(self.quadrature_degree)) + ds_b(degree=(self.quadrature_degree))) # 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) pibar_avg = Function(Vtrace) rho_avg_prb = LinearVariationalProblem(a_tr, L_tr(rhobar), rhobar_avg) pi_avg_prb = LinearVariationalProblem(a_tr, L_tr(pibar), pibar_avg) rho_avg_solver = LinearVariationalSolver(rho_avg_prb, solver_parameters=cg_ilu_parameters, options_prefix='rhobar_avg_solver') pi_avg_solver = LinearVariationalSolver(pi_avg_prb, solver_parameters=cg_ilu_parameters, options_prefix='pibar_avg_solver') with timed_region("Gusto:HybridProjectRhobar"): rho_avg_solver.solve() with timed_region("Gusto:HybridProjectPibar"): pi_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. eqn = ( # momentum equation 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_w*V(w), n=n)*pibar_avg('+')*dS_vp + beta_cp*jump(theta_w*V(w), n=n)*pibar_avg('+')*dS_hp + beta_cp*dot(theta_w*V(w), n)*pibar_avg*ds_tbp - beta_cp*div(thetabar_w*w)*pi*dxp # 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 mu is not None: eqn += dt*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.state.bcs @timed_function("Gusto:LinearSolve") def solve(self): """ Apply the solver with rhs state.xrhs and result state.dy. """ # Solve the hybridized system self.hybridized_solver.solve() broken_u, rho1, _ = self.urhol0.split() u1 = self.u_hdiv # Project broken_u into the HDiv space u1.assign(0.0) with timed_region("Gusto:HybridProjectHDiv"): par_loop(self._average_kernel, dx, {"w": (self._weight, READ), "vec_in": (broken_u, READ), "vec_out": (u1, INC)}) # Reapply bcs to ensure they are satisfied for bc in self.bcs: bc.apply(u1) # Copy back into u and rho cpts of dy u, rho, theta = self.state.dy.split() u.assign(u1) rho.assign(rho1) # Reconstruct theta with timed_region("Gusto:ThetaRecon"): self.theta_solver.solve() # Copy into theta cpt of dy theta.assign(self.theta)
def __init__(self, state): super(Condensation, self).__init__(state) # obtain our fields self.theta = state.fields('theta') self.water_v = state.fields('water_v') self.water_c = state.fields('water_c') rho = state.fields('rho') # declare function space Vt = self.theta.function_space() param = self.state.parameters # define some parameters as attributes dt = self.state.timestepping.dt R_d = param.R_d p_0 = param.p_0 kappa = param.kappa cp = param.cp cv = param.cv c_pv = param.c_pv c_pl = param.c_pl c_vv = param.c_vv R_v = param.R_v L_v0 = param.L_v0 T_0 = param.T_0 w_sat1 = param.w_sat1 w_sat2 = param.w_sat2 w_sat3 = param.w_sat3 w_sat4 = param.w_sat4 # make useful fields Pi = ((R_d * rho * self.theta / p_0)**(kappa / (1.0 - kappa))) T = Pi * self.theta * R_d / (R_d + self.water_v * R_v) p = p_0 * Pi**(1.0 / kappa) L_v = L_v0 - (c_pl - c_pv) * (T - T_0) R_m = R_d + R_v * self.water_v c_pml = cp + c_pv * self.water_v + c_pl * self.water_c c_vml = cv + c_vv * self.water_v + c_pl * self.water_c # use Teten's formula to calculate w_sat w_sat = (w_sat1 / (p * exp(w_sat2 * (T - T_0) / (T - w_sat3)) - w_sat4)) # make appropriate condensation rate dot_r_cond = ((self.water_v - w_sat) / (dt * (1.0 + ((L_v**2.0 * w_sat) / (cp * R_v * T**2.0))))) # make cond_rate function, that needs to be the same for all updates in one time step self.cond_rate = Function(Vt) # adjust cond rate so negative concentrations don't occur self.lim_cond_rate = Interpolator( conditional(dot_r_cond < 0, max_value(dot_r_cond, -self.water_c / dt), min_value(dot_r_cond, self.water_v / dt)), self.cond_rate) # tell the prognostic fields what to update to self.water_v_new = Interpolator(self.water_v - dt * self.cond_rate, Vt) self.water_c_new = Interpolator(self.water_c + dt * self.cond_rate, Vt) self.theta_new = Interpolator( self.theta * (1.0 + dt * self.cond_rate * (cv * L_v / (c_vml * cp * T) - R_v * cv * c_pml / (R_m * cp * c_vml))), Vt)
def _setup_solver(self): state = self.state # just cutting down line length a bit dt = state.timestepping.dt beta = dt*state.timestepping.alpha mu = state.mu # Split up the rhs vector (symbolically) u_in, p_in, b_in = split(state.xrhs) # Build the reduced function space for u,p M = MixedFunctionSpace((state.V[0], state.V[1])) w, phi = TestFunctions(M) u, p = TrialFunctions(M) # Get background fields bbar = state.bbar # Analytical (approximate) elimination of theta k = state.k # Upward pointing unit vector b = -dot(k,u)*dot(k,grad(bbar))*beta + b_in # vertical projection def V(u): return k*inner(u,k) eqn = ( inner(w, (u - u_in))*dx - beta*div(w)*p*dx - beta*inner(w,k)*b*dx + phi*div(u)*dx ) if mu is not None: eqn += dt*mu*inner(w,k)*inner(u,k)*dx aeqn = lhs(eqn) Leqn = rhs(eqn) # Place to put result of u p solver self.up = Function(M) # Boundary conditions (assumes extruded mesh) dim = M.sub(0).ufl_element().value_shape()[0] bc = ("0.0",)*dim bcs = [DirichletBC(M.sub(0), Expression(bc), "bottom"), DirichletBC(M.sub(0), Expression(bc), "top")] # preconditioner equation L = self.L Ap = ( inner(w,u) + L*L*div(w)*div(u) + phi*p/L/L )*dx # Solver for u, p up_problem = LinearVariationalProblem( aeqn, Leqn, self.up, bcs=bcs, aP=Ap) nullspace = MixedVectorSpaceBasis(M, [M.sub(0), VectorSpaceBasis(constant=True)]) self.up_solver = LinearVariationalSolver(up_problem, solver_parameters=self.params, nullspace=nullspace) # Reconstruction of b b = TrialFunction(state.V[2]) gamma = TestFunction(state.V[2]) u, p = self.up.split() self.b = Function(state.V[2]) b_eqn = gamma*(b - b_in + dot(k,u)*dot(k,grad(bbar))*beta)*dx b_problem = LinearVariationalProblem(lhs(b_eqn), rhs(b_eqn), self.b) self.b_solver = LinearVariationalSolver(b_problem)
def PartiallyPeriodicRectangleMesh(nx, ny, Lx, Ly, direction="x", quadrilateral=False, reorder=None, comm=COMM_WORLD): """Generates RectangleMesh that is periodic in the x or y direction. :arg nx: The number of cells in the x direction :arg ny: The number of cells in the y direction :arg Lx: The extent in the x direction :arg Ly: The extent in the y direction :kwarg direction: The direction of the periodicity (default x). :kwarg quadrilateral: (optional), creates quadrilateral mesh, defaults to False :kwarg reorder: (optional), should the mesh be reordered :kwarg comm: Optional communicator to build the mesh on (defaults to COMM_WORLD). If direction == "x" the boundary edges in this mesh are numbered as follows: * 1: plane y == 0 * 2: plane y == Ly If direction == "y" the boundary edges are: * 1: plane x == 0 * 2: plane x == Lx """ if direction not in ("x", "y"): raise ValueError("Unsupported periodic direction '%s'" % direction) # handle x/y directions: na, La are for the periodic axis na, nb, La, Lb = nx, ny, Lx, Ly if direction == "y": na, nb, La, Lb = ny, nx, Ly, Lx if na < 3: raise ValueError("2D periodic meshes with fewer than 3 \ cells in each direction are not currently supported") m = CylinderMesh(na, nb, 1.0, 1.0, longitudinal_direction="z", quadrilateral=quadrilateral, reorder=reorder, comm=comm) coord_fs = VectorFunctionSpace(m, 'DG', 1, dim=2) old_coordinates = m.coordinates new_coordinates = Function(coord_fs) # make x-periodic mesh # unravel x coordinates like in periodic interval # set y coordinates to z coordinates periodic_kernel = """double Y,pi; Y = 0.0; for(int i=0; i<old_coords.dofs; i++) { Y += old_coords[i][1]; } pi=3.141592653589793; for(int i=0;i<new_coords.dofs;i++){ new_coords[i][0] = atan2(old_coords[i][1],old_coords[i][0])/pi/2; if(new_coords[i][0]<0.) new_coords[i][0] += 1; if(new_coords[i][0]==0 && Y<0.) new_coords[i][0] = 1.0; new_coords[i][0] *= Lx[0]; new_coords[i][1] = old_coords[i][2]*Ly[0]; }""" cLx = Constant(La) cLy = Constant(Lb) par_loop(periodic_kernel, dx, {"new_coords": (new_coordinates, WRITE), "old_coords": (old_coordinates, READ), "Lx": (cLx, READ), "Ly": (cLy, READ)}) if direction == "y": # flip x and y coordinates operator = np.asarray([[0, 1], [1, 0]]) new_coordinates.dat.data[:] = np.dot(new_coordinates.dat.data, operator.T) return mesh.Mesh(new_coordinates)
def _setup_solver(self): state = self.state # just cutting down line length a bit dt = state.timestepping.dt beta = dt*state.timestepping.alpha cp = state.parameters.cp mu = state.mu # Split up the rhs vector (symbolically) u_in, rho_in, theta_in = split(state.xrhs) # Build the reduced function space for u,rho M = MixedFunctionSpace((state.V[0], state.V[1])) w, phi = TestFunctions(M) u, rho = TrialFunctions(M) n = FacetNormal(state.mesh) # Get background fields thetabar = state.thetabar rhobar = state.rhobar pibar = exner(thetabar, rhobar, state) pibar_rho = exner_rho(thetabar, rhobar, state) pibar_theta = exner_theta(thetabar, rhobar, state) # Analytical (approximate) elimination of theta k = state.k # Upward pointing unit vector theta = -dot(k,u)*dot(k,grad(thetabar))*beta + theta_in # Only include theta' (rather than pi') in the vertical # component of the gradient # the pi prime term (here, bars are for mean and no bars are # for linear perturbations) pi = pibar_theta*theta + pibar_rho*rho # vertical projection def V(u): return k*inner(u,k) eqn = ( inner(w, (u - u_in))*dx - beta*cp*div(theta*V(w))*pibar*dx # following does nothing but is preserved in the comments # to remind us why (because V(w) is purely vertical. # + beta*cp*jump(theta*V(w),n)*avg(pibar)*dS_v - beta*cp*div(thetabar*w)*pi*dx + beta*cp*jump(thetabar*w,n)*avg(pi)*dS_v + (phi*(rho - rho_in) - beta*inner(grad(phi), u)*rhobar)*dx + beta*jump(phi*u, n)*avg(rhobar)*(dS_v + dS_h) ) if mu is not None: eqn += dt*mu*inner(w,k)*inner(u,k)*dx aeqn = lhs(eqn) Leqn = rhs(eqn) # Place to put result of u rho solver self.urho = Function(M) # Boundary conditions (assumes extruded mesh) dim = M.sub(0).ufl_element().value_shape()[0] bc = ("0.0",)*dim bcs = [DirichletBC(M.sub(0), Expression(bc), "bottom"), DirichletBC(M.sub(0), Expression(bc), "top")] # Solver for u, rho urho_problem = LinearVariationalProblem( aeqn, Leqn, self.urho, bcs=bcs) self.urho_solver = LinearVariationalSolver(urho_problem, solver_parameters=self.params, options_prefix='ImplicitSolver') # Reconstruction of theta theta = TrialFunction(state.V[2]) gamma = TestFunction(state.V[2]) u, rho = self.urho.split() self.theta = Function(state.V[2]) theta_eqn = gamma*(theta - theta_in + dot(k,u)*dot(k,grad(thetabar))*beta)*dx theta_problem = LinearVariationalProblem(lhs(theta_eqn), rhs(theta_eqn), self.theta) self.theta_solver = LinearVariationalSolver(theta_problem, options_prefix='thetabacksubstitution')
class HybridizationPC(SCBase): needs_python_pmat = True """A Slate-based python preconditioner that solves a mixed H(div)-conforming problem using hybridization. Currently, this preconditioner supports the hybridization of the RT and BDM mixed methods of arbitrary degree. The forward eliminations and backwards reconstructions are performed element-local using the Slate language. """ @timed_function("HybridInit") 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)) domain = "{[i,j]: 0 <= i < %d and 0 <= j < %d}" % shapes instructions = """ for i, j w[i,j] = w[i,j] + 1 end """ self.weight = Function(V[self.vidx]) par_loop((domain, instructions), ufl.dx, {"w": (self.weight, INC)}, is_loopy_kernel=True) instructions = """ for i, j vec_out[i,j] = vec_out[i,j] + vec_in[i,j]/w[i,j] end """ self.average_kernel = (domain, instructions) # 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 _reconstruction_calls(self, split_mixed_op, split_trace_op): """This generates the reconstruction calls for the unknowns using the Lagrange multipliers. :arg split_mixed_op: a ``dict`` of split forms that make up the broken mixed operator from the original problem. :arg split_trace_op: a ``dict`` of split forms that make up the trace contribution in the hybridized mixed system. """ from firedrake.assemble import create_assembly_callable # We always eliminate the velocity block first id0, id1 = (self.vidx, self.pidx) # TODO: When PyOP2 is able to write into mixed dats, # the reconstruction expressions can simplify into # one clean expression. A = Tensor(split_mixed_op[(id0, id0)]) B = Tensor(split_mixed_op[(id0, id1)]) C = Tensor(split_mixed_op[(id1, id0)]) D = Tensor(split_mixed_op[(id1, id1)]) K_0 = Tensor(split_trace_op[(0, id0)]) K_1 = Tensor(split_trace_op[(0, id1)]) # Split functions and reconstruct each bit separately split_residual = self.broken_residual.split() split_sol = self.broken_solution.split() g = AssembledVector(split_residual[id0]) f = AssembledVector(split_residual[id1]) sigma = split_sol[id0] u = split_sol[id1] lambdar = AssembledVector(self.trace_solution) M = D - C * A.inv * B R = K_1.T - C * A.inv * K_0.T u_rec = M.solve(f - C * A.inv * g - R * lambdar, decomposition="PartialPivLU") self._sub_unknown = create_assembly_callable(u_rec, tensor=u, form_compiler_parameters=self.ctx.fc_params) sigma_rec = A.solve(g - B * AssembledVector(u) - K_0.T * lambdar, decomposition="PartialPivLU") self._elim_unknown = create_assembly_callable(sigma_rec, tensor=sigma, form_compiler_parameters=self.ctx.fc_params) @timed_function("HybridUpdate") def update(self, pc): """Update by assembling into the operator. No need to reconstruct symbolic objects. """ self._assemble_S() self.S.force_evaluation() def forward_elimination(self, pc, x): """Perform the forward elimination of fields and provide the reduced right-hand side for the condensed system. :arg pc: a Preconditioner instance. :arg x: a PETSc vector containing the incoming right-hand side. """ with timed_region("HybridBreak"): with self.unbroken_residual.dat.vec_wo as v: x.copy(v) # Transfer unbroken_rhs into broken_rhs # NOTE: Scalar space is already "broken" so no need for # any projections unbroken_scalar_data = self.unbroken_residual.split()[self.pidx] broken_scalar_data = self.broken_residual.split()[self.pidx] unbroken_scalar_data.dat.copy(broken_scalar_data.dat) # Assemble the new "broken" hdiv residual # We need a residual R' in the broken space that # gives R'[w] = R[w] when w is in the unbroken space. # We do this by splitting the residual equally between # basis functions that add together to give unbroken # basis functions. unbroken_res_hdiv = self.unbroken_residual.split()[self.vidx] broken_res_hdiv = self.broken_residual.split()[self.vidx] broken_res_hdiv.assign(0) par_loop(self.average_kernel, ufl.dx, {"w": (self.weight, READ), "vec_in": (unbroken_res_hdiv, READ), "vec_out": (broken_res_hdiv, INC)}, is_loopy_kernel=True) with timed_region("HybridRHS"): # Compute the rhs for the multiplier system self._assemble_Srhs() def sc_solve(self, pc): """Solve the condensed linear system for the condensed field. :arg pc: a Preconditioner instance. """ # Solve the system for the Lagrange multipliers with self.schur_rhs.dat.vec_ro as b: if self.trace_ksp.getInitialGuessNonzero(): acc = self.trace_solution.dat.vec else: acc = self.trace_solution.dat.vec_wo with acc as x_trace: self.trace_ksp.solve(b, x_trace) def backward_substitution(self, pc, y): """Perform the backwards recovery of eliminated fields. :arg pc: a Preconditioner instance. :arg y: a PETSc vector for placing the resulting fields. """ # We assemble the unknown which is an expression # of the first eliminated variable. self._sub_unknown() # Recover the eliminated unknown self._elim_unknown() with timed_region("HybridProject"): # Project the broken solution into non-broken spaces broken_pressure = self.broken_solution.split()[self.pidx] unbroken_pressure = self.unbroken_solution.split()[self.pidx] broken_pressure.dat.copy(unbroken_pressure.dat) # Compute the hdiv projection of the broken hdiv solution broken_hdiv = self.broken_solution.split()[self.vidx] unbroken_hdiv = self.unbroken_solution.split()[self.vidx] unbroken_hdiv.assign(0) par_loop(self.average_kernel, ufl.dx, {"w": (self.weight, READ), "vec_in": (broken_hdiv, READ), "vec_out": (unbroken_hdiv, INC)}, is_loopy_kernel=True) with self.unbroken_solution.dat.vec_ro as v: v.copy(y) def view(self, pc, viewer=None): """Viewer calls for the various configurable objects in this PC.""" super(HybridizationPC, self).view(pc, viewer) viewer.pushASCIITab() viewer.printfASCII("Solves K * P^-1 * K.T using local eliminations.\n") viewer.printfASCII("KSP solver for the multipliers:\n") viewer.pushASCIITab() self.trace_ksp.view(viewer) viewer.popASCIITab() viewer.printfASCII("Locally reconstructing the broken solutions from the multipliers.\n") viewer.pushASCIITab() viewer.printfASCII("Project the broken hdiv solution into the HDiv space.\n") viewer.popASCIITab()