Source code for dolfinx.fem.forms

# Copyright (C) 2017-2023 Chris N. Richardson, Garth N. Wells and Michal Habera
#
# This file is part of DOLFINx (https://www.fenicsproject.org)
#
# SPDX-License-Identifier:    LGPL-3.0-or-later

from __future__ import annotations

import collections
import typing

import numpy as np
import numpy.typing as npt

import ufl
from dolfinx import cpp as _cpp
from dolfinx import default_scalar_type, jit
from dolfinx.fem import IntegralType
from dolfinx.fem.function import FunctionSpaceBase

if typing.TYPE_CHECKING:
    from dolfinx.fem import function
    from dolfinx.mesh import Mesh


[docs]class Form: def __init__(self, form, ufcx_form=None, code: typing.Optional[str] = None): """A finite element form Note: Forms should normally be constructed using :func:`form` and not using this class initialiser. This class is combined with different base classes that depend on the scalar type used in the Form. Args: form: Compiled form object. ufcx_form: UFCx form code: Form C++ code """ self._code = code self._ufcx_form = ufcx_form self._cpp_object = form @property def ufcx_form(self): """The compiled ufcx_form object""" return self._ufcx_form @property def code(self) -> typing.Union[str, None]: """C code strings""" return self._code @property def rank(self) -> int: return self._cpp_object.rank @property def function_spaces(self) -> typing.List[FunctionSpaceBase]: """Function spaces on which this form is defined""" return self._cpp_object.function_spaces @property def dtype(self) -> np.dtype: """Scalar type of this form""" return self._cpp_object.dtype @property def mesh(self) -> Mesh: """Mesh on which this form is defined""" return self._cpp_object.mesh @property def integral_types(self): """Integral types in the form""" return self._cpp_object.integral_types
_ufl_to_dolfinx_domain = {"cell": IntegralType.cell, "exterior_facet": IntegralType.exterior_facet, "interior_facet": IntegralType.interior_facet, "vertex": IntegralType.vertex}
[docs]def form(form: typing.Union[ufl.Form, typing.Iterable[ufl.Form]], dtype: typing.Optional[npt.DTypeLike] = default_scalar_type, form_compiler_options: typing.Optional[dict] = None, jit_options: typing.Optional[dict] = None): """Create a Form or an array of Forms. Args: form: A UFL form or list(s) of UFL forms. dtype: Scalar type to use for the compiled form. form_compiler_options: See :func:`ffcx_jit <dolfinx.jit.ffcx_jit>` jit_options: See :func:`ffcx_jit <dolfinx.jit.ffcx_jit>`. Returns: Compiled finite element Form. Note: This function is responsible for the compilation of a UFL form (using FFCx) and attaching coefficients and domains specific data to the underlying C++ form. It dynamically create a :class:`Form` instance with an appropriate base class for the scalar type, e.g. :func:`_cpp.fem.Form_float64`. """ if form_compiler_options is None: form_compiler_options = dict() if dtype == np.float32: ftype = _cpp.fem.Form_float32 form_compiler_options["scalar_type"] = "float" elif dtype == np.float64: ftype = _cpp.fem.Form_float64 form_compiler_options["scalar_type"] = "double" elif dtype == np.complex64: ftype = _cpp.fem.Form_complex64 form_compiler_options["scalar_type"] = "float _Complex" elif dtype == np.complex128: ftype = _cpp.fem.Form_complex128 form_compiler_options["scalar_type"] = "double _Complex" else: raise NotImplementedError(f"Type {dtype} not supported.") def _form(form): """"Compile a single UFL form""" # Extract subdomain data from UFL form sd = form.subdomain_data() domain, = list(sd.keys()) # Assuming single domain # Check that subdomain data for each integral type is the same for data in sd.get(domain).values(): assert all([d is data[0] for d in data]) mesh = domain.ufl_cargo() if mesh is None: raise RuntimeError("Expecting to find a Mesh in the form.") ufcx_form, module, code = jit.ffcx_jit(mesh.comm, form, form_compiler_options=form_compiler_options, jit_options=jit_options) # For each argument in form extract its function space V = [arg.ufl_function_space()._cpp_object for arg in form.arguments()] # Prepare coefficients data. For every coefficient in form take # its C++ object. original_coeffs = form.coefficients() coeffs = [original_coeffs[ufcx_form.original_coefficient_position[i] ]._cpp_object for i in range(ufcx_form.num_coefficients)] constants = [c._cpp_object for c in form.constants()] # NOTE Could remove this and let the user convert meshtags by # calling compute_integration_domains themselves def get_integration_domains(integral_type, subdomain): """Get integration domains from subdomain data""" if subdomain is None: return [] else: try: if integral_type in (IntegralType.exterior_facet, IntegralType.interior_facet): tdim = subdomain.topology.dim subdomain._cpp_object.topology.create_connectivity(tdim - 1, tdim) subdomain._cpp_object.topology.create_connectivity(tdim, tdim - 1) return _cpp.fem.compute_integration_domains(integral_type, subdomain._cpp_object) except AttributeError: return subdomain # Subdomain markers (possibly empty list for some integral types) subdomains = {_ufl_to_dolfinx_domain[key]: get_integration_domains( _ufl_to_dolfinx_domain[key], subdomain_data[0]) for (key, subdomain_data) in sd.get(domain).items()} f = ftype(module.ffi.cast("uintptr_t", module.ffi.addressof(ufcx_form)), V, coeffs, constants, subdomains, mesh) return Form(f, ufcx_form, code) def _create_form(form): """Recursively convert ufl.Forms to dolfinx.fem.Form, otherwise return form argument""" if isinstance(form, ufl.Form): return _form(form) elif isinstance(form, collections.abc.Iterable): return list(map(lambda sub_form: _create_form(sub_form), form)) return form return _create_form(form)
[docs]def extract_function_spaces(forms: typing.Union[typing.Iterable[Form], # type: ignore [return] typing.Iterable[typing.Iterable[Form]]], index: int = 0) -> typing.Iterable[typing.Union[None, function.FunctionSpaceBase]]: """Extract common function spaces from an array of forms. If `forms` is a list of linear form, this function returns of list of the corresponding test functions. If `forms` is a 2D array of bilinear forms, for index=0 the list common test function space for each row is returned, and if index=1 the common trial function spaces for each column are returned.""" _forms = np.array(forms) if _forms.ndim == 0: raise RuntimeError("Expected an array for forms, not a single form") elif _forms.ndim == 1: assert index == 0 for form in _forms: if form is not None: assert form.rank == 1, "Expected linear form" return [form.function_spaces[0] if form is not None else None for form in forms] # type: ignore[union-attr] elif _forms.ndim == 2: assert index == 0 or index == 1 extract_spaces = np.vectorize(lambda form: form.function_spaces[index] if form is not None else None) V = extract_spaces(_forms) def unique_spaces(V): # Pick spaces from first column V0 = V[:, 0] # Iterate over each column for col in range(1, V.shape[1]): # Iterate over entry in column, updating if current # space is None, or where both spaces are not None check # that they are the same for row in range(V.shape[0]): if V0[row] is None and V[row, col] is not None: V0[row] = V[row, col] elif V0[row] is not None and V[row, col] is not None: assert V0[row] is V[row, col], "Cannot extract unique function spaces" return V0 if index == 0: return list(unique_spaces(V)) elif index == 1: return list(unique_spaces(V.transpose())) else: raise RuntimeError("Unsupported array of forms")