diff --git a/test/test_external_operator.py b/test/test_external_operator.py index 0cdebb422..0de32401b 100644 --- a/test/test_external_operator.py +++ b/test/test_external_operator.py @@ -181,7 +181,7 @@ def test_differentiation_procedure_action(V1, V2): def test_extractions(domain_2d, V1): - from ufl.algorithms.analysis import (extract_arguments, extract_arguments_and_coefficients, + from ufl.algorithms.analysis import (extract_arguments, extract_arguments_and_coefficients_and_geometric_quantities, extract_base_form_operators, extract_coefficients, extract_constants) u = Coefficient(V1) @@ -192,7 +192,7 @@ def test_extractions(domain_2d, V1): assert extract_coefficients(e) == [u] assert extract_arguments(e) == [vstar_e] - assert extract_arguments_and_coefficients(e) == ([vstar_e], [u]) + assert extract_arguments_and_coefficients_and_geometric_quantities(e) == ([vstar_e], [u], []) assert extract_constants(e) == [c] assert extract_base_form_operators(e) == [e] @@ -200,7 +200,7 @@ def test_extractions(domain_2d, V1): assert extract_coefficients(F) == [u] assert extract_arguments(e) == [vstar_e] - assert extract_arguments_and_coefficients(e) == ([vstar_e], [u]) + assert extract_arguments_and_coefficients_and_geometric_quantities(e) == ([vstar_e], [u], []) assert extract_constants(F) == [c] assert F.base_form_operators() == (e,) @@ -209,14 +209,14 @@ def test_extractions(domain_2d, V1): assert extract_coefficients(e) == [u] assert extract_arguments(e) == [vstar_e, u_hat] - assert extract_arguments_and_coefficients(e) == ([vstar_e, u_hat], [u]) + assert extract_arguments_and_coefficients_and_geometric_quantities(e) == ([vstar_e, u_hat], [u], []) assert extract_base_form_operators(e) == [e] F = e * dx assert extract_coefficients(F) == [u] assert extract_arguments(e) == [vstar_e, u_hat] - assert extract_arguments_and_coefficients(e) == ([vstar_e, u_hat], [u]) + assert extract_arguments_and_coefficients_and_geometric_quantities(e) == ([vstar_e, u_hat], [u], []) assert F.base_form_operators() == (e,) w = Coefficient(V1) @@ -225,14 +225,14 @@ def test_extractions(domain_2d, V1): assert extract_coefficients(e2) == [u, w] assert extract_arguments(e2) == [vstar_e2, u_hat] - assert extract_arguments_and_coefficients(e2) == ([vstar_e2, u_hat], [u, w]) + assert extract_arguments_and_coefficients_and_geometric_quantities(e2) == ([vstar_e2, u_hat], [u, w], []) assert extract_base_form_operators(e2) == [e, e2] F = e2 * dx assert extract_coefficients(e2) == [u, w] assert extract_arguments(e2) == [vstar_e2, u_hat] - assert extract_arguments_and_coefficients(e2) == ([vstar_e2, u_hat], [u, w]) + assert extract_arguments_and_coefficients_and_geometric_quantities(e2) == ([vstar_e2, u_hat], [u, w], []) assert F.base_form_operators() == (e, e2) diff --git a/test/test_interpolate.py b/test/test_interpolate.py index 71f3cd145..65d9eea7e 100644 --- a/test/test_interpolate.py +++ b/test/test_interpolate.py @@ -8,7 +8,7 @@ from ufl import (Action, Adjoint, Argument, Coefficient, FunctionSpace, Mesh, TestFunction, TrialFunction, action, adjoint, derivative, dx, grad, inner, replace, triangle) from ufl.algorithms.ad import expand_derivatives -from ufl.algorithms.analysis import (extract_arguments, extract_arguments_and_coefficients, extract_base_form_operators, +from ufl.algorithms.analysis import (extract_arguments, extract_arguments_and_coefficients_and_geometric_quantities, extract_base_form_operators, extract_coefficients) from ufl.algorithms.expand_indices import expand_indices from ufl.core.interpolate import Interpolate @@ -141,12 +141,12 @@ def test_extract_base_form_operators(V1, V2): # -- Interpolate(u, V2) -- # Iu = Interpolate(u, V2) assert extract_arguments(Iu) == [vstar] - assert extract_arguments_and_coefficients(Iu) == ([vstar], [u]) + assert extract_arguments_and_coefficients_and_geometric_quantities(Iu) == ([vstar], [u], []) F = Iu * dx # Form composition: Iu * dx <=> Action(v * dx, Iu(u; v*)) assert extract_arguments(F) == [] - assert extract_arguments_and_coefficients(F) == ([], [u]) + assert extract_arguments_and_coefficients_and_geometric_quantities(F) == ([], [u], []) for e in [Iu, F]: assert extract_coefficients(e) == [u] @@ -155,7 +155,7 @@ def test_extract_base_form_operators(V1, V2): # -- Interpolate(u, V2) -- # Iv = Interpolate(uhat, V2) assert extract_arguments(Iv) == [vstar, uhat] - assert extract_arguments_and_coefficients(Iv) == ([vstar, uhat], []) + assert extract_arguments_and_coefficients_and_geometric_quantities(Iv) == ([vstar, uhat], [], []) assert extract_coefficients(Iv) == [] assert extract_base_form_operators(Iv) == [Iv] diff --git a/test/test_mixed_function_space_with_mixed_mesh.py b/test/test_mixed_function_space_with_mixed_mesh.py new file mode 100644 index 000000000..aed6b123a --- /dev/null +++ b/test/test_mixed_function_space_with_mixed_mesh.py @@ -0,0 +1,110 @@ +from ufl import (triangle, Mesh, MixedMesh, FunctionSpace, TestFunction, TrialFunction, Coefficient, Constant, + Measure, SpatialCoordinate, FacetNormal, CellVolume, FacetArea, inner, grad, div, split, ) +from ufl.algorithms import compute_form_data +from ufl.finiteelement import FiniteElement, MixedElement +from ufl.pullback import identity_pullback, contravariant_piola +from ufl.sobolevspace import H1, HDiv, L2 +from ufl.domain import extract_domains + + +def test_mixed_function_space_with_mixed_mesh_basic(): + cell = triangle + elem0 = FiniteElement("Lagrange", cell, 1, (), identity_pullback, H1) + elem1 = FiniteElement("Brezzi-Douglas-Marini", cell, 1, (2, ), contravariant_piola, HDiv) + elem2 = FiniteElement("Discontinuous Lagrange", cell, 0, (), identity_pullback, L2) + elem = MixedElement([elem0, elem1, elem2]) + mesh0 = Mesh(FiniteElement("Lagrange", cell, 1, (2, ), identity_pullback, H1), ufl_id=100) + mesh1 = Mesh(FiniteElement("Lagrange", cell, 1, (2, ), identity_pullback, H1), ufl_id=101) + mesh2 = Mesh(FiniteElement("Lagrange", cell, 1, (2, ), identity_pullback, H1), ufl_id=102) + domain = MixedMesh(mesh0, mesh1, mesh2) + V = FunctionSpace(domain, elem) + u = TrialFunction(V) + v = TestFunction(V) + f = Coefficient(V, count=1000) + g = Coefficient(V, count=2000) + u0, u1, u2 = split(u) + v0, v1, v2 = split(v) + f0, f1, f2 = split(f) + g0, g1, g2 = split(g) + dx1 = Measure("dx", mesh1) + ds2 = Measure("ds", mesh2) + x = SpatialCoordinate(mesh1) + form = x[1] * f0 * inner(grad(u0), v1) * dx1(999) + div(f1) * g2 * inner(u1, grad(v2)) * ds2(888) + fd = compute_form_data(form, + do_apply_function_pullbacks=True, + do_apply_integral_scaling=True, + do_apply_geometry_lowering=True, + preserve_geometry_types=(CellVolume, FacetArea), + do_apply_restrictions=True, + do_estimate_degrees=True, + complex_mode=False) + id0, id1 = fd.integral_data + assert fd.preprocessed_form.arguments() == (v, u) + assert fd.reduced_coefficients == [f, g] + assert form.coefficients()[fd.original_coefficient_positions[0]] is f + assert form.coefficients()[fd.original_coefficient_positions[1]] is g + assert id0.domain is mesh1 + assert id0.integral_type == 'cell' + assert id0.subdomain_id == (999, ) + assert fd.original_form.domain_numbering()[id0.domain] == 0 + assert id0.integral_coefficients == set([f]) + assert id0.enabled_coefficients == [True, False] + assert id1.domain is mesh2 + assert id1.integral_type == 'exterior_facet' + assert id1.subdomain_id == (888, ) + assert fd.original_form.domain_numbering()[id1.domain] == 1 + assert id1.integral_coefficients == set([f, g]) + assert id1.enabled_coefficients == [True, True] + + +def test_mixed_function_space_with_mixed_mesh_restriction(): + cell = triangle + elem0 = FiniteElement("Lagrange", cell, 1, (), identity_pullback, H1) + elem1 = FiniteElement("Brezzi-Douglas-Marini", cell, 1, (2, ), contravariant_piola, HDiv) + elem2 = FiniteElement("Discontinuous Lagrange", cell, 0, (), identity_pullback, L2) + elem = MixedElement([elem0, elem1, elem2]) + mesh0 = Mesh(FiniteElement("Lagrange", cell, 1, (2, ), identity_pullback, H1), ufl_id=100) + mesh1 = Mesh(FiniteElement("Lagrange", cell, 1, (2, ), identity_pullback, H1), ufl_id=101) + mesh2 = Mesh(FiniteElement("Lagrange", cell, 1, (2, ), identity_pullback, H1), ufl_id=102) + domain = MixedMesh(mesh0, mesh1, mesh2) + V = FunctionSpace(domain, elem) + V0 = FunctionSpace(mesh0, elem0) + V1 = FunctionSpace(mesh1, elem1) + V2 = FunctionSpace(mesh2, elem2) + u1 = TrialFunction(V1) + v2 = TestFunction(V2) + f = Coefficient(V, count=1000) + g = Coefficient(V, count=2000) + f0, f1, f2 = split(f) + g0, g1, g2 = split(g) + dS1 = Measure("dS", mesh1) + x2 = SpatialCoordinate(mesh2) + form = inner(x2, g1) * g2 * inner(u1('-'), grad(v2('|'))) * dS1(999) + fd = compute_form_data(form, + do_apply_function_pullbacks=True, + do_apply_integral_scaling=True, + do_apply_geometry_lowering=True, + preserve_geometry_types=(CellVolume, FacetArea), + do_apply_restrictions=True, + do_estimate_degrees=True, + do_split_coefficients=(f, g), + do_assume_single_integral_type=False, + complex_mode=False) + integral_data, = fd.integral_data + assert integral_data.domain_integral_type_map[mesh1] == "interior_facet" + assert integral_data.domain_integral_type_map[mesh2] == "exterior_facet" + + +def test_mixed_function_space_with_mixed_mesh_signature(): + cell = triangle + mesh0 = Mesh(FiniteElement("Lagrange", cell, 1, (2, ), identity_pullback, H1), ufl_id=100) + mesh1 = Mesh(FiniteElement("Lagrange", cell, 1, (2, ), identity_pullback, H1), ufl_id=101) + dx0 = Measure("dx", mesh0) + dx1 = Measure("dx", mesh1) + n0 = FacetNormal(mesh0) + n1 = FacetNormal(mesh1) + form_a = inner(n1, n1) * dx0(999) + form_b = inner(n0, n0) * dx1(999) + assert form_a.signature() == form_b.signature() + assert extract_domains(form_a) == (mesh0, mesh1) + assert extract_domains(form_b) == (mesh1, mesh0) diff --git a/ufl/__init__.py b/ufl/__init__.py index 59c19e61f..7994a0bda 100644 --- a/ufl/__init__.py +++ b/ufl/__init__.py @@ -62,6 +62,7 @@ - AbstractDomain - Mesh + - MixedMesh - MeshView * Sobolev spaces:: @@ -256,7 +257,7 @@ from ufl.core.external_operator import ExternalOperator from ufl.core.interpolate import Interpolate, interpolate from ufl.core.multiindex import Index, indices -from ufl.domain import AbstractDomain, Mesh, MeshView +from ufl.domain import AbstractDomain, Mesh, MixedMesh, MeshView from ufl.finiteelement import AbstractFiniteElement from ufl.form import BaseForm, Form, FormSum, ZeroBaseForm from ufl.formoperators import (action, adjoint, derivative, energy_norm, extract_blocks, functional, lhs, replace, rhs, @@ -287,7 +288,7 @@ __all__ = [ 'product', 'as_cell', 'AbstractCell', 'Cell', 'TensorProductCell', - 'AbstractDomain', 'Mesh', 'MeshView', + 'AbstractDomain', 'Mesh', 'MixedMesh', 'MeshView', 'L2', 'H1', 'H2', 'HCurl', 'HDiv', 'HInf', 'HEin', 'HDivDiv', 'identity_pullback', 'l2_piola', 'contravariant_piola', 'covariant_piola', 'double_contravariant_piola', 'double_covariant_piola', diff --git a/ufl/action.py b/ufl/action.py index 2401e05ea..377d8cdea 100644 --- a/ufl/action.py +++ b/ufl/action.py @@ -188,8 +188,8 @@ def _get_action_form_arguments(left, right): elif isinstance(right, CoefficientDerivative): # Action differentiation pushes differentiation through # right as a consequence of Leibniz formula. - from ufl.algorithms.analysis import extract_arguments_and_coefficients - right_args, right_coeffs = extract_arguments_and_coefficients(right) + from ufl.algorithms.analysis import extract_arguments_and_coefficients_and_geometric_quantities + right_args, right_coeffs, _ = extract_arguments_and_coefficients_and_geometric_quantities(right) arguments = left_args + tuple(right_args) coefficients += tuple(right_coeffs) elif isinstance(right, (BaseCoefficient, Zero)): diff --git a/ufl/algorithms/analysis.py b/ufl/algorithms/analysis.py index ba3afe6c6..b61b0df10 100644 --- a/ufl/algorithms/analysis.py +++ b/ufl/algorithms/analysis.py @@ -12,9 +12,11 @@ from itertools import chain from ufl.algorithms.traversal import iter_expressions +from ufl.domain import Mesh from ufl.argument import BaseArgument, Coargument from ufl.coefficient import BaseCoefficient from ufl.constant import Constant +from ufl.geometry import GeometricQuantity from ufl.core.base_form_operator import BaseFormOperator from ufl.core.terminal import Terminal from ufl.corealg.traversal import traverse_unique_terminals, unique_pre_traversal @@ -187,19 +189,20 @@ def extract_base_form_operators(a): return sorted_by_count(extract_type(a, BaseFormOperator)) -def extract_arguments_and_coefficients(a): - """Build two sorted lists of all arguments and coefficients in a. +def extract_arguments_and_coefficients_and_geometric_quantities(a): + """Build three sorted lists of all arguments, coefficients, and geometric quantities in a. - This function is faster than extract_arguments + extract_coefficients + This function is faster than extract_arguments + extract_coefficients + extract_geometric_quantities for large forms, and has more validation built in. Args: a: A BaseForm, Integral or Expr """ # Extract lists of all BaseArgument and BaseCoefficient instances - base_coeff_and_args = extract_type(a, (BaseArgument, BaseCoefficient)) - arguments = [f for f in base_coeff_and_args if isinstance(f, BaseArgument)] - coefficients = [f for f in base_coeff_and_args if isinstance(f, BaseCoefficient)] + base_coeff_and_args_and_gq = extract_type(a, (BaseArgument, BaseCoefficient, GeometricQuantity)) + arguments = [f for f in base_coeff_and_args_and_gq if isinstance(f, BaseArgument)] + coefficients = [f for f in base_coeff_and_args_and_gq if isinstance(f, BaseCoefficient)] + geometric_quantities = [f for f in base_coeff_and_args_and_gq if isinstance(f, GeometricQuantity)] # Build number,part: instance mappings, should be one to one bfnp = dict((f, (f.number(), f.part())) for f in arguments) @@ -214,19 +217,30 @@ def extract_arguments_and_coefficients(a): if len(fcounts) != len(set(fcounts.values())): raise ValueError( "Found different coefficients with same counts.\n" - "The arguments found are:\n" + "\n".join(f" {c}" for c in coefficients)) + "The Coefficients found are:\n" + "\n".join(f" {c}" for c in coefficients)) + + # Build count: instance mappings, should be one to one + gqcounts = {} + for gq in geometric_quantities: + assert isinstance(gq._domain, Mesh), f"Found that {gq}._domain is {gq._domain}" + gqcounts[gq] = (type(gq).name, gq._domain._ufl_id) + if len(gqcounts) != len(set(gqcounts.values())): + raise ValueError( + "Found different geometric quantities with same (geometric_quantity_type, domain).\n" + "The GeometricQuantities found are:\n" + "\n".join(f" {gq}" for gq in geometric_quantities)) # Passed checks, so we can safely sort the instances by count arguments = _sorted_by_number_and_part(arguments) coefficients = sorted_by_count(coefficients) + geometric_quantities = list(sorted(geometric_quantities, key=lambda gq: (type(gq).name, gq._domain._ufl_id))) - return arguments, coefficients + return arguments, coefficients, geometric_quantities def extract_elements(form): """Build sorted tuple of all elements used in form.""" - args = chain(*extract_arguments_and_coefficients(form)) - return tuple(f.ufl_element() for f in args) + arguments, coefficients, _ = extract_arguments_and_coefficients_and_geometric_quantities(form) + return tuple(f.ufl_element() for f in arguments + coefficients) def extract_unique_elements(form): diff --git a/ufl/algorithms/apply_coefficient_split.py b/ufl/algorithms/apply_coefficient_split.py new file mode 100644 index 000000000..46fc45368 --- /dev/null +++ b/ufl/algorithms/apply_coefficient_split.py @@ -0,0 +1,232 @@ +"""Apply coefficient split. + +This module contains classes and functions to split coefficients defined on mixed function spaces. +""" + +import functools +import numpy +from ufl.classes import Restricted +from ufl.corealg.map_dag import map_expr_dag +from ufl.corealg.multifunction import MultiFunction +from ufl.domain import extract_unique_domain +from ufl.classes import (Coefficient, Form, ReferenceGrad, ReferenceValue, + Indexed, MultiIndex, Index, FixedIndex, + ComponentTensor, ListTensor, Zero, + NegativeRestricted, PositiveRestricted, SingleValueRestricted, ToBeRestricted) +from ufl import indices +from ufl.checks import is_cellwise_constant +from ufl.tensors import as_tensor + + +class CoefficientSplitter(MultiFunction): + + def __init__(self, coefficient_split): + MultiFunction.__init__(self) + self._coefficient_split = coefficient_split + + expr = MultiFunction.reuse_if_untouched + + def modified_terminal(self, o): + restriction = None + local_derivatives = 0 + reference_value = False + t = o + while not t._ufl_is_terminal_: + assert t._ufl_is_terminal_modifier_, f"Got {repr(t)}" + if isinstance(t, ReferenceValue): + assert not reference_value, "Got twice pulled back terminal!" + reference_value = True + t, = t.ufl_operands + elif isinstance(t, ReferenceGrad): + local_derivatives += 1 + t, = t.ufl_operands + elif isinstance(t, Restricted): + assert restriction is None, "Got twice restricted terminal!" + restriction = t._side + t, = t.ufl_operands + elif t._ufl_terminal_modifiers_: + raise ValueError("Missing handler for terminal modifier type %s, object is %s." % (type(t), repr(t))) + else: + raise ValueError("Unexpected type %s object %s." % (type(t), repr(t))) + if not isinstance(t, Coefficient): + # Only split coefficients + return o + if t not in self._coefficient_split: + # Only split mixed coefficients + return o + # Reference value expected + assert reference_value + # Derivative indices + beta = indices(local_derivatives) + components = [] + for subcoeff in self._coefficient_split[t]: + c = subcoeff + # Apply terminal modifiers onto the subcoefficient + if reference_value: + c = ReferenceValue(c) + for n in range(local_derivatives): + # Return zero if expression is trivially constant. This has to + # happen here because ReferenceGrad has no access to the + # topological dimension of a literal zero. + if is_cellwise_constant(c): + dim = extract_unique_domain(subcoeff).topological_dimension() + c = Zero(c.ufl_shape + (dim,), c.ufl_free_indices, c.ufl_index_dimensions) + else: + c = ReferenceGrad(c) + if restriction == '+': + c = PositiveRestricted(c) + elif restriction == '-': + c = NegativeRestricted(c) + elif restriction == '|': + c = SingleValueRestricted(c) + elif restriction == '?': + c = ToBeRestricted(c) + elif restriction is not None: + raise RuntimeError(f"Got unknown restriction: {restriction}") + # Collect components of the subcoefficient + for alpha in numpy.ndindex(subcoeff.ufl_element().reference_value_shape): + # New modified terminal: component[alpha + beta] + components.append(c[alpha + beta]) + # Repack derivative indices to shape + c, = indices(1) + return ComponentTensor(as_tensor(components)[c], MultiIndex((c,) + beta)) + + positive_restricted = modified_terminal + negative_restricted = modified_terminal + single_value_restricted = modified_terminal + to_be_restricted = modified_terminal + reference_grad = modified_terminal + reference_value = modified_terminal + terminal = modified_terminal + + +def apply_coefficient_split(expr, coefficient_split): + """Split mixed coefficients, so mixed elements need not be + implemented. + + :arg split: A :py:class:`dict` mapping each mixed coefficient to a + sequence of subcoefficients. If None, calling this + function is a no-op. + """ + if coefficient_split is None: + return expr + splitter = CoefficientSplitter(coefficient_split) + return map_expr_dag(splitter, expr) + + +class FixedIndexRemover(MultiFunction): + + def __init__(self, fimap): + MultiFunction.__init__(self) + self.fimap = fimap + self._object_cache = {} + + expr = MultiFunction.reuse_if_untouched + + @staticmethod + def _cached(f): + @functools.wraps(f) + def wrapper(self, *args, **kwargs): + o, = args + if o in self._object_cache: + return self._object_cache[o] + else: + return self._object_cache.setdefault(o, f(self, o)) + return wrapper + + @_cached + def zero(self, o): + free_indices = [] + index_dimensions = [] + for i, d in zip(o.ufl_free_indices, o.ufl_index_dimensions): + if Index(i) in self.fimap: + ind_j = self.fimap[Index(i)] + if not isinstance(ind_j, FixedIndex): + free_indices.append(ind_j.count()) + index_dimensions.append(d) + else: + free_indices.append(i) + index_dimensions.append(d) + return Zero(shape=o.ufl_shape, free_indices=tuple(free_indices), index_dimensions=tuple(index_dimensions)) + + @_cached + def list_tensor(self, o): + cc = [] + for o1 in o.ufl_operands: + comp = map_expr_dag(self, o1) + cc.append(comp) + return ListTensor(*cc) + + @_cached + def multi_index(self, o): + return MultiIndex(tuple(self.fimap.get(i, i) for i in o.indices())) + + +class IndexRemover(MultiFunction): + + def __init__(self): + MultiFunction.__init__(self) + self._object_cache = {} + + expr = MultiFunction.reuse_if_untouched + + @staticmethod + def _cached(f): + @functools.wraps(f) + def wrapper(self, *args, **kwargs): + o, = args + if o in self._object_cache: + return self._object_cache[o] + else: + return self._object_cache.setdefault(o, f(self, o)) + return wrapper + + @_cached + def _zero_simplify(self, o): + operand, = o.ufl_operands + operand = map_expr_dag(self, operand) + if isinstance(operand, Zero): + return Zero(shape=o.ufl_shape, free_indices=o.ufl_free_indices, index_dimensions=o.ufl_index_dimensions) + else: + return o._ufl_expr_reconstruct_(operand) + + @_cached + def indexed(self, o): + o1, i1 = o.ufl_operands + if isinstance(o1, ComponentTensor): + o2, i2 = o1.ufl_operands + fimap = dict(zip(i2.indices(), i1.indices(), strict=True)) + rule = FixedIndexRemover(fimap) + v = map_expr_dag(rule, o2) + return map_expr_dag(self, v) + elif isinstance(o1, ListTensor): + if isinstance(i1[0], FixedIndex): + o1 = o1.ufl_operands[i1[0]._value] + if len(i1) > 1: + i1 = MultiIndex(i1[1:]) + return map_expr_dag(self, Indexed(o1, i1)) + else: + return map_expr_dag(self, o1) + o1 = map_expr_dag(self, o1) + return Indexed(o1, i1) + + # Do something nicer + positive_restricted = _zero_simplify + negative_restricted = _zero_simplify + single_value_restricted = _zero_simplify + to_be_restricted = _zero_simplify + reference_grad = _zero_simplify + reference_value = _zero_simplify + + +def remove_component_and_list_tensors(o): + if isinstance(o, Form): + integrals = [] + for integral in o.integrals(): + integrand = remove_component_and_list_tensors(integral.integrand()) + if not isinstance(integrand, Zero): + integrals.append(integral.reconstruct(integrand=integrand)) + return o._ufl_expr_reconstruct_(integrals) + else: + rule = IndexRemover() + return map_expr_dag(rule, o) diff --git a/ufl/algorithms/apply_derivatives.py b/ufl/algorithms/apply_derivatives.py index 7ef8a0a5f..b220133fc 100644 --- a/ufl/algorithms/apply_derivatives.py +++ b/ufl/algorithms/apply_derivatives.py @@ -9,6 +9,7 @@ import warnings from collections import defaultdict from math import pi +import numpy from ufl.action import Action from ufl.algorithms.analysis import extract_arguments @@ -27,7 +28,7 @@ from ufl.corealg.map_dag import map_expr_dag from ufl.corealg.multifunction import MultiFunction from ufl.differentiation import BaseFormCoordinateDerivative, BaseFormOperatorDerivative, CoordinateDerivative -from ufl.domain import extract_unique_domain +from ufl.domain import extract_unique_domain, MixedMesh from ufl.form import Form, ZeroBaseForm from ufl.operators import (bessel_I, bessel_J, bessel_K, bessel_Y, cell_avg, conditional, cos, cosh, exp, facet_avg, ln, sign, sin, sinh, sqrt) @@ -41,6 +42,16 @@ # - ReferenceDivRuleset +def flatten_domain_element(domain, element): + """Return the flattened (domain, element) pairs for mixed domain problems.""" + if not isinstance(domain, MixedMesh): + return ((domain, element), ) + flattened = () + for d, e in zip(domain, element.sub_elements, strict=True): + flattened += flatten_domain_element(d, e) + return flattened + + class GenericDerivativeRuleset(MultiFunction): """A generic derivative.""" @@ -593,16 +604,48 @@ def reference_value(self, o): """Differentiate a reference_value.""" # grad(o) == grad(rv(f)) -> K_ji*rgrad(rv(f))_rj f = o.ufl_operands[0] - if isinstance(f.ufl_element().pullback, PhysicalPullback): - # TODO: Do we need to be more careful for immersed things? - return ReferenceGrad(o) - if not f._ufl_is_terminal_: raise ValueError("ReferenceValue can only wrap a terminal") - domain = extract_unique_domain(f) - K = JacobianInverse(domain) - Do = grad_to_reference_grad(o, K) - return Do + domain = extract_unique_domain(f, expand_mixed_mesh=False) + if isinstance(domain, MixedMesh): + element = f.ufl_function_space().ufl_element() + assert element.num_sub_elements == len(domain), f"Got {element.num_sub_elements} != {len(domain)}" + g = ReferenceGrad(o) + vsh = g.ufl_shape[:-1] + ref_dim = g.ufl_shape[-1] + components = [] + dofoffset = 0 + for d, e in flatten_domain_element(domain, element): + esh = e.reference_value_shape + if esh == (): + esh = (1, ) + if isinstance(e.pullback, PhysicalPullback): + if ref_dim != self._var_shape[0]: + raise NotImplementedError("""PhysicalPullback not handled for immersed domain : + reference dim ({ref_dim}) != physical dim (self._var_shape[0])""") + for idx in range(int(numpy.prod(esh))): + for i in range(ref_dim): + components.append(g[(dofoffset + idx, ) + (i, )]) + else: + K = JacobianInverse(d) + rdim, gdim = K.ufl_shape + assert rdim == ref_dim, f"{rdim} != {ref_dim}" + assert gdim == self._var_shape[0], f"{gdim} != {self._var_shape[0]}" + for idx in range(int(numpy.prod(esh))): + for i in range(gdim): + temp = Zero() + for j in range(rdim): + temp += g[(dofoffset + idx, ) + (j, )] * K[j, i] + components.append(temp) + dofoffset += int(numpy.prod(esh)) + return as_tensor(numpy.asarray(components).reshape(vsh + self._var_shape)) + else: + if isinstance(f.ufl_element().pullback, PhysicalPullback): + # TODO: Do we need to be more careful for immersed things? + return ReferenceGrad(o) + else: + K = JacobianInverse(domain) + return grad_to_reference_grad(o, K) def reference_grad(self, o): """Differentiate a reference_grad.""" @@ -612,10 +655,40 @@ def reference_grad(self, o): valid_operand = f._ufl_is_in_reference_frame_ or isinstance(f, (JacobianInverse, SpatialCoordinate)) if not valid_operand: raise ValueError("ReferenceGrad can only wrap a reference frame type!") - domain = extract_unique_domain(f) - K = JacobianInverse(domain) - Do = grad_to_reference_grad(o, K) - return Do + domain = extract_unique_domain(f, expand_mixed_mesh=False) + if isinstance(domain, MixedMesh): + if not f._ufl_is_in_reference_frame_: + raise RuntimeError("Expecting a reference frame type") + while not f._ufl_is_terminal_: + f, = f.ufl_operands + element = f.ufl_function_space().ufl_element() + assert element.num_sub_elements == len(domain), f"Got {element.num_sub_elements} != {len(domain)}" + g = ReferenceGrad(o) + vsh = g.ufl_shape[:-1] + ref_dim = g.ufl_shape[-1] + components = [] + dofoffset = 0 + for d, e in flatten_domain_element(domain, element): + esh = e.reference_value_shape + if esh == (): + esh = (1, ) + K = JacobianInverse(d) + rdim, gdim = K.ufl_shape + assert rdim == ref_dim, f"{rdim} != {ref_dim}" + assert gdim == self._var_shape[0], f"{gdim} != {self._var_shape[0]}" + for idx in range(int(numpy.prod(esh))): + for midx in numpy.ndindex(g.ufl_shape[1:-1]): + for i in range(gdim): + temp = Zero() + for j in range(rdim): + temp += g[(dofoffset + idx, ) + midx + (j, )] * K[j, i] + components.append(temp) + dofoffset += int(numpy.prod(esh)) + assert g.ufl_shape[0] == dofoffset, f"{g.ufl_shape[0]} != {dofoffset}" + return as_tensor(numpy.asarray(components).reshape(vsh + self._var_shape)) + else: + K = JacobianInverse(domain) + return grad_to_reference_grad(o, K) # --- Nesting of gradients diff --git a/ufl/algorithms/apply_restrictions.py b/ufl/algorithms/apply_restrictions.py index 8f788a009..18353306d 100644 --- a/ufl/algorithms/apply_restrictions.py +++ b/ufl/algorithms/apply_restrictions.py @@ -10,30 +10,34 @@ # # SPDX-License-Identifier: LGPL-3.0-or-later - from ufl.algorithms.map_integrands import map_integrand_dags from ufl.classes import Restricted from ufl.corealg.map_dag import map_expr_dag from ufl.corealg.multifunction import MultiFunction -from ufl.domain import extract_unique_domain +from ufl.domain import extract_unique_domain, MixedMesh from ufl.measure import integral_type_to_measure_name from ufl.sobolevspace import H1 +from ufl.classes import ReferenceGrad, ReferenceValue +from ufl.restriction import PositiveRestricted, SingleValueRestricted class RestrictionPropagator(MultiFunction): """Restriction propagator.""" - def __init__(self, side=None): + def __init__(self, side=None, assume_single_integral_type=True): """Initialise.""" MultiFunction.__init__(self) self.current_restriction = side - self.default_restriction = "+" + self.default_restriction = "+" if assume_single_integral_type else "?" # Caches for propagating the restriction with map_expr_dag - self.vcaches = {"+": {}, "-": {}} - self.rcaches = {"+": {}, "-": {}} + self.vcaches = {"+": {}, "-": {}, "|": {}, "?": {}} + self.rcaches = {"+": {}, "-": {}, "|": {}, "?": {}} if self.current_restriction is None: - self._rp = {"+": RestrictionPropagator("+"), - "-": RestrictionPropagator("-")} + self._rp = {"+": RestrictionPropagator("+", assume_single_integral_type), + "-": RestrictionPropagator("-", assume_single_integral_type), + "|": RestrictionPropagator("|", assume_single_integral_type), + "?": RestrictionPropagator("?", assume_single_integral_type)} + self.assume_single_integral_type = assume_single_integral_type def restricted(self, o): """When hitting a restricted quantity, visit child with a separate restriction algorithm.""" @@ -55,9 +59,12 @@ def _ignore_restriction(self, o): def _require_restriction(self, o): """Restrict a discontinuous quantity to current side, require a side to be set.""" - if self.current_restriction is None: + if self.current_restriction is not None: + return o(self.current_restriction) + elif not self.assume_single_integral_type: + return o + else: raise ValueError(f"Discontinuous type {o._ufl_class_.__name__} must be restricted.") - return o(self.current_restriction) def _default_restricted(self, o): """Restrict a continuous quantity to default side if no current restriction is set.""" @@ -172,11 +179,17 @@ def facet_normal(self, o): return self._require_restriction(o) -def apply_restrictions(expression): +def apply_restrictions(expression, assume_single_integral_type=True): """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() + if assume_single_integral_type: + integral_types = [k for k in integral_type_to_measure_name.keys() + if k.startswith("interior_facet")] + else: + # Integration type of the integral is not necessarily the same as + # the integral type of a given function; e.g., the former can be + # ``exterior_facet`` and the latter ``interior_facet``. + integral_types = None + rules = RestrictionPropagator(assume_single_integral_type=assume_single_integral_type) return map_integrand_dags(rules, expression, only_integral_type=integral_types) @@ -184,14 +197,15 @@ def apply_restrictions(expression): class DefaultRestrictionApplier(MultiFunction): """Default restriction applier.""" - def __init__(self, side=None): + def __init__(self, side=None, assume_single_integral_type=True): """Initialise.""" MultiFunction.__init__(self) self.current_restriction = side - self.default_restriction = "+" - if self.current_restriction is None: - self._rp = {"+": DefaultRestrictionApplier("+"), - "-": DefaultRestrictionApplier("-")} + # If multiple domains exist, the restriction on a function defined on + # a certain domain can not be determined by merely inspecting the + # local part of the DAG. "?" restrictions will be replaced with the + # appropriate ones later using ``replace_to_be_restricted`` function. + self.default_restriction = "+" if assume_single_integral_type else "?" def terminal(self, o): """Apply to terminal.""" @@ -236,13 +250,149 @@ def _default_restricted(self, o): facet_origin = _default_restricted # FIXME: Is this valid for quads? -def apply_default_restrictions(expression): +def apply_default_restrictions(expression, assume_single_integral_type=True): """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() + if assume_single_integral_type: + integral_types = [k for k in integral_type_to_measure_name.keys() + if k.startswith("interior_facet")] + else: + integral_types = None + rules = DefaultRestrictionApplier(assume_single_integral_type=assume_single_integral_type) return map_integrand_dags(rules, expression, only_integral_type=integral_types) + + +class DomainRestrictionMapMaker(MultiFunction): + """Make a map from domains to restriction(s). + + Inspect the DAG and collect domain-restrictions map. + This must be done per integral_data. + """ + + def __init__(self, domain_restriction_map): + MultiFunction.__init__(self) + self._domain_restriction_map = domain_restriction_map + + expr = MultiFunction.reuse_if_untouched + + def _modifier(self, o): + restriction = None + local_derivatives = 0 + reference_value = False + t = o + while not t._ufl_is_terminal_: + assert t._ufl_is_terminal_modifier_, f"Expecting a terminal modifier: got {repr(t)}" + if isinstance(t, ReferenceValue): + assert not reference_value, "Got twice pulled back terminal" + reference_value = True + t, = t.ufl_operands + elif isinstance(t, ReferenceGrad): + local_derivatives += 1 + t, = t.ufl_operands + elif isinstance(t, Restricted): + assert restriction is None, "Got twice restricted terminal" + restriction = t._side + t, = t.ufl_operands + elif t._ufl_terminal_modifiers_: + raise ValueError("Missing handler for terminal modifier type %s, object is %s." % (type(t), repr(t))) + else: + raise ValueError("Unexpected type %s object %s." % (type(t), repr(t))) + domain = extract_unique_domain(t, expand_mixed_mesh=False) + if isinstance(domain, MixedMesh): + raise RuntimeError(f"Not expecting a terminal object on a mixed mesh at this stage: found {repr(t)}") + if domain is not None: + if domain not in self._domain_restriction_map: + self._domain_restriction_map[domain] = set() + if restriction in ['+', '-', '|']: + self._domain_restriction_map[domain].add(restriction) + elif restriction not in ['?', None]: + raise RuntimeError + return o + + reference_value = _modifier + reference_grad = _modifier + positive_restricted = _modifier + negative_restricted = _modifier + single_value_restricted = _modifier + to_be_restricted = _modifier + terminal = _modifier + + +def make_domain_restriction_map(integral_data): + """Make domain-restriction map for the given integral_data.""" + domain_restriction_map = {} + rule = DomainRestrictionMapMaker(domain_restriction_map) + for integral in integral_data.integrals: + _ = map_expr_dag(rule, integral.integrand()) + return domain_restriction_map + + +def make_domain_integral_type_map(integral_data): + domain_restriction_map = make_domain_restriction_map(integral_data) + integration_domain = integral_data.domain + integration_type = integral_data.integral_type + domain_integral_type_dict = {} + for d, rs in domain_restriction_map.items(): + if rs in [{'+'}, {'-'}, {'+', '-'}]: + domain_integral_type_dict[d] = "interior_facet" + elif rs == {'|'}: + domain_integral_type_dict[d] = "exterior_facet" + elif rs == set(): + if d.topological_dimension() == integration_domain.topological_dimension(): + if integration_type == "cell": + domain_integral_type_dict[d] = "cell" + elif integration_type in ["exterior_facet", "interior_facet"]: + domain_integral_type_dict[d] = "exterior_facet" + else: + raise NotImplementedError + else: + raise NotImplementedError + else: + raise RuntimeError(f"Found inconsistent restrictions {rs} for domain {d}") + if integration_domain in domain_integral_type_dict: + if domain_integral_type_dict[integration_domain] != integration_type: + raise RuntimeError(f"""Found inconsistent integral types for the integration domain ({integration_domain}) : + {domain_integral_type_dict[integration_domain]} != {integration_type}""") + else: + domain_integral_type_dict[integration_domain] = integration_type + return domain_integral_type_dict + + +class ToBeRestrectedReplacer(MultiFunction): + """Replace ``?`` restrictions.""" + + def __init__(self, domain_integral_type_map): + MultiFunction.__init__(self) + self.domain_integral_type_map = domain_integral_type_map + + expr = MultiFunction.reuse_if_untouched + + def to_be_restricted(self, o): + mt, = o.ufl_operands + domain = extract_unique_domain(mt) + if isinstance(domain, MixedMesh): + raise RuntimeError(f"""Not expecting a (modified) terminal object on a mixed mesh at this stage : + got {repr(o)}""") + if domain not in self.domain_integral_type_map: + raise RuntimeError(f"Integral type on {domain} not known") + integral_type = self.domain_integral_type_map[domain] + if integral_type == "cell": + return mt + elif integral_type == "exterior_facet": + return SingleValueRestricted(mt) + elif integral_type == "interial_facet": + return PositiveRestricted(mt) + else: + raise RuntimeError(f"Unknown integral type: {integral_type}") + + +def replace_to_be_restricted(integral_data): + new_integrals = [] + rule = ToBeRestrectedReplacer(integral_data.domain_integral_type_map) + for integral in integral_data.integrals: + integrand = map_expr_dag(rule, integral.integrand()) + new_integrals.append(integral.reconstruct(integrand=integrand)) + return new_integrals diff --git a/ufl/algorithms/balancing.py b/ufl/algorithms/balancing.py index 477ec3f6f..19d1d7ba5 100644 --- a/ufl/algorithms/balancing.py +++ b/ufl/algorithms/balancing.py @@ -6,14 +6,15 @@ # # SPDX-License-Identifier: LGPL-3.0-or-later -from ufl.classes import (CellAvg, FacetAvg, Grad, Indexed, NegativeRestricted, PositiveRestricted, ReferenceGrad, - ReferenceValue) +from ufl.classes import (CellAvg, FacetAvg, Grad, Indexed, + NegativeRestricted, PositiveRestricted, SingleValueRestricted, ToBeRestricted, + ReferenceGrad, ReferenceValue) from ufl.corealg.map_dag import map_expr_dag from ufl.corealg.multifunction import MultiFunction modifier_precedence = [ ReferenceValue, ReferenceGrad, Grad, CellAvg, FacetAvg, PositiveRestricted, - NegativeRestricted, Indexed + NegativeRestricted, SingleValueRestricted, ToBeRestricted, Indexed ] modifier_precedence = { @@ -76,6 +77,8 @@ def _modifier(self, expr, *ops): facet_avg = _modifier positive_restricted = _modifier negative_restricted = _modifier + single_value_restricted = _modifier + to_be_restricted = _modifier def balance_modifiers(expr): diff --git a/ufl/algorithms/check_arities.py b/ufl/algorithms/check_arities.py index c93727ad8..1c661add6 100644 --- a/ufl/algorithms/check_arities.py +++ b/ufl/algorithms/check_arities.py @@ -101,6 +101,8 @@ def linear_operator(self, o, a): # Positive and negative restrictions behave as linear operators positive_restricted = linear_operator negative_restricted = linear_operator + single_value_restricted = linear_operator + to_be_restricted = linear_operator # Cell and facet average are linear operators cell_avg = linear_operator diff --git a/ufl/algorithms/compute_form_data.py b/ufl/algorithms/compute_form_data.py index af0ddf68d..fa9105dfe 100644 --- a/ufl/algorithms/compute_form_data.py +++ b/ufl/algorithms/compute_form_data.py @@ -18,7 +18,9 @@ from ufl.algorithms.apply_function_pullbacks import apply_function_pullbacks from ufl.algorithms.apply_geometry_lowering import apply_geometry_lowering from ufl.algorithms.apply_integral_scaling import apply_integral_scaling -from ufl.algorithms.apply_restrictions import apply_default_restrictions, apply_restrictions +from ufl.algorithms.apply_restrictions import (apply_default_restrictions, apply_restrictions, + replace_to_be_restricted, make_domain_integral_type_map) +from ufl.algorithms.apply_coefficient_split import apply_coefficient_split, remove_component_and_list_tensors from ufl.algorithms.check_arities import check_form_arity from ufl.algorithms.comparison_checker import do_comparison_check # See TODOs at the call sites of these below: @@ -28,10 +30,12 @@ from ufl.algorithms.formdata import FormData from ufl.algorithms.formtransformations import compute_form_arities from ufl.algorithms.remove_complex_nodes import remove_complex_nodes +from ufl.algorithms.replace import replace from ufl.classes import Coefficient, Form, FunctionSpace, GeometricFacetQuantity from ufl.corealg.traversal import traverse_unique_terminals -from ufl.domain import extract_unique_domain +from ufl.domain import extract_unique_domain, extract_domains, MixedMesh from ufl.utils.sequences import max_degree +from ufl.constantvalue import Zero def _auto_select_degree(elements): @@ -181,7 +185,7 @@ def _build_coefficient_replace_map(coefficients, element_mapping=None): # coefficient had a domain, the new one does too. # This should be overhauled with requirement that Expressions # always have a domain. - domain = extract_unique_domain(f) + domain = extract_unique_domain(f, expand_mixed_mesh=False) if domain is not None: new_e = FunctionSpace(domain, new_e) new_f = Coefficient(new_e, count=i) @@ -248,6 +252,8 @@ def compute_form_data( do_apply_geometry_lowering=False, preserve_geometry_types=(), do_apply_default_restrictions=True, do_apply_restrictions=True, do_estimate_degrees=True, do_append_everywhere_integrals=True, + do_assume_single_integral_type=True, + do_split_coefficients=None, complex_mode=False, ): """Compute form data. @@ -295,9 +301,17 @@ def compute_form_data( if do_apply_integral_scaling: form = apply_integral_scaling(form) + # Can allow for some simplifications if there indeed is only a single domain + if not do_assume_single_integral_type: + have_single_domain = len(extract_domains(form)) == 1 + # Apply default restriction to fully continuous terminals if do_apply_default_restrictions: - form = apply_default_restrictions(form) + if do_assume_single_integral_type: + form = apply_default_restrictions(form) + else: + # Apply '?' restrictions in general multi-domain problems + form = apply_default_restrictions(form, assume_single_integral_type=have_single_domain) # Lower abstractions for geometric quantities into a smaller set # of quantities, allowing the form compiler to deal with a smaller @@ -323,7 +337,10 @@ def compute_form_data( # Propagate restrictions to terminals if do_apply_restrictions: - form = apply_restrictions(form) + if do_assume_single_integral_type: + form = apply_restrictions(form) + else: + form = apply_restrictions(form, assume_single_integral_type=have_single_domain) # If in real mode, remove any complex nodes introduced during form processing. if not complex_mode: @@ -401,6 +418,57 @@ def compute_form_data( # compatible data structure. self.max_subdomain_ids = _compute_max_subdomain_ids(self.integral_data) + # Split coefficients that are contained in ``do_split_coefficients`` tuple + # into components and store a dict in ``self`` that maps + # each coefficient to its components + if do_split_coefficients is not None: + coefficient_split = {} + for o in self.reduced_coefficients: + c = self.function_replace_map[o] + elem = c.ufl_element() + mesh = extract_unique_domain(c, expand_mixed_mesh=False) + # Use MixedMesh as an indicator for MixedElement as + # the followings would be ambiguous: + # -- elem.num_sub_elements > 1 + # -- isinstance(elem.pullback, MixedPullback) + if isinstance(mesh, MixedMesh) and o in do_split_coefficients: + coefficient_split[c] = [Coefficient(FunctionSpace(m, e)) + for m, e in zip(mesh, elem.sub_elements, strict=True)] + self.coefficient_split = coefficient_split + for itg_data in self.integral_data: + new_integrals = [] + for integral in itg_data.integrals: + integrand = replace(integral.integrand(), self.function_replace_map) + integrand = remove_component_and_list_tensors(integrand) + integrand = apply_coefficient_split(integrand, self.coefficient_split) + integrand = remove_component_and_list_tensors(integrand) + if not isinstance(integrand, Zero): + new_integrals.append(integral.reconstruct(integrand=integrand)) + itg_data.integrals = new_integrals + else: + self.coefficient_split = {} + + # Make ``itg_data.domain_integral_type_map``; this is only significant + # when we handle general multi-domain problems + if do_assume_single_integral_type: + for itg_data in self.integral_data: + itg_data.domain_integral_type_map = {itg_data.domain: itg_data.integral_type} + else: + if have_single_domain: + # Make a short-cut; there is no '?' restrictions by construction + for itg_data in self.integral_data: + itg_data.domain_integral_type_map = {itg_data.domain: itg_data.integral_type} + else: + # Inspect the form and replacce all '?' restrictions with appropriate ones + # in general multi-domain problems; we must have split coefficients into components + # to simplify the DAG and facilitate this inspection + if do_split_coefficients is None: + raise ValueError("""Need to pass 'do_split_coefficients=tuple_of_coefficients_to_splilt' + for general multi-domain problems""") + for itg_data in self.integral_data: + itg_data.domain_integral_type_map = make_domain_integral_type_map(itg_data) + itg_data.integrals = replace_to_be_restricted(itg_data) + # --- Checks _check_elements(self) _check_facet_geometry(self.integral_data) diff --git a/ufl/algorithms/domain_analysis.py b/ufl/algorithms/domain_analysis.py index 3a11b123a..ff756e260 100644 --- a/ufl/algorithms/domain_analysis.py +++ b/ufl/algorithms/domain_analysis.py @@ -29,7 +29,8 @@ class IntegralData(object): __slots__ = ('domain', 'integral_type', 'subdomain_id', 'integrals', 'metadata', 'integral_coefficients', - 'enabled_coefficients') + 'enabled_coefficients', + 'domain_integral_type_map') def __init__(self, domain, integral_type, subdomain_id, integrals, metadata): @@ -51,6 +52,7 @@ def __init__(self, domain, integral_type, subdomain_id, integrals, # this stage: self.integral_coefficients = None self.enabled_coefficients = None + self.domain_integral_type_map = None # TODO: I think we can get rid of this with some refactoring # in ffc: diff --git a/ufl/algorithms/estimate_degrees.py b/ufl/algorithms/estimate_degrees.py index 334041d3f..f288a9364 100644 --- a/ufl/algorithms/estimate_degrees.py +++ b/ufl/algorithms/estimate_degrees.py @@ -15,7 +15,7 @@ from ufl.constantvalue import IntValue from ufl.corealg.map_dag import map_expr_dags from ufl.corealg.multifunction import MultiFunction -from ufl.domain import extract_unique_domain +from ufl.domain import extract_unique_domain, extract_domains from ufl.form import Form from ufl.integral import Integral @@ -95,7 +95,9 @@ def _reduce_degree(self, v, f): This is used when derivatives are taken. It does not reduce the degree when TensorProduct elements or quadrilateral elements are involved. """ - if isinstance(f, int) and extract_unique_domain(v).ufl_cell().cellname() not in ["quadrilateral", "hexahedron"]: + # Can have multiple domains of the same cell type. + cell, = set(d.ufl_cell() for d in extract_domains(v)) + if isinstance(f, int) and cell.cellname() not in ["quadrilateral", "hexahedron"]: return max(f - 1, 0) else: return f @@ -175,6 +177,14 @@ def negative_restricted(self, v, a): """Apply to negative_restricted.""" return a + def single_value_restricted(self, v, a): + """Apply to single_value_restricted.""" + return a + + def to_be_restricted(self, v, a): + """Apply to to_be_restricted.""" + return a + def conj(self, v, a): """Apply to conj.""" return a diff --git a/ufl/algorithms/formtransformations.py b/ufl/algorithms/formtransformations.py index 58d168c14..d5ed76dd3 100644 --- a/ufl/algorithms/formtransformations.py +++ b/ufl/algorithms/formtransformations.py @@ -240,6 +240,8 @@ def linear_operator(self, x, arg): # Positive and negative restrictions behave as linear operators positive_restricted = linear_operator negative_restricted = linear_operator + single_value_restricted = linear_operator + to_be_restricted = linear_operator # Cell and facet average are linear operators cell_avg = linear_operator diff --git a/ufl/differentiation.py b/ufl/differentiation.py index 34c003b85..c33f693c6 100644 --- a/ufl/differentiation.py +++ b/ufl/differentiation.py @@ -300,7 +300,7 @@ def __new__(cls, f): """Create a new ReferenceGrad.""" # Return zero if expression is trivially constant if is_cellwise_constant(f): - dim = extract_unique_domain(f).topological_dimension() + dim = extract_unique_domain(f, expand_mixed_mesh=False).topological_dimension() return Zero(f.ufl_shape + (dim,), f.ufl_free_indices, f.ufl_index_dimensions) return CompoundDerivative.__new__(cls) @@ -308,7 +308,7 @@ def __new__(cls, f): def __init__(self, f): """Initalise.""" CompoundDerivative.__init__(self, (f,)) - self._dim = extract_unique_domain(f).topological_dimension() + self._dim = extract_unique_domain(f, expand_mixed_mesh=False).topological_dimension() def _ufl_expr_reconstruct_(self, op): """Return a new object of the same type with new operands.""" diff --git a/ufl/domain.py b/ufl/domain.py index f438006d0..73d55616a 100644 --- a/ufl/domain.py +++ b/ufl/domain.py @@ -120,6 +120,77 @@ def _ufl_sort_key_(self): return (self.geometric_dimension(), self.topological_dimension(), "Mesh", typespecific) + def meshes(self): + return (self, ) + + +@attach_ufl_id +class MixedMesh(AbstractDomain, UFLObject): + """Symbolic representation of a mesh.""" + + def __init__(self, *meshes, ufl_id=None, cargo=None): + """Initialise.""" + self._ufl_id = self._init_ufl_id(ufl_id) + # Store reference to object that will not be used by UFL + self._ufl_cargo = cargo + if cargo is not None and cargo.ufl_id() != self._ufl_id: + raise ValueError("Expecting cargo object (e.g. dolfin.Mesh) to have the same ufl_id.") + if any(isinstance(m, MixedMesh) for m in meshes): + raise NotImplementedError("Currently component meshes can not include MixedMesh instances") + # TODO: Should make a "MixedCell" object here: + # currently only support single cell type. + self._ufl_cell, = set(m.ufl_cell() for m in meshes) + gdim, = set(m.geometric_dimension() for m in meshes) + # ReferenceGrad requires topological_dimension of the mixed mesh: + # set tdim to max topological dim for a general mixed mesh (currently not implemented). + tdim = max(m.topological_dimension() for m in meshes) + AbstractDomain.__init__(self, tdim, gdim) + self._meshes = tuple(meshes) + + def ufl_cargo(self): + """Return carried object that will not be used by UFL.""" + return self._ufl_cargo + + def ufl_cell(self): + """Get the cell.""" + return self._ufl_cell + + def __repr__(self): + """Representation.""" + r = "MixedMesh(*%s, ufl_id=%s)" % (repr(self._meshes), repr(self._ufl_id)) + return r + + def __str__(self): + """Format as a string.""" + return "" % (self._ufl_id,) + + def _ufl_hash_data_(self): + """UFL hash data.""" + return ("MixedMesh", self._ufl_id) + + def _ufl_signature_data_(self, renumbering): + """UFL signature data.""" + return ("MixedMesh", tuple(m._ufl_signature_data_(renumbering) for m in self._meshes)) + + def _ufl_sort_key_(self): + """UFL sort key.""" + typespecific = (self._ufl_id, ) + return (self.geometric_dimension(), self.topological_dimension(), + "MixedMesh", typespecific) + + def __len__(self): + return len(self._meshes) + + def __getitem__(self, i): + assert i < len(self), f"index ({i}) >= num. meshes ({len(self)})" + return self._meshes[i] + + def __iter__(self): + return iter(self._meshes) + + def meshes(self): + return self._meshes + @attach_ufl_id class MeshView(AbstractDomain, UFLObject): @@ -183,12 +254,16 @@ def as_domain(domain): """Convert any valid object to an AbstractDomain type.""" if isinstance(domain, AbstractDomain): # Modern UFL files and dolfin behaviour + if isinstance(domain, MixedMesh): + domain, = set(domain._meshes) return domain - try: return extract_unique_domain(domain) except AttributeError: - return domain.ufl_domain() + domain = domain.ufl_domain() + if isinstance(domain, MixedMesh): + domain, = set(domain._meshes) + return domain def sort_domains(domains): @@ -196,13 +271,22 @@ def sort_domains(domains): return tuple(sorted(domains, key=lambda domain: domain._ufl_sort_key_())) -def join_domains(domains): +def join_domains(domains, expand_mixed_mesh=True): """Take a list of domains and return a tuple with only unique domain objects. Checks that domains with the same id are compatible. """ # Use hashing to join domains, ignore None - domains = set(domains) - set((None,)) + domains_ = set(domains) - set((None,)) + if expand_mixed_mesh: + domains = set() + for domain in domains_: + if isinstance(domain, MixedMesh): + domains.update(domain._meshes) + else: + domains.add(domain) + else: + domains = domains_ if not domains: return () @@ -242,17 +326,23 @@ def join_domains(domains): # TODO: Move these to an analysis module? -def extract_domains(expr): +def extract_domains(expr, expand_mixed_mesh=True): """Return all domains expression is defined on.""" - domainlist = [] - for t in traverse_unique_terminals(expr): - domainlist.extend(t.ufl_domains()) - return sorted(join_domains(domainlist)) + from ufl.form import Form + + if isinstance(expr, Form): + # To be consistent with the numbering used in signature + return tuple(expr.domain_numbering().keys()) + else: + domainlist = [] + for t in traverse_unique_terminals(expr): + domainlist.extend(t.ufl_domains()) + return sort_domains(join_domains(domainlist, expand_mixed_mesh=expand_mixed_mesh)) -def extract_unique_domain(expr): +def extract_unique_domain(expr, expand_mixed_mesh=True): """Return the single unique domain expression is defined on or throw an error.""" - domains = extract_domains(expr) + domains = extract_domains(expr, expand_mixed_mesh=expand_mixed_mesh) if len(domains) == 1: return domains[0] elif domains: @@ -265,9 +355,11 @@ def find_geometric_dimension(expr): """Find the geometric dimension of an expression.""" gdims = set() for t in traverse_unique_terminals(expr): - domain = extract_unique_domain(t) - if domain is not None: - gdims.add(domain.geometric_dimension()) + # Can have multiple domains of the same cell type. + domains = extract_domains(t) + if len(domains) > 0: + gdim, = set(domain.geometric_dimension() for domain in domains) + gdims.add(gdim) if hasattr(t, "ufl_element"): element = t.ufl_element() if element is not None: diff --git a/ufl/exproperators.py b/ufl/exproperators.py index 60eb87566..590319cf6 100644 --- a/ufl/exproperators.py +++ b/ufl/exproperators.py @@ -24,7 +24,7 @@ from ufl.index_combination_utils import create_slice_indices, merge_overlapping_indices from ufl.indexed import Indexed from ufl.indexsum import IndexSum -from ufl.restriction import NegativeRestricted, PositiveRestricted +from ufl.restriction import NegativeRestricted, PositiveRestricted, SingleValueRestricted, ToBeRestricted from ufl.tensoralgebra import Inner, Transposed from ufl.tensors import ComponentTensor, as_tensor from ufl.utils.stacks import StackDict @@ -305,7 +305,7 @@ def _abs(self): Expr.__abs__ = _abs -# --- Extend Expr with restiction operators a("+"), a("-") --- +# --- Extend Expr with restiction operators a("+"), a("-"), a("|"), a("?") --- def _restrict(self, side): """Restrict.""" @@ -313,6 +313,10 @@ def _restrict(self, side): return PositiveRestricted(self) if side == "-": return NegativeRestricted(self) + if side == "|": + return SingleValueRestricted(self) + if side == "?": + return ToBeRestricted(self) raise ValueError(f"Invalid side '{side}' in restriction operator.") @@ -335,7 +339,7 @@ def _eval(self, coord, mapping=None, component=()): def _call(self, arg, mapping=None, component=()): """Take the restriction or evaluate depending on argument.""" - if arg in ("+", "-"): + if arg in ("+", "-", "|", "?"): if mapping is not None: raise ValueError("Not expecting a mapping when taking restriction.") return _restrict(self, arg) diff --git a/ufl/form.py b/ufl/form.py index 3f85b709d..059f17a17 100644 --- a/ufl/form.py +++ b/ufl/form.py @@ -21,11 +21,12 @@ from ufl.constantvalue import Zero from ufl.core.expr import Expr, ufl_err_str from ufl.core.ufl_type import UFLType, ufl_type -from ufl.domain import extract_unique_domain, sort_domains +from ufl.domain import extract_unique_domain, sort_domains, MixedMesh from ufl.equation import Equation from ufl.integral import Integral from ufl.utils.counted import Counted from ufl.utils.sorting import sorted_by_count +from ufl.corealg.traversal import traverse_unique_terminals # Export list for ufl.classes __all_classes__ = ["Form", "BaseForm", "ZeroBaseForm"] @@ -233,6 +234,7 @@ class Form(BaseForm): "_coefficient_numbering", "_constants", "_constant_numbering", + "_geometric_quantities", "_terminal_numbering", "_hash", "_signature", @@ -588,17 +590,26 @@ def __repr__(self): def _analyze_domains(self): """Analyze domains.""" + from ufl.argument import Argument + from ufl.coefficient import Coefficient + from ufl.geometry import GeometricQuantity from ufl.domain import join_domains, sort_domains - # Collect unique integration domains - integration_domains = join_domains([itg.ufl_domain() for itg in self._integrals]) - - # Make canonically ordered list of the domains - self._integration_domains = sort_domains(integration_domains) - - # TODO: Not including domains from coefficients and arguments - # here, may need that later - self._domain_numbering = dict((d, i) for i, d in enumerate(self._integration_domains)) + # Collect integration domains + self._integration_domains = sort_domains(join_domains([itg.ufl_domain() for itg in self._integrals])) + # Collect domains in integrands systematically + domains_in_integrands = [] + for integral in self._integrals: + integrand = integral.integrand() + for t in traverse_unique_terminals(integrand): + if isinstance(t, (Coefficient, Constant, Argument, GeometricQuantity)): + domain = extract_unique_domain(t, expand_mixed_mesh=False) + domains = domain._meshes if isinstance(domain, MixedMesh) else (domain, ) + for d in domains: + if d not in self._integration_domains and d not in domains_in_integrands: + domains_in_integrands.append(d) + all_domains = self._integration_domains + sort_domains(join_domains(domains_in_integrands)) + self._domain_numbering = dict((d, i) for i, d in enumerate(all_domains)) def _analyze_subdomain_data(self): """Analyze subdomain data.""" @@ -625,14 +636,16 @@ def _analyze_subdomain_data(self): def _analyze_form_arguments(self): """Analyze which Argument and Coefficient objects can be found in the form.""" - from ufl.algorithms.analysis import extract_arguments_and_coefficients - arguments, coefficients = extract_arguments_and_coefficients(self) + from ufl.algorithms.analysis import extract_arguments_and_coefficients_and_geometric_quantities + arguments, coefficients, geometric_quantities = \ + extract_arguments_and_coefficients_and_geometric_quantities(self) # Define canonical numbering of arguments and coefficients self._arguments = tuple( sorted(set(arguments), key=lambda x: x.number())) self._coefficients = tuple( sorted(set(coefficients), key=lambda x: x.count())) + self._geometric_quantities = geometric_quantities # sorted by (type, domain) def _analyze_base_form_operators(self): """Analyze which BaseFormOperator objects can be found in the form.""" @@ -642,38 +655,11 @@ def _analyze_base_form_operators(self): def _compute_renumbering(self): """Compute renumbering.""" - # Include integration domains and coefficients in renumbering dn = self.domain_numbering() tn = self.terminal_numbering() renumbering = {} renumbering.update(dn) renumbering.update(tn) - - # Add domains of coefficients, these may include domains not - # among integration domains - k = len(dn) - for c in self.coefficients(): - d = extract_unique_domain(c) - if d is not None and d not in renumbering: - renumbering[d] = k - k += 1 - - # Add domains of arguments, these may include domains not - # among integration domains - for a in self._arguments: - d = a.ufl_function_space().ufl_domain() - if d is not None and d not in renumbering: - renumbering[d] = k - k += 1 - - # Add domains of constants, these may include domains not - # among integration domains - for c in self._constants: - d = extract_unique_domain(c) - if d is not None and d not in renumbering: - renumbering[d] = k - k += 1 - return renumbering def _compute_signature(self): diff --git a/ufl/formatting/ufl2unicode.py b/ufl/formatting/ufl2unicode.py index 5c193da40..27e255f13 100644 --- a/ufl/formatting/ufl2unicode.py +++ b/ufl/formatting/ufl2unicode.py @@ -128,6 +128,8 @@ class UC: superscript_plus = u"⁺" superscript_minus = u"⁻" + superscript_vertical_bar = u"|" + superscript_question_mark = u"?" superscript_equals = u"⁼" superscript_left_paren = u"⁽" superscript_right_paren = u"⁾" @@ -745,6 +747,14 @@ def negative_restricted(self, o, f): """Format a negative_restriced.""" return f"{par(f)}{UC.superscript_minus}" + def single_value_restricted(self, o, f): + """Format a sigle_value_restriced.""" + return f"{par(f)}{UC.superscript_vertical_bar}" + + def to_be_restricted(self, o, f): + """Format a to_be_restriced.""" + return f"{par(f)}{UC.superscript_question_mark}" + def cell_avg(self, o, f): """Format a cell_avg.""" f = overline_string(f) diff --git a/ufl/index_combination_utils.py b/ufl/index_combination_utils.py index 1f7682ae0..e7741f5d6 100644 --- a/ufl/index_combination_utils.py +++ b/ufl/index_combination_utils.py @@ -161,7 +161,7 @@ def create_slice_indices(component, shape, fi): slice_indices.extend(ii) all_indices.extend(ii) else: - raise ValueError(f"Not expecting {ind}.") + raise ValueError(f"Not expecting {ind} [type {type(ind)}].") if len(all_indices) != len(shape): raise ValueError("Component and shape length don't match.") diff --git a/ufl/pullback.py b/ufl/pullback.py index dee12f0ce..6d2d8bd9f 100644 --- a/ufl/pullback.py +++ b/ufl/pullback.py @@ -15,7 +15,7 @@ from ufl.core.expr import Expr from ufl.core.multiindex import indices -from ufl.domain import extract_unique_domain +from ufl.domain import extract_unique_domain, AbstractDomain, MixedMesh from ufl.tensors import as_tensor if TYPE_CHECKING: @@ -54,11 +54,12 @@ def physical_value_shape(self, element) -> typing.Tuple[int, ...]: def is_identity(self) -> bool: """Is this pull back the identity (or the identity applied to mutliple components).""" - def apply(self, expr: Expr) -> Expr: + def apply(self, expr: Expr, domain: AbstractDomain = None) -> Expr: """Apply the pull back. Args: expr: A function on a physical cell + domain: The domain on which the function is defined Returns: The function pulled back to the reference cell """ @@ -77,11 +78,12 @@ def is_identity(self) -> bool: """Is this pull back the identity (or the identity applied to mutliple components).""" return True - def apply(self, expr): + def apply(self, expr, domain=None): """Apply the pull back. Args: expr: A function on a physical cell + domain: The domain on which the function is defined Returns: The function pulled back to the reference cell """ @@ -111,17 +113,19 @@ def is_identity(self) -> bool: """Is this pull back the identity (or the identity applied to mutliple components).""" return False - def apply(self, expr): + def apply(self, expr, domain=None): """Apply the pull back. Args: expr: A function on a physical cell + domain: The domain on which the function is defined Returns: The function pulled back to the reference cell """ from ufl.classes import Jacobian, JacobianDeterminant - domain = extract_unique_domain(expr) + if domain is None: + domain = extract_unique_domain(expr) J = Jacobian(domain) detJ = JacobianDeterminant(J) transform = (1.0 / detJ) * J @@ -155,17 +159,19 @@ def is_identity(self) -> bool: """Is this pull back the identity (or the identity applied to mutliple components).""" return False - def apply(self, expr): + def apply(self, expr, domain=None): """Apply the pull back. Args: expr: A function on a physical cell + domain: The domain on which the function is defined Returns: The function pulled back to the reference cell """ from ufl.classes import JacobianInverse - domain = extract_unique_domain(expr) + if domain is None: + domain = extract_unique_domain(expr) K = JacobianInverse(domain) # Apply transform "row-wise" to TensorElement(PiolaMapped, ...) *k, i, j = indices(len(expr.ufl_shape) + 1) @@ -197,17 +203,19 @@ def is_identity(self) -> bool: """Is this pull back the identity (or the identity applied to mutliple components).""" return False - def apply(self, expr): + def apply(self, expr, domain=None): """Apply the pull back. Args: expr: A function on a physical cell + domain: The domain on which the function is defined Returns: The function pulled back to the reference cell """ from ufl.classes import JacobianDeterminant - domain = extract_unique_domain(expr) + if domain is None: + domain = extract_unique_domain(expr) detJ = JacobianDeterminant(domain) return expr / detJ @@ -235,17 +243,19 @@ def is_identity(self) -> bool: """Is this pull back the identity (or the identity applied to mutliple components).""" return False - def apply(self, expr): + def apply(self, expr, domain=None): """Apply the pull back. Args: expr: A function on a physical cell + domain: The domain on which the function is defined Returns: The function pulled back to the reference cell """ from ufl.classes import Jacobian, JacobianDeterminant - domain = extract_unique_domain(expr) + if domain is None: + domain = extract_unique_domain(expr) J = Jacobian(domain) detJ = JacobianDeterminant(J) # Apply transform "row-wise" to TensorElement(PiolaMapped, ...) @@ -278,17 +288,19 @@ def is_identity(self) -> bool: """Is this pull back the identity (or the identity applied to mutliple components).""" return False - def apply(self, expr): + def apply(self, expr, domain=None): """Apply the pull back. Args: expr: A function on a physical cell + domain: The domain on which the function is defined Returns: The function pulled back to the reference cell """ from ufl.classes import JacobianInverse - domain = extract_unique_domain(expr) + if domain is None: + domain = extract_unique_domain(expr) K = JacobianInverse(domain) # Apply transform "row-wise" to TensorElement(PiolaMapped, ...) *k, i, j, m, n = indices(len(expr.ufl_shape) + 2) @@ -328,11 +340,12 @@ def is_identity(self) -> bool: """Is this pull back the identity (or the identity applied to mutliple components).""" return all(e.pullback.is_identity for e in self._element.sub_elements) - def apply(self, expr): + def apply(self, expr, domain=None): """Apply the pull back. Args: expr: A function on a physical cell + domain: The domain on which the function is defined Returns: The function pulled back to the reference cell """ @@ -340,11 +353,16 @@ def apply(self, expr): g_components = [] offset = 0 # For each unique piece in reference space, apply the appropriate pullback - for subelem in self._element.sub_elements: + if domain is None: + domain = extract_unique_domain(expr, expand_mixed_mesh=False) + if isinstance(domain, MixedMesh): + assert len(domain) == self._element.num_sub_elements + for i, subelem in enumerate(self._element.sub_elements): rsub = as_tensor(np.asarray( rflat[offset: offset + subelem.reference_value_size] ).reshape(subelem.reference_value_shape)) - rmapped = subelem.pullback.apply(rsub) + subdomain = domain[i] if isinstance(domain, MixedMesh) else None + rmapped = subelem.pullback.apply(rsub, domain=subdomain) # Flatten into the pulled back expression for the whole thing g_components.extend([rmapped[idx] for idx in np.ndindex(rmapped.ufl_shape)]) offset += subelem.reference_value_size @@ -397,11 +415,12 @@ def is_identity(self) -> bool: """Is this pull back the identity (or the identity applied to mutliple components).""" return all(e.pullback.is_identity for e in self._element.sub_elements) - def apply(self, expr): + def apply(self, expr, domain=None): """Apply the pull back. Args: expr: A function on a physical cell + domain: The domain on which the function is defined Returns: The function pulled back to the reference cell """ @@ -411,13 +430,18 @@ def apply(self, expr): for subelem in self._element.sub_elements: offsets.append(offsets[-1] + subelem.reference_value_size) # For each unique piece in reference space, apply the appropriate pullback + if domain is None: + domain = extract_unique_domain(expr, expand_mixed_mesh=False) + if isinstance(domain, MixedMesh): + assert len(domain) == self._element.num_sub_elements for component in np.ndindex(self._block_shape): i = self._symmetry[component] subelem = self._element.sub_elements[i] rsub = as_tensor(np.asarray( rflat[offsets[i]:offsets[i+1]] ).reshape(subelem.reference_value_shape)) - rmapped = subelem.pullback.apply(rsub) + subdomain = domain[i] if isinstance(domain, MixedMesh) else None + rmapped = subelem.pullback.apply(rsub, domain=subdomain) # Flatten into the pulled back expression for the whole thing g_components.extend([rmapped[idx] for idx in np.ndindex(rmapped.ufl_shape)]) # And reshape appropriately @@ -455,11 +479,12 @@ def is_identity(self) -> bool: """Is this pull back the identity (or the identity applied to mutliple components).""" return True - def apply(self, expr): + def apply(self, expr, domain=None): """Apply the pull back. Args: expr: A function on a physical cell + domain: The domain on which the function is defined Returns: The function pulled back to the reference cell """ @@ -492,11 +517,12 @@ def is_identity(self) -> bool: """Is this pull back the identity (or the identity applied to mutliple components).""" return True - def apply(self, expr): + def apply(self, expr, domain=None): """Apply the pull back. Args: expr: A function on a physical cell + domain: The domain on which the function is defined Returns: The function pulled back to the reference cell """ diff --git a/ufl/restriction.py b/ufl/restriction.py index 2871cd53f..05eae3f04 100644 --- a/ufl/restriction.py +++ b/ufl/restriction.py @@ -56,3 +56,19 @@ class NegativeRestricted(Restricted): __slots__ = () _side = "-" + + +@ufl_type(is_terminal_modifier=True) +class SingleValueRestricted(Restricted): + """Single value restriction.""" + + __slots__ = () + _side = "|" + + +@ufl_type(is_terminal_modifier=True) +class ToBeRestricted(Restricted): + """Single value restriction.""" + + __slots__ = () + _side = "?"