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 : #pragma once
10 :
11 : #include "FEProblem.h"
12 : #include "DomainInterface.h"
13 : #include "InputParameters.h"
14 : #include "SwiftTypes.h"
15 : #include "SwiftUtils.h"
16 : #include "TensorBuffer.h"
17 :
18 : // list tensor buffer includes here
19 : #include "PlainTensorBuffer.h"
20 : #include "NEML2TensorBuffer.h"
21 :
22 : #include "AuxiliarySystem.h"
23 : #include "libmesh/petsc_vector.h"
24 : #include "libmesh/print_trace.h"
25 :
26 : #include <cstddef>
27 : #include <list>
28 : #include <memory>
29 : #include <torch/torch.h>
30 : #include <type_traits>
31 :
32 : class UniformTensorMesh;
33 : class TensorOperatorBase;
34 : template <typename T = torch::Tensor>
35 : class TensorTimeIntegrator;
36 : class TensorOutput;
37 : class TensorSolver;
38 : class CreateTensorSolverAction;
39 :
40 : namespace Swift
41 : {
42 : struct ConstantBase
43 : {
44 : virtual ~ConstantBase() = default;
45 : virtual std::string getType() = 0;
46 : };
47 :
48 : template <typename T>
49 10 : struct Constant : public ConstantBase
50 : {
51 0 : virtual std::string getType() override { return libMesh::demangle(typeid(T).name()); }
52 : T _value;
53 : };
54 : }
55 :
56 : /**
57 : * Problem for solving eigenvalue problems
58 : */
59 : class TensorProblem : public FEProblem, public DomainInterface
60 : {
61 : public:
62 : static InputParameters validParams();
63 :
64 : TensorProblem(const InputParameters & parameters);
65 : ~TensorProblem() override;
66 :
67 : // setup stuff
68 : void init() override;
69 :
70 : // run compute objects
71 : void execute(const ExecFlagType & exec_type) override;
72 :
73 : // move tensors in time
74 : void advanceState() override;
75 :
76 : // recompute quantities on grid size change
77 : virtual void gridChanged();
78 :
79 : virtual void addTensorBuffer(const std::string & buffer_type,
80 : const std::string & buffer_name,
81 : InputParameters & parameters);
82 :
83 : template <typename T>
84 : std::shared_ptr<TensorBuffer<T>> addTensorBuffer(const std::string & buffer_name);
85 :
86 : virtual void addTensorComputeInitialize(const std::string & compute_name,
87 : const std::string & name,
88 : InputParameters & parameters);
89 : virtual void addTensorComputeSolve(const std::string & compute_name,
90 : const std::string & name,
91 : InputParameters & parameters);
92 : virtual void addTensorComputePostprocess(const std::string & compute_name,
93 : const std::string & name,
94 : InputParameters & parameters);
95 :
96 : virtual void addTensorOutput(const std::string & output_name,
97 : const std::string & name,
98 : InputParameters & parameters);
99 :
100 : /// check if a tensor buffer with the given type exists
101 : template <typename T>
102 : bool hasBuffer(const std::string & buffer_name);
103 :
104 : /// returns the current state of the tensor
105 : template <typename T = torch::Tensor>
106 : T & getBuffer(const std::string & buffer_name);
107 :
108 : /// requests a tensor regardless of type
109 : TensorBufferBase & getBufferBase(const std::string & buffer_name);
110 :
111 : /// return the old states of the tensor
112 : template <typename T = torch::Tensor>
113 : const std::vector<T> & getBufferOld(const std::string & buffer_name, unsigned int max_states);
114 :
115 : /// returns a reference to a raw torch::Tensor view of buffer_name
116 : const torch::Tensor & getRawBuffer(const std::string & buffer_name);
117 :
118 : /// returns a reference to a copy of buffer_name that is guaranteed to be contiguous and located on the CPU device
119 : const torch::Tensor & getRawCPUBuffer(const std::string & buffer_name);
120 :
121 110 : virtual Real & subDt() { return _sub_dt; }
122 1298 : virtual Real & subTime() { return _sub_time; }
123 12 : virtual Real & outputTime() { return _output_time; }
124 :
125 : /// align a 1d tensor in a specific dimension
126 : torch::Tensor align(torch::Tensor t, unsigned int dim) const;
127 :
128 : /// get the domain shape (to build tensors from scratch) TODO: make sure this is local
129 32 : const torch::IntArrayRef & getShape() { return _shape; }
130 :
131 : typedef std::vector<std::shared_ptr<TensorOperatorBase>> TensorComputeList;
132 : const TensorComputeList & getComputes() const { return _computes; }
133 :
134 : template <typename T = TensorOperatorBase>
135 : T & getCompute(const std::string & name) const;
136 :
137 : typedef std::vector<std::shared_ptr<TensorOutput>> TensorOutputList;
138 : const TensorOutputList & getOutputs() const { return _outputs; }
139 :
140 : template <typename T>
141 : const T & getConstant(const std::string & name);
142 :
143 : template <typename T>
144 : void declareConstant(const std::string & name, const T & value);
145 :
146 : /// The CreateTensorSolverAction calls this to set the active solver
147 : void setSolver(std::shared_ptr<TensorSolver> solver,
148 : const MooseTensor::Key<CreateTensorSolverAction> &);
149 :
150 : /// get a reference to the current solver
151 : template <typename T>
152 : T & getSolver() const;
153 :
154 : static TensorProblem & cast(MooseObject * moose_object, Problem & problem);
155 :
156 : protected:
157 : void updateDOFMap();
158 :
159 : template <typename FLOAT_TYPE>
160 : void mapBuffersToAux();
161 :
162 : template <typename FLOAT_TYPE>
163 : void mapAuxToBuffers();
164 :
165 : virtual void addTensorCompute(const std::string & compute_name,
166 : const std::string & name,
167 : InputParameters & parameters,
168 : TensorComputeList & list);
169 :
170 : /// execute initial conditionobjects
171 : void executeTensorInitialConditions();
172 :
173 : /// perform output tasks
174 : void executeTensorOutputs(const ExecFlagType & exec_type);
175 :
176 : /// helper to get the TensorBuffer wrapper object that holds the actual tensor data
177 : template <typename T = torch::Tensor>
178 : TensorBuffer<T> & getBufferHelper(const std::string & buffer_name);
179 :
180 : /// tensor options
181 : const torch::TensorOptions _options;
182 :
183 : /// show debug ouput
184 : const bool _debug;
185 :
186 : /// solver substeps
187 : const unsigned int _substeps;
188 :
189 : /// substepping timestep
190 : Real _sub_dt;
191 : Real _sub_time;
192 :
193 : /// simulation time for the currently running output thread
194 : Real _output_time;
195 :
196 : /// list of TensorBuffers (i.e. tensors)
197 : std::map<std::string, std::shared_ptr<TensorBufferBase>> _tensor_buffer;
198 :
199 : /// old timesteps
200 : std::vector<Real> _old_dt;
201 :
202 : const unsigned int & _dim;
203 :
204 : /// grid spacing
205 : const RealVectorValue & _grid_spacing;
206 :
207 : /// global grid size
208 : const std::array<int64_t, 3> & _n;
209 :
210 : /// domain shape
211 : const torch::IntArrayRef & _shape;
212 :
213 : /// solve objects
214 : TensorComputeList _computes;
215 :
216 : /// initialization objects
217 : TensorComputeList _ics;
218 :
219 : /// postprocessing objects
220 : TensorComputeList _pps;
221 :
222 : /// on demand objects that are explicitly triggered by other objects
223 : TensorComputeList _on_demand;
224 :
225 : /// time integrator objects
226 : std::vector<std::shared_ptr<TensorTimeIntegrator<>>> _time_integrators;
227 :
228 : std::vector<std::shared_ptr<TensorOutput>> _outputs;
229 :
230 : /// map from buffer name to variable name
231 : std::map<std::string, AuxVariableName> _buffer_to_var_name;
232 :
233 : /// buffers to solution vector indices
234 : std::map<std::string, std::tuple<const MooseVariableFieldBase *, std::vector<std::size_t>, bool>>
235 : _buffer_to_var;
236 : std::map<std::string, std::tuple<const MooseVariableFieldBase *, std::vector<std::size_t>, bool>>
237 : _var_to_buffer;
238 :
239 : /// The [TensorSolver]
240 : std::shared_ptr<TensorSolver> _solver;
241 :
242 : /// parameters
243 : std::map<std::string, std::unique_ptr<Swift::ConstantBase>> _constants;
244 : std::set<std::string> _fetched_constants;
245 : std::list<Real> _literal_constants;
246 : bool _can_fetch_constants;
247 :
248 : private:
249 : std::list<InputParameters> _buffer_params;
250 : };
251 :
252 : template <typename T>
253 : T &
254 0 : TensorProblem::getSolver() const
255 : {
256 0 : if (_solver)
257 : {
258 0 : const auto specialized_solver = dynamic_cast<T *>(_solver.get());
259 0 : if (specialized_solver)
260 0 : return *specialized_solver;
261 0 : mooseError(
262 0 : "No TensorSolver supporting the requested type '", typeid(T).name(), "' has been set up.");
263 : }
264 0 : mooseError("No TensorSolver has been set up.");
265 : }
266 :
267 : template <typename T>
268 : bool
269 : TensorProblem::hasBuffer(const std::string & buffer_name)
270 : {
271 : auto it = _tensor_buffer.find(buffer_name);
272 : return (it != _tensor_buffer.end() && !dynamic_cast<TensorBuffer<T> *>(it->second.get()));
273 : }
274 :
275 : template <typename T>
276 : TensorBuffer<T> &
277 0 : TensorProblem::getBufferHelper(const std::string & buffer_name)
278 : {
279 : auto it = _tensor_buffer.find(buffer_name);
280 :
281 0 : if (it == _tensor_buffer.end())
282 0 : return *addTensorBuffer<T>(buffer_name);
283 : else
284 : {
285 0 : auto tensor_buffer = dynamic_cast<TensorBuffer<T> *>(it->second.get());
286 0 : if (!tensor_buffer)
287 0 : mooseError("TensorBuffer '",
288 : buffer_name,
289 : "' of the requested type '",
290 : libMesh::demangle(typeid(T).name()),
291 : "' was previously declared as '",
292 : it->second->type(),
293 : "'.");
294 : return *tensor_buffer;
295 : }
296 : }
297 :
298 : template <typename T>
299 : T &
300 : TensorProblem::getBuffer(const std::string & buffer_name)
301 : {
302 1422 : return getBufferHelper<T>(buffer_name).getTensor();
303 : }
304 :
305 : template <typename T>
306 : const std::vector<T> &
307 130 : TensorProblem::getBufferOld(const std::string & buffer_name, unsigned int max_states)
308 : {
309 130 : return getBufferHelper<T>(buffer_name).getOldTensor(max_states);
310 : }
311 :
312 : template <typename T>
313 : std::shared_ptr<TensorBuffer<T>>
314 0 : TensorProblem::addTensorBuffer(const std::string & buffer_name)
315 : {
316 0 : if (_debug)
317 0 : mooseInfoRepeated("Automatically adding tensor '",
318 : buffer_name,
319 : "' of type '",
320 : libMesh::demangle(typeid(T).name()),
321 : "'");
322 :
323 0 : _buffer_params.push_back(TensorBufferSpecialization<T>::type::validParams());
324 : auto & params = _buffer_params.back();
325 0 : params.template set<std::string>("_object_name") = buffer_name;
326 0 : params.template set<FEProblem *>("_fe_problem") = this;
327 0 : params.template set<FEProblemBase *>("_fe_problem_base") = this;
328 0 : params.template set<THREAD_ID>("_tid") = 0;
329 0 : params.template set<std::string>("_type") = "TensorBufferBase";
330 0 : params.template set<MooseApp *>("_moose_app") = &getMooseApp();
331 0 : params.finalize(buffer_name);
332 : auto tensor_buffer = std::make_shared<typename TensorBufferSpecialization<T>::type>(params);
333 :
334 0 : _tensor_buffer.try_emplace(buffer_name, tensor_buffer);
335 0 : return tensor_buffer;
336 : }
337 :
338 : template <typename T>
339 : const T &
340 332 : TensorProblem::getConstant(const std::string & name)
341 : {
342 332 : if (!_can_fetch_constants)
343 0 : mooseError("Constants must be gotten during object construction.");
344 :
345 : // deal with literal scalars
346 : if constexpr (std::is_same_v<T, Real>)
347 : {
348 : // check if we are able to convert the name into a Real
349 : Real real_value;
350 332 : std::istringstream ss(name);
351 332 : if (ss >> real_value && ss.eof())
352 : {
353 322 : _literal_constants.push_back(real_value);
354 : return _literal_constants.back();
355 : }
356 332 : }
357 :
358 : const auto it = _constants.find(name);
359 10 : if (it != _constants.end())
360 : {
361 4 : auto * t_param = dynamic_cast<Swift::Constant<T> *>(it->second.get());
362 4 : if (t_param == nullptr)
363 0 : mooseError("Fetching constant '",
364 : name,
365 : "' (which is of type ",
366 0 : it->second->getType(),
367 : ") with the wrong type (",
368 : libMesh::demangle(typeid(T).name()),
369 : ").");
370 4 : return t_param->_value;
371 : }
372 : else
373 : {
374 6 : auto param = std::make_unique<Swift::Constant<T>>();
375 6 : const auto & ref = param->_value;
376 6 : _constants[name] = std::move(param);
377 6 : _fetched_constants.insert(name);
378 : return ref;
379 6 : }
380 : }
381 :
382 : template <typename T>
383 : void
384 4 : TensorProblem::declareConstant(const std::string & name, const T & value)
385 : {
386 4 : if (!_can_fetch_constants)
387 0 : mooseError("Constants must be declared during object construction.");
388 :
389 : const auto it = _constants.find(name);
390 4 : if (it != _constants.end())
391 : {
392 0 : auto * t_param = dynamic_cast<Swift::Constant<T> *>(it->second.get());
393 0 : if (t_param == nullptr)
394 0 : mooseError("Declaring constant '",
395 : name,
396 : "' (which was already fetched as type ",
397 0 : it->second->getType(),
398 : ") with the wrong type (",
399 : libMesh::demangle(typeid(T).name()),
400 : ").");
401 :
402 : if (_fetched_constants.count(name))
403 : {
404 : // placeholder parameter created by getParameter
405 0 : t_param->_value = value;
406 : _fetched_constants.erase(name);
407 : }
408 : else
409 0 : mooseError("Parameter '", name, "' has already been declared.");
410 : }
411 : else
412 : {
413 4 : auto param = std::make_unique<Swift::Constant<T>>();
414 4 : param->_value = value;
415 4 : _constants[name] = std::move(param);
416 4 : }
417 4 : }
418 :
419 : template <typename T>
420 : T &
421 4 : TensorProblem::getCompute(const std::string & name) const
422 : {
423 4 : for (const auto & tcb : _computes)
424 8 : if (std::dynamic_pointer_cast<MooseObject>(tcb)->name() == name)
425 4 : if (const auto ptr = std::dynamic_pointer_cast<T>(tcb); ptr)
426 4 : return *ptr;
427 :
428 0 : mooseError("No compute with name '",
429 : name,
430 : "' of type ",
431 : libMesh::demangle(typeid(T).name()),
432 : " found.");
433 : }
|
