Coverage for dragonfly_sim/utils/derived_fields.py: 89%
184 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-17 03:37 +0000
1"""
2Derived fields and the chain rule through them.
4A derived field is a Function the residual (or a force functional) reads
5that is not one of the unknowns, but is itself a function of the state U,
6the angle of attack alpha and/or the mesh coordinates x:
8 the Green-Gauss gradient G = M^-1 b(U, alpha, x) (core/reconstruction.py)
9 the cell centres c = C x (core/reconstruction.py)
10 the wall distance d = d(x) (core/wall_distance.py)
12For R(U, f(U, alpha, x), alpha, x) every total derivative picks up one
13chain term per field,
15 dR/dU = dR/dU|_f + sum_f (dR/df) (df/dU)
16 dR/dalpha = dR/dalpha|_f + sum_f (dR/df) (df/dalpha)
17 dR/dx = dR/dx|_f + sum_f (dR/df) (df/dx)
19and the same holds for a scalar functional J. A DerivedField supplies the
20three right-hand factors; DerivedFieldChainRule assembles the partial
21derivatives with UFL and combines them. With no derived fields nothing here
22is used, and the model's derivative assemblers are the plain UFL ones.
23"""
24import numpy as np
25import ufl
26import dolfinx
27import dolfinx.fem.petsc
28from mpi4py import MPI
29from petsc4py import PETSc
32class DerivedField:
33 """
34 Interface. `function` is the dolfinx Function the forms read.
36 update() recompute the field at the current state and mesh
37 update_geometry() refresh whatever depends on the mesh only (called
38 after mesh motion, before update())
39 state_jacobian() PETSc Mat d(field)/dU (only if depends_on_state)
40 alpha_vector() PETSc Vec d(field)/dalpha, or None
41 mesh_jacobian() PETSc Mat d(field)/dx, columns in the mesh's
42 coordinate_space, or None
43 """
44 function = None
45 # True if the field is a function of the state (then dR/dU is condensed)
46 depends_on_state = False
48 def update(self):
49 pass
51 def update_geometry(self):
52 pass
54 def state_jacobian(self):
55 return None
57 def alpha_vector(self):
58 return None
60 def mesh_jacobian(self):
61 return None
63 def discard(self, obj):
64 """Free a matrix or vector returned by state_jacobian / alpha_vector /
65 mesh_jacobian once the caller is done with it. A field that caches
66 them keeps its own (override); a python-context operator also frees
67 what its context holds."""
68 destroy_operator(obj)
70 def release_cache(self):
71 """Free whatever the field caches between calls (nothing by default)."""
74def destroy_operator(obj):
75 """Destroy a PETSc Mat or Vec. A python-context Mat's context frees what it
76 owns in its destroy(mat) method, which PETSc calls with the Mat."""
77 obj.destroy()
80def _nonempty(form):
81 """The form with derivatives expanded, or None if it is identically zero."""
82 form = ufl.algorithms.expand_derivatives(form)
83 return None if form.empty() else form
86def _assemble_vec(form):
87 vec = dolfinx.fem.petsc.assemble_vector(form)
88 vec.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
89 return vec
92def _assemble_mat(form):
93 A = dolfinx.fem.petsc.assemble_matrix(form)
94 A.assemble()
95 return A
98def mesh_coordinate_argument(form, coordinate_space):
99 """The Argument a mesh-coordinate derivative of `form` is taken along."""
100 args = form.arguments()
101 number = max(a.number() for a in args) + 1 if args else 0
102 return ufl.Argument(coordinate_space, number)
105class DerivedFieldChainRule:
106 """
107 Total derivatives of a residual form F (rank 1) or a functional J
108 (rank 0) through a list of DerivedFields. Compiled forms are cached per
109 (form id, field) on this object, so reuse one instance per model.
110 """
112 def __init__(self, mesh_obj, u, fields, alpha_var=None, jit_options=None):
113 self.mesh_obj = mesh_obj
114 self.u = u
115 self.fields = list(fields)
116 self.alpha_var = alpha_var
117 self.jit_options = jit_options or {}
118 self._forms = {}
120 # ----------------------------------------------------------------
121 def _compiled(self, key, factory):
122 if key not in self._forms:
123 ufl_form = factory()
124 self._forms[key] = None if ufl_form is None else dolfinx.fem.form(
125 ufl_form, jit_options=self.jit_options)
126 return self._forms[key]
128 def _partial_field(self, form, key, field):
129 """dform/dfield: a matrix form (rank-1 form) or a vector form (rank 0)."""
130 V_f = field.function.function_space
131 rank = len(form.arguments())
132 arg = ufl.TrialFunction(V_f) if rank == 1 else ufl.TestFunction(V_f)
133 return self._compiled((key, "field", id(field)),
134 lambda: _nonempty(ufl.derivative(form, field.function, arg)))
136 def update(self):
137 for f in self.fields:
138 f.update()
140 # ----------------------------------------------------------------
141 # Residual (rank 1)
142 # ----------------------------------------------------------------
143 def residual_alpha(self, F, key, out):
144 """out <- dR/dalpha (total); `out` is a Vec in the residual's layout."""
145 dF = self._compiled((key, "alpha"), lambda: _nonempty(ufl.diff(F, self.alpha_var)))
146 with out.localForm() as loc:
147 loc.set(0.0)
148 if dF is not None:
149 dolfinx.fem.petsc.assemble_vector(out, dF)
150 out.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
151 for field in self.fields:
152 a = field.alpha_vector()
153 if a is None:
154 continue
155 dFdf = self._partial_field(F, key, field)
156 if dFdf is None:
157 field.discard(a)
158 continue
159 A = _assemble_mat(dFdf)
160 A.multAdd(a, out, out)
161 A.destroy()
162 field.discard(a)
163 return out
165 def residual_mesh_operator(self, F, key):
166 """dR/dx (total) as a PETSc python Mat with mult / multTranspose."""
167 Vx = self.mesh_obj.coordinate_space
168 spatial = ufl.SpatialCoordinate(self.mesh_obj.mesh)
169 # coordinate derivatives are left to the form compiler (expanding
170 # them early is not supported on interior-facet integrals)
171 a_x = self._compiled((key, "mesh"), lambda: ufl.derivative(
172 F, spatial, mesh_coordinate_argument(F, Vx)))
173 A_x = _assemble_mat(a_x)
174 chains = []
175 for field in self.fields:
176 X = field.mesh_jacobian()
177 if X is None:
178 continue
179 dFdf = self._partial_field(F, key, field)
180 if dFdf is None:
181 field.discard(X)
182 continue
183 chains.append((_assemble_mat(dFdf), X, field))
184 return _shell_matrix(A_x, chains)
186 # ----------------------------------------------------------------
187 # Functionals (rank 0)
188 # ----------------------------------------------------------------
189 def functional_state(self, J, key):
190 """dJ/dU (total) as a Vec in the state layout."""
191 V = self.u.function_space
192 dJ = self._compiled((key, "state"), lambda: _nonempty(
193 ufl.derivative(J, self.u, ufl.TestFunction(V))))
194 g = _assemble_vec(dJ)
195 for field in self.fields:
196 B = field.state_jacobian()
197 if B is None:
198 continue
199 dJdf = self._partial_field(J, key, field)
200 if dJdf is not None:
201 v = _assemble_vec(dJdf)
202 B.multTransposeAdd(v, g, g)
203 v.destroy()
204 # a fresh matrix per call: free it now, not at the next
205 # PETSc.garbage_cleanup (several GB for a 3D Green-Gauss gradient);
206 # a field that caches it keeps it
207 field.discard(B)
208 return g
210 def functional_alpha(self, J, key):
211 """dJ/dalpha (total) as a number. Collective."""
212 dJ = self._compiled((key, "alpha"), lambda: None if self.alpha_var is None
213 else _nonempty(ufl.diff(J, self.alpha_var)))
214 total = 0.0
215 if dJ is not None:
216 total = self.mesh_obj.mesh.comm.allreduce(dolfinx.fem.assemble_scalar(dJ), op=MPI.SUM)
217 for field in self.fields:
218 a = field.alpha_vector()
219 if a is None:
220 continue
221 dJdf = self._partial_field(J, key, field)
222 if dJdf is not None:
223 v = _assemble_vec(dJdf)
224 total += v.dot(a)
225 v.destroy()
226 field.discard(a)
227 return total
229 def functional_mesh(self, J, key):
230 """dJ/dx (total) as a Vec on the mesh's coordinate_space."""
231 Vx = self.mesh_obj.coordinate_space
232 spatial = ufl.SpatialCoordinate(self.mesh_obj.mesh)
233 dJ = self._compiled((key, "mesh"), lambda: ufl.derivative(
234 J, spatial, ufl.TestFunction(Vx)))
235 g = _assemble_vec(dJ)
236 for field in self.fields:
237 X = field.mesh_jacobian()
238 if X is None:
239 continue
240 dJdf = self._partial_field(J, key, field)
241 if dJdf is not None:
242 v = _assemble_vec(dJdf)
243 X.multTransposeAdd(v, g, g)
244 v.destroy()
245 field.discard(X)
246 return g
249class _ChainedMeshOperator:
250 """Python-Mat context: y = A_x v + sum_k A_k (X_k v), and its transpose."""
252 def __init__(self, A_x, chains):
253 self.A_x = A_x
254 self.chains = chains
255 self._work = [(A.createVecRight(), A.createVecLeft()) for A, _, _ in chains]
257 def mult(self, mat, v, y):
258 self.A_x.mult(v, y)
259 for (A, X, _), (w_f, _) in zip(self.chains, self._work):
260 X.mult(v, w_f)
261 A.multAdd(w_f, y, y)
263 def multTranspose(self, mat, w, y):
264 self.A_x.multTranspose(w, y)
265 for (A, X, _), (w_f, _) in zip(self.chains, self._work):
266 A.multTranspose(w, w_f)
267 X.multTransposeAdd(w_f, y, y)
270def _shell_matrix(A_x, chains):
271 ctx = _ChainedMeshOperator(A_x, chains)
272 M = PETSc.Mat().createPython(A_x.getSizes(), comm=A_x.getComm())
273 M.setPythonContext(ctx)
274 M.setUp()
275 return M
278def destroy_chained_operator(M):
279 """Free a residual_mesh_operator result and the matrices it holds now.
280 Left to the garbage collector, a parallel petsc4py object is only queued,
281 and freed at the next PETSc.garbage_cleanup; for a 3D RANS operator that
282 is several GB held into the next adjoint."""
283 ctx = M.getPythonContext()
284 M.destroy()
285 ctx.A_x.destroy()
286 for (A, X, field), (w_f, w_r) in zip(ctx.chains, ctx._work):
287 A.destroy()
288 field.discard(X) # a field that caches its mesh Jacobian keeps it
289 w_f.destroy()
290 w_r.destroy()
291 ctx.chains, ctx._work = [], []
294def diagonal_mass_inverse(V, dx):
295 """1/|K| per dof of a DG0 (blocked) space V: the inverse of its diagonal
296 mass matrix, as a Vec in V's layout."""
297 tau = ufl.TestFunction(V)
298 shape = V.ufl_element().reference_value_shape
299 ones = ufl.as_tensor(np.ones(shape).tolist()) if shape else 1.0
300 m = dolfinx.fem.petsc.assemble_vector(dolfinx.fem.form(ufl.inner(ones, tau) * dx))
301 m.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
302 inv = m.duplicate()
303 inv.setArray(1.0 / m.getArray(readonly=True))
304 m.destroy()
305 return inv