23427573198

Committed 10 Mar 2026 06:22PM UTC coverage: 77.13% (+0.1%) from 76.999%

Build # 23427573198

Build Type

push

github

Committed by

web-flow

Commit Message

Update `compute_form_data` and `FormData` to use a proper class and remove cell and facet average from balance modifiers. (#469)

* Cell and facet averages is not terminal modifiers.

* Move form data creation into form data

* Mypy fixes ones typing has been activated

* Fix typehint

* Make some consistency  across Integral and IntegralData

* Fixes related to @schnellerhase's comments

Coverage Stats

222 of 236 new or added lines in 7 files covered. (94.07%)

3 existing lines in 2 files now uncovered.

9342 of 12112 relevant lines covered (77.13%)

0.77 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

93.75

/ufl/algorithms/formdata.py

"""FormData class easy for collecting of various data about a form."""

# Copyright (C) 2008-2016 Martin Sandve Alnæs
#
# This file is part of UFL (https://www.fenicsproject.org)
#
# SPDX-License-Identifier:    LGPL-3.0-or-later
#
# Modified by Anders Logg, 2008.
# Modified by Jørgen S. Dokken, 2026.

from functools import cached_property
from itertools import chain
from typing import Any

from ufl.algorithms.analysis import extract_coefficients, extract_sub_elements, unique_tuple
from ufl.algorithms.apply_coefficient_split import CoefficientSplitter
from ufl.algorithms.apply_restrictions import apply_restrictions, default_restriction_map
from ufl.algorithms.check_arities import check_integrand_arity
from ufl.algorithms.domain_analysis import IntegralData, reconstruct_form_from_integral_data
from ufl.algorithms.replace import replace
from ufl.classes import Argument, Coefficient, FunctionSpace, GeometricFacetQuantity
from ufl.corealg.traversal import traverse_unique_terminals
from ufl.domain import AbstractDomain, MeshSequence, extract_domains, extract_unique_domain
from ufl.form import Form, Zero
from ufl.utils.formatting import estr, lstr, tstr
from ufl.utils.sequences import max_degree


def _check_form_arity(
    integral_data: list[IntegralData], arguments: tuple[Argument], complex_mode: bool
):
    """Check if form arity is valid."""
    for integral in integral_data:
        for itg in integral.integrals:
            check_integrand_arity(itg.integrand(), arguments, complex_mode)


def _auto_select_degree(elements):
    """Automatically select degree for all elements of the form.

    This is be used in cases where the degree has not been specified by the user.
    This feature is used by DOLFIN to allow the specification of Expressions with
    undefined degrees.
    """
    # Use max degree of all elements, at least 1 (to work with
    # Lagrange elements)
    return max_degree({e.embedded_superdegree for e in elements} - {None} | {1})


def _compute_element_mapping(form: Form):
    """Compute element mapping for element replacement."""
    # The element mapping is a slightly messy concept with two use
    # cases:
    # - Expression with missing cell or element TODO: Implement proper
    #   Expression handling in UFL and get rid of this
    # - Constant with missing cell TODO: Fix anything that needs to be
    #   worked around to drop this requirement

    # Extract all elements and include subelements of mixed elements
    elements = [obj.ufl_element() for obj in chain(form.arguments(), form.coefficients())]
    elements = extract_sub_elements(elements)

    # Try to find a common degree for elements
    common_degree = _auto_select_degree(elements)

    # Compute element map
    element_mapping = {}
    for element in elements:
        # Flag for whether element needs to be reconstructed
        reconstruct = False

        # Set cell
        cell = element.cell
        if cell is None:
            domains = form.ufl_domains()
            if not all(domains[0].ufl_cell() == d.ufl_cell() for d in domains):
                raise ValueError(
                    "Cannot replace unknown element cell without unique common cell in form."
                )
            cell = domains[0].ufl_cell()
            reconstruct = True

        # Set degree
        degree = element.embedded_superdegree
        if degree is None:
            degree = common_degree
            reconstruct = True

        # Reconstruct element and add to map
        if reconstruct:
            element_mapping[element] = element.reconstruct(cell=cell, degree=degree)
        else:
            element_mapping[element] = element

    return element_mapping


def _compute_max_subdomain_ids(integral_data: list[IntegralData]) -> dict[str, int]:
    """Compute the maximum subdomain ids."""
    max_subdomain_ids: dict[str, int] = {}
    for itg_data in integral_data:
        it = itg_data.integral_type
        for integral in itg_data.integrals:
            # Convert string for default integral to -1
            sids = (-1 if isinstance(si, str) else si for si in integral.subdomain_id())
            newmax = max(sids) + 1
            prevmax = max_subdomain_ids.get(it, 0)
            max_subdomain_ids[it] = max(prevmax, newmax)
    return max_subdomain_ids


def _check_elements(form_data):
    """Check elements."""
    for element in chain(form_data.unique_elements, form_data.unique_sub_elements):
        if element.cell is None:
            raise ValueError(f"Found element with undefined cell: {element}")


def _check_facet_geometry(integral_data):
    """Check facet geometry."""
    for itg_data in integral_data:
        for itg in itg_data.integrals:
            for expr in traverse_unique_terminals(itg.integrand()):
                cls = expr._ufl_class_
                if issubclass(cls, GeometricFacetQuantity):
                    domain = extract_unique_domain(expr, expand_mesh_sequence=False)
                    if isinstance(domain, MeshSequence):
                        raise RuntimeError(
                            f"Not expecting a terminal object on a "
                            f"mesh sequence at this stage: found {expr!r}"
                        )
                    it = itg_data.domain_integral_type_map[domain]
                    # Facet geometry is only valid in facet integrals.
                    # Allowing custom integrals to pass as well, although
                    # that's not really strict enough.
                    if not ("facet" in it or "custom" in it or "interface" in it):
                        # Not a facet integral
                        raise ValueError(f"Integral of type {it} cannot contain a {cls.__name__}.")


def _build_coefficient_replace_map(
    coefficients: list[Coefficient], element_mapping=None
) -> tuple[list[Coefficient], dict[Coefficient, Coefficient]]:
    """Create new Coefficient objects with count starting at 0.

    Returns:
        lists of the new objects and mapping from old to new objects
    """
    if element_mapping is None:
        element_mapping = {}

    new_coefficients = []
    replace_map = {}
    for i, f in enumerate(coefficients):
        old_e = f.ufl_element()
        new_e = element_mapping.get(old_e, old_e)
        # XXX: This is a hack to ensure that if the original
        # 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, expand_mesh_sequence=False)
        if domain is not None:
            new_e = FunctionSpace(domain, new_e)
        new_f = Coefficient(new_e, count=i)
        new_coefficients.append(new_f)
        replace_map[f] = new_f

    return new_coefficients, replace_map


class FormData:
    """Class collecting various information extracted from a Form by calling preprocess."""

    _original_form: Form
    _integral_data: list[IntegralData]
    _reduced_coefficients: list[Coefficient]
    _original_coefficient_positions: list[int]
    _function_replace_map: dict[Coefficient, Coefficient]
    _coefficient_elements: tuple[Any, ...]
    _coefficient_split: dict[Coefficient, list[Coefficient]]

    def __init__(
        self,
        original_form: Form,
        integral_data: list[IntegralData],
        do_apply_default_restrictions: bool = True,
        do_apply_restrictions: bool = True,
        do_replace_functions: bool = False,
        coefficients_to_split: tuple[Coefficient, ...] | None = None,
        complex_mode: bool = False,
    ):
        """Create form-data for a form that has been processed.

        Args:
            original_form: The form.
            integral_data: List of integral data objects corresponding to
                the form.
            do_apply_default_restrictions: Apply default restrictions, defined in
                {py:mod}`ufl.algorithms.apply_restrictions` to integrals if no
                restriction has been set.
            do_apply_restrictions: Apply restrictions towards terminal nodes.
            do_replace_functions: Replace functions with with its cannonically numbered
                function or thos provided in coefficients_to_split.
            coefficients_to_split: Sequence of coefficients to split over a MeshSequence
            complex_mode: If form has been processed as complex or not.
        """
        self._original_form = original_form
        self._integral_data = integral_data

        # --- Create replacements for arguments and coefficients

        # Figure out which form coefficients each integral should enable
        for itg_data in self.integral_data:
            itg_coeffs = set()
            # Get all coefficients in integrand
            for itg in itg_data.integrals:
                itg_coeffs.update(extract_coefficients(itg.integrand()))
            # Store with IntegralData object
            itg_data.integral_coefficients = itg_coeffs

        # Figure out which coefficients from the original form are
        # actually used in any integral (Differentiation may reduce the
        # set of coefficients w.r.t. the original form)
        reduced_coefficients_set: set[Coefficient] = set()
        for itg_data in self.integral_data:
            assert itg_data.integral_coefficients is not None
            reduced_coefficients_set.update(itg_data.integral_coefficients)
        self._reduced_coefficients = sorted(reduced_coefficients_set, key=lambda c: c.count())
        self._original_coefficient_positions = [
            i
            for i, c in enumerate(self.original_form.coefficients())
            if c in self.reduced_coefficients
        ]

        # Store back into integral data which form coefficients are used
        # by each integral
        for itg_data in self.integral_data:
            assert itg_data.integral_coefficients is not None
            itg_data.enabled_coefficients = [
                bool(coeff in itg_data.integral_coefficients) for coeff in self.reduced_coefficients
            ]

        # Mappings from elements and coefficients that reside in form to
        # objects with canonical numbering as well as completed cells and
        # elements
        renumbered_coefficients, function_replace_map = _build_coefficient_replace_map(
            self.reduced_coefficients, self.element_replace_map
        )
        self._function_replace_map = function_replace_map

        self._coefficient_elements = tuple(f.ufl_element() for f in renumbered_coefficients)

        # --- Apply replace(integrand, self.function_replace_map)
        if do_replace_functions:
            for itg_data in self.integral_data:
                new_integrals = []
                for integral in itg_data.integrals:
                    integrand = replace(integral.integrand(), self.function_replace_map)
                    new_integrals.append(integral.reconstruct(integrand=integrand))
                itg_data.integrals = new_integrals

        # --- Split mixed coefficients with their components
        if coefficients_to_split is None:
            self._coefficient_split = {}
        else:
            if not do_replace_functions:
                raise ValueError("Must call with do_replace_functions=True")
            for itg_data in self.integral_data:
                new_integrals = []
                for integral in itg_data.integrals:
                    # Propagate restrictions as required by CoefficientSplitter.
                    # Can not yet apply default restrictions at this point
                    # as restrictions can only be applied to the components.
                    new_integral = apply_restrictions(integral)
                    new_integrals.append(new_integral)
                itg_data.integrals = new_integrals
            # Split coefficients that are contained in ``coefficients_to_split``
            # into components, and store a dict in ``self`` that maps
            # each coefficient to its components.
            coefficient_split = {}
            for o in self.reduced_coefficients:
                if o in coefficients_to_split:
                    c = self.function_replace_map[o]
                    mesh = extract_unique_domain(c, expand_mesh_sequence=False)
                    elem = c.ufl_element()
                    coefficient_split[c] = [
                        Coefficient(FunctionSpace(m, e))
                        for m, e in zip(mesh.iterable_like(elem), elem.sub_elements)  # type: ignore
                    ]
            self._coefficient_split = coefficient_split
            coeff_splitter = CoefficientSplitter(self.coefficient_split)
            for itg_data in self.integral_data:
                new_integrals = []
                for integral in itg_data.integrals:
                    integrand = coeff_splitter(integral.integrand())
                    # Potentially need to call `remove_component_tensors()` here, but
                    # early-simplifications of Indexed objects seem sufficient.
                    if not isinstance(integrand, Zero):
                        new_integrals.append(integral.reconstruct(integrand=integrand))
                itg_data.integrals = new_integrals

        # Propagate restrictions to terminals
        if do_apply_restrictions:
            for itg_data in self.integral_data:
                # Need the following if block in case not all participating domains
                # have been included in the Measure (backwards compat).
                if all(
                    not integral_type.startswith("interior_facet")
                    for _, integral_type in itg_data.domain_integral_type_map.items()
                ):
                    continue
                default_restrictions: dict[AbstractDomain, str | None] | None
                if do_apply_default_restrictions:
                    default_restrictions = {
                        domain: default_restriction_map[integral_type]
                        for domain, integral_type in itg_data.domain_integral_type_map.items()
                    }
                    # Need the following dict update in case not all participating domains
                    # have been included in the Measure (backwards compat).
                    extra = {
                        domain: default_restriction_map[itg_data.integral_type]
                        for integral in itg_data.integrals
                        for domain in extract_domains(integral)
                        if domain not in default_restrictions
                    }
                    default_restrictions.update(extra)
                else:
                    default_restrictions = None
                new_integrals = []
                for integral in itg_data.integrals:
                    new_integral = apply_restrictions(
                        integral,
                        default_restrictions=default_restrictions,
                    )
                    new_integrals.append(new_integral)
                itg_data.integrals = new_integrals

        _check_elements(self)
        _check_facet_geometry(self.integral_data)
        _check_form_arity(self.integral_data, self.original_form.arguments(), complex_mode)

    def __str__(self):
        """Return formatted summary of form data."""
        types = sorted(self.max_subdomain_ids.keys())
        geometry = (("Geometric dimension", self.geometric_dimension),)
        subdomains = tuple(
            (f"Number of {integral_type} subdomains", self.max_subdomain_ids[integral_type])
            for integral_type in types
        )
        functions = (
            # Arguments
            ("Rank", self.rank),
            ("Arguments", lstr(self.original_form.arguments())),
            # Coefficients
            ("Number of coefficients", self.num_coefficients),
            ("Coefficients", lstr(self.reduced_coefficients)),
            # Elements
            ("Unique elements", estr(self.unique_elements)),
            ("Unique sub elements", estr(self.unique_sub_elements)),
        )
        return tstr(geometry + subdomains + functions)

    @property
    def integral_data(self) -> list[IntegralData]:
        """The integral data of the form."""
        return self._integral_data

    @property
    def reduced_coefficients(self) -> list[Coefficient]:
        """Set of active coeffcient in the form."""
        return self._reduced_coefficients

    @property
    def num_coefficients(self):
        """Number of active coefficients in the form."""
        return len(self.reduced_coefficients)

    @property
    def rank(self):
        """Rank of the form."""
        return len(self.original_form.arguments())

    @cached_property
    def geometric_dimension(self):
        """Geometric dimension of the form."""
        domains = extract_domains(self.original_form())
        gdims = set([domain.geometric_dimension for domain in domains])
        assert len(gdims) == 1
        return gdims.pop()

    @property
    def function_replace_map(self) -> dict[Coefficient, Coefficient]:
        """Map from coefficients in form to those used in IntegralData."""
        return self._function_replace_map

    @cached_property
    def element_replace_map(self):
        """Mapping from incomplete elements to new well-defined elements.

        This is to support the Expression construct in dolfin which
        subclasses Coefficient but doesn't provide an element,
        and the Constant construct that doesn't provide the domain that
        a Coefficient is supposed to have. A future design iteration in
        UFL/UFC/FFC/DOLFIN may allow removal of this mapping with the
        introduction of UFL types for Expression-like functions that can
        be evaluated in quadrature points.

        Note:
            This property should likely be removed.
        """
        return _compute_element_mapping(self.original_form)

    @cached_property
    def max_subdomain_ids(self) -> dict[str, int]:
        """For each integral type, return the maximum subdomain id."""
        return _compute_max_subdomain_ids(self.integral_data)

    @cached_property
    def argument_elements(self) -> tuple[Any, ...]:
        """The set of elements the arguments in the form."""
        return tuple(f.ufl_element() for f in self.original_form.arguments())

    @property
    def coefficient_elements(self) -> tuple[Any, ...]:
        """The set of elements used for coefficients in the form."""
        return self._coefficient_elements

    @cached_property
    def coordinate_elements(self):
        """The set of coordinate elements in the form."""
        return tuple(domain.ufl_coordinate_element() for domain in self.original_form.ufl_domains())

    @cached_property
    def unique_elements(self):
        """Set of unique elements (not expanded for sub elements) in the form."""
        return unique_tuple(
            self.argument_elements + self.coefficient_elements + self.coordinate_elements
        )

    @cached_property
    def unique_sub_elements(self):
        """Set of unique elements (expanded by sub elements) in the form."""
        all_elements = self.argument_elements + self.coefficient_elements + self.coordinate_elements
        all_sub_elements = extract_sub_elements(all_elements)
        return unique_tuple(all_sub_elements)

    @property
    def coefficient_split(self) -> dict[Coefficient, list[Coefficient]]:
        """Map from coefficient to its split counterparts."""
        return self._coefficient_split

    @cached_property
    def preprocessed_form(self):
        """This is used in tests and is rather slow."""
        return reconstruct_form_from_integral_data(self.integral_data)

    @property
    def original_form(self):
        """The input UFL form."""
        return self._original_form

    @property
    def original_coefficient_positions(self):
        """Original placement of coefficients in form."""
        return self._original_coefficient_positions

1	"""FormData class easy for collecting of various data about a form."""
2
3	# Copyright (C) 2008-2016 Martin Sandve Alnæs
4	#
5	# This file is part of UFL (https://www.fenicsproject.org)
6	#
7	# SPDX-License-Identifier: LGPL-3.0-or-later
8	#
9	# Modified by Anders Logg, 2008.
10	# Modified by Jørgen S. Dokken, 2026.
11
12	from functools import cached_property	1✔
13	from itertools import chain	1✔
14	from typing import Any	1✔
15
16	from ufl.algorithms.analysis import extract_coefficients, extract_sub_elements, unique_tuple	1✔
17	from ufl.algorithms.apply_coefficient_split import CoefficientSplitter	1✔
18	from ufl.algorithms.apply_restrictions import apply_restrictions, default_restriction_map	1✔
19	from ufl.algorithms.check_arities import check_integrand_arity	1✔
20	from ufl.algorithms.domain_analysis import IntegralData, reconstruct_form_from_integral_data	1✔
21	from ufl.algorithms.replace import replace	1✔
22	from ufl.classes import Argument, Coefficient, FunctionSpace, GeometricFacetQuantity	1✔
23	from ufl.corealg.traversal import traverse_unique_terminals	1✔
24	from ufl.domain import AbstractDomain, MeshSequence, extract_domains, extract_unique_domain	1✔
25	from ufl.form import Form, Zero	1✔
26	from ufl.utils.formatting import estr, lstr, tstr	1✔
27	from ufl.utils.sequences import max_degree	1✔
28
29
30	def _check_form_arity(	1✔
31	integral_data: list[IntegralData], arguments: tuple[Argument], complex_mode: bool
32	):
33	"""Check if form arity is valid."""
34	for integral in integral_data:	1✔
35	for itg in integral.integrals:	1✔
36	check_integrand_arity(itg.integrand(), arguments, complex_mode)	1✔
37
38
39	def _auto_select_degree(elements):	1✔
40	"""Automatically select degree for all elements of the form.
41
42	This is be used in cases where the degree has not been specified by the user.
43	This feature is used by DOLFIN to allow the specification of Expressions with
44	undefined degrees.
45	"""
46	# Use max degree of all elements, at least 1 (to work with
47	# Lagrange elements)
48	return max_degree({e.embedded_superdegree for e in elements} - {None} \| {1})	1✔
49
50
51	def _compute_element_mapping(form: Form):	1✔
52	"""Compute element mapping for element replacement."""
53	# The element mapping is a slightly messy concept with two use
54	# cases:
55	# - Expression with missing cell or element TODO: Implement proper
56	# Expression handling in UFL and get rid of this
57	# - Constant with missing cell TODO: Fix anything that needs to be
58	# worked around to drop this requirement
59
60	# Extract all elements and include subelements of mixed elements
61	elements = [obj.ufl_element() for obj in chain(form.arguments(), form.coefficients())]	1✔
62	elements = extract_sub_elements(elements)	1✔
63
64	# Try to find a common degree for elements
65	common_degree = _auto_select_degree(elements)	1✔
66
67	# Compute element map
68	element_mapping = {}	1✔
69	for element in elements:	1✔
70	# Flag for whether element needs to be reconstructed
71	reconstruct = False	1✔
72
73	# Set cell
74	cell = element.cell	1✔
75	if cell is None:	1✔
NEW 76	domains = form.ufl_domains()	×
NEW 77	if not all(domains[0].ufl_cell() == d.ufl_cell() for d in domains):	×
NEW 78	raise ValueError(	×
79	"Cannot replace unknown element cell without unique common cell in form."
80	)
NEW 81	cell = domains[0].ufl_cell()	×
NEW 82	reconstruct = True	×
83
84	# Set degree
85	degree = element.embedded_superdegree	1✔
86	if degree is None:	1✔
NEW 87	degree = common_degree	×
NEW 88	reconstruct = True	×
89
90	# Reconstruct element and add to map
91	if reconstruct:	1✔
NEW 92	element_mapping[element] = element.reconstruct(cell=cell, degree=degree)	×
93	else:
94	element_mapping[element] = element	1✔
95
96	return element_mapping	1✔
97
98
99	def _compute_max_subdomain_ids(integral_data: list[IntegralData]) -> dict[str, int]:	1✔
100	"""Compute the maximum subdomain ids."""
101	max_subdomain_ids: dict[str, int] = {}	1✔
102	for itg_data in integral_data:	1✔
103	it = itg_data.integral_type	1✔
104	for integral in itg_data.integrals:	1✔
105	# Convert string for default integral to -1
106	sids = (-1 if isinstance(si, str) else si for si in integral.subdomain_id())	1✔
107	newmax = max(sids) + 1	1✔
108	prevmax = max_subdomain_ids.get(it, 0)	1✔
109	max_subdomain_ids[it] = max(prevmax, newmax)	1✔
110	return max_subdomain_ids	1✔
111
112
113	def _check_elements(form_data):	1✔
114	"""Check elements."""
115	for element in chain(form_data.unique_elements, form_data.unique_sub_elements):	1✔
116	if element.cell is None:	1✔
NEW 117	raise ValueError(f"Found element with undefined cell: {element}")	×
118
119
120	def _check_facet_geometry(integral_data):	1✔
121	"""Check facet geometry."""
122	for itg_data in integral_data:	1✔
123	for itg in itg_data.integrals:	1✔
124	for expr in traverse_unique_terminals(itg.integrand()):	1✔
125	cls = expr._ufl_class_	1✔
126	if issubclass(cls, GeometricFacetQuantity):	1✔
127	domain = extract_unique_domain(expr, expand_mesh_sequence=False)	1✔
128	if isinstance(domain, MeshSequence):	1✔
NEW 129	raise RuntimeError(	×
130	f"Not expecting a terminal object on a "
131	f"mesh sequence at this stage: found {expr!r}"
132	)
133	it = itg_data.domain_integral_type_map[domain]	1✔
134	# Facet geometry is only valid in facet integrals.
135	# Allowing custom integrals to pass as well, although
136	# that's not really strict enough.
137	if not ("facet" in it or "custom" in it or "interface" in it):	1✔
138	# Not a facet integral
NEW 139	raise ValueError(f"Integral of type {it} cannot contain a {cls.__name__}.")	×
140
141
142	def _build_coefficient_replace_map(	1✔
143	coefficients: list[Coefficient], element_mapping=None
144	) -> tuple[list[Coefficient], dict[Coefficient, Coefficient]]:
145	"""Create new Coefficient objects with count starting at 0.
146
147	Returns:
148	lists of the new objects and mapping from old to new objects
149	"""
150	if element_mapping is None:	1✔
NEW 151	element_mapping = {}	×
152
153	new_coefficients = []	1✔
154	replace_map = {}	1✔
155	for i, f in enumerate(coefficients):	1✔
156	old_e = f.ufl_element()	1✔
157	new_e = element_mapping.get(old_e, old_e)	1✔
158	# XXX: This is a hack to ensure that if the original
159	# coefficient had a domain, the new one does too.
160	# This should be overhauled with requirement that Expressions
161	# always have a domain.
162	domain = extract_unique_domain(f, expand_mesh_sequence=False)	1✔
163	if domain is not None:	1✔
164	new_e = FunctionSpace(domain, new_e)	1✔
165	new_f = Coefficient(new_e, count=i)	1✔
166	new_coefficients.append(new_f)	1✔
167	replace_map[f] = new_f	1✔
168
169	return new_coefficients, replace_map	1✔
170
171
172	class FormData:	1✔
173	"""Class collecting various information extracted from a Form by calling preprocess."""
174
175	_original_form: Form	1✔
176	_integral_data: list[IntegralData]	1✔
177	_reduced_coefficients: list[Coefficient]	1✔
178	_original_coefficient_positions: list[int]	1✔
179	_function_replace_map: dict[Coefficient, Coefficient]	1✔
180	_coefficient_elements: tuple[Any, ...]	1✔
181	_coefficient_split: dict[Coefficient, list[Coefficient]]	1✔
182
183	def __init__(	1✔
184	self,
185	original_form: Form,
186	integral_data: list[IntegralData],
187	do_apply_default_restrictions: bool = True,
188	do_apply_restrictions: bool = True,
189	do_replace_functions: bool = False,
190	coefficients_to_split: tuple[Coefficient, ...] \| None = None,
191	complex_mode: bool = False,
192	):
193	"""Create form-data for a form that has been processed.
194
195	Args:
196	original_form: The form.
197	integral_data: List of integral data objects corresponding to
198	the form.
199	do_apply_default_restrictions: Apply default restrictions, defined in
200	{py:mod}`ufl.algorithms.apply_restrictions` to integrals if no
201	restriction has been set.
202	do_apply_restrictions: Apply restrictions towards terminal nodes.
203	do_replace_functions: Replace functions with with its cannonically numbered
204	function or thos provided in coefficients_to_split.
205	coefficients_to_split: Sequence of coefficients to split over a MeshSequence
206	complex_mode: If form has been processed as complex or not.
207	"""
208	self._original_form = original_form	1✔
209	self._integral_data = integral_data	1✔
210
211	# --- Create replacements for arguments and coefficients
212
213	# Figure out which form coefficients each integral should enable
214	for itg_data in self.integral_data:	1✔
215	itg_coeffs = set()	1✔
216	# Get all coefficients in integrand
217	for itg in itg_data.integrals:	1✔
218	itg_coeffs.update(extract_coefficients(itg.integrand()))	1✔
219	# Store with IntegralData object
220	itg_data.integral_coefficients = itg_coeffs	1✔
221
222	# Figure out which coefficients from the original form are
223	# actually used in any integral (Differentiation may reduce the
224	# set of coefficients w.r.t. the original form)
225	reduced_coefficients_set: set[Coefficient] = set()	1✔
226	for itg_data in self.integral_data:	1✔
227	assert itg_data.integral_coefficients is not None	1✔
228	reduced_coefficients_set.update(itg_data.integral_coefficients)	1✔
229	self._reduced_coefficients = sorted(reduced_coefficients_set, key=lambda c: c.count())	1✔
230	self._original_coefficient_positions = [	1✔
231	i
232	for i, c in enumerate(self.original_form.coefficients())
233	if c in self.reduced_coefficients
234	]
235
236	# Store back into integral data which form coefficients are used
237	# by each integral
238	for itg_data in self.integral_data:	1✔
239	assert itg_data.integral_coefficients is not None	1✔
240	itg_data.enabled_coefficients = [	1✔
241	bool(coeff in itg_data.integral_coefficients) for coeff in self.reduced_coefficients
242	]
243
244	# Mappings from elements and coefficients that reside in form to
245	# objects with canonical numbering as well as completed cells and
246	# elements
247	renumbered_coefficients, function_replace_map = _build_coefficient_replace_map(	1✔
248	self.reduced_coefficients, self.element_replace_map
249	)
250	self._function_replace_map = function_replace_map	1✔
251
252	self._coefficient_elements = tuple(f.ufl_element() for f in renumbered_coefficients)	1✔
253
254	# --- Apply replace(integrand, self.function_replace_map)
255	if do_replace_functions:	1✔
256	for itg_data in self.integral_data:	1✔
257	new_integrals = []	1✔
258	for integral in itg_data.integrals:	1✔
259	integrand = replace(integral.integrand(), self.function_replace_map)	1✔
260	new_integrals.append(integral.reconstruct(integrand=integrand))	1✔
261	itg_data.integrals = new_integrals	1✔
262
263	# --- Split mixed coefficients with their components
264	if coefficients_to_split is None:	1✔
265	self._coefficient_split = {}	1✔
266	else:
267	if not do_replace_functions:	1✔
NEW 268	raise ValueError("Must call with do_replace_functions=True")	×
269	for itg_data in self.integral_data:	1✔
270	new_integrals = []	1✔
271	for integral in itg_data.integrals:	1✔
272	# Propagate restrictions as required by CoefficientSplitter.
273	# Can not yet apply default restrictions at this point
274	# as restrictions can only be applied to the components.
275	new_integral = apply_restrictions(integral)	1✔
276	new_integrals.append(new_integral)	1✔
277	itg_data.integrals = new_integrals	1✔
278	# Split coefficients that are contained in ``coefficients_to_split``
279	# into components, and store a dict in ``self`` that maps
280	# each coefficient to its components.
281	coefficient_split = {}	1✔
282	for o in self.reduced_coefficients:	1✔
283	if o in coefficients_to_split:	1✔
284	c = self.function_replace_map[o]	1✔
285	mesh = extract_unique_domain(c, expand_mesh_sequence=False)	1✔
286	elem = c.ufl_element()	1✔
287	coefficient_split[c] = [	1✔
288	Coefficient(FunctionSpace(m, e))
289	for m, e in zip(mesh.iterable_like(elem), elem.sub_elements) # type: ignore
290	]
291	self._coefficient_split = coefficient_split	1✔
292	coeff_splitter = CoefficientSplitter(self.coefficient_split)	1✔
293	for itg_data in self.integral_data:	1✔
294	new_integrals = []	1✔
295	for integral in itg_data.integrals:	1✔
296	integrand = coeff_splitter(integral.integrand())	1✔
297	# Potentially need to call `remove_component_tensors()` here, but
298	# early-simplifications of Indexed objects seem sufficient.
299	if not isinstance(integrand, Zero):	1✔
300	new_integrals.append(integral.reconstruct(integrand=integrand))	1✔
301	itg_data.integrals = new_integrals	1✔
302
303	# Propagate restrictions to terminals
304	if do_apply_restrictions:	1✔
305	for itg_data in self.integral_data:	1✔
306	# Need the following if block in case not all participating domains
307	# have been included in the Measure (backwards compat).
308	if all(	1✔
309	not integral_type.startswith("interior_facet")
310	for _, integral_type in itg_data.domain_integral_type_map.items()
311	):
312	continue	1✔
313	default_restrictions: dict[AbstractDomain, str \| None] \| None
314	if do_apply_default_restrictions:	1✔
315	default_restrictions = {	1✔
316	domain: default_restriction_map[integral_type]
317	for domain, integral_type in itg_data.domain_integral_type_map.items()
318	}
319	# Need the following dict update in case not all participating domains
320	# have been included in the Measure (backwards compat).
321	extra = {	1✔
322	domain: default_restriction_map[itg_data.integral_type]
323	for integral in itg_data.integrals
324	for domain in extract_domains(integral)
325	if domain not in default_restrictions
326	}
327	default_restrictions.update(extra)	1✔
328	else:
NEW 329	default_restrictions = None	×
330	new_integrals = []	1✔
331	for integral in itg_data.integrals:	1✔
332	new_integral = apply_restrictions(	1✔
333	integral,
334	default_restrictions=default_restrictions,
335	)
336	new_integrals.append(new_integral)	1✔
337	itg_data.integrals = new_integrals	1✔
338
339	_check_elements(self)	1✔
340	_check_facet_geometry(self.integral_data)	1✔
341	_check_form_arity(self.integral_data, self.original_form.arguments(), complex_mode)	1✔
342
343	def __str__(self):	1✔
344	"""Return formatted summary of form data."""
345	types = sorted(self.max_subdomain_ids.keys())	1✔
346	geometry = (("Geometric dimension", self.geometric_dimension),)	1✔
347	subdomains = tuple(	1✔
348	(f"Number of {integral_type} subdomains", self.max_subdomain_ids[integral_type])
349	for integral_type in types
350	)
351	functions = (	1✔
352	# Arguments
353	("Rank", self.rank),
354	("Arguments", lstr(self.original_form.arguments())),
355	# Coefficients
356	("Number of coefficients", self.num_coefficients),
357	("Coefficients", lstr(self.reduced_coefficients)),
358	# Elements
359	("Unique elements", estr(self.unique_elements)),
360	("Unique sub elements", estr(self.unique_sub_elements)),
361	)
362	return tstr(geometry + subdomains + functions)	1✔
363
364	@property	1✔
365	def integral_data(self) -> list[IntegralData]:	1✔
366	"""The integral data of the form."""
367	return self._integral_data	1✔
368
369	@property	1✔
370	def reduced_coefficients(self) -> list[Coefficient]:	1✔
371	"""Set of active coeffcient in the form."""
372	return self._reduced_coefficients	1✔
373
374	@property	1✔
375	def num_coefficients(self):	1✔
376	"""Number of active coefficients in the form."""
377	return len(self.reduced_coefficients)	1✔
378
379	@property	1✔
380	def rank(self):	1✔
381	"""Rank of the form."""
382	return len(self.original_form.arguments())	1✔
383
384	@cached_property	1✔
385	def geometric_dimension(self):	1✔
386	"""Geometric dimension of the form."""
387	domains = extract_domains(self.original_form())	1✔
388	gdims = set([domain.geometric_dimension for domain in domains])	1✔
389	assert len(gdims) == 1	1✔
390	return gdims.pop()	1✔
391
392	@property	1✔
393	def function_replace_map(self) -> dict[Coefficient, Coefficient]:	1✔
394	"""Map from coefficients in form to those used in IntegralData."""
395	return self._function_replace_map	1✔
396
397	@cached_property	1✔
398	def element_replace_map(self):	1✔
399	"""Mapping from incomplete elements to new well-defined elements.
400
401	This is to support the Expression construct in dolfin which
402	subclasses Coefficient but doesn't provide an element,
403	and the Constant construct that doesn't provide the domain that
404	a Coefficient is supposed to have. A future design iteration in
405	UFL/UFC/FFC/DOLFIN may allow removal of this mapping with the
406	introduction of UFL types for Expression-like functions that can
407	be evaluated in quadrature points.
408
409	Note:
410	This property should likely be removed.
411	"""
412	return _compute_element_mapping(self.original_form)	1✔
413
414	@cached_property	1✔
415	def max_subdomain_ids(self) -> dict[str, int]:	1✔
416	"""For each integral type, return the maximum subdomain id."""
417	return _compute_max_subdomain_ids(self.integral_data)	1✔
418
419	@cached_property	1✔
420	def argument_elements(self) -> tuple[Any, ...]:	1✔
421	"""The set of elements the arguments in the form."""
422	return tuple(f.ufl_element() for f in self.original_form.arguments())	1✔
423
424	@property	1✔
425	def coefficient_elements(self) -> tuple[Any, ...]:	1✔
426	"""The set of elements used for coefficients in the form."""
427	return self._coefficient_elements	1✔
428
429	@cached_property	1✔
430	def coordinate_elements(self):	1✔
431	"""The set of coordinate elements in the form."""
432	return tuple(domain.ufl_coordinate_element() for domain in self.original_form.ufl_domains())	1✔
433
434	@cached_property	1✔
435	def unique_elements(self):	1✔
436	"""Set of unique elements (not expanded for sub elements) in the form."""
437	return unique_tuple(	1✔
438	self.argument_elements + self.coefficient_elements + self.coordinate_elements
439	)
440
441	@cached_property	1✔
442	def unique_sub_elements(self):	1✔
443	"""Set of unique elements (expanded by sub elements) in the form."""
444	all_elements = self.argument_elements + self.coefficient_elements + self.coordinate_elements	1✔
445	all_sub_elements = extract_sub_elements(all_elements)	1✔
446	return unique_tuple(all_sub_elements)	1✔
447
448	@property	1✔
449	def coefficient_split(self) -> dict[Coefficient, list[Coefficient]]:	1✔
450	"""Map from coefficient to its split counterparts."""
451	return self._coefficient_split	1✔
452
453	@cached_property	1✔
454	def preprocessed_form(self):	1✔
455	"""This is used in tests and is rather slow."""
456	return reconstruct_form_from_integral_data(self.integral_data)	1✔
457
458	@property	1✔
459	def original_form(self):	1✔
460	"""The input UFL form."""
461	return self._original_form	1✔
462
463	@property	1✔
464	def original_coefficient_positions(self):	1✔
465	"""Original placement of coefficients in form."""
466	return self._original_coefficient_positions	1✔

FEniCS / ufl / 23427573198

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous