def apply_restrictions(expression): "Propagate restriction nodes to wrap differential terminals directly." integral_types = [k for k in integral_type_to_measure_name.keys() if k.startswith("interior_facet")] rules = RestrictionPropagator() return map_integrand_dags(rules, expression, only_integral_type=integral_types)
def replace(e, mapping): """Replace subexpressions in expression. @param e: An Expr or Form. @param mapping: A dict with from:to replacements to perform. """ mapping2 = dict((k, as_ufl(v)) for (k, v) in mapping.items()) # Workaround for problem with delayed derivative evaluation # The problem is that J = derivative(f(g, h), g) does not evaluate immediately # So if we subsequently do replace(J, {g: h}) we end up with an expression: # derivative(f(h, h), h) # rather than what were were probably thinking of: # replace(derivative(f(g, h), g), {g: h}) # # To fix this would require one to expand derivatives early (which # is not attractive), or make replace lazy too. if has_exact_type(e, CoefficientDerivative): # Hack to avoid circular dependencies from ufl.algorithms.ad import expand_derivatives e = expand_derivatives(e) return map_integrand_dags(Replacer(mapping2), e)
def coarsen_form(form): """Return a coarse mesh version of a form :arg form: The :class:`~ufl.classes.Form` to coarsen. This maps over the form and replaces coefficients and arguments with their coarse mesh equivalents.""" if form is None: return None assert isinstance(form, ufl.Form), \ "Don't know how to coarsen %r" % type(form) mapper = CoarsenIntegrand() integrals = sum_integrands(form).integrals() forms = [] # Ugh, visitors can't deal with measures (they're not actual # Exprs) so we need to map the transformer over the integrand and # reconstruct the integral by building the measure by hand. for it in integrals: integrand = map_integrand_dags(mapper, it.integrand()) mesh = it.ufl_domain() hierarchy, level = utils.get_level(mesh) new_mesh = hierarchy[level - 1] measure = ufl.Measure(it.integral_type(), domain=new_mesh, subdomain_id=it.subdomain_id(), subdomain_data=it.subdomain_data(), metadata=it.metadata()) forms.append(integrand * measure) return reduce(add, forms)
def coarsen_form(form): """Return a coarse mesh version of a form :arg form: The :class:`ufl.Form` to coarsen. This maps over the form and replaces coefficients and arguments with their coarse mesh equivalents.""" if form is None: return None assert isinstance(form, ufl.Form), \ "Don't know how to coarsen %r" % type(form) mapper = CoarsenIntegrand() integrals = sum_integrands(form).integrals() forms = [] # Ugh, visitors can't deal with measures (they're not actual # Exprs) so we need to map the transformer over the integrand and # reconstruct the integral by building the measure by hand. for it in integrals: integrand = map_integrand_dags(mapper, it.integrand()) mesh = it.ufl_domain() hierarchy, level = utils.get_level(mesh) new_mesh = hierarchy[level-1] measure = ufl.Measure(it.integral_type(), domain=new_mesh, subdomain_id=it.subdomain_id(), subdomain_data=it.subdomain_data(), metadata=it.metadata()) forms.append(integrand * measure) return reduce(add, forms)
def remove_complex_nodes(expr): """Replaces complex operator nodes with their children. This is called during compute_form_data if the compiler wishes to compile real-valued forms. In essence this strips all trace of complex support from the preprocessed form. """ return map_integrand_dags(ComplexNodeRemoval(), expr)
def split(self, form): """Split the form. :arg form: the form to split. This is a no-op if none of the arguments in the form are defined on :class:`~.MixedFunctionSpace`\s. The return-value is a tuple for which each entry is. .. code-block:: python (argument_indices, form) Where ``argument_indices`` is a tuple indicating which part of the mixed space the form belongs to, it has length equal to the number of arguments in the form. Hence functionals have a 0-tuple, 1-forms have a 1-tuple and 2-forms a 2-tuple of indices. For example, consider the following code: .. code-block:: python V = FunctionSpace(m, 'CG', 1) W = V*V*V u, v, w = TrialFunctions(W) p, q, r = TestFunctions(W) a = q*u*dx + p*w*dx Then splitting the form returns a tuple of two forms. .. code-block:: python ((0, 2), w*p*dx), (1, 0), q*u*dx)) """ args = form.arguments() if all(len(a.function_space()) == 1 for a in args): # No mixed spaces, just return the form directly. idx = tuple([0]*len(form.arguments())) return (SplitForm(indices=idx, form=form), ) forms = [] # How many subspaces do we have for each argument? shape = tuple(len(a.function_space()) for a in args) # Walk over all the indices of the spaces for idx in numpy.ndindex(shape): # Which subspace are we currently interested in? self.idx = dict(enumerate(idx)) # Cache for the arguments we construct self._args = {} # Visit the form f = map_integrand_dags(self, form) # Zero-simplification may result in an empty form, only # collect those that are non-zero. if len(f.integrals()) > 0: forms.append(SplitForm(indices=idx, form=f)) return tuple(forms)
def applyRestrictions(form): integral_types = [ k for k in integral_type_to_measure_name.keys() if k.startswith("interior_facet") ] return map_integrand_dags(RestrictionPropagator(), form, only_integral_type=integral_types)
def __init__(self, form): """Constructor for the Tensor class.""" if not isinstance(form, Form): raise ValueError("Only UFL forms are acceptable inputs for creating terminal tensors.") r = len(form.arguments()) if r not in (0, 1, 2): raise NotImplementedError("Currently don't support tensors of rank %d." % r) # Checks for positive restrictions on integrals integrals = form.integrals() mapper = CheckRestrictions() for it in integrals: map_integrand_dags(mapper, it) super(Tensor, self).__init__() self.form = form
def __init__(self, form): """Constructor for the Tensor class.""" if not isinstance(form, Form): raise ValueError("Only UFL forms are acceptable inputs.") r = len(form.arguments()) if r not in (0, 1, 2): raise NotImplementedError("No support for tensors of rank %d." % r) # Checks for positive restrictions on integrals integrals = form.integrals() mapper = CheckRestrictions() for it in integrals: map_integrand_dags(mapper, it) super(Tensor, self).__init__() self.form = form
def apply_default_restrictions(expression): """Some terminals can be restricted from either side. This applies a default restriction to such terminals if unrestricted.""" integral_types = [k for k in integral_type_to_measure_name.keys() if k.startswith("interior_facet")] rules = DefaultRestrictionApplier() return map_integrand_dags(rules, expression, only_integral_type=integral_types)
def apply_function_pullbacks(expr): """Change representation of coefficients and arguments in expression by applying Piola mappings where applicable and representing all form arguments in reference value. @param expr: An Expr. """ return map_integrand_dags(FunctionPullbackApplier(), expr)
def component_tensor(self, o, *dops): assert o not in self._mapping assert len(dops) == 2 assert isinstance(dops[1], MultiIndex) if dops[0] in self._mapping: replaced_ufl_operand_0 = self._mapping[dops[0]] else: replaced_ufl_operand_0 = map_integrand_dags(self, dops[0]) return ComponentTensor(replaced_ufl_operand_0, dops[1])
def list_tensor(self, o, *dops): assert o not in self._mapping replaced_ufl_operands = list() for ufl_operand in dops: if ufl_operand in self._mapping: replaced_ufl_operands.append(self._mapping[ufl_operand]) else: replaced_ufl_operands.append( map_integrand_dags(self, ufl_operand)) return ListTensor(*replaced_ufl_operands)
def apply_restrictions(expression): "Propagate restriction nodes to wrap differential terminals directly." integral_types = [ k for k in integral_type_to_measure_name.keys() if k.startswith("interior_facet") ] rules = RestrictionPropagator() return map_integrand_dags(rules, expression, only_integral_type=integral_types)
def indexed(self, o, *dops): if o in self._mapping: return self._mapping[o] else: assert len(dops) == 2 assert isinstance(dops[1], MultiIndex) if dops[0] in self._mapping: replaced_ufl_operand_0 = self._mapping[dops[0]] else: replaced_ufl_operand_0 = map_integrand_dags(self, dops[0]) return Indexed(replaced_ufl_operand_0, dops[1])
def apply_default_restrictions(expression): """Some terminals can be restricted from either side. This applies a default restriction to such terminals if unrestricted.""" integral_types = [ k for k in integral_type_to_measure_name.keys() if k.startswith("interior_facet") ] rules = DefaultRestrictionApplier() return map_integrand_dags(rules, expression, only_integral_type=integral_types)
def _transform_and_attach_multi_index(self, Class, o, *dops): if o not in self.ufl_to_replaced_ufl: self._store_sympy_symbol(o) assert len(dops) == 2 assert isinstance(dops[1], MultiIndex) transformed_ufl_operand_0 = list() split_sum(map_integrand_dags(self, dops[0]), transformed_ufl_operand_0) new_o = sum([Class(operand, dops[1]) for operand in transformed_ufl_operand_0]) self._update_sympy_symbol(o, new_o) self.ufl_to_replaced_ufl[o] = new_o return new_o else: return self.ufl_to_replaced_ufl[o]
def evaluate(e, u, u0): """Evaluate a UFL expression (or list of expressions). Basically replaces function(s) u by function(s) u0. Parameters ---------- e: UFL Expr or list of UFL expressions/forms u: UFL Function or list of UFL functions u0: UFL Function or list of UFL functions Returns ------- Expression (or list of expressions) with function(s) u replaced by function(s) u0. """ from ufl.algorithms.map_integrands import map_integrand_dags from ufl.algorithms.replace import Replacer if isinstance(u, list) and isinstance(u0, list): repmap = {v: v0 for v, v0 in zip(u, u0)} elif not isinstance(u, list) and not isinstance(u0, list): repmap = {u: u0} else: raise RuntimeError( "Incompatible functions (list-of-functions and function) provided." ) replacer = Replacer(repmap) if isinstance(e, list): e0 = [] for e_ in e: e0.append(map_integrand_dags(replacer, e_)) else: e0 = map_integrand_dags(replacer, e) return e0
def replace(e, mapping): """Replace objects in expression. @param e: An Expr or Form. @param mapping: A dict with from:to replacements to perform. """ mapping2 = dict((k, as_ufl(v)) for (k, v) in mapping.items()) # We have expanded derivative evaluation in ParametrizedTensorFactory assert not has_exact_type(e, CoefficientDerivative) return map_integrand_dags(Replacer(mapping2), e)
def split(self, form, ix=0, iy=0): # Remember which block to extract self.idx = [ix, iy] args = form.arguments() if len(args) == 0: # Functional can't be split return form if all(len(a.function_space()) == 1 for a in args): # not mixed, just return self # SHOULD/CAN THESE BE RESTORED? # WHAT EXACTLY WERE THEY CHECKING? # assert (len(idx) == 1 for idx in self.blocks.values()) # assert (idx[0] == 0 for idx in self.blocks.values()) return form return map_integrand_dags(self, form)
def replace(e, mapping): """Replace terminal objects in expression. @param e: An Expr or Form. @param mapping: A dict with from:to replacements to perform. """ mapping2 = dict((k, as_ufl(v)) for (k, v) in mapping.items()) # Workaround for problem with delayed derivative evaluation if has_exact_type(e, CoefficientDerivative): # Hack to avoid circular dependencies from ufl.algorithms.ad import expand_derivatives e = expand_derivatives(e) return map_integrand_dags(Replacer(mapping2), e)
def replace(e, mapping): """Replace terminal objects in expression. @param e: An Expr or Form. @param mapping: A dict with from:to replacements to perform. """ mapping2 = dict((k, as_ufl(v)) for (k, v) in mapping.items()) # Workaround for problem with delayed derivative evaluation if has_exact_type(e, CoefficientDerivative): # Hack to avoid circular dependencies from ufl.algorithms.ad import expand_derivatives e = expand_derivatives(e) return map_integrand_dags(Replacer(mapping2), e)
def expand_sum_product(form): form = remove_complex_nodes(form) # TODO support forms in the complex field. This is currently needed otherwise conj((a+b)*c) does not get expanded. # Patch Expr.__mul__ and Expr.__rmul__ patch_expr_mul() # Call sympy replacer expanded_form = map_integrand_dags(SympyExpander(), form) # Split sums expanded_split_form_integrals = list() for integral in expanded_form.integrals(): expanded_split_form_integrands = list() split_sum(integral.integrand(), expanded_split_form_integrands) expanded_split_form_integrals.extend([integral.reconstruct(integrand=integrand) for integrand in expanded_split_form_integrands]) expanded_split_form = Form(expanded_split_form_integrals) # Undo patch to Expr.__mul__ and Expr.__rmul__ unpatch_expr_mul() # Return return expanded_split_form
def __init__(self, form): """Constructor for the Tensor class.""" if not isinstance(form, Form): if isinstance(form, Function): raise TypeError("Use AssembledVector instead of Tensor.") raise TypeError("Only UFL forms are acceptable inputs.") r = len(form.arguments()) if r not in (0, 1, 2): raise NotImplementedError("No support for tensors of rank %d." % r) # Remove any negative restrictions and replace with zero form = map_integrand_dags(RemoveNegativeRestrictions(), form) super(Tensor, self).__init__() self.form = form
def __init__(self, form): """Constructor for the Tensor class.""" if not isinstance(form, Form): if isinstance(form, Function): raise TypeError("Use AssembledVector instead of Tensor.") raise TypeError("Only UFL forms are acceptable inputs.") r = len(form.arguments()) if r not in (0, 1, 2): raise NotImplementedError("No support for tensors of rank %d." % r) # Remove any negative restrictions and replace with zero form = map_integrand_dags(RemoveNegativeRestrictions(), form) super(Tensor, self).__init__() self.form = form
def expand_sum_product(form): # Patch Expr.__mul__ and Expr.__rmul__ patch_expr_mul() # Call sympy replacer expanded_form = map_integrand_dags(SympyExpander(), form) # Split sums expanded_split_form_integrals = list() for integral in expanded_form.integrals(): expanded_split_form_integrands = list() split_sum(integral.integrand(), expanded_split_form_integrands) expanded_split_form_integrals.extend([ integral.reconstruct(integrand=integrand) for integrand in expanded_split_form_integrands ]) expanded_split_form = Form(expanded_split_form_integrals) # Undo patch to Expr.__mul__ and Expr.__rmul__ unpatch_expr_mul() # Return return expanded_split_form
def coarsen_bc(bc): new_V = coarsen_thing(bc.function_space()) val = bc._original_val zeroed = bc._currently_zeroed subdomain = bc.sub_domain method = bc.method new_val = val if isinstance(val, firedrake.Expression): new_val = val if isinstance(val, (firedrake.Constant, firedrake.Function)): mapper = CoarsenIntegrand() new_val = map_integrand_dags(mapper, val) new_bc = firedrake.DirichletBC(new_V, new_val, subdomain, method=method) if zeroed: new_bc.homogenize() return new_bc
def coarsen_bc(bc): new_V = coarsen_thing(bc.function_space()) val = bc._original_val zeroed = bc._currently_zeroed subdomain = bc.sub_domain method = bc.method new_val = val if isinstance(val, firedrake.Expression): new_val = val if isinstance(val, (firedrake.Constant, firedrake.Function)): mapper = CoarsenIntegrand() new_val = map_integrand_dags(mapper, val) new_bc = firedrake.DirichletBC(new_V, new_val, subdomain, method=method) if zeroed: new_bc.homogenize() return new_bc
def lower_form_order(form): # get test and trial functions of (possibly mixed) form testfunc, trialfunc = form.arguments() mesh = testfunc.ufl_domain() testelem = testfunc.ufl_element() trialelem = trialfunc.ufl_element() assert(testelem == trialelem) elem_lowest = lower_element_order(testelem) ismixed = False if isinstance(testelem, MixedElement): ismixed = True if ismixed: spacelist = [] for elem in elem_lowest: spacelist.append(FunctionSpace(mesh, elem)) space_lowest = MixedFunctionSpace(spacelist) testfuncs_lowest = TestFunctions(space_lowest) testfuncs = TestFunctions(testfunc.ufl_function_space()) trialfuncs = TrialFunctions(trialfunc.ufl_function_space()) rdict = {} for trialfunc, trialfunc_lowest in zip(trialfuncs, testfuncs_lowest): rdict[trialfunc] = trialfunc_lowest for testfunc, testfunc_lowest in zip(testfuncs, testfuncs_lowest): rdict[testfunc] = testfunc_lowest replacer = Replacer(rdict) else: space_lowest = FunctionSpace(mesh, elem_lowest) trialfunc_lowest = TrialFunction(space_lowest) testfunc_lowest = TestFunction(space_lowest) replacer = Replacer({testfunc: testfunc_lowest, trialfunc: trialfunc_lowest}) newform = map_integrand_dags(replacer, form) return newform
def split(self, form, argument_indices): """Split a form. :arg form: the form to split. :arg argument_indices: indices of test and trial spaces to extract. This should be 0-, 1-, or 2-tuple (whose length is the same as the number of arguments as the ``form``) whose entries are either an integer index, or else an iterable of indices. Returns a new :class:`ufl.classes.Form` on the selected subspace. """ args = form.arguments() self._arg_cache = {} self.blocks = dict(zip((0, 1), argument_indices)) if len(args) == 0: # Functional can't be split return form if all(len(a.function_space()) == 1 for a in args): assert (len(idx) == 1 for idx in self.blocks.values()) assert (idx[0] == 0 for idx in self.blocks.values()) return form f = map_integrand_dags(self, form) return f
def split(self, form, argument_indices): """Split a form. :arg form: the form to split. :arg argument_indices: indices of test and trial spaces to extract. This should be 0-, 1-, or 2-tuple (whose length is the same as the number of arguments as the ``form``) whose entries are either an integer index, or else an iterable of indices. Returns a new :class:`ufl.classes.Form` on the selected subspace. """ args = form.arguments() self._arg_cache = {} self.blocks = dict(enumerate(argument_indices)) if len(args) == 0: # Functional can't be split return form if all(len(a.function_space()) == 1 for a in args): assert (len(idx) == 1 for idx in self.blocks.values()) assert (idx[0] == 0 for idx in self.blocks.values()) return form f = map_integrand_dags(self, form) return f
def split(self, form): """Split the form. :arg form: the form to split. This is a no-op if none of the arguments in the form are defined on :class:`~.MixedFunctionSpace`\s. The return-value is a tuple for which each entry is. .. code-block:: python (argument_indices, form) Where ``argument_indices`` is a tuple indicating which part of the mixed space the form belongs to, it has length equal to the number of arguments in the form. Hence functionals have a 0-tuple, 1-forms have a 1-tuple and 2-forms a 2-tuple of indices. For example, consider the following code: .. code-block:: python V = FunctionSpace(m, 'CG', 1) W = V*V*V u, v, w = TrialFunctions(W) p, q, r = TestFunctions(W) a = q*u*dx + p*w*dx Then splitting the form returns a tuple of two forms. .. code-block:: python ((0, 2), w*p*dx), (1, 0), q*u*dx)) """ from ufl.algorithms.map_integrands import map_integrand_dags from numpy import ndindex args = form.arguments() if all(a.function_space().num_sub_spaces() == 0 for a in args): # No mixed spaces, just return the form directly. idx = tuple([0]*len(form.arguments())) return (SplitForm(indices=idx, form=form), ) forms = [] # How many subspaces do we have for each argument? shape = tuple(max(1, a.function_space().num_sub_spaces()) for a in args) # Walk over all the indices of the spaces for idx in ndindex(shape): # Which subspace are we currently interested in? self.idx = dict(enumerate(idx)) # Cache for the arguments we construct self._args = {} # Visit the form f = map_integrand_dags(self, form) # Zero-simplification may result in an empty form, only # collect those that are non-zero. if len(f.integrals()) > 0: forms.append(SplitForm(indices=idx, form=f)) return tuple(forms)
def apply_time_derivatives(expression, t, timedep_coeffs=[]): rules = TimeDerivativeRuleDispatcher(t, timedep_coeffs) return map_integrand_dags(rules, expression)
def do_comparison_check(form): """Raises an error if invalid comparison nodes exist""" return map_integrand_dags(CheckComparisons(), form)
def split(self, form, ix, iy=0): # Remember which block to extract self.idx = [ix, iy] return map_integrand_dags(self, form)
def apply_coordinate_derivatives(expression): rules = CoordinateDerivativeRuleDispatcher() return map_integrand_dags(rules, expression)
def split(self, form, ix, iy=0): # Remember which block to extract self.idx = [ix, iy] return map_integrand_dags(self, form)
def apply_algebra_lowering(expr): """Expands high level compound operators (e.g. inner) to equivalent representations using basic operators (e.g. index notation).""" return map_integrand_dags(LowerCompoundAlgebra(), expr)
def do_comparison_check(form): """Raises an error if invalid comparison nodes exist""" return map_integrand_dags(CheckComparisons(), form)
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 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)
def apply_algebra_lowering(expr): """Expands high level compound operators (e.g. inner) to equivalent representations using basic operators (e.g. index notation).""" return map_integrand_dags(LowerCompoundAlgebra(), expr)
def apply_coordinate_derivatives(expression): rules = CoordinateDerivativeRuleDispatcher() return map_integrand_dags(rules, expression)
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 ufl.algorithms.map_integrands import map_integrand_dags from firedrake import (FunctionSpace, TrialFunction, TrialFunctions, TestFunction, Function, BrokenElement, MixedElement, FacetNormal, Constant, DirichletBC, Projector) from firedrake.assemble import (allocate_matrix, create_assembly_callable) from firedrake.formmanipulation import ArgumentReplacer, split_form # Extract the problem context prefix = pc.getOptionsPrefix() _, P = pc.getOperators() context = P.getPythonContext() test, trial = context.a.arguments() V = test.function_space() if V.mesh().cell_set._extruded: # TODO: Merge FIAT branch to support TPC trace elements raise NotImplementedError("Not implemented on extruded meshes.") # Break the function spaces and define fully discontinuous spaces broken_elements = [BrokenElement(Vi.ufl_element()) for Vi in V] elem = MixedElement(broken_elements) V_d = FunctionSpace(V.mesh(), elem) arg_map = {test: TestFunction(V_d), trial: TrialFunction(V_d)} # Replace the problems arguments with arguments defined # on the new discontinuous spaces replacer = ArgumentReplacer(arg_map) new_form = map_integrand_dags(replacer, context.a) # Create the space of approximate traces. # The vector function space will have a non-empty value_shape W = next(v for v in V if bool(v.ufl_element().value_shape())) if W.ufl_element().family() in ["Raviart-Thomas", "RTCF"]: tdegree = W.ufl_element().degree() - 1 else: tdegree = W.ufl_element().degree() # NOTE: Once extruded is ready, we will need to be aware of this # and construct the appropriate trace space for the HDiv element TraceSpace = FunctionSpace(V.mesh(), "HDiv Trace", tdegree) # NOTE: For extruded, we will need to add "on_top" and "on_bottom" trace_conditions = [ DirichletBC(TraceSpace, Constant(0.0), "on_boundary") ] # Set up the functions for the original, hybridized # and schur complement systems self.broken_solution = Function(V_d) self.broken_rhs = Function(V_d) self.trace_solution = Function(TraceSpace) self.unbroken_solution = Function(V) self.unbroken_rhs = Function(V) # Create the symbolic Schur-reduction Atilde = Tensor(new_form) gammar = TestFunction(TraceSpace) n = FacetNormal(V.mesh()) # Vector trial function will have a non-empty ufl_shape sigma = next(f for f in TrialFunctions(V_d) if bool(f.ufl_shape)) # NOTE: Once extruded is ready, this will change slightly # to include both horizontal and vertical interior facets K = Tensor(gammar('+') * ufl.dot(sigma, n) * ufl.dS) # Assemble the Schur complement operator and right-hand side self.schur_rhs = Function(TraceSpace) self._assemble_Srhs = create_assembly_callable( K * Atilde.inv * self.broken_rhs, tensor=self.schur_rhs, form_compiler_parameters=context.fc_params) schur_comp = K * Atilde.inv * K.T self.S = allocate_matrix(schur_comp, bcs=trace_conditions, form_compiler_parameters=context.fc_params) self._assemble_S = create_assembly_callable( schur_comp, tensor=self.S, bcs=trace_conditions, form_compiler_parameters=context.fc_params) self._assemble_S() self.S.force_evaluation() Smat = self.S.petscmat # Nullspace for the multiplier problem nullsp = P.getNullSpace() if nullsp.handle != 0: new_vecs = get_trace_nullspace_vecs(K * Atilde.inv, nullsp, V, V_d, TraceSpace) tr_nullsp = PETSc.NullSpace().create(vectors=new_vecs, comm=pc.comm) Smat.setNullSpace(tr_nullsp) # Set up the KSP for the system of Lagrange multipliers ksp = PETSc.KSP().create(comm=pc.comm) ksp.setOptionsPrefix(prefix + "trace_") ksp.setTolerances(rtol=1e-13) ksp.setOperators(Smat) ksp.setUp() ksp.setFromOptions() self.ksp = ksp # Now we construct the local tensors for the reconstruction stage # TODO: Add support for mixed tensors and these variables # become unnecessary split_forms = split_form(new_form) A = Tensor(next(sf.form for sf in split_forms if sf.indices == (0, 0))) B = Tensor(next(sf.form for sf in split_forms if sf.indices == (1, 0))) C = Tensor(next(sf.form for sf in split_forms if sf.indices == (1, 1))) trial = TrialFunction( FunctionSpace(V.mesh(), BrokenElement(W.ufl_element()))) K_local = Tensor(gammar('+') * ufl.dot(trial, n) * ufl.dS) # Split functions and reconstruct each bit separately sigma_h, u_h = self.broken_solution.split() g, f = self.broken_rhs.split() # Pressure reconstruction M = B * A.inv * B.T + C u_sol = M.inv * f + M.inv * ( B * A.inv * K_local.T * self.trace_solution - B * A.inv * g) self._assemble_pressure = create_assembly_callable( u_sol, tensor=u_h, form_compiler_parameters=context.fc_params) # Velocity reconstruction sigma_sol = A.inv * g + A.inv * (B.T * u_h - K_local.T * self.trace_solution) self._assemble_velocity = create_assembly_callable( sigma_sol, tensor=sigma_h, form_compiler_parameters=context.fc_params) # Set up the projector for projecting the broken solution # into the unbroken finite element spaces # NOTE: Tolerance here matters! sigma_b, _ = self.broken_solution.split() sigma_u, _ = self.unbroken_solution.split() self.projector = Projector(sigma_b, sigma_u, solver_parameters={ "ksp_type": "cg", "ksp_rtol": 1e-13 })