Coverage for dragonfly_sim/utils/filewriter.py: 98%
163 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-17 03:37 +0000
1import weakref
3import dolfinx
4import numpy as np
5import ufl
6from mpi4py import MPI
9# Rank-0 serial copies of each parent mesh, keyed by mesh_deformation: the
10# copy shared by deforming writers follows the parent geometry, the other one
11# keeps the geometry the parent had when it was built.
12_serial_meshes = weakref.WeakKeyDictionary()
15def _gather_rows(comm, rows):
16 """Gather the (n_local, k) arrays of all ranks to rank 0, stacked in rank order."""
17 rows = np.ascontiguousarray(rows)
18 ncols = rows.shape[1]
19 counts = comm.gather(rows.size, root=0)
20 recv = None
21 if comm.rank == 0:
22 recv = (np.empty(sum(counts), dtype=rows.dtype), counts)
23 comm.Gatherv(rows, recv, root=0)
24 return recv[0].reshape(-1, ncols) if comm.rank == 0 else None
27def _gather_geometry(mesh):
28 """Owned geometry nodes of all ranks on rank 0, in global node order. Collective."""
29 return _gather_rows(mesh.comm, mesh.geometry.x[:mesh.geometry.index_map().size_local])
32def _serial_mesh(mesh, mesh_deformation):
33 """Return the ghost-free serial copy of `mesh` on rank 0 (None elsewhere).
35 Owned cells and owned geometry nodes are numbered contiguously per rank in
36 global numbering, so gathering them in rank order yields the cells and
37 nodes in global order, each exactly once. Collective.
38 """
39 copies = _serial_meshes.setdefault(mesh, {})
40 if mesh_deformation in copies:
41 return copies[mesh_deformation]
42 comm = mesh.comm
43 tdim = mesh.topology.dim
44 gdim = mesh.geometry.dim
45 num_cells_owned = mesh.topology.index_map(tdim).size_local
46 x_map = mesh.geometry.index_map()
48 cell_nodes = mesh.geometry.dofmaps[0][:num_cells_owned]
49 cells_global = x_map.local_to_global(cell_nodes.ravel()).reshape(cell_nodes.shape)
51 x_all = _gather_geometry(mesh)
52 cells_all = _gather_rows(comm, cells_global.astype(np.int64))
53 serial = None
54 if comm.rank == 0:
55 domain = ufl.Mesh(mesh.ufl_domain().ufl_coordinate_element())
56 serial = dolfinx.mesh.create_mesh(MPI.COMM_SELF, cells_all, domain, x_all[:, :gdim])
57 copies[mesh_deformation] = serial
58 return serial
61def _cell_dof_indices(V, cells):
62 """(len(cells), dofs_per_cell * bs) dof-array indices of `cells`, cell-major."""
63 bs = V.dofmap.bs
64 cell_dofs = V.dofmap.list[cells]
65 return (cell_dofs[:, :, None] * bs + np.arange(bs)).reshape(
66 len(cells), cell_dofs.shape[1] * bs).astype(np.int32)
69class _Field:
70 """One named field of a FileWriter: its parent space and output function."""
72 def __init__(self, name, V):
73 self.name = name
74 self.V = V
75 tdim = V.mesh.topology.dim
76 owned_cells = np.arange(V.mesh.topology.index_map(tdim).size_local, dtype=np.int32)
77 # Parent dof-array indices of each owned cell, one row per cell
78 self.parent_idx = _cell_dof_indices(V, owned_cells)
79 # Target for UFL expressions, which are interpolated on the parent mesh
80 self.parent_function = dolfinx.fem.Function(V)
81 self.gather = not (V.ufl_element().discontinuous and V.dofmap.dof_layout.num_dofs > 1)
82 self.target_funcspace = None
83 self.function = None
84 self.target_idx = None
87class FileWriter:
88 """Writes dolfinx Functions to VTX (.bp) files for ParaView.
90 One writer holds one or more named fields on the same mesh, written to a
91 single file: ``func_space`` is either a function space (one field, named
92 "solution") or a dict ``{name: function space}``, e.g. density, momentum
93 (a vector field), energy and pressure. All fields of a writer must share
94 the mesh and be either all gathered or all per-rank (see below); the mesh
95 is then stored once per frame for all of them.
97 ParaView's VTX reader concatenates the per-rank blocks of a parallel write
98 without merging the points the blocks share. Where neighbouring cells
99 share points within a rank -- DG0 fields, written as cell data on the
100 mesh geometry, and continuous fields -- the MPI partition boundaries then
101 show up as cracks (exterior faces, boundary edges, jumps in
102 point-interpolated cell data). Those spaces are gathered to rank 0 and
103 written as a single block on a serial copy of the mesh.
105 Note on the DG0 gather: rank 0 holds the whole mesh (built once per parent
106 mesh and shared by the writers on it) plus one frame of the fields, and
107 every frame passes through rank 0. Per-frame data is small (one value per
108 cell and component); the serial mesh is what bounds the mesh size. Past
109 that point, the parallel alternatives are writing DG0 as per-cell-node DG
110 data of the geometry degree (flat cells, no smoothing across cells) or
111 dolfinx's XDMFFile.
113 Discontinuous spaces with more than one dof per cell (DG p >= 1) are
114 written by VTX with separate points for every cell, so every cell face is
115 a crack in serial as well and partition boundaries add none. Each rank
116 writes its own owned cells in parallel; ghost cells are left out because
117 they would be written twice.
119 Keywords, set per writer to match the simulation and the fields:
121 mesh_deformation
122 False: frames are written on the geometry the mesh had when the
123 writer was created. True: each frame is written on the mesh geometry
124 at the time of the write (e.g. a shape-optimisation design
125 evaluation, or a moving mesh).
126 time_dependent
127 False: frames are labelled by a counter (``write_counter``, e.g. the
128 design-evaluation index). True: frames are labelled by the physical
129 time passed as ``time=``, which must increase from frame to frame.
130 """
132 def __init__(self, filename, comm_obj, func_space, mesh_deformation=False,
133 time_dependent=False):
134 self.filename = filename
135 self.comm = comm_obj
136 self.single_field = not isinstance(func_space, dict)
137 spaces = {"solution": func_space} if self.single_field else dict(func_space)
138 if not spaces:
139 raise ValueError("FileWriter {}: no fields given".format(filename))
140 self.fields = [_Field(name, V) for name, V in spaces.items()]
141 # the first field's space; the mesh and the gather mode are shared
142 self.functionspace = self.fields[0].V
143 self.mesh_deformation = mesh_deformation
144 self.time_dependent = time_dependent
146 self.construct_function_and_writer()
148 @property
149 def field_names(self):
150 return [f.name for f in self.fields]
152 # The single-field attributes (the first field's), kept for callers and
153 # tests that write one field
154 @property
155 def function(self):
156 return self.fields[0].function
158 @property
159 def target_funcspace(self):
160 return self.fields[0].target_funcspace
162 @property
163 def parent_idx(self):
164 return self.fields[0].parent_idx
166 def construct_function_and_writer(self):
167 mesh = self.functionspace.mesh
168 tdim = mesh.topology.dim
169 if any(f.V.mesh is not mesh for f in self.fields):
170 raise ValueError("FileWriter {}: all fields must be on the same mesh".format(
171 self.filename))
172 gather = {f.gather for f in self.fields}
173 if len(gather) != 1:
174 raise ValueError(
175 "FileWriter {}: fields mix gathered (DG0, continuous) and per-rank (DG p >= 1) "
176 "spaces; write them to separate files".format(self.filename))
177 self.gather = gather.pop()
179 if self.gather:
180 # Serial cell s is input cell original_cell_index[s], which is
181 # global parent cell original_cell_index[s] (the gather order);
182 # likewise for the geometry nodes and input_global_indices
183 target_mesh = _serial_mesh(mesh, self.mesh_deformation)
184 writer_comm = MPI.COMM_SELF
185 if target_mesh is not None:
186 self.target_rows = np.asarray(target_mesh.topology.original_cell_index)
187 self.target_nodes = np.asarray(target_mesh.geometry.input_global_indices)
188 else:
189 owned_cells = np.arange(mesh.topology.index_map(tdim).size_local, dtype=np.int32)
190 target_mesh, entity_map, _, self.target_nodes = dolfinx.mesh.create_submesh(
191 mesh, tdim, owned_cells)
192 writer_comm = self.comm
193 sub_cells = target_mesh.topology.index_map(tdim)
194 self.target_rows = entity_map.sub_topology_to_topology(
195 np.arange(sub_cells.size_local + sub_cells.num_ghosts, dtype=np.int32), False)
196 self.target_mesh = target_mesh
198 self.writer = None
199 if target_mesh is not None:
200 for f in self.fields:
201 f.target_funcspace = dolfinx.fem.functionspace(target_mesh, f.V.ufl_element())
202 f.function = dolfinx.fem.Function(f.target_funcspace, name=f.name)
203 f.target_idx = _cell_dof_indices(
204 f.target_funcspace, np.arange(len(self.target_rows), dtype=np.int32))
205 self.writer = dolfinx.io.VTXWriter(
206 writer_comm,
207 "{}.bp".format(self.filename),
208 [f.function for f in self.fields],
209 engine="BP5",
210 )
211 self._check_dof_mapping()
212 self.write_counter = 0
213 self.last_time = None
215 def _copy_cell_values(self, array, field=None):
216 """Copy a parent dof array into a field's output function. Collective."""
217 field = self.fields[0] if field is None else field
218 values = array[field.parent_idx]
219 if self.gather:
220 values = _gather_rows(self.comm, values)
221 if self.writer is not None:
222 field.function.x.array[field.target_idx] = values[self.target_rows]
224 def _sync_geometry(self):
225 """Copy the current parent node coordinates onto the output mesh. Collective."""
226 x_parent = self.functionspace.mesh.geometry.x
227 if self.gather:
228 x_parent = _gather_geometry(self.functionspace.mesh)
229 if self.writer is not None:
230 self.target_mesh.geometry.x[:] = x_parent[self.target_nodes]
232 def _check_dof_mapping(self):
233 """Verify that mapped dofs sit at the same coordinates on both meshes.
235 The cell-local dof order is assumed equal on the parent and output
236 meshes; this holds for the Lagrange/DG spaces written here but would
237 break for spaces whose entity dofs depend on cell orientation.
238 Collective.
239 """
240 ok = True
241 for f in self.fields:
242 bs = f.V.dofmap.bs
243 coords = np.repeat(f.V.tabulate_dof_coordinates(), bs, axis=0)
244 for d in range(f.V.mesh.geometry.dim):
245 self._copy_cell_values(coords[:, d], f)
246 if self.writer is not None:
247 target_coords = np.repeat(
248 f.target_funcspace.tabulate_dof_coordinates(), bs, axis=0)[:, d]
249 ok &= np.allclose(f.function.x.array, target_coords, rtol=0, atol=1e-10)
250 if not self.comm.allreduce(ok, op=MPI.LAND):
251 raise RuntimeError(
252 "FileWriter: dof mapping onto the output mesh failed for {}".format(
253 self.filename))
255 def _parent_array(self, field, u_interp):
256 if isinstance(u_interp, dolfinx.fem.Function):
257 return u_interp.x.array
258 # UFL expression - interpolate on the parent mesh, then copy
259 field.parent_function.interpolate(
260 dolfinx.fem.Expression(u_interp, field.V.element.interpolation_points))
261 return field.parent_function.x.array
263 def interpolate_and_write(self, u_interp, write_counter=None, time=None):
264 """Write one frame. ``u_interp`` is, for a single-field writer, a
265 Function on the writer's space or a UFL expression; for a multi-field
266 writer, a dict ``{name: Function or UFL expression}`` with one entry
267 per field.
269 A counter-labelled writer (time_dependent=False) takes an optional
270 ``write_counter``; a time-dependent writer requires ``time``.
271 Collective.
272 """
273 if self.time_dependent:
274 if time is None or write_counter is not None:
275 raise ValueError(
276 "FileWriter {}: time-dependent output is labelled by time=, "
277 "not write_counter".format(self.filename))
278 if self.last_time is not None and not time > self.last_time:
279 raise ValueError(
280 "FileWriter {}: time {} does not increase past the last written "
281 "time {}".format(self.filename, time, self.last_time))
282 elif time is not None:
283 raise ValueError(
284 "FileWriter {}: time= needs time_dependent=True".format(self.filename))
286 if self.single_field:
287 values = {self.fields[0].name: u_interp}
288 else:
289 values = dict(u_interp)
290 if set(values) != set(self.field_names):
291 raise ValueError("FileWriter {}: expected the fields {}, got {}".format(
292 self.filename, self.field_names, sorted(values)))
293 for f in self.fields:
294 self._copy_cell_values(self._parent_array(f, values[f.name]), f)
295 if self.mesh_deformation:
296 self._sync_geometry()
298 if self.time_dependent:
299 label = self.last_time = float(time)
300 else:
301 # Only advance the internal counter when the caller did not supply an
302 # index. Bumping it unconditionally meant an explicit write_counter
303 # still perturbed writer state, so callers that pin the index (the
304 # aligned per-design-evaluation write) could not actually pin it.
305 if write_counter is None:
306 write_counter = self.write_counter
307 self.write_counter += 1
308 label = write_counter
310 if self.writer is not None:
311 self.writer.write(label)