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 "Moose.h"
10 : #include "MooseError.h"
11 : #include "TensorProblem.h"
12 : #include "TensorSolver.h"
13 : #include "UniformTensorMesh.h"
14 :
15 : #include "TensorOperatorBase.h"
16 : #include "TensorTimeIntegrator.h"
17 : #include "TensorOutput.h"
18 : #include "DomainAction.h"
19 :
20 : #include "SwiftUtils.h"
21 : #include "DependencyResolverInterface.h"
22 : #include <memory>
23 :
24 : registerMooseObject("SwiftApp", TensorProblem);
25 :
26 : InputParameters
27 256 : TensorProblem::validParams()
28 : {
29 256 : InputParameters params = FEProblem::validParams();
30 256 : params.addClassDescription(
31 : "A normal Problem object that adds the ability to perform spectral solves.");
32 256 : params.set<bool>("skip_nl_system_check") = true;
33 512 : params.addParam<bool>("print_debug_output", false, "Show Tensor specific debug outputs");
34 512 : params.addParam<unsigned int>(
35 : "spectral_solve_substeps",
36 512 : 1,
37 : "How many substeps to divide the spectral solve for each MOOSE timestep into.");
38 512 : params.addParam<std::vector<std::string>>("scalar_constant_names", "Scalar constant names");
39 512 : params.addParam<std::vector<Real>>("scalar_constant_values", "Scalar constant values");
40 256 : return params;
41 0 : }
42 :
43 128 : TensorProblem::TensorProblem(const InputParameters & parameters)
44 : : FEProblem(parameters),
45 : DomainInterface(this),
46 128 : _options(MooseTensor::floatTensorOptions()),
47 256 : _debug(getParam<bool>("print_debug_output")),
48 256 : _substeps(getParam<unsigned int>("spectral_solve_substeps")),
49 128 : _dim(_domain.getDim()),
50 128 : _grid_spacing(_domain.getGridSpacing()),
51 128 : _n((_domain.getGridSize())),
52 128 : _shape(_domain.getShape()),
53 : _solver(nullptr),
54 128 : _can_fetch_constants(true)
55 : {
56 : // get constants (for scalar constants we provide a shortcut in the problem block)
57 128 : for (const auto & [name, value] :
58 260 : getParam<std::string, Real>("scalar_constant_names", "scalar_constant_values"))
59 132 : declareConstant<Real>(name, value);
60 :
61 : // make sure AuxVariables are contiguous in the solution vector
62 128 : getAuxiliarySystem().sys().identify_variable_groups(false);
63 128 : }
64 :
65 244 : TensorProblem::~TensorProblem()
66 : {
67 : // wait for outputs to be completed (otherwise resources might get freed that the output thread
68 : // depends on)
69 130 : for (auto & output : _outputs)
70 8 : output->waitForCompletion();
71 366 : }
72 :
73 : void
74 126 : TensorProblem::init()
75 : {
76 126 : unsigned int n_threads = libMesh::n_threads();
77 126 : if (n_threads != 1)
78 : {
79 63 : mooseInfo("Setting libTorch to use ", n_threads, " threads on the CPU.");
80 63 : torch::set_num_threads(n_threads);
81 : }
82 :
83 : // initialize tensors
84 986 : for (auto pair : _tensor_buffer)
85 860 : pair.second->init();
86 :
87 : // compute grid dependent quantities
88 126 : gridChanged();
89 :
90 : // init computes (must happen before dependency update)
91 630 : for (auto & initializer : _ics)
92 504 : initializer->init();
93 :
94 : // init computes (must happen before dependency update)
95 652 : for (auto & cmp : _computes)
96 526 : cmp->init();
97 :
98 : // update dependencies
99 126 : if (_solver)
100 110 : _solver->updateDependencies();
101 : else
102 : {
103 : // dependency resolution of TensorComputes
104 16 : DependencyResolverInterface::sort(_computes);
105 : }
106 :
107 : // dependency resolution of TensorICs
108 126 : DependencyResolverInterface::sort(_ics);
109 :
110 : // dependency resolution of Tensor Postprocessors
111 126 : DependencyResolverInterface::sort(_pps);
112 :
113 : // show computes
114 126 : if (_debug)
115 : {
116 0 : _console << COLOR_CYAN << "Compute object execution order:\n" << COLOR_DEFAULT;
117 0 : for (auto & cmp : _computes)
118 : {
119 0 : _console << " " << cmp->name() << '\n' << COLOR_YELLOW;
120 0 : for (const auto & ri : cmp->getRequestedItems())
121 0 : _console << " <- " << ri << '\n';
122 0 : _console << COLOR_GREEN;
123 0 : for (const auto & si : cmp->getSuppliedItems())
124 0 : _console << " -> " << si << '\n';
125 0 : _console << COLOR_DEFAULT;
126 : }
127 : }
128 :
129 : // call base class init
130 126 : FEProblem::init();
131 :
132 : // init outputs
133 134 : for (auto & output : _outputs)
134 8 : output->init();
135 :
136 : // check computes (after dependency update)
137 652 : for (auto & cmp : _computes)
138 526 : cmp->check();
139 :
140 126 : updateDOFMap();
141 :
142 : // debug output
143 : std::string variable_mapping;
144 166 : for (const auto & [buffer_name, tuple] : _buffer_to_var)
145 100 : variable_mapping += (std::get<bool>(tuple) ? "NODAL " : "ELEMENTAL ") + buffer_name + '\n';
146 126 : if (!variable_mapping.empty())
147 20 : mooseInfo("Direct buffer to solution vector mappings:\n", variable_mapping);
148 126 : }
149 :
150 : void
151 8494 : TensorProblem::execute(const ExecFlagType & exec_type)
152 : {
153 8494 : if (exec_type == EXEC_INITIAL)
154 : {
155 : // check for constants
156 126 : if (_fetched_constants.size() == 1)
157 2 : mooseError(
158 2 : "Constant ", Moose::stringify(_fetched_constants), " was requested but never declared.");
159 124 : if (_fetched_constants.size() > 1)
160 2 : mooseError("Constants ",
161 2 : Moose::stringify(_fetched_constants),
162 : " were requested but never declared.");
163 122 : _can_fetch_constants = false;
164 :
165 : // update time
166 122 : _sub_time = FEProblem::time();
167 :
168 122 : executeTensorInitialConditions();
169 :
170 122 : executeTensorOutputs(EXEC_INITIAL);
171 : }
172 :
173 8490 : if (exec_type == EXEC_TIMESTEP_BEGIN)
174 : {
175 : // update time
176 1624 : _sub_time = FEProblem::timeOld();
177 :
178 : // run solver
179 1624 : if (_solver)
180 1572 : _solver->computeBuffer();
181 : else
182 92 : for (auto & cmp : _computes)
183 40 : cmp->computeBuffer();
184 : }
185 :
186 8490 : if (exec_type == EXEC_TIMESTEP_END)
187 : {
188 : // run outputs
189 1624 : executeTensorOutputs(EXEC_TIMESTEP_END);
190 : }
191 :
192 8490 : FEProblem::execute(exec_type);
193 8490 : }
194 :
195 : void
196 122 : TensorProblem::executeTensorInitialConditions()
197 : {
198 : // run ICs
199 610 : for (auto & ic : _ics)
200 488 : ic->computeBuffer();
201 :
202 : // compile ist of compute output tensors
203 : std::set<std::string> _is_output;
204 636 : for (auto & cmp : _computes)
205 1028 : _is_output.insert(cmp->getSuppliedItems().begin(), cmp->getSuppliedItems().end());
206 :
207 : // // check for uninitialized tensors
208 : // for (auto & [name, t] : _tensor_buffer)
209 : // if (!t.defined() && _is_output.count(name) == 0)
210 : // mooseWarning(name, " is not initialized and not an output of any [Solve] compute.");
211 122 : }
212 :
213 : /// perform output tasks
214 : void
215 1746 : TensorProblem::executeTensorOutputs(const ExecFlagType & exec_flag)
216 : {
217 : // run postprocessing before output
218 1876 : for (auto & pp : _pps)
219 130 : pp->computeBuffer();
220 :
221 : // wait for prior asynchronous activity on CPU buffers to complete
222 : // (this is a synchronization barrier for the threaded CPU activity)
223 1834 : for (auto & output : _outputs)
224 88 : output->waitForCompletion();
225 :
226 : // update output time
227 1746 : _output_time = _time;
228 :
229 : // prepare CPU buffers (this is a synchronization barrier for the GPU)
230 16130 : for (const auto & pair : _tensor_buffer)
231 14384 : pair.second->makeCPUCopy();
232 :
233 : // run direct buffer outputs (asynchronous in threads)
234 1834 : for (auto & output : _outputs)
235 88 : if (output->shouldRun(exec_flag))
236 88 : output->startOutput();
237 :
238 1746 : if (_options.dtype() == torch::kFloat64)
239 1746 : mapBuffersToAux<double>();
240 0 : else if (_options.dtype() == torch::kFloat32)
241 0 : mapBuffersToAux<float>();
242 : else
243 0 : mooseError("torch::Dtype unsupported by mapBuffersToAux.");
244 1746 : }
245 :
246 : void
247 126 : TensorProblem::updateDOFMap()
248 : {
249 252 : TIME_SECTION("update", 3, "Updating Tensor DOF Map", true);
250 126 : const auto & min_global = _domain.getDomainMin();
251 :
252 : // variable mapping
253 : const auto & aux = getAuxiliarySystem();
254 126 : if (!const_cast<libMesh::System &>(aux.system()).is_initialized())
255 0 : mooseError("Aux system is not initialized :(");
256 :
257 126 : auto sys_num = aux.number();
258 166 : for (auto & [buffer_name, tuple] : _buffer_to_var)
259 : {
260 : auto & [var, dofs, is_nodal] = tuple;
261 40 : if (var->isArray() || var->isVector() || var->isFV())
262 0 : mooseError("Unsupported variable type for mapping");
263 40 : auto var_num = var->number();
264 :
265 65940 : auto compute_iteration_index = [this](Point p, long int n0, long int n1)
266 : {
267 131880 : return static_cast<long int>(p(0) / _grid_spacing(0)) +
268 65940 : (_dim > 1 ? static_cast<long int>(p(1) / _grid_spacing(1)) * n0 : 0) +
269 65940 : (_dim > 2 ? static_cast<long int>(p(2) / _grid_spacing(2)) * n0 * n1 : 0);
270 40 : };
271 :
272 40 : if (is_nodal)
273 : {
274 20 : long int n0 = _n[0] + 1;
275 20 : long int n1 = _n[1] + 1;
276 20 : long int n2 = _n[2] + 1;
277 36 : dofs.resize(n0 * (_dim > 1 ? n1 : 1) * (_dim > 2 ? n2 : 1));
278 :
279 : // loop over nodes
280 20 : const static Point shift = _grid_spacing / 2.0 - min_global;
281 67720 : for (const auto & node : _mesh.getMesh().node_ptr_range())
282 : {
283 33840 : const auto dof_index = node->dof_number(sys_num, var_num, 0);
284 33840 : const auto iteration_index = compute_iteration_index(*node + shift, n0, n1);
285 33840 : dofs[iteration_index] = dof_index;
286 20 : }
287 : }
288 : else
289 : {
290 20 : long int n0 = _n[0];
291 20 : long int n1 = _n[1];
292 20 : long int n2 = _n[2];
293 20 : dofs.resize(n0 * n1 * n2);
294 :
295 : // loop over elements
296 20 : const static Point shift = -min_global;
297 64240 : for (const auto & elem : _mesh.getMesh().element_ptr_range())
298 : {
299 32100 : const auto dof_index = elem->dof_number(sys_num, var_num, 0);
300 : const auto iteration_index =
301 32100 : compute_iteration_index(elem->vertex_average() + shift, n0, n1);
302 32100 : dofs[iteration_index] = dof_index;
303 20 : }
304 : }
305 : }
306 126 : }
307 :
308 : template <typename FLOAT_TYPE>
309 : void
310 1746 : TensorProblem::mapBuffersToAux()
311 : {
312 : // nothing to map?
313 1746 : if (_buffer_to_var.empty())
314 : return;
315 :
316 1320 : TIME_SECTION("update", 3, "Mapping Tensor buffers to Variables", true);
317 :
318 660 : auto * current_solution = &getAuxiliarySystem().solution();
319 660 : auto * solution_vector = dynamic_cast<PetscVector<Number> *>(current_solution);
320 660 : if (!solution_vector)
321 0 : mooseError(
322 : "Cannot map directly to the solution vector because NumericVector is not a PetscVector!");
323 :
324 : auto value = solution_vector->get_array();
325 :
326 : // const monomial variables
327 1980 : for (const auto & [buffer_name, tuple] : _buffer_to_var)
328 : {
329 : const auto & [var, dofs, is_nodal] = tuple;
330 : libmesh_ignore(var);
331 1320 : const long int n0 = is_nodal ? _n[0] + 1 : _n[0];
332 1320 : const long int n1 = is_nodal ? _n[1] + 1 : _n[1];
333 1320 : const long int n2 = is_nodal ? _n[2] + 1 : _n[2];
334 :
335 : // TODO: better design that works for NEML2 tensors as well
336 1320 : const auto buffer = getRawCPUBuffer(buffer_name);
337 1320 : if (buffer.sizes().size() != _dim)
338 0 : mooseError("Buffer '",
339 : buffer_name,
340 : "' is not a scalar tensor field and is not yet supported for AuxVariable mapping");
341 : std::size_t idx = 0;
342 1320 : switch (_dim)
343 : {
344 : {
345 0 : case 1:
346 0 : const auto b = buffer.template accessor<FLOAT_TYPE, 1>();
347 0 : for (const auto i : make_range(n0))
348 0 : value[dofs[idx++]] = b[i % _n[0]];
349 : break;
350 : }
351 1232 : case 2:
352 : {
353 1232 : const auto b = buffer.template accessor<FLOAT_TYPE, 2>();
354 60808 : for (const auto j : make_range(n1))
355 3014352 : for (const auto i : make_range(n0))
356 2954776 : value[dofs[idx++]] = b[i % _n[0]][j % _n[1]];
357 : break;
358 : }
359 88 : case 3:
360 : {
361 88 : const auto b = buffer.template accessor<FLOAT_TYPE, 3>();
362 572 : for (const auto k : make_range(n2))
363 3168 : for (const auto j : make_range(n1))
364 17688 : for (const auto i : make_range(n0))
365 15004 : value[dofs[idx++]] = b[i % _n[0]][j % _n[1]][k % _n[2]];
366 : break;
367 : }
368 0 : default:
369 0 : mooseError("Unsupported dimension");
370 : }
371 : }
372 :
373 : solution_vector->restore_array();
374 660 : getAuxiliarySystem().sys().update();
375 : }
376 :
377 : template <typename FLOAT_TYPE>
378 : void
379 : TensorProblem::mapAuxToBuffers()
380 : {
381 : // nothing to map?
382 : if (_var_to_buffer.empty())
383 : return;
384 :
385 : TIME_SECTION("update", 3, "Mapping Variables to Tensor buffers", true);
386 :
387 : const auto * current_solution = &getAuxiliarySystem().solution();
388 : const auto * solution_vector = dynamic_cast<const PetscVector<Number> *>(current_solution);
389 : if (!solution_vector)
390 : mooseError(
391 : "Cannot map directly to the solution vector because NumericVector is not a PetscVector!");
392 :
393 : const auto value = solution_vector->get_array_read();
394 :
395 : // const monomial variables
396 : for (const auto & [buffer_name, tuple] : _var_to_buffer)
397 : {
398 : const auto & [var, dofs, is_nodal] = tuple;
399 : libmesh_ignore(var);
400 : const auto buffer = getBufferBase(buffer_name).getRawCPUTensor();
401 : std::size_t idx = 0;
402 : switch (_dim)
403 : {
404 : {
405 : case 1:
406 : auto b = buffer.template accessor<FLOAT_TYPE, 1>();
407 : for (const auto i : make_range(_n[0]))
408 : b[i % _n[0]] = value[dofs[idx++]];
409 : break;
410 : }
411 : case 2:
412 : {
413 : auto b = buffer.template accessor<FLOAT_TYPE, 2>();
414 : for (const auto j : make_range(_n[1]))
415 : {
416 : for (const auto i : make_range(_n[0]))
417 : b[i % _n[0]][j % _n[1]] = value[dofs[idx++]];
418 : if (is_nodal)
419 : idx++;
420 : }
421 : break;
422 : }
423 : case 3:
424 : {
425 : auto b = buffer.template accessor<FLOAT_TYPE, 3>();
426 : for (const auto k : make_range(_n[2]))
427 : {
428 : for (const auto j : make_range(_n[1]))
429 : {
430 : for (const auto i : make_range(_n[0]))
431 : b[i % _n[0]][j % _n[1]][k % _n[2]] = value[dofs[idx++]];
432 : if (is_nodal)
433 : idx++;
434 : }
435 : if (is_nodal)
436 : idx += _n[0] + 1;
437 : }
438 : break;
439 : }
440 : default:
441 : mooseError("Unsupported dimension");
442 : }
443 : }
444 : }
445 :
446 : void
447 77836 : TensorProblem::advanceState()
448 : {
449 77836 : FEProblem::advanceState();
450 :
451 77836 : if (timeStep() <= 1)
452 5162 : return;
453 :
454 : // move buffers in time
455 72674 : std::size_t total_max = 0;
456 899618 : for (auto & pair : _tensor_buffer)
457 869298 : total_max = std::max(total_max, pair.second->advanceState());
458 :
459 : // move dt in time (UGH, we need the _substep_dt!!!!)
460 72674 : if (_old_dt.size() < total_max)
461 82 : _old_dt.push_back(0.0);
462 72674 : if (!_old_dt.empty())
463 : {
464 83134 : for (std::size_t i = _old_dt.size() - 1; i > 0; --i)
465 40780 : _old_dt[i] = _old_dt[i - 1];
466 42354 : _old_dt[0] = _dt;
467 : }
468 : }
469 :
470 : void
471 126 : TensorProblem::gridChanged()
472 : {
473 : // _domain.gridChanged();
474 126 : }
475 :
476 : void
477 340 : TensorProblem::addTensorBuffer(const std::string & buffer_type,
478 : const std::string & buffer_name,
479 : InputParameters & parameters)
480 : {
481 : // add buffer
482 340 : if (_tensor_buffer.find(buffer_name) != _tensor_buffer.end())
483 0 : mooseError("TensorBuffer '", buffer_name, "' already exists in the system");
484 :
485 : // Add a pointer to the TensorProblem and the Domain
486 340 : parameters.addPrivateParam<TensorProblem *>("_tensor_problem", this);
487 : // parameters.addPrivateParam<const DomainAction *>("_domain", &_domain);
488 :
489 : // Create the object
490 340 : auto tensor_buffer = _factory.create<TensorBufferBase>(buffer_type, buffer_name, parameters, 0);
491 340 : logAdd("TensorBufferBase", buffer_name, buffer_type, parameters);
492 :
493 340 : _tensor_buffer.try_emplace(buffer_name, tensor_buffer);
494 :
495 : // store variable mapping
496 340 : const auto & var_names = parameters.get<std::vector<AuxVariableName>>("map_to_aux_variable");
497 340 : if (!var_names.empty())
498 : {
499 : const auto & aux = getAuxiliarySystem();
500 : const auto var_name = var_names[0];
501 40 : if (!aux.hasVariable(var_name))
502 0 : mooseError("AuxVariable '", var_name, "' does not exist in the system.");
503 :
504 : bool is_nodal;
505 40 : const auto & var = aux.getVariable(0, var_name);
506 : if (var.feType() == FEType(FIRST, LAGRANGE))
507 : is_nodal = true;
508 : else if (var.feType() == FEType(CONSTANT, MONOMIAL))
509 : is_nodal = false;
510 : else
511 0 : mooseError("Only first order lagrange and constant monomial variables are supported for "
512 : "direct transfer. Try using the ProjectTensorAux kernel to transfer buffers to "
513 : "variables of any other type.");
514 :
515 80 : _buffer_to_var[buffer_name] = std::make_tuple(&var, std::vector<std::size_t>{}, is_nodal);
516 :
517 : // call this to mark the CPU copy as requested
518 40 : getRawCPUBuffer(buffer_name);
519 : }
520 340 : }
521 :
522 : void
523 536 : TensorProblem::addTensorComputeSolve(const std::string & compute_type,
524 : const std::string & compute_name,
525 : InputParameters & parameters)
526 : {
527 536 : addTensorCompute(compute_type, compute_name, parameters, _computes);
528 536 : }
529 :
530 : void
531 512 : TensorProblem::addTensorComputeInitialize(const std::string & compute_type,
532 : const std::string & compute_name,
533 : InputParameters & parameters)
534 : {
535 512 : addTensorCompute(compute_type, compute_name, parameters, _ics);
536 512 : }
537 :
538 : void
539 30 : TensorProblem::addTensorComputePostprocess(const std::string & compute_name,
540 : const std::string & name,
541 : InputParameters & parameters)
542 : {
543 30 : addTensorCompute(compute_name, name, parameters, _pps);
544 30 : }
545 :
546 : void
547 1078 : TensorProblem::addTensorCompute(const std::string & compute_type,
548 : const std::string & compute_name,
549 : InputParameters & parameters,
550 : TensorComputeList & list)
551 : {
552 : // Add a pointer to the TensorProblem and the Domain
553 1078 : parameters.addPrivateParam<TensorProblem *>("_tensor_problem", this);
554 1078 : parameters.addPrivateParam<const DomainAction *>("_domain", &_domain);
555 :
556 : // Create the object
557 : auto compute_object =
558 1078 : _factory.create<TensorOperatorBase>(compute_type, compute_name, parameters, 0);
559 1078 : logAdd("TensorOperatorBase", compute_name, compute_type, parameters);
560 1078 : list.push_back(compute_object);
561 1078 : }
562 :
563 : void
564 12 : TensorProblem::addTensorOutput(const std::string & output_type,
565 : const std::string & output_name,
566 : InputParameters & parameters)
567 : {
568 : // Add a pointer to the TensorProblem and the Domain
569 12 : parameters.addPrivateParam<TensorProblem *>("_tensor_problem", this);
570 12 : parameters.addPrivateParam<const DomainAction *>("_domain", &_domain);
571 :
572 : // Create the object
573 12 : auto output_object = _factory.create<TensorOutput>(output_type, output_name, parameters, 0);
574 10 : logAdd("TensorInitialCondition", output_name, output_type, parameters);
575 10 : _outputs.push_back(output_object);
576 10 : }
577 :
578 : void
579 110 : TensorProblem::setSolver(std::shared_ptr<TensorSolver> solver,
580 : const MooseTensor::Key<CreateTensorSolverAction> &)
581 : {
582 110 : if (_solver)
583 0 : mooseError("A solver has already been set up.");
584 :
585 : _solver = solver;
586 110 : }
587 :
588 : TensorBufferBase &
589 1990 : TensorProblem::getBufferBase(const std::string & buffer_name)
590 : {
591 : auto it = _tensor_buffer.find(buffer_name);
592 1990 : if (it == _tensor_buffer.end())
593 0 : mooseError("TensorBuffer '", buffer_name, " does not exist in the system.");
594 1990 : return *it->second.get();
595 : }
596 :
597 : const torch::Tensor &
598 588 : TensorProblem::getRawBuffer(const std::string & buffer_name)
599 : {
600 588 : return getBufferBase(buffer_name).getRawTensor();
601 : }
602 :
603 : const torch::Tensor &
604 1372 : TensorProblem::getRawCPUBuffer(const std::string & buffer_name)
605 : {
606 1372 : return getBufferBase(buffer_name).getRawCPUTensor();
607 : }
608 :
609 : TensorProblem &
610 298 : TensorProblem::cast(MooseObject * moose_object, Problem & problem)
611 : {
612 298 : if (auto tensor_problem = dynamic_cast<TensorProblem *>(&problem); tensor_problem)
613 298 : return *tensor_problem;
614 0 : moose_object->mooseError("Object requires a TensorProblem.");
615 : }
|
