Line data Source code
1 : /**********************************************************************/
2 : /* DO NOT MODIFY THIS HEADER */
3 : /* Swift, a Fourier spectral solver for MOOSE */
4 : /* */
5 : /* Copyright 2024 Battelle Energy Alliance, LLC */
6 : /* ALL RIGHTS RESERVED */
7 : /**********************************************************************/
8 :
9 : #include "XDMFTensorOutput.h"
10 : #include "TensorProblem.h"
11 : #include "Conversion.h"
12 :
13 : #include <ATen/core/TensorBody.h>
14 : #include <cstddef>
15 : #include <cstdint>
16 : #include <filesystem>
17 :
18 : #ifdef LIBMESH_HAVE_HDF5
19 : namespace
20 : {
21 : void addDataToHDF5(hid_t file_id,
22 : const std::string & dataset_name,
23 : const char * data,
24 : std::vector<std::size_t> & ndim,
25 : hid_t type);
26 : }
27 : #endif
28 :
29 : registerMooseObject("SwiftApp", XDMFTensorOutput);
30 :
31 : InputParameters
32 22 : XDMFTensorOutput::validParams()
33 : {
34 22 : auto params = TensorOutput::validParams();
35 22 : params.addClassDescription("Output a tensor in XDMF format.");
36 : #ifdef LIBMESH_HAVE_HDF5
37 44 : params.addParam<bool>("enable_hdf5", false, "Use HDF5 for binary data storage.");
38 : #endif
39 22 : MultiMooseEnum outputMode("CELL NODE OVERSIZED_NODAL");
40 44 : outputMode.addDocumentation("CELL", "Output as discontinuous elemental fields.");
41 44 : outputMode.addDocumentation(
42 : "NODE",
43 : "Output as nodal fields, increasing each box dimension by one and duplicating values from "
44 : "opposing surfaces to create a periodic continuation.");
45 44 : outputMode.addDocumentation("OVERSIZED_NODAL",
46 : "Output oversized tensors as nodal fields without forcing "
47 : "periodicity (suitable for displacement variables).");
48 :
49 44 : params.addParam<MultiMooseEnum>("output_mode", outputMode, "Output as cell or node data");
50 44 : params.addParam<bool>("transpose",
51 44 : true,
52 : "The Paraview XDMF reader swaps x-y (x-z in 3d), so we transpose the "
53 : "tensors before we output to make the data look right in Paraview.");
54 22 : return params;
55 22 : }
56 :
57 12 : XDMFTensorOutput::XDMFTensorOutput(const InputParameters & parameters)
58 : : TensorOutput(parameters),
59 12 : _dim(_domain.getDim()),
60 12 : _frame(0),
61 24 : _transpose(getParam<bool>("transpose"))
62 : #ifdef LIBMESH_HAVE_HDF5
63 : ,
64 24 : _enable_hdf5(getParam<bool>("enable_hdf5")),
65 24 : _hdf5_name(_file_base + ".h5")
66 : #endif
67 : {
68 36 : const auto output_mode = getParam<MultiMooseEnum>("output_mode").getSetValueIDs<OutputMode>();
69 : const auto nbuffers = _out_buffers.size();
70 :
71 12 : if (output_mode.size() == 0)
72 : // default all to Cell
73 20 : for (const auto & pair : _out_buffers)
74 16 : _output_mode[pair.first] = OutputMode::CELL;
75 8 : else if (output_mode.size() != nbuffers)
76 0 : paramError(
77 : "output_mode", "Specify one output mode per buffer.", output_mode.size(), " != ", nbuffers);
78 : else
79 : {
80 8 : const auto & buffer_name = getParam<std::vector<TensorInputBufferName>>("buffer");
81 22 : for (const auto i : make_range(nbuffers))
82 14 : _output_mode[buffer_name[i]] = output_mode[i];
83 : }
84 :
85 : #ifdef LIBMESH_HAVE_HDF5
86 : // Check if the library is thread-safe
87 : hbool_t is_threadsafe;
88 12 : H5is_library_threadsafe(&is_threadsafe);
89 12 : if (!is_threadsafe)
90 : {
91 12 : for (const auto & output : _tensor_problem.getOutputs())
92 2 : if (output.get() != this && dynamic_cast<XDMFTensorOutput *>(output.get()))
93 2 : mooseError(
94 : "Using an hdf5 library that is not threadsafe and multiple XDMF output objects. "
95 : "Consolidate the XDMF outputs or build Swift with a thread safe build of libhdf5.");
96 10 : mooseWarning("Using an hdf5 library that is not threadsafe.");
97 : }
98 : #endif
99 10 : }
100 :
101 16 : XDMFTensorOutput::~XDMFTensorOutput()
102 : {
103 : #ifdef LIBMESH_HAVE_HDF5
104 8 : if (_enable_hdf5)
105 8 : H5Fclose(_hdf5_file_id);
106 : #endif
107 16 : }
108 :
109 : void
110 8 : XDMFTensorOutput::init()
111 : {
112 : // get mesh metadata
113 8 : auto sdim = Moose::stringify(_dim);
114 : std::vector<Real> origin;
115 : std::vector<Real> dgrid;
116 24 : for (const auto i : make_range(_dim))
117 : {
118 : // we need to transpose the tensor because of
119 : // https://discourse.paraview.org/t/axis-swapped-with-xdmf-topologytype-3dcorectmesh/3059/4
120 16 : const auto j = _transpose ? _dim - i - 1 : i;
121 16 : _ndata[0].push_back(_domain.getGridSize()[j]);
122 16 : _ndata[1].push_back(_domain.getGridSize()[j] + 1);
123 16 : _nnode.push_back(_domain.getGridSize()[j] + 1);
124 16 : dgrid.push_back(_domain.getGridSpacing()(j));
125 16 : origin.push_back(_domain.getDomainMin()(j));
126 : }
127 16 : _data_grid[0] = Moose::stringify(_ndata[0], " ");
128 16 : _data_grid[1] = Moose::stringify(_ndata[1], " ");
129 16 : _node_grid = Moose::stringify(_nnode, " ");
130 :
131 : //
132 : // setup XDMF skeleton
133 : //
134 :
135 : // Top level xdmf block
136 8 : auto xdmf = _doc.append_child("Xdmf");
137 8 : xdmf.append_attribute("xmlns:xi") = "http://www.w3.org/2003/XInclude";
138 8 : xdmf.append_attribute("Version") = "2.2";
139 :
140 : // Domain
141 8 : auto domain = xdmf.append_child("Domain");
142 :
143 : // - Topology
144 8 : auto topology = domain.append_child("Topology");
145 8 : topology.append_attribute("TopologyType") = (sdim + "DCoRectMesh").c_str();
146 8 : topology.append_attribute("Dimensions").set_value(_node_grid.c_str());
147 :
148 : // - Geometry
149 8 : auto geometry = domain.append_child("Geometry");
150 8 : std::string type = "ORIGIN_";
151 8 : const char * dxyz[] = {"DX", "DY", "DZ"};
152 24 : for (const auto i : make_range(_dim))
153 16 : type += dxyz[i];
154 8 : geometry.append_attribute("Type") = type.c_str();
155 :
156 : // -- Origin
157 : {
158 8 : auto data = geometry.append_child("DataItem");
159 8 : data.append_attribute("Format").set_value("XML");
160 8 : data.append_attribute("Dimensions") = sdim.c_str();
161 16 : data.append_child(pugi::node_pcdata).set_value(Moose::stringify(origin, " ").c_str());
162 : }
163 :
164 : // -- Grid spacing
165 : {
166 8 : auto data = geometry.append_child("DataItem");
167 8 : data.append_attribute("Format") = "XML";
168 8 : data.append_attribute("Dimensions") = sdim.c_str();
169 16 : data.append_child(pugi::node_pcdata).set_value(Moose::stringify(dgrid, " ").c_str());
170 : }
171 :
172 : // - TimeSeries Grid
173 8 : _tgrid = domain.append_child("Grid");
174 8 : _tgrid.append_attribute("Name") = "TimeSeries";
175 8 : _tgrid.append_attribute("GridType") = "Collection";
176 8 : _tgrid.append_attribute("CollectionType") = "Temporal";
177 :
178 : // write XDMF file
179 8 : _doc.save_file((_file_base + ".xmf").c_str());
180 : #ifdef LIBMESH_HAVE_HDF5
181 : // delete HDF5 file
182 8 : if (_enable_hdf5)
183 : {
184 8 : std::filesystem::remove(_hdf5_name);
185 : // open new file
186 8 : _hdf5_file_id = H5Fcreate(_hdf5_name.c_str(), H5F_ACC_TRUNC, H5P_DEFAULT, H5P_DEFAULT);
187 8 : if (_hdf5_file_id < 0)
188 : {
189 0 : H5Eprint(H5E_DEFAULT, stderr);
190 0 : mooseError("Error opening HDF5 file '", _hdf5_name, "'.");
191 : }
192 : }
193 : #endif
194 16 : }
195 :
196 : void
197 88 : XDMFTensorOutput::output()
198 : {
199 : // add grid for new timestep
200 88 : auto grid = _tgrid.append_child("Grid");
201 88 : grid.append_attribute("Name") = ("T" + Moose::stringify(_frame)).c_str();
202 88 : grid.append_attribute("GridType") = "Uniform";
203 :
204 : // time
205 88 : auto time = grid.append_child("Time");
206 88 : time.append_attribute("Value") = _time;
207 :
208 : // add references
209 88 : grid.append_child("xi:include").append_attribute("xpointer") = "xpointer(//Xdmf/Domain/Topology)";
210 88 : grid.append_child("xi:include").append_attribute("xpointer") = "xpointer(//Xdmf/Domain/Geometry)";
211 :
212 : // loop over buffers
213 352 : for (const auto & [buffer_name, original_buffer] : _out_buffers)
214 : {
215 : // skip empty tensors (no device)
216 264 : if (!original_buffer->defined())
217 0 : continue;
218 :
219 264 : const auto output_mode = _output_mode[buffer_name];
220 : torch::Tensor buffer;
221 :
222 : // we need to transpose the tensor because of
223 : // https://discourse.paraview.org/t/axis-swapped-with-xdmf-topologytype-3dcorectmesh/3059/4
224 264 : switch (output_mode)
225 : {
226 44 : case OutputMode::NODE:
227 44 : if (_transpose)
228 : {
229 0 : if (_dim == 2)
230 : {
231 0 : buffer = buffer = torch::transpose(extendTensor(*original_buffer), 0, 1).contiguous();
232 0 : break;
233 : }
234 0 : else if (_dim == 3)
235 : {
236 0 : buffer = buffer = torch::transpose(extendTensor(*original_buffer), 0, 2).contiguous();
237 0 : break;
238 : }
239 : }
240 44 : buffer = extendTensor(*original_buffer);
241 44 : break;
242 :
243 220 : case OutputMode::CELL:
244 : case OutputMode::OVERSIZED_NODAL:
245 220 : if (_transpose)
246 : {
247 0 : if (_dim == 2)
248 : {
249 0 : buffer = torch::transpose(*original_buffer, 0, 1).contiguous();
250 0 : break;
251 : }
252 0 : else if (_dim == 3)
253 : {
254 0 : buffer = torch::transpose(*original_buffer, 0, 2).contiguous();
255 : }
256 : break;
257 : }
258 :
259 220 : buffer = *original_buffer;
260 : break;
261 : }
262 :
263 : const auto is_cell = output_mode == OutputMode::CELL;
264 :
265 : const auto sizes = buffer.sizes();
266 : const auto total_dims = sizes.size();
267 :
268 264 : if (total_dims < _dim)
269 0 : mooseError("Tensor has fewer dimensions than specified spatial dimension.");
270 :
271 : int64_t num_grid_fields = 1;
272 792 : for (const auto i : make_range(_dim))
273 528 : num_grid_fields *= sizes[i];
274 :
275 : int64_t num_scalar_fields = 1;
276 264 : for (std::size_t i = _dim; i < total_dims; ++i)
277 0 : num_scalar_fields *= sizes[i];
278 :
279 528 : std::vector<int64_t> reshape_sizes = {num_grid_fields, num_scalar_fields};
280 264 : const auto reshaped = buffer.reshape(reshape_sizes);
281 :
282 : // now loop over scalar components
283 264 : const std::array<std::string, 3> xyz = {"x", "y", "z"};
284 528 : for (const auto index : make_range(num_scalar_fields))
285 : {
286 : auto name =
287 : buffer_name + (num_scalar_fields > 1
288 528 : ? "_" + (num_scalar_fields <= 3 ? xyz[index] : Moose::stringify(index))
289 264 : : "");
290 :
291 264 : buffer = reshaped.select(1, index).contiguous();
292 :
293 264 : auto attr = grid.append_child("Attribute");
294 264 : attr.append_attribute("Name") = name.c_str();
295 308 : attr.append_attribute("Center") = output_mode == OutputMode::CELL ? "Cell" : "Node";
296 264 : auto data = attr.append_child("DataItem");
297 264 : data.append_attribute("DataType") = "Float";
298 :
299 : // TODO: recompute data grid from sizes[0:_dim]
300 308 : data.append_attribute("Dimensions") = _data_grid[is_cell ? 0 : 1].c_str();
301 :
302 : // save file
303 528 : const auto setname = name + "." + Moose::stringify(_frame);
304 :
305 : char * raw_ptr = static_cast<char *>(buffer.data_ptr());
306 264 : std::size_t raw_size = buffer.numel();
307 :
308 : #ifdef LIBMESH_HAVE_HDF5
309 264 : if (_enable_hdf5)
310 : {
311 264 : if (buffer.dtype() == torch::kFloat32)
312 0 : addDataToHDF5(_hdf5_file_id, setname, raw_ptr, _ndata[is_cell ? 0 : 1], H5T_NATIVE_FLOAT);
313 264 : else if (buffer.dtype() == torch::kFloat64)
314 264 : addDataToHDF5(
315 264 : _hdf5_file_id, setname, raw_ptr, _ndata[is_cell ? 0 : 1], H5T_NATIVE_DOUBLE);
316 0 : else if (buffer.dtype() == torch::kInt32)
317 0 : addDataToHDF5(_hdf5_file_id, setname, raw_ptr, _ndata[is_cell ? 0 : 1], H5T_NATIVE_INT32);
318 0 : else if (buffer.dtype() == torch::kInt64)
319 0 : addDataToHDF5(_hdf5_file_id, setname, raw_ptr, _ndata[is_cell ? 0 : 1], H5T_NATIVE_INT64);
320 : else
321 0 : mooseError("Unsupported output type");
322 :
323 264 : data.append_attribute("Format") = "HDF";
324 264 : const auto h5path = _hdf5_name + ":/" + setname;
325 264 : data.append_child(pugi::node_pcdata).set_value(h5path.c_str());
326 : }
327 : else
328 : #endif
329 : {
330 0 : if (buffer.dtype() == torch::kFloat32)
331 : {
332 0 : data.append_attribute("Precision") = "4";
333 0 : raw_size *= 4;
334 : }
335 0 : else if (buffer.dtype() == torch::kFloat64)
336 : {
337 0 : data.append_attribute("Precision") = "8";
338 0 : raw_size *= 8;
339 : }
340 0 : else if (buffer.dtype() == torch::kInt32 || buffer.dtype() == torch::kInt64)
341 : {
342 0 : data.append_attribute("Precision") = "1";
343 : raw_size *= 1;
344 : }
345 : else
346 0 : mooseError("Unsupported output type");
347 :
348 0 : const auto fname = _file_base + "." + setname + ".bin";
349 0 : auto file = std::fstream(fname.c_str(), std::ios::out | std::ios::binary);
350 0 : file.write(raw_ptr, raw_size);
351 0 : file.close();
352 :
353 0 : data.append_attribute("Format") = "Binary";
354 0 : data.append_attribute("Endian") = "Little";
355 0 : data.append_child(pugi::node_pcdata).set_value(fname.c_str());
356 0 : }
357 : }
358 : }
359 :
360 : // write XDMF file
361 88 : _doc.save_file((_file_base + ".xmf").c_str());
362 :
363 : #ifdef LIBMESH_HAVE_HDF5
364 : // flush hdf5 file contents to disk
365 88 : if (_enable_hdf5)
366 88 : H5Fflush(_hdf5_file_id, H5F_SCOPE_GLOBAL);
367 : #endif
368 :
369 : // increment frame
370 88 : _frame++;
371 88 : }
372 :
373 : torch::Tensor
374 44 : XDMFTensorOutput::extendTensor(torch::Tensor tensor)
375 : {
376 : // for nodal data we increase each dimension by one and fill in a copy of the slice at 0
377 : torch::Tensor first;
378 : using torch::indexing::Slice;
379 :
380 44 : if (_dim == 3)
381 : {
382 0 : first = tensor.index({0, Slice(), Slice()}).unsqueeze(0);
383 0 : tensor = torch::cat({tensor, first}, 0);
384 0 : first = tensor.index({Slice(), 0, Slice()}).unsqueeze(1);
385 0 : tensor = torch::cat({tensor, first}, 1);
386 0 : first = tensor.index({Slice(), Slice(), 0}).unsqueeze(2);
387 0 : tensor = torch::cat({tensor, first}, 2);
388 : }
389 :
390 44 : else if (_dim == 2)
391 : {
392 132 : first = tensor.index({0}).unsqueeze(0);
393 132 : tensor = torch::cat({tensor, first}, 0);
394 264 : first = tensor.index({Slice(), 0}).unsqueeze(1);
395 132 : tensor = torch::cat({tensor, first}, 1);
396 : }
397 : else
398 0 : mooseError("Unsupported tensor dimension");
399 :
400 88 : return tensor.contiguous();
401 88 : }
402 :
403 : torch::Tensor
404 0 : XDMFTensorOutput::upsampleTensor(torch::Tensor tensor)
405 : {
406 : // For nodal nonperiodic data we transform into reciprocal space, add one additional K-vector on
407 : // each dimension, and transform back. This should amount to interpolating the spectral "shape
408 : // functions" at the nodes, rather than at the cell centers.
409 0 : std::vector<int64_t> shape(tensor.sizes().begin(), tensor.sizes().end());
410 0 : for (const auto i : make_range(_dim))
411 0 : shape[i]++;
412 :
413 : // return back transform with frequency padding
414 : return torch::fft::irfftn(
415 0 : _domain.fft(tensor), torch::IntArrayRef(shape.data(), _dim), _domain.getDimIndices());
416 0 : }
417 :
418 : #ifdef LIBMESH_HAVE_HDF5
419 : namespace
420 : {
421 : void
422 264 : addDataToHDF5(hid_t file_id,
423 : const std::string & dataset_name,
424 : const char * data,
425 : std::vector<std::size_t> & ndim,
426 : hid_t type)
427 : {
428 : hid_t dataset_id, dataspace_id, plist_id;
429 : herr_t status;
430 :
431 : // Open the file in read/write mode, create if it doesn't exist
432 :
433 : // hsize_t chunk_dims[RANK];
434 264 : std::vector<hsize_t> dims(ndim.begin(), ndim.end());
435 :
436 : // Check if the dataset already exists
437 264 : if (H5Lexists(file_id, dataset_name.c_str(), H5P_DEFAULT) > 0)
438 0 : mooseError("Dataset '", dataset_name, "' already exists in HDF5 file.");
439 :
440 : // Create a new dataset
441 264 : dataspace_id = H5Screate_simple(dims.size(), dims.data(), nullptr);
442 264 : if (dataspace_id < 0)
443 0 : mooseError("Error creating dataspace");
444 :
445 264 : plist_id = H5Pcreate(H5P_DATASET_CREATE);
446 264 : if (plist_id < 0)
447 0 : mooseError("Error creating property list");
448 :
449 264 : status = H5Pset_chunk(plist_id, dims.size(), dims.data());
450 264 : if (status < 0)
451 0 : mooseError("Error setting chunking");
452 :
453 264 : if (H5Zfilter_avail(H5Z_FILTER_DEFLATE))
454 : {
455 : unsigned filter_info;
456 264 : H5Zget_filter_info(H5Z_FILTER_DEFLATE, &filter_info);
457 264 : if (filter_info & H5Z_FILTER_CONFIG_ENCODE_ENABLED)
458 : {
459 264 : status = H5Pset_deflate(plist_id, 9);
460 264 : if (status < 0)
461 0 : mooseError("Error setting compression filter");
462 : }
463 : }
464 :
465 264 : dataset_id = H5Dcreate(
466 : file_id, dataset_name.c_str(), type, dataspace_id, H5P_DEFAULT, plist_id, H5P_DEFAULT);
467 264 : if (dataset_id < 0)
468 : {
469 : mooseInfo(dataset_id,
470 0 : ' ',
471 : file_id,
472 0 : ' ',
473 0 : dataset_name.c_str(),
474 0 : ' ',
475 : type,
476 0 : ' ',
477 : dataspace_id,
478 0 : ' ',
479 0 : H5P_DEFAULT,
480 0 : ' ',
481 : plist_id,
482 0 : ' ',
483 0 : H5P_DEFAULT);
484 0 : mooseError("Error creating dataset");
485 : }
486 :
487 : // Write data to the dataset
488 264 : status = H5Dwrite(dataset_id, type, H5S_ALL, dataspace_id, H5P_DEFAULT, data);
489 :
490 : // Close resources
491 264 : H5Pclose(plist_id);
492 264 : H5Dclose(dataset_id);
493 264 : H5Sclose(dataspace_id);
494 264 : }
495 : }
496 : #endif
|
