18532473541

Committed 15 Oct 2025 02:32PM UTC coverage: 76.622% (+0.7%) from 75.964%

Build # 18532473541

Build Type

push

github

Committed by

web-flow

Commit Message

Handle restrictions on MeshSequence (#264)

* simplify Indexed(ComponentTensor(Indexed(ListTensor(...), ...), ...), ...) on construction

Co-authored-by: Pablo Brubeck <brubeck@protonmail.com>

* handle restrictions on MeshSequence

* Suggestions (#67)

* Add some type hints and simplify _extract_and_check_domain

* Use future annotations and set ruff target

* IntegralData: make domain_integral_type_map positional

* compute_form_data: remove do_assume_single_integral_type

* Add more words

* DROP AFTER 428

* try removing skip 3.14

---------

Co-authored-by: Pablo Brubeck <brubeck@protonmail.com>
Co-authored-by: Jørgen Schartum Dokken <dokken92@gmail.com>

Run Details

181 of 209 new or added lines in 12 files covered. (86.6%)

1 existing line in 1 file now uncovered.

9131 of 11917 relevant lines covered (76.62%)

0.77 hits per line

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

92.89

/ufl/algorithms/compute_form_data.py

"""This module provides the compute_form_data function.

Form compilers will typically call compute_form_dataprior to code
generation to preprocess/simplify a raw input form given by a user.
"""
# 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

from itertools import chain

from ufl.algorithms.analysis import extract_coefficients, extract_sub_elements, unique_tuple
from ufl.algorithms.apply_algebra_lowering import apply_algebra_lowering

# These are the main symbolic processing steps:
from ufl.algorithms.apply_coefficient_split import CoefficientSplitter
from ufl.algorithms.apply_derivatives import apply_coordinate_derivatives, apply_derivatives
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_restrictions, default_restriction_map
from ufl.algorithms.check_arities import check_form_arity
from ufl.algorithms.comparison_checker import do_comparison_check

# See TODOs at the call sites of these below:
from ufl.algorithms.domain_analysis import (
    build_integral_data,
    group_form_integrals,
    reconstruct_form_from_integral_data,
)
from ufl.algorithms.estimate_degrees import estimate_total_polynomial_degree
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.remove_component_tensors import remove_component_tensors
from ufl.algorithms.replace import replace
from ufl.classes import Coefficient, Form, FunctionSpace, GeometricFacetQuantity
from ufl.constantvalue import Zero
from ufl.corealg.traversal import traverse_unique_terminals
from ufl.domain import extract_domains, extract_unique_domain
from ufl.utils.sequences import max_degree


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):
    """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):
    """Compute the maximum subdomain ids."""
    max_subdomain_ids = {}
    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 _compute_form_data_elements(self, arguments, coefficients, domains):
    """Compute form data elements."""
    self.argument_elements = tuple(f.ufl_element() for f in arguments)
    self.coefficient_elements = tuple(f.ufl_element() for f in coefficients)
    self.coordinate_elements = tuple(domain.ufl_coordinate_element() for domain in domains)

    # TODO: Include coordinate elements from argument and coefficient
    # domains as well? Can they differ?

    # Note: Removed self.elements and self.sub_elements to make sure
    #       code that depends on the selection of argument +
    #       coefficient elements blow up, as opposed to silently
    #       almost working, with the introduction of the coordinate
    #       elements here.

    all_elements = self.argument_elements + self.coefficient_elements + self.coordinate_elements
    all_sub_elements = extract_sub_elements(all_elements)

    self.unique_elements = unique_tuple(all_elements)
    self.unique_sub_elements = unique_tuple(all_sub_elements)


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:
            it = itg_data.integral_type
            # 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
                for expr in traverse_unique_terminals(itg.integrand()):
                    cls = expr._ufl_class_
                    if issubclass(cls, GeometricFacetQuantity):
                        raise ValueError(f"Integral of type {it} cannot contain a {cls.__name__}.")


def _check_form_arity(preprocessed_form):
    """Check that we don't have a mixed linear/bilinear form or anything like that."""
    # FIXME: This is slooow and should be moved to form compiler
    # and/or replaced with something faster
    if 1 != len(compute_form_arities(preprocessed_form)):
        raise ValueError("All terms in form must have same rank.")


def _build_coefficient_replace_map(coefficients, element_mapping=None):
    """Create new Coefficient objects with count starting at 0.

    Returns:
        mapping from old to new objects, and lists of the 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


def attach_estimated_degrees(form):
    """Attach estimated polynomial degree to a form's integrals.

    Args:
        form: The Form` to inspect.

    Returns:
        A new Form with estimate degrees attached.
    """
    integrals = form.integrals()

    new_integrals = []
    for integral in integrals:
        md = {}
        md.update(integral.metadata())
        degree = estimate_total_polynomial_degree(integral.integrand())
        md["estimated_polynomial_degree"] = degree
        new_integrals.append(integral.reconstruct(metadata=md))
    return Form(new_integrals)


def preprocess_form(form, complex_mode):
    """Preprocess a form."""
    # Note: Default behaviour here will process form the way that is
    # currently expected by vanilla FFC

    # Check that the form does not try to compare complex quantities:
    # if the quantites being compared are 'provably' real, wrap them
    # with Real, otherwise throw an error.
    if complex_mode:
        form = do_comparison_check(form)

    # Lower abstractions for tensor-algebra types into index notation,
    # reducing the number of operators later algorithms and form
    # compilers need to handle
    form = apply_algebra_lowering(form)

    # After lowering to index notation, remove any complex nodes that
    # have been introduced but are not wanted when working in real mode,
    # allowing for purely real forms to be written
    if not complex_mode:
        form = remove_complex_nodes(form)

    # Apply differentiation before function pullbacks, because for
    # example coefficient derivatives are more complicated to derive
    # after coefficients are rewritten, and in particular for
    # user-defined coefficient relations it just gets too messy
    form = apply_derivatives(form)

    return form


def compute_form_data(
    form,
    do_apply_function_pullbacks=False,
    do_apply_integral_scaling=False,
    do_apply_geometry_lowering=False,
    preserve_geometry_types=(),
    do_apply_default_restrictions=True,
    do_apply_restrictions=True,
    do_estimate_degrees=True,
    do_append_everywhere_integrals=True,
    do_replace_functions=False,
    coefficients_to_split=None,
    complex_mode=False,
    do_remove_component_tensors=False,
):
    """Compute form data.

    The default arguments configured to behave the way old FFC expects.
    """
    # TODO: Move this to the constructor instead
    self = FormData()

    # --- Store untouched form for reference.
    # The user of FormData may get original arguments,
    # original coefficients, and form signature from this object.
    # But be aware that the set of original coefficients are not
    # the same as the ones used in the final UFC form.
    # See 'reduced_coefficients' below.
    self.original_form = form

    # --- Pass form integrands through some symbolic manipulation

    form = preprocess_form(form, complex_mode)

    # --- Group form integrals
    # TODO: Refactor this, it's rather opaque what this does
    # TODO: Is self.original_form.ufl_domains() right here?
    #       It will matter when we start including 'num_domains' in ufc form.
    form = group_form_integrals(
        form,
        self.original_form.ufl_domains(),
        do_append_everywhere_integrals=do_append_everywhere_integrals,
    )

    # Estimate polynomial degree of integrands now, before applying
    # any pullbacks and geometric lowering.  Otherwise quad degrees
    # blow up horrifically.
    if do_estimate_degrees:
        form = attach_estimated_degrees(form)

    if do_apply_function_pullbacks:
        # Rewrite coefficients and arguments in terms of their
        # reference cell values with Piola transforms and symmetry
        # transforms injected where needed.
        # Decision: Not supporting grad(dolfin.Expression) without a
        #           Domain.  Current dolfin works if Expression has a
        #           cell but this should be changed to a mesh.
        form = apply_function_pullbacks(form)

    # Scale integrals to reference cell frames
    if do_apply_integral_scaling:
        form = apply_integral_scaling(form)

    # Lower abstractions for geometric quantities into a smaller set
    # of quantities, allowing the form compiler to deal with a smaller
    # set of types and treating geometric quantities like any other
    # expressions w.r.t. loop-invariant code motion etc.
    if do_apply_geometry_lowering:
        form = apply_geometry_lowering(form, preserve_geometry_types)

    # Apply differentiation again, because the algorithms above can
    # generate new derivatives or rewrite expressions inside
    # derivatives
    if do_apply_function_pullbacks or do_apply_geometry_lowering:
        form = apply_derivatives(form)

        # Neverending story: apply_derivatives introduces new Jinvs,
        # which needs more geometry lowering
        if do_apply_geometry_lowering:
            form = apply_geometry_lowering(form, preserve_geometry_types)
            # Lower derivatives that may have appeared
            form = apply_derivatives(form)

    form = apply_coordinate_derivatives(form)

    # If in real mode, remove any complex nodes introduced during form processing.
    if not complex_mode:
        form = remove_complex_nodes(form)

    # Remove component tensors
    if do_remove_component_tensors:
        form = remove_component_tensors(form)

    # --- Group integrals into IntegralData objects
    # Most of the heavy lifting is done above in group_form_integrals.
    self.integral_data = build_integral_data(form.integrals())

    # --- 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()
    for itg_data in self.integral_data:
        reduced_coefficients_set.update(itg_data.integral_coefficients)
    self.reduced_coefficients = sorted(reduced_coefficients_set, key=lambda c: c.count())
    self.num_coefficients = len(self.reduced_coefficients)
    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:
        itg_data.enabled_coefficients = [
            bool(coeff in itg_data.integral_coefficients) for coeff in self.reduced_coefficients
        ]

    # --- Collect some trivial data

    # Get rank of form from argument list (assuming not a mixed arity form)
    self.rank = len(self.original_form.arguments())

    # Extract common geometric dimension (topological is not common!)
    self.geometric_dimension = self.original_form.integrals()[0].ufl_domain().geometric_dimension()

    # --- Build mapping from old incomplete element objects 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.
    self.element_replace_map = _compute_element_mapping(self.original_form)

    # 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

    # --- Store various lists of elements and sub elements (adds
    #     members to self)
    _compute_form_data_elements(
        self,
        self.original_form.arguments(),
        renumbered_coefficients,
        self.original_form.ufl_domains(),
    )

    # --- Store number of domains for integral types
    # TODO: Group this by domain first. For now keep a backwards
    # compatible data structure.
    self.max_subdomain_ids = _compute_max_subdomain_ids(self.integral_data)

    # --- 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)
                ]
        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
            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

    # --- Checks
    _check_elements(self)
    _check_facet_geometry(self.integral_data)

    # TODO: This is a very expensive check... Replace with something
    # faster!
    preprocessed_form = reconstruct_form_from_integral_data(self.integral_data)

    # TODO: Test how fast this is
    check_form_arity(preprocessed_form, self.original_form.arguments(), complex_mode)

    # TODO: This member is used by unit tests, change the tests to
    # remove this!
    self.preprocessed_form = preprocessed_form

    return self

FEniCS / ufl / 18532473541

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