Source code for ffcx.analysis

# Copyright (C) 2007-2020 Anders Logg, Martin Alnaes, Kristian B. Oelgaard,
#                         Michal Habera and others
#
# This file is part of FFCx. (https://www.fenicsproject.org)
#
# SPDX-License-Identifier:    LGPL-3.0-or-later
"""Compiler stage 1: Analysis.

This module implements the analysis/preprocessing of variational forms,
including automatic selection of elements, degrees and form
representation type.
"""

import logging
import typing

import numpy
import numpy.typing

import basix.ufl_wrapper
import ufl
from ffcx.element_interface import convert_element, QuadratureElement
from warnings import warn

logger = logging.getLogger("ffcx")


[docs]class UFLData(typing.NamedTuple): form_data: typing.Tuple[ufl.algorithms.formdata.FormData, ...] # Tuple of ufl form data unique_elements: typing.List[basix.ufl_wrapper._BasixElementBase] # List of unique elements # Lookup table from each unique element to its index in `unique_elements` element_numbers: typing.Dict[basix.ufl_wrapper._BasixElementBase, int] unique_coordinate_elements: typing.List[basix.ufl_wrapper._BasixElementBase] # List of unique coordinate elements # List of ufl Expressions as tuples (expression, points, original_expression) expressions: typing.List[typing.Tuple[ufl.core.expr.Expr, numpy.typing.NDArray[numpy.float64], ufl.core.expr.Expr]]
[docs]def analyze_ufl_objects(ufl_objects: typing.List, options: typing.Dict) -> UFLData: """Analyze ufl object(s). Options ---------- ufl_objects options FFCx options. These options take priority over all other set options. Returns a data structure holding ------- form_datas Form_data objects unique_elements Unique elements across all forms and expressions element_numbers Mapping to unique numbers for all elements unique_coordinate_elements Unique coordinate elements across all forms and expressions expressions List of all expressions after post-processing, with its evaluation points and the original expression """ logger.info(79 * "*") logger.info("Compiler stage 1: Analyzing UFL objects") logger.info(79 * "*") elements = [] coordinate_elements = [] # Group objects by types forms = [] expressions = [] processed_expressions = [] for ufl_object in ufl_objects: if isinstance(ufl_object, ufl.form.Form): forms.append(ufl_object) elif isinstance(ufl_object, ufl.FiniteElementBase): elements.append(convert_element(ufl_object)) elif isinstance(ufl_object, ufl.Mesh): coordinate_elements.append(convert_element(ufl_object.ufl_coordinate_element())) elif isinstance(ufl_object[0], ufl.core.expr.Expr): original_expression = ufl_object[0] points = numpy.asarray(ufl_object[1]) expressions.append((original_expression, points)) else: raise TypeError("UFL objects not recognised.") form_data = tuple(_analyze_form(form, options) for form in forms) for data in form_data: elements += [convert_element(e) for e in data.unique_sub_elements] coordinate_elements += [convert_element(e) for e in data.coordinate_elements] for original_expression, points in expressions: elements += [convert_element(e) for e in ufl.algorithms.extract_elements(original_expression)] processed_expression = _analyze_expression(original_expression, options) processed_expressions += [(processed_expression, points, original_expression)] elements += ufl.algorithms.analysis.extract_sub_elements(elements) # Sort elements so sub-elements come before mixed elements unique_elements = ufl.algorithms.sort_elements(set(elements)) unique_coordinate_element_list = sorted(set(coordinate_elements), key=lambda x: repr(x)) for e in unique_elements: assert isinstance(e, basix.ufl_wrapper._BasixElementBase) # Compute dict (map) from element to index element_numbers = {element: i for i, element in enumerate(unique_elements)} return UFLData(form_data=form_data, unique_elements=unique_elements, element_numbers=element_numbers, unique_coordinate_elements=unique_coordinate_element_list, expressions=processed_expressions)
def _analyze_expression(expression: ufl.core.expr.Expr, options: typing.Dict): """Analyzes and preprocesses expressions.""" preserve_geometry_types = (ufl.classes.Jacobian, ) expression = ufl.algorithms.apply_algebra_lowering.apply_algebra_lowering(expression) expression = ufl.algorithms.apply_derivatives.apply_derivatives(expression) expression = ufl.algorithms.apply_function_pullbacks.apply_function_pullbacks(expression) expression = ufl.algorithms.apply_geometry_lowering.apply_geometry_lowering(expression, preserve_geometry_types) expression = ufl.algorithms.apply_derivatives.apply_derivatives(expression) expression = ufl.algorithms.apply_geometry_lowering.apply_geometry_lowering(expression, preserve_geometry_types) expression = ufl.algorithms.apply_derivatives.apply_derivatives(expression) complex_mode = "_Complex" in options["scalar_type"] if not complex_mode: expression = ufl.algorithms.remove_complex_nodes.remove_complex_nodes(expression) return expression def _analyze_form(form: ufl.form.Form, options: typing.Dict) -> ufl.algorithms.formdata.FormData: """Analyzes UFL form and attaches metadata. Options ---------- form options Returns ------- form_data - Form data computed by UFL with metadata attached Note ---- The main workload of this function is extraction of unique/default metadata from options, integral metadata or inherited from UFL (in case of quadrature degree) """ if form.empty(): raise RuntimeError(f"Form ({form}) seems to be zero: cannot compile it.") if _has_custom_integrals(form): raise RuntimeError(f"Form ({form}) contains unsupported custom integrals.") # Set default spacing for coordinate elements to be equispaced for n, i in enumerate(form._integrals): element = i._ufl_domain._ufl_coordinate_element if not isinstance(element, basix.ufl_wrapper._BasixElementBase) and element.degree() > 2: warn("UFL coordinate elements using elements not created via Basix may not work with DOLFINx") # Check for complex mode complex_mode = "_Complex" in options["scalar_type"] # Compute form metadata form_data = ufl.algorithms.compute_form_data( form, do_apply_function_pullbacks=True, do_apply_integral_scaling=True, do_apply_geometry_lowering=True, preserve_geometry_types=(ufl.classes.Jacobian,), do_apply_restrictions=True, do_append_everywhere_integrals=False, # do not add dx integrals to dx(i) in UFL complex_mode=complex_mode) # If form contains a quadrature element, use the custom quadrature scheme custom_q = None for e in form_data.unique_elements: e = convert_element(e) if isinstance(e, QuadratureElement): if custom_q is None: custom_q = e._points, e._weights else: assert numpy.allclose(e._points, custom_q[0]) assert numpy.allclose(e._weights, custom_q[1]) # Determine unique quadrature degree, quadrature scheme and # precision per each integral data for id, integral_data in enumerate(form_data.integral_data): # Iterate through groups of integral data. There is one integral # data for all integrals with same domain, itype, subdomain_id # (but possibly different metadata). # # Quadrature degree and quadrature scheme must be the same for # all integrals in this integral data group, i.e. must be the # same for for the same (domain, itype, subdomain_id) # Extract precision p_default = -1 precisions = set([integral.metadata().get("precision", p_default) for integral in integral_data.integrals]) precisions.discard(p_default) if len(precisions) == 1: p = precisions.pop() elif len(precisions) == 0: # Default precision p = numpy.finfo("double").precision + 1 # == 16 else: raise RuntimeError("Only one precision allowed within integrals grouped by subdomain.") integral_data.metadata["precision"] = p qd_default = -1 qr_default = "default" for i, integral in enumerate(integral_data.integrals): metadata = integral.metadata() if custom_q is None: # Extract quadrature degree qd_metadata = integral.metadata().get("quadrature_degree", qd_default) pd_estimated = numpy.max(integral.metadata()["estimated_polynomial_degree"]) if qd_metadata != qd_default: qd = qd_metadata else: qd = pd_estimated # Extract quadrature rule qr = integral.metadata().get("quadrature_rule", qr_default) logger.info(f"Integral {i}, integral group {id}:") logger.info(f"--- quadrature rule: {qr}") logger.info(f"--- quadrature degree: {qd}") logger.info(f"--- precision: {p}") metadata.update({"quadrature_degree": qd, "quadrature_rule": qr, "precision": p}) else: metadata.update({"quadrature_points": custom_q[0], "quadrature_weights": custom_q[1], "quadrature_rule": "custom", "precision": p}) integral_data.integrals[i] = integral.reconstruct(metadata=metadata) return form_data def _has_custom_integrals(o: typing.Union[ufl.integral.Integral, ufl.classes.Form, list, tuple]) -> bool: """Check for custom integrals.""" if isinstance(o, ufl.integral.Integral): return o.integral_type() in ufl.custom_integral_types elif isinstance(o, ufl.classes.Form): return any(_has_custom_integrals(itg) for itg in o.integrals()) elif isinstance(o, (list, tuple)): return any(_has_custom_integrals(itg) for itg in o) else: raise NotImplementedError