def apply_geometry_lowering(form, preserve_types=()):
"""Change GeometricQuantity objects in expression to the lowest level GeometricQuantity objects.
Assumes the expression is preprocessed or at least that derivatives have been expanded.
@param form:
An Expr or Form.
"""
if isinstance(form, Form):
newintegrals = [
apply_geometry_lowering(integral, preserve_types)
for integral in form.integrals()
]
return Form(newintegrals)
elif isinstance(form, Integral):
integral = form
if integral.integral_type() in (custom_integral_types +
point_integral_types):
automatic_preserve_types = [SpatialCoordinate, Jacobian]
else:
automatic_preserve_types = [CellCoordinate]
preserve_types = set(preserve_types) | set(automatic_preserve_types)
mf = GeometryLoweringApplier(preserve_types)
newintegrand = map_expr_dag(mf, integral.integrand())
return integral.reconstruct(integrand=newintegrand)
elif isinstance(form, Expr):
expr = form
mf = GeometryLoweringApplier(preserve_types)
return map_expr_dag(mf, expr)
else:
error("Invalid type %s" % (form.__class__.__name__, ))
def apply_integral_scaling(form):
"Multiply integrands by a factor to scale the integral to reference frame."
# TODO: Consider adding an in_reference_frame property to Integral
# and checking it here and setting it in the returned form
if isinstance(form, Form):
newintegrals = [
apply_integral_scaling(integral) for integral in form.integrals()
]
return Form(newintegrals)
elif isinstance(form, Integral):
integral = form
integrand = integral.integrand()
# Compute and apply integration scaling factor since we want to compute
# coordinate derivatives at the end, the scaling factor has to be moved
# inside those
scale = compute_integrand_scaling_factor(integral)
def scale_coordinate_derivative(o, scale):
o_ = o.ufl_operands
if isinstance(o, CoordinateDerivative):
return CoordinateDerivative(
scale_coordinate_derivative(o_[0], scale), o_[1], o_[2],
o_[3])
else:
return scale * o
newintegrand = scale_coordinate_derivative(integrand, scale)
return integral.reconstruct(integrand=newintegrand)
else:
error("Invalid type %s" % (form.__class__.__name__, ))
def attach_estimated_degrees(form):
"""Attach estimated polynomial degree to a form's integrals.
:arg form: The :class:`~.Form` to inspect.
:returns: A new Form with estimate degrees attached.
"""
integrals = form.integrals()
new_integrals = []
for integral in integrals:
md = {}
md.update(integral.metadata())
degree = estimate_total_polynomial_degree(integral.integrand())
md["estimated_polynomial_degree"] = degree
new_integrals.append(integral.reconstruct(metadata=md))
return Form(new_integrals)
def apply_integral_scaling(form):
"Multiply integrands by a factor to scale the integral to reference frame."
# TODO: Consider adding an in_reference_frame property to Integral
# and checking it here and setting it in the returned form
if isinstance(form, Form):
newintegrals = [
apply_integral_scaling(integral) for integral in form.integrals()
]
return Form(newintegrals)
elif isinstance(form, Integral):
integral = form
integrand = integral.integrand()
# Compute and apply integration scaling factor since we want to compute
# coordinate derivatives at the end, the scaling factor has to be moved
# inside those
scale, degree = compute_integrand_scaling_factor(integral)
md = {}
md.update(integral.metadata())
new_degree = degree
cur_degree = md.get("estimated_polynomial_degree")
if cur_degree is not None:
if isinstance(cur_degree, tuple) and isinstance(degree, tuple):
new_degree = tuple(d[0] + d[1]
for d in zip(cur_degree, degree))
elif isinstance(cur_degree, tuple):
new_degree = tuple(d + degree for d in cur_degree)
elif isinstance(degree, tuple):
new_degree = tuple(cur_degree + d for d in degree)
else:
new_degree = cur_degree + degree
md["estimated_polynomial_degree"] = new_degree
def scale_coordinate_derivative(o, scale):
o_ = o.ufl_operands
if isinstance(o, CoordinateDerivative):
return CoordinateDerivative(
scale_coordinate_derivative(o_[0], scale), o_[1], o_[2],
o_[3])
else:
return scale * o
newintegrand = scale_coordinate_derivative(integrand, scale)
return integral.reconstruct(integrand=newintegrand, metadata=md)
else:
error("Invalid type %s" % (form.__class__.__name__, ))
def attach_coordinate_derivatives(form, coordinate_derivatives):
if coordinate_derivatives is None:
return form
if isinstance(form, Form):
cds = coordinate_derivatives
new_integrals = [attach_coordinate_derivatives(integ, cds) for integ in form.integrals()]
return Form(new_integrals)
elif isinstance(form, Integral):
integral = form
integrand = integral.integrand()
# apply the stored coordinate derivatives back onto the integrand
for tup in reversed(coordinate_derivatives):
integrand = CoordinateDerivative(integrand, tup[0], tup[1], tup[2])
return integral.reconstruct(integrand=integrand)
else:
error("Invalid type %s" % (form.__class__.__name__,))
def replace_terminal_data(o, mapping):
"""Return a new form where the terminals have been replaced using the
provided mapping.
:arg o: The object to have its terminals replaced. This must either be a
:class:`~.Form` or :class:`~.Integral`.
:arg mapping: A mapping suitable for reconstructing the form such as the one
returned by :func:`strip_terminal_data`.
:returns: The new form.
"""
if isinstance(o, Form):
return Form(
[replace_terminal_data(itg, mapping) for itg in o.integrals()])
elif isinstance(o, Integral):
expr_map, domain_map = mapping
integrand = replace(o.integrand(), expr_map)
return o.reconstruct(integrand, domain=domain_map[o.ufl_domain()])
else:
raise ValueError("Only Form or Integral inputs expected")
def apply_integral_scaling(form):
"Multiply integrands by a factor to scale the integral to reference frame."
# TODO: Consider adding an in_reference_frame property to Integral
# and checking it here and setting it in the returned form
if isinstance(form, Form):
newintegrals = [
apply_integral_scaling(integral) for integral in form.integrals()
]
return Form(newintegrals)
elif isinstance(form, Integral):
integral = form
# Compute and apply integration scaling factor
scale = compute_integrand_scaling_factor(integral)
newintegrand = integral.integrand() * scale
return integral.reconstruct(integrand=newintegrand)
else:
error("Invalid type %s" % (form.__class__.__name__, ))
def strip_coordinate_derivatives(form):
if isinstance(form, Form):
if len(form.integrals()) == 0:
return form, None
stripped_integrals = []
coordinate_derivatives = []
for integral in form.integrals():
(si, cd) = strip_coordinate_derivatives(integral)
stripped_integrals.append(si)
coordinate_derivatives.append(cd)
assert_that_coordinate_derivatives_are_the_same(coordinate_derivatives)
# now that we have checked that the coordinate derivative that we apply is
# the same for all integrals, we only have to return them for the first integral
return (Form(stripped_integrals), coordinate_derivatives[0])
elif isinstance(form, Integral):
integral = form
integrand = integral.integrand()
checker = CoordinateDerivativeIsOutermostChecker()
map_expr_dags(checker, [integrand])
coordinate_derivatives = []
# grab all coordinate derivatives and store them, so that we can apply
# them later again
def take_top_coordinate_derivatives(o):
o_ = o.ufl_operands
if isinstance(o, CoordinateDerivative):
coordinate_derivatives.append((o_[1], o_[2], o_[3]))
return take_top_coordinate_derivatives(o_[0])
else:
return o
newintegrand = take_top_coordinate_derivatives(integrand)
return (integral.reconstruct(integrand=newintegrand), coordinate_derivatives)
else:
error("Invalid type %s" % (form.__class__.__name__,))
def strip_terminal_data(o):
"""Return a new form where all terminals have been replaced by UFL-only
equivalents.
:arg o: The object to be stripped. This must either be a :class:`~.Form`
or :class:`~.Integral`.
:returns: A 2-tuple containing an equivalent UFL-only object and a mapping
allowing the original form to be reconstructed using
:func:`replace_terminal_data`.
This function is useful for forms containing augmented UFL objects that
hold references to large data structures. These objects are be extracted
into the mapping allowing the form to be cached without leaking memory.
"""
# We need to keep track of two maps because integrals store references to the
# domain and ``replace`` expects only a mapping containing ``Expr`` objects.
if isinstance(o, Form):
integrals = []
expr_map = {}
domain_map = {}
for integral in o.integrals():
itg, (emap, dmap) = strip_terminal_data(integral)
integrals.append(itg)
expr_map.update(emap)
domain_map.update(dmap)
return Form(integrals), (expr_map, domain_map)
elif isinstance(o, Integral):
handler = TerminalStripper()
integrand = map_expr_dag(handler, o.integrand())
domain = strip_domain(o.ufl_domain())
# invert the mapping so it can be passed straight into replace_terminal_data
expr_map = {v: k for k, v in handler.mapping.items()}
domain_map = {domain: o.ufl_domain()}
return o.reconstruct(integrand, domain=domain), (expr_map, domain_map)
else:
raise ValueError("Only Form or Integral inputs expected")