-
-
Notifications
You must be signed in to change notification settings - Fork 172
/
forms.py
213 lines (172 loc) · 8.14 KB
/
forms.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
# Copyright (C) 2017-2021 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 collections.abc
import typing
from dolfinx.fem.function import FunctionSpace
if typing.TYPE_CHECKING:
from dolfinx.fem import function
from dolfinx.mesh import Mesh
import numpy as np
import ufl
from dolfinx import cpp as _cpp
from dolfinx import jit
from petsc4py import PETSc
class FormMetaClass:
def __init__(self, form, V: list[_cpp.fem.FunctionSpace], coeffs, constants,
subdomains: dict[_cpp.mesh.MeshTags_int32, typing.Union[None, typing.Any]], mesh: _cpp.mesh.Mesh,
ffi, code):
"""A finite element form
Notes:
Forms should normally be constructed using
:func:`forms.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 UFC form
V: The argument function spaces
coeffs: Finite element coefficients that appear in the form
constants: Constants appearing in the form
subdomains: Subdomains for integrals
mesh: The mesh that the form is defined on
"""
self._code = code
self._ufcx_form = form
super().__init__(ffi.cast("uintptr_t", ffi.addressof(self._ufcx_form)),
V, coeffs, constants, subdomains, mesh) # type: ignore
@property
def ufcx_form(self):
"""The compiled ufcx_form object"""
return self._ufcx_form
@property
def code(self) -> str:
"""C code strings"""
return self._code
@property
def function_spaces(self) -> typing.List[FunctionSpace]:
"""Function spaces on which this form is defined"""
return super().function_spaces # type: ignore
@property
def dtype(self) -> np.dtype:
"""dtype of this form"""
return super().dtype # type: ignore
@property
def mesh(self) -> Mesh:
"""Mesh on which this form is defined"""
return super().mesh # type: ignore
@property
def integral_types(self):
"""Integral types in the form"""
return super().integral_types # type: ignore
form_types = typing.Union[FormMetaClass, _cpp.fem.Form_float32, _cpp.fem.Form_float64,
_cpp.fem.Form_complex64, _cpp.fem.Form_complex128]
def form(form: typing.Union[ufl.Form, typing.Iterable[ufl.Form]], dtype: np.dtype = PETSc.ScalarType,
form_compiler_options: dict = {}, jit_options: dict = {}):
"""Create a DOLFINx 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
Notes:
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. `_cpp.fem.Form_float64`.
"""
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.complex128:
ftype = _cpp.fem.Form_complex128
form_compiler_options["scalar_type"] = "double _Complex"
else:
raise NotImplementedError(f"Type {dtype} not supported.")
formcls = type("Form", (FormMetaClass, ftype), {})
def _form(form):
""""Compile a single UFL form"""
# Extract subdomain data from UFL form
sd = form.subdomain_data()
subdomains, = list(sd.values()) # Assuming single domain
domain, = list(sd.keys()) # Assuming single domain
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_coefficients = form.coefficients()
coeffs = [original_coefficients[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()]
# Subdomain markers (possibly None for some dimensions)
subdomains = {_cpp.fem.IntegralType.cell: subdomains.get("cell"),
_cpp.fem.IntegralType.exterior_facet: subdomains.get("exterior_facet"),
_cpp.fem.IntegralType.interior_facet: subdomains.get("interior_facet"),
_cpp.fem.IntegralType.vertex: subdomains.get("vertex")}
return formcls(ufcx_form, V, coeffs, constants, subdomains, mesh, module.ffi, 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)
def extract_function_spaces(forms: typing.Union[typing.Iterable[FormMetaClass], # type: ignore [return]
typing.Iterable[typing.Iterable[FormMetaClass]]],
index: int = 0) -> typing.Iterable[typing.Union[None, function.FunctionSpace]]:
"""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")