diff --git a/test/test_external_operator.py b/test/test_external_operator.py index ab996cce1..0169b3f94 100644 --- a/test/test_external_operator.py +++ b/test/test_external_operator.py @@ -205,10 +205,10 @@ def test_differentiation_procedure_action(V1, V2): def test_extractions(domain_2d, V1): from ufl.algorithms.analysis import ( extract_arguments, - extract_arguments_and_coefficients, extract_base_form_operators, extract_coefficients, extract_constants, + extract_terminals_with_domain, ) u = Coefficient(V1) @@ -219,7 +219,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_terminals_with_domain(e) == ([vstar_e], [u], []) assert extract_constants(e) == [c] assert extract_base_form_operators(e) == [e] @@ -227,7 +227,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_terminals_with_domain(e) == ([vstar_e], [u], []) assert extract_constants(F) == [c] assert F.base_form_operators() == (e,) @@ -236,14 +236,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_terminals_with_domain(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_terminals_with_domain(e) == ([vstar_e, u_hat], [u], []) assert F.base_form_operators() == (e,) w = Coefficient(V1) @@ -252,14 +252,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_terminals_with_domain(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_terminals_with_domain(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 877f55a35..cb6f38099 100644 --- a/test/test_interpolate.py +++ b/test/test_interpolate.py @@ -26,9 +26,9 @@ from ufl.algorithms.ad import expand_derivatives from ufl.algorithms.analysis import ( extract_arguments, - extract_arguments_and_coefficients, extract_base_form_operators, extract_coefficients, + extract_terminals_with_domain, ) from ufl.algorithms.expand_indices import expand_indices from ufl.core.interpolate import Interpolate @@ -157,12 +157,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_terminals_with_domain(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_terminals_with_domain(F) == ([], [u], []) for e in [Iu, F]: assert extract_coefficients(e) == [u] @@ -171,7 +171,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_terminals_with_domain(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..42e0d4cb1 --- /dev/null +++ b/test/test_mixed_function_space_with_mixed_mesh.py @@ -0,0 +1,109 @@ +from ufl import (triangle, Mesh, MixedMesh, FunctionSpace, TestFunction, TrialFunction, Coefficient, + 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) + 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 0782ebad7..f52a96d1f 100644 --- a/ufl/__init__.py +++ b/ufl/__init__.py @@ -62,6 +62,7 @@ -AbstractDomain -Mesh + -MixedMesh -MeshView * Sobolev spaces:: @@ -265,7 +266,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, MeshView, MixedMesh from ufl.finiteelement import AbstractFiniteElement from ufl.form import BaseForm, Form, FormSum, ZeroBaseForm from ufl.formoperators import ( @@ -442,6 +443,7 @@ "TensorProductCell", "AbstractDomain", "Mesh", + "MixedMesh", "MeshView", "L2", "H1", diff --git a/ufl/action.py b/ufl/action.py index fb6416865..c81bcabfe 100644 --- a/ufl/action.py +++ b/ufl/action.py @@ -196,9 +196,9 @@ 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 + from ufl.algorithms.analysis import extract_terminals_with_domain - right_args, right_coeffs = extract_arguments_and_coefficients(right) + right_args, right_coeffs, _ = extract_terminals_with_domain(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 1d0464a95..f6dc8d9cd 100644 --- a/ufl/algorithms/analysis.py +++ b/ufl/algorithms/analysis.py @@ -18,7 +18,9 @@ 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 +from ufl.domain import Mesh from ufl.form import BaseForm, Form +from ufl.geometry import GeometricQuantity from ufl.utils.sorting import sorted_by_count, topological_sorting # TODO: Some of these can possibly be optimised by implementing @@ -198,19 +200,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_terminals_with_domain(a): + """Build three sorted lists of all arguments, coefficients, and geometric quantities in `a`. - This function is faster than extract_arguments + extract_coefficients - for large forms, and has more validation built in. + This function is faster than extracting each type of terminal + separately 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)] + # Extract lists of all BaseArgument, BaseCoefficient, and GeometricQuantity instances + terminals = extract_type(a, (BaseArgument, BaseCoefficient, GeometricQuantity)) + arguments = [f for f in terminals if isinstance(f, BaseArgument)] + coefficients = [f for f in terminals if isinstance(f, BaseCoefficient)] + geometric_quantities = [f for f in terminals 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) @@ -226,20 +229,36 @@ 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: + if not isinstance(gq._domain, Mesh): + raise TypeError(f"{gq}._domain must be a Mesh: got {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_terminals_with_domain(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..4e7d1200d --- /dev/null +++ b/ufl/algorithms/apply_coefficient_split.py @@ -0,0 +1,210 @@ +"""Apply coefficient split. + +This module contains classes and functions to split coefficients defined on mixed function spaces. +""" + +import numpy +from ufl.classes import Restricted +from ufl.corealg.map_dag import map_expr_dag +from ufl.corealg.multifunction import MultiFunction, memoized_handler +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 + + @memoized_handler + 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)) + + @memoized_handler + def list_tensor(self, o): + cc = [] + for o1 in o.ufl_operands: + comp = map_expr_dag(self, o1) + cc.append(comp) + return ListTensor(*cc) + + @memoized_handler + 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 + + @memoized_handler + 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) + + @memoized_handler + def indexed(self, o): + o1, i1 = o.ufl_operands + if isinstance(o1, ComponentTensor): + o2, i2 = o1.ufl_operands + assert len(i2.indices()) == len(i1.indices()) + fimap = dict(zip(i2.indices(), i1.indices())) + 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 da7b61da1..6e610ecae 100644 --- a/ufl/algorithms/apply_derivatives.py +++ b/ufl/algorithms/apply_derivatives.py @@ -10,6 +10,8 @@ from collections import defaultdict from math import pi +import numpy as np + from ufl.action import Action from ufl.algorithms.analysis import extract_arguments from ufl.algorithms.map_integrands import map_integrand_dags @@ -55,7 +57,7 @@ BaseFormOperatorDerivative, CoordinateDerivative, ) -from ufl.domain import extract_unique_domain +from ufl.domain import MixedMesh, extract_unique_domain from ufl.form import Form, ZeroBaseForm from ufl.operators import ( bessel_I, @@ -84,6 +86,17 @@ # - 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 = () + assert len(domain) == len(element.sub_elements) + for d, e in zip(domain, element.sub_elements): + flattened += flatten_domain_element(d, e) + return flattened + + class GenericDerivativeRuleset(MultiFunction): """A generic derivative.""" @@ -657,16 +670,51 @@ 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() + if element.num_sub_elements != len(domain): + raise RuntimeError(f"{element.num_sub_elements} != {len(domain)}") + g = ReferenceGrad(o) + 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(np.prod(esh))): + for i in range(ref_dim): + components.append(g[(dofoffset + idx,) + (i,)]) + else: + K = JacobianInverse(d) + rdim, gdim = K.ufl_shape + if rdim != ref_dim: + raise RuntimeError(f"{rdim} != {ref_dim}") + if gdim != self._var_shape[0]: + raise RuntimeError(f"{gdim} != {self._var_shape[0]}") + for idx in range(int(np.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(np.prod(esh)) + return as_tensor(np.asarray(components).reshape(g.ufl_shape[:-1] + 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.""" @@ -678,10 +726,43 @@ def reference_grad(self, o): ) 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() + if element.num_sub_elements != len(domain): + raise RuntimeError(f"{element.num_sub_elements} != {len(domain)}") + g = ReferenceGrad(o) + 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 + if rdim != ref_dim: + raise RuntimeError(f"{rdim} != {ref_dim}") + if gdim != self._var_shape[0]: + raise RuntimeError(f"{gdim} != {self._var_shape[0]}") + for idx in range(int(np.prod(esh))): + for midx in np.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(np.prod(esh)) + if g.ufl_shape[0] != dofoffset: + raise RuntimeError(f"{g.ufl_shape[0]} != {dofoffset}") + return as_tensor(np.asarray(components).reshape(g.ufl_shape[:-1] + 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 a7550d722..03f10e0ad 100644 --- a/ufl/algorithms/apply_restrictions.py +++ b/ufl/algorithms/apply_restrictions.py @@ -10,28 +10,36 @@ # # SPDX-License-Identifier: LGPL-3.0-or-later +from ufl.algorithms.formdata import FormData 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.form import Form 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.""" @@ -57,9 +65,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.""" @@ -174,25 +185,41 @@ 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() - return map_integrand_dags(rules, expression, only_integral_type=integral_types) + 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) + if isinstance(expression, FormData): + for integral_data in expression.integral_data: + integral_data.integrals = tuple( + map_integrand_dags(rules, integral, only_integral_type=integral_types) + for integral in integral_data.integrals + ) + return expression + else: + return map_integrand_dags(rules, expression, only_integral_type=integral_types) 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.""" @@ -237,13 +264,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 e671d01d0..8a0b5fb18 100644 --- a/ufl/algorithms/balancing.py +++ b/ufl/algorithms/balancing.py @@ -13,6 +13,8 @@ Indexed, NegativeRestricted, PositiveRestricted, + SingleValueRestricted, + ToBeRestricted, ReferenceGrad, ReferenceValue, ) @@ -27,6 +29,8 @@ FacetAvg, PositiveRestricted, NegativeRestricted, + SingleValueRestricted, + ToBeRestricted, Indexed, ] @@ -86,6 +90,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 e1d9b366b..a5e51cbce 100644 --- a/ufl/algorithms/check_arities.py +++ b/ufl/algorithms/check_arities.py @@ -113,6 +113,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 687e8edc4..9683a4b73 100644 --- a/ufl/algorithms/compute_form_data.py +++ b/ufl/algorithms/compute_form_data.py @@ -19,7 +19,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 @@ -33,10 +35,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 MixedMesh, extract_domains, extract_unique_domain from ufl.utils.sequences import max_degree +from ufl.constantvalue import Zero def _auto_select_degree(elements): @@ -184,7 +188,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) @@ -256,6 +260,8 @@ def compute_form_data( 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. @@ -306,9 +312,9 @@ def compute_form_data( if do_apply_integral_scaling: form = apply_integral_scaling(form) - # Apply default restriction to fully continuous terminals - if do_apply_default_restrictions: - form = apply_default_restrictions(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 # Lower abstractions for geometric quantities into a smaller set # of quantities, allowing the form compiler to deal with a smaller @@ -332,9 +338,13 @@ def compute_form_data( form = apply_coordinate_derivatives(form) - # Propagate restrictions to terminals - if do_apply_restrictions: - form = apply_restrictions(form) + # Apply default restriction to fully continuous terminals + if do_apply_default_restrictions: + 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) # If in real mode, remove any complex nodes introduced during form processing. if not complex_mode: @@ -344,6 +354,13 @@ def compute_form_data( # Most of the heavy lifting is done above in group_form_integrals. self.integral_data = build_integral_data(form.integrals()) + # Propagate restrictions to terminals + if do_apply_restrictions: + if do_assume_single_integral_type: + apply_restrictions(self) + else: + apply_restrictions(self, assume_single_integral_type=have_single_domain) + # --- Create replacements for arguments and coefficients # Figure out which form coefficients each integral should enable @@ -415,6 +432,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: + assert len(mesh) == len(elem.sub_elements) + coefficient_split[c] = [Coefficient(FunctionSpace(m, e)) + for m, e in zip(mesh, elem.sub_elements)] + 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 = 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 19e0f4e3c..63cc454aa 100644 --- a/ufl/algorithms/domain_analysis.py +++ b/ufl/algorithms/domain_analysis.py @@ -39,6 +39,7 @@ class IntegralData(object): "metadata", "integral_coefficients", "enabled_coefficients", + "domain_integral_type_map", ) def __init__(self, domain, integral_type, subdomain_id, integrals, metadata): @@ -60,6 +61,7 @@ def __init__(self, domain, integral_type, subdomain_id, integrals, metadata): # 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 da550116c..ee8cfe7dc 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_domains, extract_unique_domain from ufl.form import Form from ufl.integral import Integral @@ -98,10 +98,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 @@ -181,6 +180,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 20a6a1743..679d5dc94 100644 --- a/ufl/algorithms/formtransformations.py +++ b/ufl/algorithms/formtransformations.py @@ -241,6 +241,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 74e33617d..ee5844e41 100644 --- a/ufl/differentiation.py +++ b/ufl/differentiation.py @@ -307,14 +307,16 @@ 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() + # TODO: Use max topological dimension if there are multiple topological dimensions. + 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) def __init__(self, f): """Initalise.""" CompoundDerivative.__init__(self, (f,)) - self._dim = extract_unique_domain(f).topological_dimension() + # TODO: Use max topological dimension if there are multiple topological dimensions. + 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 8f4ae5138..961b0ba93 100644 --- a/ufl/domain.py +++ b/ufl/domain.py @@ -50,6 +50,33 @@ def topological_dimension(self): """Return the dimension of the topology of this domain.""" return self._topological_dimension + @property + def meshes(self): + """Return the component meshes.""" + raise NotImplementedError("meshes() method not implemented") + + def __len__(self): + """Return number of component meshes.""" + return len(self.meshes) + + def __getitem__(self, i): + """Return i-th component mesh.""" + if i >= len(self): + raise ValueError(f"index ({i}) >= num. component meshes ({len(self)})") + return self.meshes[i] + + def __iter__(self): + """Return iterable component meshes.""" + return iter(self.meshes) + + def iterable_like(self, element): + """Return iterable object that is iterable like ``element``.""" + raise NotImplementedError("iterable_like() method not implemented") + + def can_make_function_space(self, element): + """Check whether this mesh can make a function space with ``element``.""" + raise NotImplementedError("check_compatibility() method not implemented") + # TODO: Would it be useful to have a domain representing R^d? E.g. for # Expression. @@ -126,6 +153,128 @@ def _ufl_sort_key_(self): typespecific = (self._ufl_id, self._ufl_coordinate_element) return (self.geometric_dimension(), self.topological_dimension(), "Mesh", typespecific) + @property + def meshes(self): + """Return the component meshes.""" + return (self,) + + def iterable_like(self, element): + """Return iterable object that is iterable like ``element``.""" + return iter(self for _ in element.sub_elements) + + def can_make_function_space(self, element): + """Check whether this mesh can make a function space with ``element``.""" + # Can use with any element. + return True + + def iterable_like(self, element): + """Return iterable object that is iterable like ``element``.""" + return iter(self for _ in element.sub_elements) + + def can_make_function_space(self, element): + """Check whether this mesh can make a function space with ``element``.""" + # Can use with any element. + return True + + +@attach_ufl_id +class MixedMesh(AbstractDomain, UFLObject): + """Symbolic representation of a mixed mesh. + + This class represents a collection of meshes that, along with + a :class:`MixedElement`, represent a mixed function space defined on + multiple domains. This abstraction allows for defining the + mixed function space with the conventional :class:`FunctionSpace` + class and integrating multi-domain problems seamlessly. + + Currently, all component meshes must have the same cell type (and + thus the same topological dimension). + + Currently, one can only perform cell integrations when + :class:`MixedMesh`es are used. + + .. code-block:: python3 + + cell = triangle + mesh0 = Mesh(FiniteElement("Lagrange", cell, 1, (2,), identity_pullback, H1)) + mesh1 = Mesh(FiniteElement("Lagrange", cell, 1, (2,), identity_pullback, H1)) + domain = MixedMesh([mesh0, mesh1]) + elem0 = FiniteElement("Lagrange", cell, 1, (), identity_pullback, H1) + elem1 = FiniteElement("Lagrange", cell, 2, (), identity_pullback, H1) + elem = MixedElement([elem0, elem1]) + V = FunctionSpace(domain, elem) + v = TestFunction(V) + v0, v1 = split(v) + + """ + + 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""") + # 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) + # TODO: Need to change for more general mixed meshes. + (tdim,) = set(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.""" + # TODO: Might need MixedCell class for more general mixed meshes. + return self._ufl_cell + + def __repr__(self): + """Representation.""" + return "MixedMesh(%s, ufl_id=%s)" % (repr(self._meshes), repr(self._ufl_id)) + + 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) + + @property + def meshes(self): + """Return the component meshes.""" + return self._meshes + + def iterable_like(self, element): + """Return iterable object that is iterable like ``element``.""" + if len(self) != element.num_sub_elements: + raise RuntimeError(f"""len(self) ({len(self)}) != + element.num_sub_elements ({element.num_sub_elements})""") + return self + + def can_make_function_space(self, element): + """Check whether this mesh can make a function space with ``element``.""" + if len(self) != element.num_sub_elements: + return False + else: + return all(d.can_make_function_space(e) for d, e in zip(self, element.sub_elements)) + @attach_ufl_id class MeshView(AbstractDomain, UFLObject): @@ -190,12 +339,14 @@ def as_domain(domain): """Convert any valid object to an AbstractDomain type.""" if isinstance(domain, AbstractDomain): # Modern UFL files and dolfin behaviour + (domain,) = set(domain.meshes) return domain - try: return extract_unique_domain(domain) except AttributeError: - return domain.ufl_domain() + domain = domain.ufl_domain() + (domain,) = set(domain.meshes) + return domain def sort_domains(domains): @@ -203,13 +354,19 @@ 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_: + domains.update(domain.meshes) + else: + domains = domains_ if not domains: return () @@ -226,20 +383,26 @@ 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), - key=lambda D: (D.topological_dimension(), D.ufl_cell(), D.ufl_id()), - ) + from ufl.form import Form + + if isinstance(expr, Form): + if not expand_mixed_mesh: + raise NotImplementedError(""" + Currently, can only extract domains from a Form with expand_mixed_mesh=True""") + # 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: @@ -252,9 +415,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 len(gdims) != 1: raise ValueError("Cannot determine geometric dimension from expression.") diff --git a/ufl/exproperators.py b/ufl/exproperators.py index ff532e843..963a14422 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 @@ -310,7 +310,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): @@ -319,6 +319,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.") @@ -342,7 +346,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 4fc7c2165..2dd4041f0 100644 --- a/ufl/form.py +++ b/ufl/form.py @@ -108,6 +108,12 @@ def coefficients(self): self._analyze_form_arguments() return self._coefficients + def geometric_quantities(self): + """Return all ``GeometricQuantity`` objects found in form.""" + if self._geometric_quantities is None: + self._analyze_form_arguments() + return self._geometric_quantities + def ufl_domain(self): """Return the single geometric integration domain occuring in the base form. @@ -240,6 +246,7 @@ class Form(BaseForm): "_coefficient_numbering", "_constants", "_constant_numbering", + "_geometric_quantities", "_terminal_numbering", "_hash", "_signature", @@ -578,15 +585,22 @@ def _analyze_domains(self): """Analyze domains.""" 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. + domains_in_integrands = set() + for o in chain( + self.arguments(), self.coefficients(), self.constants(), self.geometric_quantities() + ): + domain = extract_unique_domain(o, expand_mixed_mesh=False) + domains_in_integrands.update(domain.meshes) + domains_in_integrands -= set(self._integration_domains) + all_domains = self._integration_domains + sort_domains(join_domains(domains_in_integrands)) + # Let problem solving environments access all domains via + # self._domain_numbering.keys() (wrapped in extract_domains()). + self._domain_numbering = dict((d, i) for i, d in enumerate(all_domains)) def _analyze_subdomain_data(self): """Analyze subdomain data.""" @@ -613,13 +627,14 @@ 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 + from ufl.algorithms.analysis import extract_terminals_with_domain - arguments, coefficients = extract_arguments_and_coefficients(self) + arguments, coefficients, geometric_quantities = extract_terminals_with_domain(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.""" @@ -630,38 +645,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 5d8d5f0c5..e827f8d14 100644 --- a/ufl/formatting/ufl2unicode.py +++ b/ufl/formatting/ufl2unicode.py @@ -138,6 +138,8 @@ class UC: superscript_plus = "⁺" superscript_minus = "⁻" + superscript_vertical_bar = "|" + superscript_question_mark = "?" superscript_equals = "⁼" superscript_left_paren = "⁽" superscript_right_paren = "⁾" @@ -749,6 +751,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/functionspace.py b/ufl/functionspace.py index 834734274..558b6bd73 100644 --- a/ufl/functionspace.py +++ b/ufl/functionspace.py @@ -57,6 +57,8 @@ def __init__(self, domain, element, label=""): else: if element.cell != domain_cell: raise ValueError("Non-matching cell of finite element and domain.") + if not domain.can_make_function_space(element): + raise ValueError(f"Mismatching domain ({domain}) and element ({element}).") AbstractFunctionSpace.__init__(self) self._label = label self._ufl_domain = domain diff --git a/ufl/geometry.py b/ufl/geometry.py index 1acaf40a2..1a5e0d24a 100644 --- a/ufl/geometry.py +++ b/ufl/geometry.py @@ -8,7 +8,7 @@ from ufl.core.terminal import Terminal from ufl.core.ufl_type import ufl_type -from ufl.domain import as_domain, extract_unique_domain +from ufl.domain import MixedMesh, as_domain, extract_unique_domain from ufl.sobolevspace import H1 """ @@ -84,6 +84,9 @@ class GeometricQuantity(Terminal): def __init__(self, domain): """Initialise.""" Terminal.__init__(self) + if isinstance(domain, MixedMesh) and len(set(domain)) > 1: + # Can not make GeometricQuantity if multiple domains exist. + raise TypeError(f"Can not create a GeometricQuantity on {domain}") self._domain = as_domain(domain) def ufl_domains(self): diff --git a/ufl/index_combination_utils.py b/ufl/index_combination_utils.py index 8bd5087a8..cb3266dad 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 0dc780e52..b0bde60a6 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 AbstractDomain, MixedMesh, extract_unique_domain from ufl.functionspace import FunctionSpace from ufl.tensors import as_tensor @@ -69,11 +69,12 @@ def physical_value_shape(self, element, domain) -> 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 """ @@ -92,11 +93,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 """ @@ -127,17 +129,18 @@ 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) + domain = domain or extract_unique_domain(expr) J = Jacobian(domain) detJ = JacobianDeterminant(J) transform = (1.0 / detJ) * J @@ -172,17 +175,18 @@ 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) + domain = domain or extract_unique_domain(expr) K = JacobianInverse(domain) # Apply transform "row-wise" to TensorElement(PiolaMapped, ...) *k, i, j = indices(len(expr.ufl_shape) + 1) @@ -215,17 +219,18 @@ 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) + domain = domain or extract_unique_domain(expr) detJ = JacobianDeterminant(domain) return expr / detJ @@ -254,17 +259,18 @@ 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) + domain = domain or extract_unique_domain(expr) J = Jacobian(domain) detJ = JacobianDeterminant(J) # Apply transform "row-wise" to TensorElement(PiolaMapped, ...) @@ -298,17 +304,18 @@ 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) + domain = domain or 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) @@ -394,31 +401,37 @@ 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 """ - domain = extract_unique_domain(expr) - space = FunctionSpace(domain, self._element) rflat = [expr[idx] for idx in np.ndindex(expr.ufl_shape)] g_components = [] offset = 0 # For each unique piece in reference space, apply the appropriate pullback - for subelem in self._element.sub_elements: + domain = domain or extract_unique_domain(expr, expand_mixed_mesh=False) + if isinstance(domain, MixedMesh): + if len(domain) != self._element.num_sub_elements: + raise ValueError(f"""num. component meshes ({len(domain)}) != + num. sub elements ({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 # And reshape appropriately + space = FunctionSpace(domain, self._element) f = as_tensor(np.asarray(g_components).reshape(space.value_shape)) if f.ufl_shape != space.value_shape: raise ValueError( @@ -438,7 +451,10 @@ def physical_value_shape(self, element, domain) -> typing.Tuple[int, ...]: The value shape when the pull back is applied to the given element """ assert element == self._element - dim = sum(FunctionSpace(domain, e).value_size for e in self._element.sub_elements) + domains = domain.iterable_like(element) + dim = sum( + FunctionSpace(d, e).value_size for d, e in zip(domains, self._element.sub_elements) + ) return (dim,) @@ -473,11 +489,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 """ @@ -489,6 +506,11 @@ 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 + domain = domain or extract_unique_domain(expr, expand_mixed_mesh=False) + if isinstance(domain, MixedMesh): + if len(domain) != self._element.num_sub_elements: + raise ValueError(f"""num. component meshes ({len(domain)}) != + num. sub elements ({self._element.num_sub_elements})""") for component in np.ndindex(self._block_shape): i = self._symmetry[component] subelem = self._element.sub_elements[i] @@ -497,7 +519,8 @@ def apply(self, expr): 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 @@ -540,11 +563,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 """ @@ -578,11 +602,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 430fefa41..99f481a00 100644 --- a/ufl/restriction.py +++ b/ufl/restriction.py @@ -57,3 +57,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 = "?" diff --git a/ufl/split_functions.py b/ufl/split_functions.py index 4f4cb4aaf..c015f1376 100644 --- a/ufl/split_functions.py +++ b/ufl/split_functions.py @@ -22,8 +22,6 @@ def split(v): If v is a Coefficient or Argument in a mixed space, returns a tuple with the function components corresponding to the subelements. """ - domain = extract_unique_domain(v) - # Default range is all of v begin = 0 end = None @@ -62,6 +60,8 @@ def split(v): "Don't know how to split tensor valued mixed functions without flattened index space." ) + domain = extract_unique_domain(v, expand_mixed_mesh=False) + # Compute value size and set default range end value_size = v.ufl_function_space().value_size if end is None: @@ -71,11 +71,13 @@ def split(v): # corresponding to beginning of range j = begin while True: - for e in element.sub_elements: - if j < FunctionSpace(domain, e).value_size: + domains = domain.iterable_like(element) + for d, e in zip(domains, element.sub_elements): + if j < FunctionSpace(d, e).value_size: + domain = d element = e break - j -= FunctionSpace(domain, e).value_size + j -= FunctionSpace(d, e).value_size # Then break when we find the subelement that covers the whole range if FunctionSpace(domain, element).value_size == (end - begin): break @@ -83,10 +85,11 @@ def split(v): # Build expressions representing the subfunction of v for each subelement offset = begin sub_functions = [] - for i, e in enumerate(element.sub_elements): + domains = domain.iterable_like(element) + for i, (d, e) in enumerate(zip(domains, element.sub_elements)): # Get shape, size, indices, and v components # corresponding to subelement value - shape = FunctionSpace(domain, e).value_shape + shape = FunctionSpace(d, e).value_shape strides = shape_to_strides(shape) rank = len(shape) sub_size = product(shape)