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
})
