From 3512b14be7c6cc4c1f21d36867396413eb1f01bb Mon Sep 17 00:00:00 2001 From: ksagiyam Date: Wed, 21 Aug 2024 02:41:28 +0100 Subject: [PATCH] introduce mixed mesh --- test/test_external_operator.py | 14 +- test/test_interpolate.py | 8 +- ...st_mixed_function_space_with_mixed_mesh.py | 81 ++++++++ ufl/__init__.py | 4 +- ufl/action.py | 4 +- ufl/algorithms/analysis.py | 43 ++-- ufl/algorithms/apply_derivatives.py | 107 ++++++++-- ufl/algorithms/compute_form_data.py | 12 +- ufl/algorithms/estimate_degrees.py | 9 +- ufl/differentiation.py | 6 +- ufl/domain.py | 190 ++++++++++++++++-- ufl/form.py | 64 +++--- ufl/functionspace.py | 2 + ufl/geometry.py | 5 +- ufl/index_combination_utils.py | 2 +- ufl/pullback.py | 71 ++++--- ufl/split_functions.py | 17 +- 17 files changed, 504 insertions(+), 135 deletions(-) create mode 100644 test/test_mixed_function_space_with_mixed_mesh.py 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..9750142fa --- /dev/null +++ b/test/test_mixed_function_space_with_mixed_mesh.py @@ -0,0 +1,81 @@ +from ufl import ( + CellVolume, + Coefficient, + FacetArea, + FacetNormal, + FunctionSpace, + Measure, + Mesh, + MixedMesh, + SpatialCoordinate, + TestFunction, + TrialFunction, + grad, + inner, + split, + triangle, +) +from ufl.algorithms import compute_form_data +from ufl.domain import extract_domains +from ufl.finiteelement import FiniteElement, MixedElement +from ufl.pullback import contravariant_piola, identity_pullback +from ufl.sobolevspace import H1, L2, HDiv + + +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) + x = SpatialCoordinate(mesh1) + form = x[1] * f0 * inner(grad(u0), v1) * dx1(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, + complex_mode=False, + ) + (id0,) = fd.integral_data + assert fd.preprocessed_form.arguments() == (v, u) + assert fd.reduced_coefficients == [f] + assert form.coefficients()[fd.original_coefficient_positions[0]] is f + 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] + + +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 9d255211a..73c0b6829 100644 --- a/ufl/__init__.py +++ b/ufl/__init__.py @@ -62,6 +62,7 @@ -AbstractDomain -Mesh + -MixedMesh -MeshView * Sobolev spaces:: @@ -263,7 +264,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 ( @@ -439,6 +440,7 @@ "TensorProductCell", "AbstractDomain", "Mesh", + "MixedMesh", "MeshView", "L2", "H1", diff --git a/ufl/action.py b/ufl/action.py index 78489ff6b..7dc221a80 100644 --- a/ufl/action.py +++ b/ufl/action.py @@ -188,9 +188,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_derivatives.py b/ufl/algorithms/apply_derivatives.py index 6fcbb1b5f..3baba3597 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 @@ -53,7 +55,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, @@ -82,6 +84,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.""" @@ -655,16 +668,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.""" @@ -676,10 +724,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/compute_form_data.py b/ufl/algorithms/compute_form_data.py index 687e8edc4..3089783e1 100644 --- a/ufl/algorithms/compute_form_data.py +++ b/ufl/algorithms/compute_form_data.py @@ -35,7 +35,7 @@ from ufl.algorithms.remove_complex_nodes import remove_complex_nodes 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 @@ -184,7 +184,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) @@ -262,6 +262,14 @@ def compute_form_data( The default arguments configured to behave the way old FFC expects. """ + # Currently, only integral_type="cell" can be used with MixedMesh. + for integral in form.integrals(): + if integral.integral_type() != "cell": + all_domains = extract_domains(integral.integrand(), expand_mixed_mesh=False) + if any(isinstance(m, MixedMesh) for m in all_domains): + raise NotImplementedError(""" + Only integral_type="cell" can be used with MixedMesh""") + # TODO: Move this to the constructor instead self = FormData() diff --git a/ufl/algorithms/estimate_degrees.py b/ufl/algorithms/estimate_degrees.py index da550116c..1cc42e998 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 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..eab56b129 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,119 @@ 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 + + +@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 +330,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 +345,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 +374,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 +406,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/form.py b/ufl/form.py index 620fb4e1e..af3dbf34f 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. @@ -241,6 +247,7 @@ class Form(BaseForm): "_coefficient_numbering", "_constants", "_constant_numbering", + "_geometric_quantities", "_terminal_numbering", "_hash", "_signature", @@ -599,15 +606,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.""" @@ -634,13 +648,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.""" @@ -651,38 +666,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/functionspace.py b/ufl/functionspace.py index cc047523c..4e85e8924 100644 --- a/ufl/functionspace.py +++ b/ufl/functionspace.py @@ -55,6 +55,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 0357c6d9b..d263128bc 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 @@ -68,11 +68,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 """ @@ -91,11 +92,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 """ @@ -126,17 +128,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 @@ -171,17 +174,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) @@ -214,17 +218,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 @@ -253,17 +258,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, ...) @@ -297,17 +303,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) @@ -348,31 +355,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( @@ -392,7 +405,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,) @@ -427,11 +443,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 """ @@ -443,6 +460,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] @@ -451,7 +473,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 @@ -492,11 +515,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 """ @@ -530,11 +554,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/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)