def residual(self, test, trial, trial_lagged, fields, bcs): u_adv = trial_lagged phi = test n = self.n u = trial F = -dot(u, div(outer(phi, u_adv)))*self.dx for id, bc in bcs.items(): if 'u' in bc: u_in = bc['u'] elif 'un' in bc: u_in = bc['un'] * n # this implies u_t=0 on the inflow else: u_in = zero(self.dim) F += conditional(dot(u_adv, n) < 0, dot(phi, u_in)*dot(u_adv, n), dot(phi, u)*dot(u_adv, n)) * self.ds(id) if not (is_continuous(self.trial_space) and normal_is_continuous(u_adv)): # s=0: u.n(-)<0 => flow goes from '+' to '-' => '+' is upwind # s=1: u.n(-)>0 => flow goes from '-' to '+' => '-' is upwind s = 0.5*(sign(dot(avg(u), n('-'))) + 1.0) u_up = u('-')*s + u('+')*(1-s) F += dot(u_up, (dot(u_adv('+'), n('+'))*phi('+') + dot(u_adv('-'), n('-'))*phi('-'))) * self.dS return -F
def incompressible_hydrostatic_balance(state, b0, p0, top=False, params=None): # get F Vu = state.spaces("HDiv") Vv = FunctionSpace(state.mesh, Vu.ufl_element()._elements[-1]) v = TrialFunction(Vv) w = TestFunction(Vv) if top: bstring = "top" else: bstring = "bottom" bcs = [DirichletBC(Vv, 0.0, bstring)] a = inner(w, v) * dx L = inner(state.k, w) * b0 * dx F = Function(Vv) solve(a == L, F, bcs=bcs) # define mixed function space VDG = state.spaces("DG") WV = (Vv) * (VDG) # get pprime v, pprime = TrialFunctions(WV) w, phi = TestFunctions(WV) bcs = [DirichletBC(WV[0], zero(), bstring)] a = (inner(w, v) + div(w) * pprime + div(v) * phi) * dx L = phi * div(F) * dx w1 = Function(WV) if params is None: params = { 'ksp_type': 'gmres', 'pc_type': 'fieldsplit', 'pc_fieldsplit_type': 'schur', 'pc_fieldsplit_schur_fact_type': 'full', 'pc_fieldsplit_schur_precondition': 'selfp', 'fieldsplit_1_ksp_type': 'preonly', 'fieldsplit_1_pc_type': 'gamg', 'fieldsplit_1_mg_levels_pc_type': 'bjacobi', 'fieldsplit_1_mg_levels_sub_pc_type': 'ilu', 'fieldsplit_0_ksp_type': 'richardson', 'fieldsplit_0_ksp_max_it': 4, 'ksp_atol': 1.e-08, 'ksp_rtol': 1.e-08 } solve(a == L, w1, bcs=bcs, solver_parameters=params) v, pprime = w1.split() p0.project(pprime)
def initialize(self, pc): _, P = pc.getOperators() assert P.type == "python" context = P.getPythonContext() (self.J, self.bcs) = (context.a, context.row_bcs) test, trial = self.J.arguments() if test.function_space() != trial.function_space(): raise NotImplementedError("test and trial spaces must be the same") Pk = test.function_space() element = Pk.ufl_element() shape = element.value_shape() mesh = Pk.ufl_domain() if len(shape) == 0: P1 = firedrake.FunctionSpace(mesh, "CG", 1) elif len(shape) == 2: P1 = firedrake.VectorFunctionSpace(mesh, "CG", 1, dim=shape[0]) else: P1 = firedrake.TensorFunctionSpace(mesh, "CG", 1, shape=shape, symmetry=element.symmetry()) # TODO: A smarter low-order operator would also interpolate # any coefficients to the coarse space. mapper = ArgumentReplacer({ test: firedrake.TestFunction(P1), trial: firedrake.TrialFunction(P1) }) self.lo_J = map_integrands.map_integrand_dags(mapper, self.J) lo_bcs = [] for bc in self.bcs: # Don't actually need the value, since it's only used for # killing parts of the restriction matrix. lo_bcs.append( firedrake.DirichletBC(P1, firedrake.zero(P1.shape), bc.sub_domain, method=bc.method)) self.lo_bcs = tuple(lo_bcs) mat_type = PETSc.Options().getString( pc.getOptionsPrefix() + "lo_mat_type", firedrake.parameters["default_matrix_type"]) self.lo_op = firedrake.assemble(self.lo_J, bcs=self.lo_bcs, mat_type=mat_type) self.lo_op.force_evaluation() A, P = pc.getOperators() nearnullsp = P.getNearNullSpace() if nearnullsp.handle != 0: # Actually have a near nullspace tmp = firedrake.Function(Pk) low = firedrake.Function(P1) vecs = [] for vec in nearnullsp.getVecs(): with tmp.dat.vec as v: vec.copy(v) low.interpolate(tmp) with low.dat.vec_ro as v: vecs.append(v.copy()) nullsp = PETSc.NullSpace().create(vectors=vecs, comm=pc.comm) self.lo_op.petscmat.setNearNullSpace(nullsp) lo = PETSc.PC().create(comm=pc.comm) lo.incrementTabLevel(1, parent=pc) lo.setOperators(self.lo_op.petscmat, self.lo_op.petscmat) lo.setOptionsPrefix(pc.getOptionsPrefix() + "lo_") lo.setFromOptions() self.lo = lo self.restriction = restriction_matrix(Pk, P1, self.bcs, self.lo_bcs) self.work = self.lo_op.petscmat.createVecs() if len(self.bcs) > 0: bc_nodes = numpy.unique( numpy.concatenate([bc.nodes for bc in self.bcs])) bc_nodes = bc_nodes[bc_nodes < Pk.dof_dset.size] bc_iset = PETSc.IS().createBlock(numpy.prod(shape), bc_nodes, comm=PETSc.COMM_SELF) self.bc_indices = bc_iset.getIndices() bc_iset.destroy() else: self.bc_indices = numpy.empty(0, dtype=numpy.int32)
def coarsen(self, fdm, comm): fctx = get_appctx(fdm) test, trial = fctx.J.arguments() fV = test.function_space() fu = fctx._problem.u cele = self.coarsen_element(fV.ufl_element()) cV = firedrake.FunctionSpace(fV.mesh(), cele) cdm = cV.dm cu = firedrake.Function(cV) interpolators = tuple( firedrake.Interpolator(fus, cus) for fus, cus in zip(fu.split(), cu.split())) def inject_state(interpolators): for interpolator in interpolators: interpolator.interpolate() parent = get_parent(fdm) assert parent is not None add_hook(parent, setup=partial(push_parent, cdm, parent), teardown=partial(pop_parent, cdm, parent), call_setup=True) replace_d = { fu: cu, test: firedrake.TestFunction(cV), trial: firedrake.TrialFunction(cV) } cJ = replace(fctx.J, replace_d) cF = replace(fctx.F, replace_d) if fctx.Jp is not None: cJp = replace(fctx.Jp, replace_d) else: cJp = None cbcs = [] for bc in fctx._problem.bcs: # Don't actually need the value, since it's only used for # killing parts of the matrix. This should be generalised # for p-FAS, if anyone ever wants to do that cV_ = cV for index in bc._indices: cV_ = cV_.sub(index) cbcs.append( firedrake.DirichletBC(cV_, firedrake.zero(cV_.shape), bc.sub_domain, method=bc.method)) fcp = fctx._problem.form_compiler_parameters cproblem = firedrake.NonlinearVariationalProblem( cF, cu, cbcs, cJ, Jp=cJp, form_compiler_parameters=fcp, is_linear=fctx._problem.is_linear) cctx = _SNESContext( cproblem, fctx.mat_type, fctx.pmat_type, appctx=fctx.appctx, pre_jacobian_callback=fctx._pre_jacobian_callback, pre_function_callback=fctx._pre_function_callback, post_jacobian_callback=fctx._post_jacobian_callback, post_function_callback=fctx._post_function_callback, options_prefix=fctx.options_prefix, transfer_manager=fctx.transfer_manager) add_hook(parent, setup=partial(push_appctx, cdm, cctx), teardown=partial(pop_appctx, cdm, cctx), call_setup=True) add_hook(parent, setup=partial(inject_state, interpolators), call_setup=True) cdm.setKSPComputeOperators(_SNESContext.compute_operators) cdm.setCreateInterpolation(self.create_interpolation) cdm.setOptionsPrefix(fdm.getOptionsPrefix()) # If we're the coarsest grid of the p-hierarchy, don't # overwrite the coarsen routine; this is so that you can # use geometric multigrid for the p-coarse problem try: self.coarsen_element(cele) cdm.setCoarsen(self.coarsen) except ValueError: pass return cdm
def compressible_hydrostatic_balance(state, theta0, rho0, exner0=None, top=False, exner_boundary=Constant(1.0), mr_t=None, solve_for_rho=False, params=None): """ Compute a hydrostatically balanced density given a potential temperature profile. By default, this uses a vertically-oriented hybridization procedure for solving the resulting discrete systems. :arg state: The :class:`State` object. :arg theta0: :class:`.Function`containing the potential temperature. :arg rho0: :class:`.Function` to write the initial density into. :arg top: If True, set a boundary condition at the top. Otherwise, set it at the bottom. :arg exner_boundary: a field or expression to use as boundary data for exner on the top or bottom as specified. :arg mr_t: the initial total water mixing ratio field. """ # Calculate hydrostatic Pi VDG = state.spaces("DG") Vu = state.spaces("HDiv") Vv = FunctionSpace(state.mesh, Vu.ufl_element()._elements[-1]) W = MixedFunctionSpace((Vv, VDG)) v, exner = TrialFunctions(W) dv, dexner = TestFunctions(W) n = FacetNormal(state.mesh) cp = state.parameters.cp # add effect of density of water upon theta theta = theta0 if mr_t is not None: theta = theta0 / (1 + mr_t) alhs = ((cp * inner(v, dv) - cp * div(dv * theta) * exner) * dx + dexner * div(theta * v) * dx) if top: bmeasure = ds_t bstring = "bottom" else: bmeasure = ds_b bstring = "top" arhs = -cp * inner(dv, n) * theta * exner_boundary * bmeasure # Possibly make g vary with spatial coordinates? g = state.parameters.g arhs -= g * inner(dv, state.k) * dx bcs = [DirichletBC(W.sub(0), zero(), bstring)] w = Function(W) exner_problem = LinearVariationalProblem(alhs, arhs, w, bcs=bcs) if params is None: params = { 'ksp_type': 'preonly', 'pc_type': 'python', 'mat_type': 'matfree', 'pc_python_type': 'gusto.VerticalHybridizationPC', # Vertical trace system is only coupled vertically in columns # block ILU is a direct solver! 'vert_hybridization': { 'ksp_type': 'preonly', 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu' } } exner_solver = LinearVariationalSolver(exner_problem, solver_parameters=params, options_prefix="exner_solver") exner_solver.solve() v, exner = w.split() if exner0 is not None: exner0.assign(exner) if solve_for_rho: w1 = Function(W) v, rho = w1.split() rho.interpolate(thermodynamics.rho(state.parameters, theta0, exner)) v, rho = split(w1) dv, dexner = TestFunctions(W) exner = thermodynamics.exner_pressure(state.parameters, rho, theta0) F = ((cp * inner(v, dv) - cp * div(dv * theta) * exner) * dx + dexner * div(theta0 * v) * dx + cp * inner(dv, n) * theta * exner_boundary * bmeasure) F += g * inner(dv, state.k) * dx rhoproblem = NonlinearVariationalProblem(F, w1, bcs=bcs) rhosolver = NonlinearVariationalSolver(rhoproblem, solver_parameters=params, options_prefix="rhosolver") rhosolver.solve() v, rho_ = w1.split() rho0.assign(rho_) else: rho0.interpolate(thermodynamics.rho(state.parameters, theta0, exner))
def coarsen_bc_value(self, bc, cV): return firedrake.zero(cV.shape)
def initialize(self, pc): _, P = pc.getOperators() assert P.type == "python" context = P.getPythonContext() (self.J, self.bcs) = (context.a, context.row_bcs) test, trial = self.J.arguments() if test.function_space() != trial.function_space(): raise NotImplementedError("test and trial spaces must be the same") Pk = test.function_space() element = Pk.ufl_element() shape = element.value_shape() mesh = Pk.ufl_domain() if len(shape) == 0: P1 = firedrake.FunctionSpace(mesh, "CG", 1) elif len(shape) == 1: P1 = firedrake.VectorFunctionSpace(mesh, "CG", 1, dim=shape[0]) else: P1 = firedrake.TensorFunctionSpace(mesh, "CG", 1, shape=shape, symmetry=element.symmetry()) # TODO: A smarter low-order operator would also interpolate # any coefficients to the coarse space. mapper = ArgumentReplacer({test: firedrake.TestFunction(P1), trial: firedrake.TrialFunction(P1)}) self.lo_J = map_integrands.map_integrand_dags(mapper, self.J) lo_bcs = [] for bc in self.bcs: # Don't actually need the value, since it's only used for # killing parts of the restriction matrix. lo_bcs.append(firedrake.DirichletBC(P1, firedrake.zero(P1.shape), bc.sub_domain, method=bc.method)) self.lo_bcs = tuple(lo_bcs) mat_type = PETSc.Options().getString(pc.getOptionsPrefix() + "lo_mat_type", firedrake.parameters["default_matrix_type"]) self.lo_op = firedrake.assemble(self.lo_J, bcs=self.lo_bcs, mat_type=mat_type) self.lo_op.force_evaluation() A, P = pc.getOperators() nearnullsp = P.getNearNullSpace() if nearnullsp.handle != 0: # Actually have a near nullspace tmp = firedrake.Function(Pk) low = firedrake.Function(P1) vecs = [] for vec in nearnullsp.getVecs(): with tmp.dat.vec as v: vec.copy(v) low.interpolate(tmp) with low.dat.vec_ro as v: vecs.append(v.copy()) nullsp = PETSc.NullSpace().create(vectors=vecs, comm=pc.comm) self.lo_op.petscmat.setNearNullSpace(nullsp) lo = PETSc.PC().create(comm=pc.comm) lo.incrementTabLevel(1, parent=pc) lo.setOperators(self.lo_op.petscmat, self.lo_op.petscmat) lo.setOptionsPrefix(pc.getOptionsPrefix() + "lo_") lo.setFromOptions() self.lo = lo self.restriction = restriction_matrix(Pk, P1, self.bcs, self.lo_bcs) self.work = self.lo_op.petscmat.createVecs() if len(self.bcs) > 0: bc_nodes = numpy.unique(numpy.concatenate([bc.nodes for bc in self.bcs])) bc_nodes = bc_nodes[bc_nodes < Pk.dof_dset.size] bc_iset = PETSc.IS().createBlock(numpy.prod(shape), bc_nodes, comm=PETSc.COMM_SELF) self.bc_indices = bc_iset.getIndices() bc_iset.destroy() else: self.bc_indices = numpy.empty(0, dtype=numpy.int32)