Coverage for dragonfly_sim/core/reconstruction.py: 64%
364 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-17 03:37 +0000
1"""
2Green-Gauss gradient reconstruction for the p = 0 discretization.
4At p = 0, grad U vanishes inside every cell, which leaves no viscous volume
5term and no vorticity for the turbulence sources, and gives only
6first-order face states. As a cell-centred finite-volume code does, a DG0
7gradient G is reconstructed by Green-Gauss over each cell's facets, with
8facet averages {U} inside and the registered boundary state U_b on the
9boundary,
11 |K| G_K = sum_f |f| U_f (x) n_f.
13G is used in three places:
15 * MUSCL: the interior-facet traces of the inviscid flux are reconstructed
16 linearly, U_K + G_K . (x_f - x_K), optionally with a smooth van Albada
17 limiter (see muscl_traces).
18 * The viscous facet gradient, {G} with its component along the line
19 between the two cell centres replaced by the direct difference (which
20 keeps the two-point coupling that suppresses odd-even modes), or on a
21 boundary facet with the normal component replaced by (U_b - U)/d_n.
22 * The turbulence sources (vorticity, gradients of the transported scalars).
24G is eliminated from the unknowns exactly (utils/gradient_condensation.py),
25so Newton sees the true Jacobian. For adjoints and shape derivatives G and
26the cell centres are DerivedFields (utils/derived_fields.py), with exact
27derivatives with respect to U, alpha and the mesh coordinates.
28"""
29import numpy as np
30from mpi4py import MPI
31from petsc4py import PETSc
32import ufl
33import dolfinx
34import dolfinx.fem.petsc
36from dragonfly_sim.utils.derived_fields import (DerivedField, diagonal_mass_inverse,
37 mesh_coordinate_argument, destroy_operator)
38from dragonfly_sim.utils.gradient_condensation import CondensedGradientProblem
40_JIT = {"cffi_extra_compile_args": ["-O3", "-march=native", "-ffast-math"]}
43class CellCentres(DerivedField):
44 """
45 DG0 cell centres, each cell's reference midpoint mapped to the cell. For
46 the degree-1 geometry used here that is the plain average of the cell's
47 vertices, c = C x, so dc/dx is the constant averaging matrix C.
48 """
50 def __init__(self, mesh_obj):
51 self.mesh_obj = mesh_obj
52 # the reconstruction (set by its install()) decides whether dc/dx is built
53 # vectorized and once -- it is a constant matrix -- or per call in a Python loop
54 self.rec = None
55 self._mesh_jac = None
56 msh = mesh_obj.mesh
57 gdim = msh.geometry.dim
58 self.function = dolfinx.fem.Function(
59 dolfinx.fem.functionspace(msh, ("DG", 0, (gdim,))), name="cell_centres")
60 self.update_geometry()
62 def update_geometry(self):
63 gdim = self.mesh_obj.mesh.geometry.dim
64 self.function.interpolate(lambda x: x[:gdim])
65 self.function.x.scatter_forward()
67 def mesh_jacobian(self):
68 if self.rec is None or not self.rec.fast_cell_centre_jacobian:
69 return self._mesh_jacobian_loop()
70 if self._mesh_jac is None:
71 self._mesh_jac = self._mesh_jacobian_vectorized()
72 return self._mesh_jac
74 def discard(self, obj):
75 if obj is not self._mesh_jac:
76 destroy_operator(obj)
78 def release_cache(self):
79 # the matrix is constant: kept until the reconstruction's cache is released
80 if self._mesh_jac is not None:
81 self._mesh_jac.destroy()
82 self._mesh_jac = None
84 def _mesh_jacobian_vectorized(self):
85 """The same matrix as _mesh_jacobian_loop, as one CSR build: row (cell,
86 component i) holds 1/nv at the cell's vertex dofs' component i."""
87 msh = self.mesh_obj.mesh
88 Vc = self.function.function_space
89 Vx = self.mesh_obj.coordinate_space
90 gdim = msh.geometry.dim
91 n_own = msh.topology.index_map(msh.topology.dim).size_local
92 rows_map, cols_map = Vc.dofmap.index_map, Vx.dofmap.index_map
93 IT = PETSc.IntType
94 # the rows of a CSR are in the order of the owned row dofs, not of the cells
95 cell_dofs = np.asarray(Vc.dofmap.list[:n_own, 0])
96 order = np.argsort(cell_dofs)
97 assert np.array_equal(cell_dofs[order], np.arange(n_own)), "cell centre dofs are not the owned blocks"
98 vertex_dofs = np.asarray(Vx.dofmap.list[:n_own])[order]
99 nv = vertex_dofs.shape[1] if n_own else 1
100 cols = cols_map.local_to_global(vertex_dofs.ravel()).reshape(n_own, nv).astype(np.int64)
101 idx = cols[:, None, :] * gdim + np.arange(gdim)[None, :, None] # (cell, component, vertex)
102 idx = np.sort(idx.reshape(n_own * gdim, nv), axis=1)
103 indptr = np.arange(n_own * gdim + 1, dtype=np.int64) * nv
104 data = np.full(n_own * gdim * nv, 1.0 / nv)
105 X = PETSc.Mat().createAIJ(
106 ((rows_map.size_local * gdim, rows_map.size_global * gdim),
107 (cols_map.size_local * gdim, cols_map.size_global * gdim)),
108 csr=(indptr.astype(IT), idx.ravel().astype(IT), data), comm=msh.comm)
109 X.assemble()
110 return X
112 def _mesh_jacobian_loop(self):
113 msh = self.mesh_obj.mesh
114 Vc = self.function.function_space
115 Vx = self.mesh_obj.coordinate_space
116 gdim = msh.geometry.dim
117 n_own = msh.topology.index_map(msh.topology.dim).size_local
118 rows_map, cols_map = Vc.dofmap.index_map, Vx.dofmap.index_map
119 X = PETSc.Mat().createAIJ(
120 ((rows_map.size_local * gdim, rows_map.size_global * gdim),
121 (cols_map.size_local * gdim, cols_map.size_global * gdim)),
122 comm=msh.comm)
123 nv = Vx.dofmap.cell_dofs(0).size if n_own else 1
124 X.setPreallocationNNZ((nv, nv))
125 X.setOption(PETSc.Mat.Option.NEW_NONZERO_ALLOCATION_ERR, False)
126 IT = PETSc.IntType
127 for cell in range(n_own):
128 r = int(rows_map.local_to_global(Vc.dofmap.cell_dofs(cell))[0])
129 cols = cols_map.local_to_global(Vx.dofmap.cell_dofs(cell)).astype(IT)
130 for i in range(gdim):
131 X.setValues(np.array([r * gdim + i], dtype=IT), cols * gdim + i,
132 np.full(cols.size, 1.0 / cols.size))
133 X.assemble()
134 return X
137class GreenGaussGradient(DerivedField):
138 """
139 G = M^-1 b(U, alpha, x), as a DerivedField (see the module docstring).
141 Options (attributes of the reconstruction):
142 cache_jacobians keep dG/dU, dG/dalpha and dG/dx between calls
143 while the state, alpha and the mesh are unchanged
144 (several functionals and products of one totals
145 evaluation share them); freed by release_cache
146 structured_state_jacobian dG/dU as the operator W (x) I plus a small
147 boundary part (GreenGaussStateOperator) instead
148 of an assembled matrix
149 """
151 def __init__(self, reconstruction):
152 self.rec = reconstruction
153 self.function = reconstruction.G
154 self.depends_on_state = True
155 self._forms = {}
156 self._cache = None
158 def _form(self, key, factory):
159 if key not in self._forms:
160 self._forms[key] = dolfinx.fem.form(factory(), jit_options=_JIT)
161 return self._forms[key]
163 def update(self):
164 self.rec.update_gradient()
166 def update_geometry(self):
167 self.rec.update_mass()
169 # ---- content-keyed cache -------------------------------------------
170 def _key(self):
171 m = self.rec.model
172 n = m.u_vec.x.petsc_vec.local_size
173 return (np.array(m.u_vec.x.array[:n], copy=True), np.array(m.mesh_obj.mesh.geometry.x, copy=True),
174 float(m.farfield_conditions['alpha']))
176 def _same_key(self, key):
177 a, b = self._cache["key"], key
178 same = a[2] == b[2] and np.array_equal(a[0], b[0]) and np.array_equal(a[1], b[1])
179 return bool(self.rec.mesh_obj.mesh.comm.allreduce(int(same), op=MPI.MIN))
181 def _cached(self, name, build):
182 """build() once per (state, alpha, mesh), or every call if caching is off."""
183 if not self.rec.cache_jacobians:
184 return build()
185 key = self._key()
186 if self._cache is None or not self._same_key(key):
187 self.release_cache()
188 self._cache = {"key": key}
189 if name not in self._cache:
190 self._cache[name] = build()
191 return self._cache[name]
193 def release_cache(self):
194 cache, self._cache = self._cache, None
195 if cache is not None:
196 for name, obj in cache.items():
197 if name != "key" and obj is not None:
198 destroy_operator(obj)
200 def discard(self, obj):
201 if self._cache is not None and any(obj is v for k, v in self._cache.items() if k != "key"):
202 return
203 destroy_operator(obj)
205 # ---- the Jacobians -------------------------------------------------
206 def state_jacobian(self):
207 return self._cached("state", self._build_state_jacobian)
209 def _build_state_jacobian(self):
210 rec = self.rec
211 if rec.structured_state_jacobian:
212 return rec.state_operator()
213 a = self._form("U", lambda: ufl.derivative(
214 rec.gg_rhs, rec.model.u_vec, ufl.TrialFunction(rec.model.u_vec.function_space)))
215 B = dolfinx.fem.petsc.assemble_matrix(a)
216 B.assemble()
217 B.diagonalScale(L=rec.inv_mass)
218 return B
220 def alpha_vector(self):
221 return self._cached("alpha", self._build_alpha_vector)
223 def _build_alpha_vector(self):
224 rec = self.rec
225 if getattr(rec.model, "alpha_var", None) is None:
226 return None
227 key = "alpha"
228 if key not in self._forms:
229 dform = ufl.algorithms.expand_derivatives(ufl.diff(rec.gg_rhs, rec.model.alpha_var))
230 self._forms[key] = None if dform.empty() else dolfinx.fem.form(dform, jit_options=_JIT)
231 if self._forms[key] is None:
232 return None
233 v = dolfinx.fem.petsc.assemble_vector(self._forms[key])
234 v.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
235 v.pointwiseMult(v, rec.inv_mass)
236 return v
238 def mesh_jacobian(self):
239 return self._cached("mesh", self._build_mesh_jacobian)
241 def _build_mesh_jacobian(self):
242 # r_G(G, U, x) = int G : tau dx - b(U; tau) = 0, so dG/dx = -M^-1 dr_G/dx
243 rec = self.rec
244 G = rec.G
245 tau = ufl.TestFunction(G.function_space)
246 r_G = ufl.inner(G, tau) * rec.model.forms.dx - rec.gg_rhs
247 Vx = rec.mesh_obj.coordinate_space
248 a = self._form("x", lambda: ufl.derivative(
249 r_G, ufl.SpatialCoordinate(rec.mesh_obj.mesh), mesh_coordinate_argument(r_G, Vx)))
250 X = dolfinx.fem.petsc.assemble_matrix(a)
251 X.assemble()
252 X.diagonalScale(L=rec.inv_mass)
253 X.scale(-1.0)
254 return X
257class GreenGaussStateOperator:
258 """
259 dG/dU = M_G^-1 db/dU as a python-context Mat (G layout x state layout)
260 with mult and multTranspose, without assembling the 18-by-6 blocks.
262 The interior-facet part of b is a facet average of U times the area
263 vector, component by component: its derivative is W (x) I_n, with W the
264 scalar operator (3 area-vector components per cell pair, assembled once
265 per mesh) and I_n the identity over the n state components. Only the
266 boundary-facet part, U_b(U) of the registered boundary states, couples
267 components; it is assembled per state (boundary cells only).
268 """
270 def __init__(self, rec, W, B_bdry):
271 self.rec, self.W, self.B_bdry = rec, W, B_bdry
272 V = rec.model.u_vec.function_space
273 self.n = V.dofmap.index_map.size_local
274 self.ns, self.gdim = rec.model.n_state, rec.model.dimensions
275 self._xk, self._vk = W.createVecRight(), W.createVecLeft()
276 self._t = rec.inv_mass.duplicate()
278 def mult(self, mat, x, y):
279 n, ns, gdim = self.n, self.ns, self.gdim
280 xa = x.getArray(readonly=True)[:n * ns].reshape(n, ns)
281 ya = np.empty((n, ns, gdim))
282 for k in range(ns):
283 self._xk.setArray(np.ascontiguousarray(xa[:, k]))
284 self.W.mult(self._xk, self._vk)
285 ya[:, k, :] = self._vk.getArray(readonly=True).reshape(n, gdim)
286 y.setArray(ya.reshape(-1))
287 if self.B_bdry is not None:
288 self.B_bdry.multAdd(x, y, y)
289 y.pointwiseMult(y, self.rec.inv_mass)
291 def multTranspose(self, mat, v, y):
292 n, ns, gdim = self.n, self.ns, self.gdim
293 t = self._t
294 t.pointwiseMult(v, self.rec.inv_mass)
295 ta = t.getArray(readonly=True)[:n * ns * gdim].reshape(n, ns, gdim)
296 out = np.empty((n, ns))
297 for k in range(ns):
298 self._vk.setArray(np.ascontiguousarray(ta[:, k, :]).reshape(-1))
299 self.W.multTranspose(self._vk, self._xk)
300 out[:, k] = self._xk.getArray(readonly=True)
301 y.setArray(out.reshape(-1))
302 if self.B_bdry is not None:
303 self.B_bdry.multTransposeAdd(t, y, y)
305 def destroy(self, mat=None):
306 # called by PETSc when the python Mat is destroyed; W belongs to the
307 # reconstruction (one per mesh)
308 for obj in (self.B_bdry, self._xk, self._vk, self._t):
309 if obj is not None:
310 obj.destroy()
311 self.B_bdry = self._xk = self._vk = self._t = None
314class GreenGaussReconstruction:
315 """
316 p = 0 gradient reconstruction attached to a CompressibleEulerModel host.
318 install() plugs it in: the MUSCL face traces (if muscl), the derived
319 fields (cell centres and G), and the condensed nonlinear problem.
320 Call it after the function spaces exist and before the boundary
321 conditions are defined; the Green-Gauss forms are built lazily, once
322 every boundary condition has registered its state.
324 limiter: None (unlimited MUSCL) or "vanalbada". limiter_scales() returns
325 one magnitude per state component for the limiter's smoothing parameter.
326 """
328 def __init__(self, model, muscl=True, limiter=None, limiter_eps=1e-2, limiter_scales=None):
329 if limiter not in (None, "vanalbada"):
330 raise ValueError("limiter must be None or 'vanalbada'")
331 self.model = model
332 self.mesh_obj = model.mesh_obj
333 self.muscl = muscl
334 self.limiter = limiter
335 self.limiter_eps = limiter_eps
336 self.limiter_scales = limiter_scales
337 msh = self.mesh_obj.mesh
338 V_G = dolfinx.fem.functionspace(msh, ("DG", 0, (model.n_state, model.dimensions)))
339 self.G = dolfinx.fem.Function(V_G, name="green_gauss_gradient")
340 self.centres = CellCentres(self.mesh_obj)
341 self.xc = self.centres.function
342 self.gradient = GreenGaussGradient(self)
343 # options of the adjoint's derived-field Jacobians (see GreenGaussGradient)
344 self.fast_cell_centre_jacobian = False
345 self.cache_jacobians = False
346 self.structured_state_jacobian = False
347 self._W = None
348 self._W_key = None
349 self.gg_rhs = None
350 self._gg_form = None
351 self.inv_mass = None
353 def install(self):
354 m = self.model
355 m.reconstruction = self
356 if self.muscl:
357 m.face_traces = self.muscl_traces
358 self.centres.rec = self
359 m.derived_fields += [self.centres, self.gradient]
360 m.problem_factory = self.make_problem
362 # ------------------------------------------------------------------
363 # Face states and gradients (UFL)
364 # ------------------------------------------------------------------
365 def muscl_traces(self):
366 """
367 Linear reconstruction of both interior-facet traces from G at the
368 facet's quadrature points: U_K + G_K . (x_f - x_K) unlimited, or with
369 limiter="vanalbada" the edge-based van Albada form (Blazek,
370 Computational Fluid Dynamics, ch. 5) generalised to an off-centre
371 facet. For the cell K and its neighbour N across the facet, with
372 s = |x_f - x_K| / |x_N - x_K|,
374 D+ = s (U_N - U_K) (central increment to the facet)
375 D- = 2 G_K . (x_f - x_K) - D+ (upwind increment)
376 U_f = U_K + Psi(D-, D+),
377 Psi(a, b) = ((a^2 + e) b + (b^2 + e) a) / (a^2 + b^2 + 2e),
379 componentwise. Psi(a, a) = a, so smooth regions keep the unlimited
380 reconstruction; at a shock D- ~ 0 and the trace falls back to first
381 order. Psi is smooth, so the exact Jacobian survives.
382 e = (limiter_eps * S_k)^2 with S_k the component's scale.
383 """
384 U = self.model.u_vec
385 G, xc = self.G, self.xc
386 # x is continuous, so x('+') = x('-'); the explicit restriction is
387 # what lets the form's mesh-coordinate derivative compile on dS
388 x = ufl.SpatialCoordinate(self.mesh_obj.mesh)
389 dp, dm = x('+') - xc('+'), x('-') - xc('-')
390 if self.limiter is None:
391 return U('+') + ufl.dot(G('+'), dp), U('-') + ufl.dot(G('-'), dm)
392 r = xc('-') - xc('+')
393 rr = ufl.sqrt(ufl.dot(r, r))
394 eps = [(self.limiter_eps * s)**2 for s in self.limiter_scales()]
395 n_state = self.model.n_state
397 def trace(UK, UN, GK, d):
398 s = ufl.sqrt(ufl.dot(d, d)) / rr
399 Dp = s * (UN - UK)
400 Dm = 2.0 * ufl.dot(GK, d) - Dp
401 return ufl.as_vector([
402 UK[k] + ((Dm[k]**2 + eps[k]) * Dp[k] + (Dp[k]**2 + eps[k]) * Dm[k])
403 / (Dm[k]**2 + Dp[k]**2 + 2.0 * eps[k]) for k in range(n_state)])
405 return trace(U('+'), U('-'), G('+'), dp), trace(U('-'), U('+'), G('-'), dm)
407 def face_gradient_interior(self):
408 U = self.model.u_vec
409 G, xc = self.G, self.xc
410 d = xc('-') - xc('+')
411 dist = ufl.sqrt(ufl.dot(d, d))
412 e = d / dist
413 G_avg = 0.5 * (G('+') + G('-'))
414 return G_avg + ufl.outer((U('-') - U('+')) / dist - ufl.dot(G_avg, e), e)
416 def face_gradient_boundary(self, U_b, U=None):
417 """The cell's G with its normal component replaced by (U_b - U)/d_n,
418 d_n the cell centre's normal distance from the facet."""
419 U = self.model.u_vec if U is None else U
420 n = self.mesh_obj.n
421 d_n = ufl.dot(ufl.SpatialCoordinate(self.mesh_obj.mesh) - self.xc, n)
422 return self.G + ufl.outer((U_b - U) / d_n - ufl.dot(self.G, n), n)
424 def green_gauss_rhs(self):
425 """b(U; tau) with |K| G_K = b: facet averages inside, the registered
426 boundary states on the boundary."""
427 tau = ufl.TestFunction(self.G.function_space)
428 U, n = self.model.u_vec, self.mesh_obj.n
429 forms = self.model.forms
430 b = ufl.inner(ufl.outer(0.5 * (U('+') + U('-')), n('+')), tau('+') - tau('-')) * forms.dS
431 b_bdry = None
432 for U_b, tag, _ in self.model.boundary_states:
433 term = ufl.inner(ufl.outer(U_b, n), tau) * forms.ds(tag)
434 b_bdry = term if b_bdry is None else b_bdry + term
435 self.gg_rhs_boundary = b_bdry
436 return b if b_bdry is None else b + b_bdry
438 # ------------------------------------------------------------------
439 # Assembly
440 # ------------------------------------------------------------------
441 def set_up(self):
442 if self._gg_form is not None:
443 return
444 msh = self.mesh_obj.mesh
445 comm = msh.comm
446 # Every boundary facet needs a facet state, or its cell's gradient
447 # silently loses a term. Counted, not integrated: facet areas depend
448 # on the quadrature on warped (bilinear) hex faces.
449 tags = sorted({t for _, t, _ in self.model.boundary_states})
450 fdim = msh.topology.dim - 1
451 n_own = msh.topology.index_map(fdim).size_local
452 mt = self.mesh_obj.meshtags
453 tagged = np.unique(mt.indices[np.isin(mt.values, tags)])
454 ext = dolfinx.mesh.exterior_facet_indices(msh.topology)
455 missing = comm.allreduce(int(np.sum(~np.isin(ext[ext < n_own], tagged))), op=MPI.SUM)
456 if missing:
457 raise RuntimeError("Green-Gauss: {} boundary facets have no registered boundary "
458 "state (tags {})".format(missing, tags))
459 self.gg_rhs = self.green_gauss_rhs()
460 self._gg_form = dolfinx.fem.form(self.gg_rhs, jit_options=_JIT)
461 self._gg_vec = dolfinx.fem.petsc.create_vector(self.G.function_space)
462 self.inv_mass = diagonal_mass_inverse(self.G.function_space, self.model.forms.dx)
464 # ------------------------------------------------------------------
465 # dG/dU without assembling its 18-by-6 blocks (GreenGaussStateOperator)
466 # ------------------------------------------------------------------
467 def _weights(self):
468 """W, the scalar operator of the interior-facet part of dG/dU
469 before M_G^-1: (3 n_cells) x n_cells, assembled once per mesh."""
470 msh = self.mesh_obj.mesh
471 x = msh.geometry.x
472 if self._W is not None:
473 same = self._W_key.shape == x.shape and np.array_equal(self._W_key, x)
474 if bool(msh.comm.allreduce(int(same), op=MPI.MIN)):
475 return self._W
476 self._W.destroy()
477 self._W = None
478 gdim = msh.geometry.dim
479 V1 = dolfinx.fem.functionspace(msh, ("DG", 0))
480 Vg = dolfinx.fem.functionspace(msh, ("DG", 0, (gdim,)))
481 V = self.model.u_vec.function_space
482 n = V.dofmap.index_map.size_local
483 for W_ in (V1, Vg, self.G.function_space):
484 if not np.array_equal(np.asarray(W_.dofmap.list[:n]), np.asarray(V.dofmap.list[:n])):
485 raise RuntimeError("the DG0 dof numbering differs between the state and gradient spaces")
486 u1, t3, nn = ufl.TrialFunction(V1), ufl.TestFunction(Vg), self.mesh_obj.n
487 a = ufl.inner(0.5 * (u1('+') + u1('-')) * nn('+'), t3('+') - t3('-')) * self.model.forms.dS
488 W = dolfinx.fem.petsc.assemble_matrix(dolfinx.fem.form(a, jit_options=_JIT))
489 W.assemble()
490 self._W, self._W_key = W, np.array(x, copy=True)
491 return W
493 def state_operator(self):
494 """dG/dU at the current state, as a python-context Mat."""
495 self.set_up()
496 W = self._weights()
497 B_bdry = None
498 if self.gg_rhs_boundary is not None:
499 if getattr(self, "_a_bdry", None) is None:
500 V = self.model.u_vec.function_space
501 d = ufl.algorithms.expand_derivatives(
502 ufl.derivative(self.gg_rhs_boundary, self.model.u_vec, ufl.TrialFunction(V)))
503 self._a_bdry = False if d.empty() else dolfinx.fem.form(d, jit_options=_JIT)
504 if self._a_bdry is not False:
505 B_bdry = dolfinx.fem.petsc.assemble_matrix(self._a_bdry)
506 B_bdry.assemble()
507 ctx = GreenGaussStateOperator(self, W, B_bdry)
508 nG = self.G.function_space.dofmap.index_map
509 nU = self.model.u_vec.function_space.dofmap.index_map
510 bsG = self.G.function_space.dofmap.index_map_bs
511 bsU = self.model.u_vec.function_space.dofmap.index_map_bs
512 M = PETSc.Mat().createPython(
513 ((nG.size_local * bsG, nG.size_global * bsG), (nU.size_local * bsU, nU.size_global * bsU)),
514 comm=self.mesh_obj.mesh.comm)
515 M.setPythonContext(ctx)
516 M.setUp()
517 return M
519 def release_cache(self):
520 self.gradient.release_cache()
521 self.centres.release_cache()
523 def update_mass(self):
524 """Refresh M_G^-1 in place after mesh motion."""
525 if self.inv_mass is None:
526 return
527 new = diagonal_mass_inverse(self.G.function_space, self.model.forms.dx)
528 new.copy(self.inv_mass)
529 new.destroy()
531 def update_gradient(self):
532 """G <- M_G^-1 b(U) at the current state and mesh. Collective."""
533 self.set_up()
534 b = self._gg_vec
535 with b.localForm() as loc:
536 loc.set(0.0)
537 dolfinx.fem.petsc.assemble_vector(b, self._gg_form)
538 b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
539 self.G.x.petsc_vec.pointwiseMult(b, self.inv_mass)
540 self.G.x.scatter_forward()
542 def make_problem(self, F_newton, J_newton):
543 self.set_up()
544 return CondensedGradientProblem(F_newton, J_newton, self.model.u_vec, self.G, self.gg_rhs,
545 self.update_gradient, self.inv_mass, jit_options=_JIT)
547 def compact_forms(self, F):
548 """(a_UU, a_UG): the compiled forms of dR_U/dU and dR_U/dG, cached for F."""
549 if getattr(self, "_compact_forms", None) is None or self._compact_forms[0] is not F:
550 self.set_up()
551 a_UU = dolfinx.fem.form(ufl.derivative(F, self.model.u_vec), jit_options=_JIT)
552 a_UG = dolfinx.fem.form(ufl.derivative(F, self.G, ufl.TrialFunction(self.G.function_space)),
553 jit_options=_JIT)
554 self._compact_forms = (F, (a_UU, a_UG))
555 return self._compact_forms[1]
557 def condensed_jacobian(self, F):
558 """The exact dR/dU of the residual form F, as a CondensedJacobian."""
559 from dragonfly_sim.utils.gradient_condensation import CondensedJacobian
560 self.set_up()
561 return CondensedJacobian(F, self.model.u_vec, self.G, self.gg_rhs, self.inv_mass,
562 jit_options=_JIT)