Coverage for dragonfly_sim/core/meshwarping.py: 79%
219 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-17 03:37 +0000
1from mpi4py import MPI
2import numpy as np
3from petsc4py import PETSc
5import resource
6import datetime
7import os
8from time import perf_counter
10# These only take effect if jax's backend has not been initialised yet; export
11# them before starting Python to be sure. Note that XLA:CPU sizes its own
12# thread pool from NPROC (else the core count), which neither setting bounds:
13# measured 2026-09-30 on jax 0.11.2, each rank uses ~3.4 cores in the warp
14# kernels without NPROC. A one-worker pool (NPROC=1, or one-core pinning) is
15# NOT set here: under one-core pinning, nested XLA executions from inside csdl
16# host callbacks deadlocked.
17os.environ.setdefault("OMP_NUM_THREADS", "1")
18os.environ.setdefault("XLA_FLAGS", "--xla_cpu_multi_thread_eigen=false")
20import jax
21# Mandatory, and before any idwarp_jax call: the kd-tree warper takes its host
22# dtype from jnp.asarray(volume_points), so without x64 it silently runs in
23# float32 -- which cannot difference O(300) coordinates against O(1e-5) motions.
24jax.config.update("jax_enable_x64", True)
26from idwarp_jax import build_volume_pts_func
27from idwarp_jax.warp import compute_idwarp_reference_length
29import csdl_alpha as csdl
30# ordered_callbacks needs csdl_alpha with ordered custom-operation callbacks (branch
31# mpi_callback_issue); older versions ignore it, and MPI runs give wrong gradients.
32assert hasattr(csdl.CustomExplicitOperation, "ordered_callbacks"), \
33 "csdl_alpha is too old: custom-operation callbacks are not ordered under MPI"
35# build_volume_pts_func applies a rigid pitch rotation to its output; pitch 0
36# makes it an exact identity (cos 0 = 1, sin 0 = 0).
37_NO_PITCH = {"pitch": 0.0}
40# Propagate movement of the wall nodes to the volume nodes by inverse-distance weighting
41class IDWarp_jax(csdl.CustomExplicitOperation):
42 """IDW mesh deformation via the ``idwarp-jax`` package's kd-tree warper.
44 Maps the replicated ``(M_global, dim)`` wall-node motions
45 (``evaluate(bdry_motions, shape_parameters)``, usually
46 ``shape_design.FFDShapeParameterization.wall_displacement``) to rank-local volume
47 node motions, in geometry-node order. Because the map is closed form rather
48 than a PDE solve, this is a ``CustomExplicitOperation`` -- CSDL calls
49 ``compute`` and ``compute_jacvec_product``.
51 Everything that is a property of the mesh rather than of the warp lives on
52 the ``Mesh`` (mesh_manager): the wall and anchor node sets
53 and their global numbering (``Mesh.boundary_node_set``), the tagged facets
54 and their orientation, applying/resetting node motions, mesh quality, and
55 the output files. The wall numbering here is the mesh's wall node set's.
57 The warp itself is ``idwarp_jax.build_volume_pts_func``, a kd-tree IDWarp
58 that batches its forward and reverse passes internally, so its memory stays
59 bounded without any chunking here. ``err_tol`` is its pruning tolerance,
60 relative to the IDW denominator; ``err_tol=0`` disables pruning and gives
61 the exact dense IDW sum. This class only turns the mesh's wall and anchor
62 node sets into the package's surface and moves arrays between CSDL and it.
64 MPI
65 ---
66 Volume points are mutually independent and each one needs the *whole*
67 surface -- which every rank already has, since ``bdry_motions`` is a
68 replicated ``(M_global, dim)`` array. So each rank builds its own warper
69 over its own interior nodes against the global surface. The forward pass
70 needs no collectives at all; the reverse pass needs exactly one
71 ``Allreduce``, in ``compute_jacvec_product``, to sum every rank's
72 contribution to the gradient with respect to the replicated wall motions.
74 Those collectives run inside csdl's JAX callbacks, so ``ordered_callbacks``
75 must stay True: otherwise XLA may run two of them concurrently (e.g. the
76 drag and lift reverse passes) and cross-pair the Allreduces across ranks.
78 Anchor nodes
79 ------------
80 IDW is a *normalized* weighted average, so unlike an elasticity warp it
81 has no farfield decay -- it reproduces rigid modes everywhere. A rigid wall
82 translation of 0.1 moves a point 300 units away by 0.1. The non-moving
83 boundaries are therefore made part of the prescribed surface, held at zero
84 displacement; ``anchor_meshtag_vals`` selects them. Together with the
85 package's mirror about y = 0 this gives the usual mesh-deformation
86 boundary conditions:
88 farfield u = 0 -> anchor nodes, prescribed zero
89 symmetry u_normal = 0 -> mirror about y = 0
90 wall u = prescribed -> prescribed from bdry_motions
92 Symmetry
93 --------
94 The package always mirrors the surface about y = 0, and pins surface nodes
95 within ``symmetry_tolerance`` of that plane to it. In 3D that is the
96 model's symmetry plane, which windtunnel_model also hard-codes at
97 y = 0. The pin is applied to the wall motions here too (``_project``), so
98 the wall nodes move exactly as the warp assumed they do.
100 2D
101 --
102 The package is intrinsically three-dimensional (cross products, 3x3
103 Rodrigues matrices) and gives zero area to faces with fewer than three
104 nodes -- and a 2D boundary facet is a two-node line. So each boundary edge
105 is extruded into a quad (``extrusion_width``). The 2D plane is embedded as
106 the **x-z** plane -- not x-y, which the y = 0 mirror would reflect about the
107 chord line -- with the ribbon spanning ``y = -w/2 .. +w/2`` and the volume
108 nodes at y = 0. Every ribbon normal then has n_y = 0 exactly, so each
109 rotation is about y, so every surface node contributes exactly zero y
110 displacement; and each volume node is its own mirror image, so the mirrored
111 contributions duplicate the real ones and change nothing but the cost.
113 Parameters
114 ----------
115 mesh_obj : mesh_manager.Mesh, with its boundary facets
116 tagged. Its ``output_suffix`` (Mesh.configure_output) names this
117 class's memory logs.
118 wall_meshtag_val : tag of the moving (design) surface.
119 anchor_meshtag_vals : tags of the boundaries held at zero displacement.
120 Pass ``()`` to disable anchoring -- but read the note above first.
121 extrusion_width : 2D only; width of the extruded boundary ribbon. ``None``
122 derives it from the smallest boundary edge. See
123 ``_resolve_extrusion_width`` for why the default is what it is -- the
124 choice trades IDW distance fidelity against the conditioning of Newell's
125 method, and neither limit is safe.
126 LdefFact : IDWarp's dimensionless deformation length:
127 ``Ldef = LdefFact * Ldef0``, with ``Ldef0`` the largest distance from
128 the wall nodes' centroid to a wall node. The package takes ``Ldef0``
129 over every prescribed node, anchors included -- the farfield radius --
130 so the factor is rescaled before it is passed on.
131 reuse_tol : the forward warp is not repeated when the wall displacement is
132 within this absolute tolerance (largest entry) of the one the last
133 warp used; see ``compute``. Finite-difference steps must exceed it.
134 err_tol : kd-tree pruning tolerance, see above. Measured against
135 ``err_tol=0`` on the 2D NACA0012 mesh: 1e-6 gives 1.2e-3 relative
136 motion error at half the cost; 1e-10 prunes almost nothing (cost of
137 the exact sum). On the 3D wing mesh 1e-6 keeps ~25% of the dense
138 interactions, with a ~47 GiB host route cache summed over all ranks.
139 **idwarp_kwargs : forwarded to ``build_volume_pts_func``, over these
140 defaults:
142 - ``normal_eps = warp_eps = 1e-30``. The package defaults (1e-6) are
143 dimensional: ``normal_eps`` is added to |2A|^2 of every face, which
144 is ~1e-14 for a 2D ribbon quad and ~1e-10 for a small 3D wall face,
145 and ``warp_eps`` floors every IDW distance at 1e-3. Either would
146 wreck the warp.
147 - ``symmetry_tolerance``: 0 in 2D, where no ribbon node may be pinned
148 to y = 0; 1e-6 in 3D, the tolerance the model tags symmetry facets
149 with.
150 - ``target_batch_interactions = 1_000_000`` (package: 6e6). Each rank
151 holds one batch of volume-point x source interactions in flight,
152 and the reverse pass's working arrays scale with it: on the 3D wing
153 mesh at err_tol=1e-6 they add ~0.9 GiB/rank at 1e6 vs ~2.4 GiB/rank
154 at 6e6, with identical results and no measurable slowdown.
155 - ``progress = print_timings = False``; this class reports its own
156 timings.
158 Everything else (``zeroCornerRotations``, ``cornerAngle``,
159 ``useRotations``, ``rotation_eps``, ``bucket_size``, ``route_cache``,
160 ...) is the package's default.
161 """
162 ordered_callbacks = True # MPI collectives in compute and compute_jacvec_product
164 def __init__(
165 self,
166 mesh_obj,
167 wall_meshtag_val: int = 2,
168 anchor_meshtag_vals=(6,), # the windtunnel's FARFIELD_TAG
169 extrusion_width: float = None,
170 LdefFact: float = 1.0,
171 err_tol: float = 1.0e-6,
172 reuse_tol: float = 1.0e-12,
173 **idwarp_kwargs,
174 ):
175 super().__init__()
177 self.mesh = mesh_obj
178 comm = self.mesh.mesh.comm
179 # the last forward warp: its input and output, reused when CSDL
180 # re-evaluates the same design (see compute)
181 self.reuse_tol = float(reuse_tol)
182 self._last_bdry_motions = None
183 self._last_motions = None
185 self.dimensions = self.mesh.mesh.geometry.dim
186 self.wall_meshtag_val = wall_meshtag_val
187 self.anchor_meshtag_vals = tuple(anchor_meshtag_vals or ())
188 self._extrusion_width_arg = extrusion_width
189 self.LdefFact = float(LdefFact)
190 self.err_tol = float(err_tol)
191 self.N_local = self.mesh.n_owned_nodes
193 # Columns of the package's 3D arrays that hold the physical coordinates.
194 # In 2D, column 1 is the extrusion direction (see the class docstring).
195 self._cols = [0, 2] if self.dimensions == 2 else [0, 1, 2]
197 self._idwarp_kwargs = dict(
198 normal_eps=1.0e-30, warp_eps=1.0e-30,
199 symmetry_tolerance=0.0 if self.dimensions == 2 else 1.0e-6,
200 target_batch_interactions=1_000_000,
201 progress=False, print_timings=False)
202 self._idwarp_kwargs.update(idwarp_kwargs)
204 # ---- the prescribed surface: the mesh's wall and anchor node sets ---
205 # A node shared by the wall and an anchor boundary stays with the wall:
206 # it has to move.
207 self.wall = self.mesh.boundary_node_set((wall_meshtag_val,))
208 self.anchors = self.mesh.boundary_node_set(self.anchor_meshtag_vals,
209 exclude=self.wall)
210 self.M_global = self.wall.n_global
211 prescribed = np.zeros(self.N_local, dtype=bool)
212 prescribed[self.wall.local_idx] = True
213 prescribed[self.anchors.local_idx] = True
214 self._interior_local_idx = np.flatnonzero(~prescribed).astype(np.int64)
215 self._n_interior = self._interior_local_idx.size
216 # Surface node ids: the wall block first, so a wall node id equals its
217 # wall global id, then the anchor block.
218 self._n_surface_nodes = self.wall.n_global + self.anchors.n_global
220 self._build_global_surface_topology()
221 self._build_warper()
223 PETSc.Sys.syncPrint(
224 "[IDWarp_jax] rank {}: N_local={}, M_local={}, anchor_local={}, "
225 "interior={}".format(
226 comm.Get_rank(), self.N_local, self.wall.n_local,
227 self.anchors.n_local, self._n_interior))
228 PETSc.Sys.syncFlush()
229 PETSc.Sys.Print(
230 "[IDWarp_jax] M_global={}, anchor_global={}, n_surface_rows={}, "
231 "n_faces={}, LdefFact={}, err_tol={:.1e}".format(
232 self.M_global, self.anchors.n_global, self._n_surface_rows,
233 len(self._mesh_dict["faces"]["type"]), self.LdefFact, self.err_tol))
236 # ==================================================================
237 # construction helpers
238 # ==================================================================
240 def _build_global_surface_topology(self):
241 """Assemble the replicated prescribed surface as the package's mesh dict.
243 ``self._mesh_dict`` holds ``points`` (the surface rows, ``self._Xs0``)
244 and ``faces`` (every face with type 1, so wall and anchor faces are all
245 prescribed). Every rank ends up with byte-identical arrays, because
246 every rank has to evaluate the IDW sum against the whole surface.
247 """
248 comm = self.mesh.mesh.comm
249 dim = self.dimensions
251 node_coords = np.zeros((self._n_surface_nodes, 3), dtype=np.float64)
252 node_coords[:self.M_global, self._cols] = self.wall.coords
253 node_coords[self.M_global:, self._cols] = self.anchors.coords
255 all_tags = (self.wall_meshtag_val,) + self.anchor_meshtag_vals
256 local_faces, local_facet_idx = self.mesh.tagged_facets(
257 all_tags, np.concatenate([self.wall.geom_gids, self.anchors.geom_gids]))
259 if self.dimensions == 2:
260 self._n_layers = 2
261 self._extrusion_width = self._resolve_extrusion_width(node_coords, local_faces)
262 half_w = 0.5 * self._extrusion_width
263 # Layer-major row numbering: row = layer * n_nodes + node_id, so the
264 # wall rows of layer 0 are [0, M_global) and _expand/_fold are plain
265 # reshapes. The ribbon runs along y, out of the embedded x-z plane.
266 Xs0 = np.vstack([node_coords, node_coords])
267 Xs0[:self._n_surface_nodes, 1] = -half_w
268 Xs0[self._n_surface_nodes:, 1] = +half_w
269 if local_faces.size:
270 a = local_faces[:, 0]
271 b = local_faces[:, 1]
272 conn = np.column_stack([
273 a, b,
274 b + self._n_surface_nodes,
275 a + self._n_surface_nodes,
276 ])
277 else:
278 conn = np.zeros((0, 4), dtype=np.int64)
279 else:
280 self._n_layers = 1
281 self._extrusion_width = None
282 Xs0 = node_coords
283 conn = local_faces # quads already in cyclic order
285 n_flipped = 0
286 if conn.shape[0]:
287 cell_mids = np.zeros((local_facet_idx.size, 3), dtype=np.float64)
288 cell_mids[:, self._cols] = \
289 self.mesh.facet_cell_midpoints(local_facet_idx)[:, :dim]
290 conn, n_flipped = self.mesh.orient_facets(conn, Xs0, cell_mids)
292 # Rank-order concatenation, so every rank sees the same face ordering.
293 gathered = comm.allgather(np.ascontiguousarray(conn, dtype=np.int32))
294 gathered = [g for g in gathered if g.size]
295 if gathered:
296 global_conn = np.concatenate(gathered, axis=0)
297 else:
298 raise RuntimeError(
299 "IDWarp_jax: no tagged boundary facets found for tags {}.".format(
300 (self.wall_meshtag_val,) + self.anchor_meshtag_vals))
302 self._n_surface_rows = Xs0.shape[0]
303 self._Xs0 = np.ascontiguousarray(Xs0, dtype=np.float64)
304 self._mesh_dict = {
305 "points": self._Xs0,
306 "faces": {
307 "points": np.ascontiguousarray(global_conn, dtype=np.int32),
308 "type": np.ones(global_conn.shape[0], dtype=np.int32),
309 },
310 }
312 self._check_surface_topology(n_flipped)
315 def _resolve_extrusion_width(self, node_coords, local_faces):
316 """Pick the 2D ribbon width: a fraction of the shortest boundary edge.
318 Two effects pull in opposite directions, and the factor below is where
319 they balance.
321 Smaller w is better for the IDW geometry: w enters the distance as
322 sqrt(d_xy^2 + w^2/4), so a width comparable to the near-wall spacing
323 would blunt exactly the weight concentration that makes the warp follow
324 the surface. At w = 0.1 * h_min the nearest volume node sees its
325 distance inflated by 0.125%, which is nothing.
327 Larger w is better for conditioning. Newell's method in idwarp_jax's
328 normal.py sums cross products of *absolute* coordinates, whose
329 magnitudes are set by the farfield, and those terms have to cancel down to twice the face area --
330 which for a ribbon quad is only edge_length * w. The relative error in
331 the resulting normal therefore goes as 1/w. Measured on the 2D mesh, as
332 the error in reproducing a rigid surface translation:
334 w = 1e-5 -> 6.5e-08 w = 1e-2 -> 6.7e-11
335 w = 1e-4 -> 6.7e-09 w = 1e-1 -> 6.7e-12
336 w = 1e-3 -> 6.7e-10 w = 1e+0 -> 6.6e-13
338 0.1 * h_min puts that error around 1e-9 for an O(1) displacement, four
339 orders below the smallest cell, at a distance distortion of 0.1%.
340 (3D is unaffected: real surface faces have areas many orders larger than
341 an extruded ribbon's, so nothing cancels catastrophically.)
343 The minimum is taken over every prescribed boundary edge, so in practice
344 it is set by the wall, where the edges are finest.
345 """
346 if self._extrusion_width_arg is not None:
347 return float(self._extrusion_width_arg)
348 comm = self.mesh.mesh.comm
349 if local_faces.size:
350 p0 = node_coords[local_faces[:, 0]]
351 p1 = node_coords[local_faces[:, 1]]
352 local_min = float(np.linalg.norm(p1 - p0, axis=1).min())
353 else:
354 local_min = np.inf
355 global_min = comm.allreduce(local_min, op=MPI.MIN)
356 if not np.isfinite(global_min) or global_min <= 0.0:
357 global_min = 1.0
358 return 1.0e-1 * global_min
360 def _check_surface_topology(self, n_flipped):
361 """Fail loudly if the ranks assembled different surfaces.
363 Identical on every rank, byte for byte: otherwise ranks evaluate
364 different surfaces and the gradients they reduce are incoherent. A
365 content hash rather than a sum, so a permutation cannot slip through.
366 """
367 comm = self.mesh.mesh.comm
368 import hashlib
369 faces = self._mesh_dict["faces"]["points"]
370 digest = (hashlib.sha1(faces.tobytes()).hexdigest(),
371 hashlib.sha1(self._Xs0.tobytes()).hexdigest())
372 digests = comm.allgather(digest)
373 if any(d != digests[0] for d in digests):
374 raise RuntimeError(
375 "IDWarp_jax: the assembled prescribed surface differs between "
376 "ranks. Per-rank (faces, Xs0) hashes: {}".format(digests))
378 PETSc.Sys.Print(
379 "[IDWarp_jax] surface: {} faces of {} nodes, {} rows, "
380 "{} reversed for outward winding".format(
381 faces.shape[0], faces.shape[1], self._n_surface_rows,
382 comm.allreduce(n_flipped, op=MPI.SUM)))
384 def _build_warper(self):
385 """Build this rank's idwarp_jax warper over its own interior nodes.
387 The volume points are the interior nodes only: wall motions are
388 prescribed and anchors are held at zero, so neither needs warping. A
389 rank that owns no interior node gets one dummy point, because the
390 package cannot preprocess an empty volume.
391 """
392 msh = self.mesh.mesh
393 comm = msh.comm
394 t0 = perf_counter()
396 n_vol = max(self._n_interior, 1)
397 self._Xv0 = np.zeros((n_vol, 3), dtype=np.float64)
398 if self._n_interior:
399 self._Xv0[:, self._cols] = \
400 msh.geometry.x[:self.N_local, :self.dimensions][self._interior_local_idx]
401 else:
402 self._Xv0[0] = self._Xs0.mean(axis=0)
404 # The package takes Ldef0 over every prescribed row, anchors included;
405 # rescale so that Ldef = LdefFact * Ldef0 over the wall rows, as in IDWarp.
406 wall_rows = (np.arange(self._n_layers)[:, None] * self._n_surface_nodes
407 + np.arange(self.M_global)).reshape(-1)
408 ldef0_wall = float(compute_idwarp_reference_length(self._Xs0[wall_rows]))
409 ldef0_surface = float(compute_idwarp_reference_length(self._Xs0))
410 ldef_fact = self.LdefFact * ldef0_wall / ldef0_surface
412 # Pinned by the package to y = 0, so pinned here too (see _project).
413 tol = self._idwarp_kwargs["symmetry_tolerance"]
414 self._on_symmetry = np.abs(self._Xs0[:self.M_global, 1]) < tol
416 self._warp, self._warp_vjp, aux = build_volume_pts_func(
417 self._mesh_dict, self._Xv0, jax.devices("cpu")[0], return_aux=True,
418 LdefFact=ldef_fact, err_tol=self.err_tol, **self._idwarp_kwargs)
419 self.Ldef = float(aux["kdtree"]["ldef"])
421 # The package numbers its surface by np.unique over the face nodes; our
422 # row indexing only holds if that is every row, in order. An unreferenced
423 # row would also carry zero area -- a silent hole in the IDW sum.
424 if not np.array_equal(aux["wall_ids"], np.arange(self._n_surface_rows)):
425 raise RuntimeError(
426 "IDWarp_jax: {} of {} prescribed surface rows are not referenced "
427 "by any face.".format(
428 self._n_surface_rows - len(aux["wall_ids"]), self._n_surface_rows))
429 # Strictly positive area weights -- catches degenerate or bowtie faces
430 # that survive the orientation pass.
431 Ai = np.asarray(aux["kdtree"]["surface_area"])
432 if not np.all(Ai > 0.0):
433 raise RuntimeError(
434 "IDWarp_jax: {} prescribed surface nodes have non-positive area "
435 "weight (min {:.3e}); the surface faces are degenerate or wound "
436 "inconsistently.".format(int((Ai <= 0.0).sum()), float(Ai.min())))
438 # Compile every batch shape now, not inside csdl's jax.pure_callback:
439 # compilation, with its thread spawning under a held GIL, is the risky
440 # part of re-entering JAX from a callback.
441 self._warp(self._Xs0, _NO_PITCH)
442 self._warp_vjp(self._Xs0, _NO_PITCH, np.zeros_like(self._Xv0))
444 routes_gib = aux["kdtree"]["routes"].nbytes / 2**30
445 cache_gib = aux["kdtree"]["route_cache_nbytes"] / 2**30
446 PETSc.Sys.Print(
447 "[IDWarp_jax] warper: Ldef = {:.4e} (LdefFact {} x wall Ldef0 {:.4e}), "
448 "Ai in [{:.3e}, {:.3e}], routes max {:.3f} GiB/rank + host cache max "
449 "{:.3f} GiB/rank, setup incl. compilation {:.2f} s".format(
450 aux["kdtree"]["ldef"], self.LdefFact, ldef0_wall,
451 float(Ai.min()), float(Ai.max()),
452 comm.allreduce(routes_gib, op=MPI.MAX),
453 comm.allreduce(cache_gib, op=MPI.MAX),
454 comm.allreduce(perf_counter() - t0, op=MPI.MAX)))
456 # ==================================================================
457 # array plumbing between the CSDL variables and the idwarp_jax arrays
458 # ==================================================================
460 def _project(self, wall_vals):
461 """Zero the y component of wall nodes on the symmetry plane.
463 Takes and returns a full (M_global, dim) array. The package pins those
464 nodes to y = 0 before warping, so this is the wall motion the warp
465 actually assumed. A diagonal projection, hence its own adjoint.
466 Identity in 2D, where no node is pinned.
467 """
468 if not self._on_symmetry.any():
469 return wall_vals
470 wall_vals = np.array(wall_vals, dtype=np.float64)
471 wall_vals[self._on_symmetry, 1] = 0.0
472 return wall_vals
474 def _expand(self, wall_vals):
475 """(M_global, dim) wall values -> (n_surface_rows, 3), in every layer.
477 The anchor rows stay zero.
478 """
479 out = np.zeros((self._n_layers, self._n_surface_nodes, 3), dtype=np.float64)
480 out[:, :self.M_global][..., self._cols] = wall_vals[:, :self.dimensions]
481 return out.reshape(-1, 3)
483 def _fold(self, rows):
484 """Adjoint of ``_expand``: (n_surface_rows, 3) -> (M_global, dim).
486 The layers are summed because the forward writes the same wall motion
487 into every layer. The anchor rows are dropped: they are sensitivities
488 with respect to a constant.
489 """
490 rows = np.asarray(rows).reshape(self._n_layers, self._n_surface_nodes, 3)
491 return np.ascontiguousarray(
492 rows[:, :self.M_global].sum(axis=0)[:, self._cols], dtype=np.float64)
494 def _assert_planar(self, motion_3d):
495 """In 2D the out-of-plane (y) motion must vanish identically.
497 It does so exactly, per surface node: both nodal normals are in-plane, so
498 the rotation is about y, so the y component of
499 M (x_v - Xs0_i) + Xs_i - x_v is (+-w/2) + (-+w/2) = 0. A nonzero value
500 here means the ribbon construction is wrong, not that a tolerance is
501 tight.
502 """
503 if self.dimensions != 2 or motion_3d.size == 0:
504 return
505 max_dy = float(np.abs(motion_3d[:, 1]).max())
506 scale = max(1.0, float(np.abs(motion_3d[:, self._cols]).max()))
507 if max_dy > 1.0e-9 * scale:
508 raise RuntimeError(
509 "IDWarp_jax: 2D warp produced an out-of-plane motion of {:.3e} "
510 "(in-plane scale {:.3e}). The extruded ribbon surface is not "
511 "mirror-symmetric about y = 0.".format(max_dy, scale))
513 def _assemble_motions(self, wall_motions, interior_3d):
514 """(N_local, dim) mesh-node motions in geometry-node order.
516 Owned wall nodes take their rows of the replicated wall array, interior
517 nodes the warp output; anchor nodes stay zero.
518 """
519 motions = np.zeros((self.N_local, self.dimensions), dtype=np.float64)
520 motions[self.wall.local_idx] = \
521 wall_motions[self.wall.global_idx, :self.dimensions]
522 if self._n_interior:
523 self._assert_planar(interior_3d)
524 motions[self._interior_local_idx] = \
525 interior_3d[:self._n_interior][:, self._cols]
526 return motions
529 # ==================================================================
530 # CSDL interface
531 # ==================================================================
533 def evaluate(self, bdry_motions: "csdl.Variable", shape_parameters: "csdl.Variable"):
534 """Register the operation.
536 ``shape_parameters`` carries no derivative -- it is declared only to keep
537 every rank's reverse chain attached to the design variables, including
538 ranks that own no wall nodes.
539 """
540 self.name = "IDWarp_jax"
542 self.declare_input("bdry_motions", bdry_motions)
543 self.declare_input("shape_parameters", shape_parameters)
545 mesh_node_motions = self.create_output(
546 "mesh_node_motions", shape=(self.N_local, self.dimensions))
548 self.declare_derivative_parameters(
549 of="mesh_node_motions", wrt="shape_parameters", dependent=False)
551 return mesh_node_motions
553 def compute(self, input_vals, output_vals):
554 """Forward warp. The only collective is the equality check and the log line:
555 every rank already has the whole surface.
557 CSDL runs the model again for each derivative evaluation of a design,
558 so the same wall displacement arrives several times; the output of the
559 last warp is returned when no entry of the wall displacement differs
560 from the one that warp used by more than ``reuse_tol`` (absolute,
561 default 1e-12). The reference stays the displacement actually warped,
562 so small changes cannot add up unnoticed: once the displacement has
563 moved more than ``reuse_tol`` from it, it is warped again.
564 """
565 t0 = perf_counter()
566 bdry_motions = np.asarray(input_vals["bdry_motions"])
568 # bdry_motions is replicated, but reduce the comparison anyway: the
569 # ranks must agree on whether to warp
570 last = self._last_bdry_motions
571 same = (last is not None and last.shape == bdry_motions.shape
572 and bool(np.all(np.abs(bdry_motions - last) <= self.reuse_tol)))
573 if self.mesh.mesh.comm.allreduce(same, op=MPI.MIN):
574 output_vals["mesh_node_motions"] = self._last_motions.copy()
575 PETSc.Sys.Print("[IDWarp_jax] forward warp: wall displacement within {:.0e} of the last "
576 "warp's, reusing it ({:.2f} s)".format(self.reuse_tol, perf_counter() - t0))
577 return
579 self._log_memory("before compute")
580 # unknown until the new warp is complete
581 self._last_bdry_motions = self._last_motions = None
583 wall = self._project(bdry_motions)
584 Xv = self._warp(self._Xs0 + self._expand(wall), _NO_PITCH)
585 motions = self._assemble_motions(wall, Xv - self._Xv0)
586 output_vals["mesh_node_motions"] = motions
587 self._last_bdry_motions = bdry_motions.copy()
588 self._last_motions = motions.copy()
590 local_max = float(np.abs(motions).max()) if motions.size else 0.0
591 PETSc.Sys.Print(
592 "[IDWarp_jax] forward warp: max |motion| = {:.6e}, {:.2f} s".format(
593 self.mesh.mesh.comm.allreduce(local_max, op=MPI.MAX), perf_counter() - t0))
594 self._log_memory("after compute")
596 def compute_jacvec_product(self, input_vals, output_vals, d_inputs, d_outputs, mode):
597 """Matrix-free VJP. Only reverse mode is supported.
599 Each rank seeds only the
600 rows it owns -- its own wall nodes for the identity block, its own
601 interior nodes for the IDW block -- so the per-rank gradients are
602 disjoint contributions to one (M_global, dim) array and summing them is
603 exact.
604 """
605 t0 = perf_counter()
606 self._log_memory("before compute_jacvec_product")
607 Xs = self._Xs0 + self._expand(self._project(input_vals["bdry_motions"]))
609 if mode == "rev":
610 d_out = d_outputs.get(
611 "mesh_node_motions",
612 np.zeros((self.N_local, self.dimensions), dtype=np.float64))
614 seed = np.zeros_like(self._Xv0)
615 seed[:self._n_interior][:, self._cols] = d_out[self._interior_local_idx]
616 d_Xs, _ = self._warp_vjp(Xs, _NO_PITCH, seed)
618 # Identity block: the owned wall rows of the output are the wall
619 # motions themselves.
620 identity = np.zeros((self.M_global, self.dimensions), dtype=np.float64)
621 identity[self.wall.global_idx] = d_out[self.wall.local_idx]
622 grad = self._project(self._fold(d_Xs) + identity)
624 # Reduce here, unconditionally, and with the buffer-mode collective.
625 # Every rank reaches this with a real data dependence, which is what
626 # keeps it from being scheduled out of step -- the reverse chain on a
627 # rank owning no wall nodes would otherwise fold to a compile-time
628 # constant and its collective could run at any point relative to the
629 # others. See shape_parameterization.WallFFD's docstring for the
630 # full account.
631 # The lowercase (pickle) form is avoided because this array is
632 # (M_global, dim) and that channel is shared by every lowercase
633 # collective on the communicator.
634 grad_combined = np.empty_like(grad)
635 self.mesh.mesh.comm.Allreduce(grad, grad_combined, op=MPI.SUM)
636 d_inputs["bdry_motions"] = grad_combined
638 PETSc.Sys.Print(
639 "[IDWarp_jax] reverse: |grad| = {:.6e}, {:.2f} s".format(
640 float(np.linalg.norm(grad_combined)), perf_counter() - t0))
642 else:
643 raise ValueError("Reverse mode required")
645 self._log_memory("after compute_jacvec_product")
647 # ==================================================================
648 # diagnostics
649 # ==================================================================
651 def _log_memory(self, note):
652 if self.mesh.output_suffix is None:
653 return
654 mem = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
655 rank = MPI.COMM_WORLD.Get_rank()
656 fname = 'memory_logs/r{}_{}.log'.format(rank, self.mesh.output_suffix)
657 try:
658 with open(fname, 'a') as f:
659 f.write(' {} {} {} {}\n'.format(
660 datetime.datetime.now(), os.getpid(), mem,
661 "IDWarp_jax, " + note))
662 except OSError:
663 pass