Line data Source code
1 : //* This file is part of the MOOSE framework
2 : //* https://mooseframework.inl.gov
3 : //*
4 : //* All rights reserved, see COPYRIGHT for full restrictions
5 : //* https://github.com/idaholab/moose/blob/master/COPYRIGHT
6 : //*
7 : //* Licensed under LGPL 2.1, please see LICENSE for details
8 : //* https://www.gnu.org/licenses/lgpl-2.1.html
9 :
10 : #include "NEML2ModelExecutor.h"
11 : #include "MOOSEToNEML2.h"
12 : #include "NEML2Utils.h"
13 : #include <string>
14 :
15 : #ifdef NEML2_ENABLED
16 : #include "libmesh/id_types.h"
17 : #include "neml2/tensors/functions/jacrev.h"
18 : #include "neml2/dispatchers/ValueMapLoader.h"
19 : #endif
20 :
21 : registerMooseObject("MooseApp", NEML2ModelExecutor);
22 :
23 : InputParameters
24 14739 : NEML2ModelExecutor::actionParams()
25 : {
26 14739 : auto params = emptyInputParameters();
27 : // allow user to explicit skip required input variables
28 14739 : params.addParam<std::vector<std::string>>(
29 : "skip_inputs",
30 : {},
31 29478 : NEML2Utils::docstring(
32 : "List of NEML2 variables to skip error checking when setting up the model input. If an "
33 : "input variable is skipped, its value will stay zero. If a required input variable is "
34 : "not skipped, an error will be raised."));
35 14739 : return params;
36 0 : }
37 :
38 : InputParameters
39 14299 : NEML2ModelExecutor::validParams()
40 : {
41 14299 : auto params = NEML2ModelInterface<GeneralUserObject>::validParams();
42 14299 : params += NEML2ModelExecutor::actionParams();
43 14299 : params.addClassDescription(NEML2Utils::docstring("Execute the specified NEML2 model"));
44 :
45 14299 : params.addRequiredParam<UserObjectName>(
46 : "batch_index_generator",
47 : "The NEML2BatchIndexGenerator used to generate the element-to-batch-index map.");
48 14299 : params.addParam<std::vector<UserObjectName>>(
49 : "gatherers",
50 : {},
51 28598 : NEML2Utils::docstring(
52 : "List of MOOSE*ToNEML2 user objects gathering MOOSE data as NEML2 input variables"));
53 14299 : params.addParam<std::vector<UserObjectName>>(
54 : "param_gatherers",
55 : {},
56 28598 : NEML2Utils::docstring(
57 : "List of MOOSE*ToNEML2 user objects gathering MOOSE data as NEML2 model parameters"));
58 :
59 : // Since we use the NEML2 model to evaluate the residual AND the Jacobian at the same time, we
60 : // want to execute this user object only at execute_on = LINEAR (i.e. during residual evaluation).
61 : // The NONLINEAR exec flag below is for computing Jacobian during automatic scaling.
62 14299 : ExecFlagEnum execute_options = MooseUtils::getDefaultExecFlagEnum();
63 57196 : execute_options = {EXEC_INITIAL, EXEC_LINEAR, EXEC_NONLINEAR};
64 14299 : params.set<ExecFlagEnum>("execute_on") = execute_options;
65 :
66 28598 : return params;
67 28598 : }
68 :
69 17 : NEML2ModelExecutor::NEML2ModelExecutor(const InputParameters & params)
70 0 : : NEML2ModelInterface<GeneralUserObject>(params)
71 : #ifdef NEML2_ENABLED
72 : ,
73 17 : _batch_index_generator(getUserObject<NEML2BatchIndexGenerator>("batch_index_generator")),
74 17 : _output_ready(false),
75 34 : _error_message("")
76 : #endif
77 : {
78 : #ifdef NEML2_ENABLED
79 17 : validateModel();
80 :
81 : // add user object dependencies by name (the UOs do not need to exist yet for this)
82 47 : for (const auto & gatherer_name : getParam<std::vector<UserObjectName>>("gatherers"))
83 30 : _depend_uo.insert(gatherer_name);
84 19 : for (const auto & gatherer_name : getParam<std::vector<UserObjectName>>("param_gatherers"))
85 2 : _depend_uo.insert(gatherer_name);
86 :
87 : // variables to skip error checking (converting vector to set to prevent duplicate checks)
88 17 : for (const auto & var_name : getParam<std::vector<std::string>>("skip_inputs"))
89 0 : _skip_vars.insert(NEML2Utils::parseVariableName(var_name));
90 : #endif
91 17 : }
92 :
93 : #ifdef NEML2_ENABLED
94 : void
95 17 : NEML2ModelExecutor::initialSetup()
96 : {
97 : // deal with user object provided inputs
98 47 : for (const auto & gatherer_name : getParam<std::vector<UserObjectName>>("gatherers"))
99 : {
100 : // gather coupled user objects late to ensure they are constructed. Do not add them as
101 : // dependencies (that's already done in the constructor).
102 30 : const auto & uo = getUserObjectByName<MOOSEToNEML2>(gatherer_name, /*is_dependency=*/false);
103 :
104 : // the target neml2 variable must exist on the input axis
105 30 : if (!model().input_axis().has_variable(NEML2Utils::parseVariableName(uo.NEML2Name())))
106 0 : mooseError("The MOOSEToNEML2 gatherer named '",
107 : gatherer_name,
108 : "' is gathering MOOSE data for a non-existent NEML2 input variable named '",
109 0 : uo.NEML2Name(),
110 : "'.");
111 :
112 : // tell the gatherer to gather for a model input variable
113 30 : const auto varname = NEML2Utils::parseVariableName(uo.NEML2Name());
114 30 : if (varname.is_old_force() || varname.is_old_state())
115 0 : uo.setMode(MOOSEToNEML2::Mode::OLD_VARIABLE);
116 : else
117 30 : uo.setMode(MOOSEToNEML2::Mode::VARIABLE);
118 :
119 30 : addGatheredVariable(gatherer_name, uo.NEML2VariableName());
120 30 : _gatherers.push_back(&uo);
121 30 : }
122 :
123 : // deal with user object provided model parameters
124 19 : for (const auto & gatherer_name : getParam<std::vector<UserObjectName>>("param_gatherers"))
125 : {
126 : // gather coupled user objects late to ensure they are constructed. Do not add them as
127 : // dependencies (that's already done in the constructor).
128 2 : const auto & uo = getUserObjectByName<MOOSEToNEML2>(gatherer_name, /*is_dependency=*/false);
129 :
130 : // introspect the NEML2 model to figure out if the gatherer UO is gathering for a NEML2 input
131 : // variable or for a NEML2 model parameter
132 2 : if (model().named_parameters().count(uo.NEML2Name()) != 1)
133 0 : mooseError("The MOOSEToNEML2 gatherer named '",
134 : gatherer_name,
135 : "' is gathering MOOSE data for a non-existent NEML2 model parameter named '",
136 0 : uo.NEML2Name(),
137 : "'.");
138 :
139 : // tell the gatherer to gather for a model parameter
140 2 : uo.setMode(MOOSEToNEML2::Mode::PARAMETER);
141 :
142 2 : addGatheredParameter(gatherer_name, uo.NEML2ParameterName());
143 2 : _gatherers.push_back(&uo);
144 : }
145 :
146 : // iterate over set of required inputs and error out if we find one that is not provided
147 17 : std::vector<neml2::VariableName> required_inputs = model().input_axis().variable_names();
148 47 : for (const auto & input : required_inputs)
149 : {
150 30 : if (_skip_vars.count(input))
151 0 : continue;
152 : // skip input state variables because they are "initial guesses" to the nonlinear system
153 30 : if (input.is_state())
154 0 : continue;
155 30 : if (!_gathered_variable_names.count(input))
156 0 : paramError("gatherers", "The required model input `", input, "` is not gathered");
157 : }
158 :
159 : // If a variable is stateful, then it'd better been retrieved by someone! In theory that's not
160 : // sufficient for stateful data management, but that's the best we can do here without being
161 : // overly restrictive.
162 47 : for (const auto & input : required_inputs)
163 30 : if (input.is_old_state() && !_retrieved_outputs.count(input.current()))
164 0 : mooseError(
165 : "The NEML2 model requires a stateful input variable `",
166 : input,
167 : "`, but its state counterpart on the output axis has not been retrieved by any object. "
168 : "Therefore, there is no way to properly propagate the corresponding stateful data in "
169 : "time. The common solution to this problem is to add a NEML2ToMOOSE retriever such as "
170 : "those called `NEML2To*MOOSEMaterialProperty`.");
171 17 : }
172 :
173 : std::size_t
174 43580 : NEML2ModelExecutor::getBatchIndex(dof_id_type elem_id) const
175 : {
176 43580 : return _batch_index_generator.getBatchIndex(elem_id);
177 : }
178 :
179 : void
180 30 : NEML2ModelExecutor::addGatheredVariable(const UserObjectName & gatherer_name,
181 : const neml2::VariableName & var)
182 : {
183 30 : if (_gathered_variable_names.count(var))
184 0 : paramError("gatherers",
185 : "The NEML2 input variable `",
186 : var,
187 : "` gathered by UO '",
188 : gatherer_name,
189 : "' is already gathered by another gatherer.");
190 30 : _gathered_variable_names.insert(var);
191 30 : }
192 :
193 : void
194 2 : NEML2ModelExecutor::addGatheredParameter(const UserObjectName & gatherer_name,
195 : const std::string & param)
196 : {
197 2 : if (_gathered_parameter_names.count(param))
198 0 : paramError("gatherers",
199 : "The NEML2 model parameter `",
200 : param,
201 : "` gathered by UO '",
202 : gatherer_name,
203 : "' is already gathered by another gatherer.");
204 2 : _gathered_parameter_names.insert(param);
205 2 : }
206 :
207 : void
208 1121 : NEML2ModelExecutor::initialize()
209 : {
210 1121 : if (!NEML2Utils::shouldCompute(_fe_problem))
211 460 : return;
212 :
213 661 : _output_ready = false;
214 661 : _error = false;
215 661 : _error_message.clear();
216 : }
217 :
218 : void
219 8 : NEML2ModelExecutor::meshChanged()
220 : {
221 8 : if (!NEML2Utils::shouldCompute(_fe_problem))
222 0 : return;
223 :
224 8 : _output_ready = false;
225 : }
226 :
227 : void
228 1121 : NEML2ModelExecutor::execute()
229 : {
230 1121 : if (!NEML2Utils::shouldCompute(_fe_problem))
231 460 : return;
232 :
233 661 : fillInputs();
234 :
235 661 : if (_t_step > 0)
236 : {
237 645 : applyPredictor();
238 645 : auto success = solve();
239 645 : if (success)
240 644 : extractOutputs();
241 : }
242 : }
243 :
244 : void
245 661 : NEML2ModelExecutor::fillInputs()
246 : {
247 : try
248 : {
249 1382 : for (const auto & uo : _gatherers)
250 721 : uo->insertInto(_in, _model_params);
251 :
252 : // Send input variables and parameters to device
253 1376 : for (auto & [var, val] : _in)
254 715 : val = val.to(device());
255 667 : for (auto & [param, pval] : _model_params)
256 6 : pval = pval.to(device());
257 :
258 : // Update model parameters
259 661 : model().set_parameters(_model_params);
260 661 : _model_params.clear();
261 :
262 : // Request gradient for the model parameters that we request AD for
263 661 : for (const auto & [y, dy] : _retrieved_parameter_derivatives)
264 0 : for (const auto & [p, tensor] : dy)
265 0 : model().get_parameter(p).requires_grad_(true);
266 : }
267 0 : catch (std::exception & e)
268 : {
269 0 : mooseError("An error occurred while filling inputs for the NEML2 model. Error message:\n",
270 0 : e.what(),
271 : NEML2Utils::NEML2_help_message);
272 0 : }
273 661 : }
274 :
275 : void
276 645 : NEML2ModelExecutor::applyPredictor()
277 : {
278 : try
279 : {
280 645 : if (!model().input_axis().has_state())
281 645 : return;
282 0 : if (!model().input_axis().has_old_state())
283 0 : return;
284 :
285 : // Set trial state variables (i.e., initial guesses).
286 : // Right now we hard-code to use the old state as the trial state.
287 : // TODO: implement other predictors
288 0 : const auto & input_state = model().input_axis().subaxis(neml2::STATE);
289 0 : const auto & input_old_state = model().input_axis().subaxis(neml2::OLD_STATE);
290 0 : for (const auto & var : input_state.variable_names())
291 0 : if (input_old_state.has_variable(var))
292 0 : _in[var.prepend(neml2::STATE)] = _in[var.prepend(neml2::OLD_STATE)];
293 : }
294 0 : catch (std::exception & e)
295 : {
296 0 : mooseError("An error occurred while applying predictor for the NEML2 model. Error message:\n",
297 0 : e.what(),
298 : NEML2Utils::NEML2_help_message);
299 0 : }
300 : }
301 :
302 : void
303 4 : NEML2ModelExecutor::expandInputs()
304 : {
305 : // Figure out what our batch size is
306 4 : std::vector<neml2::Tensor> defined;
307 12 : for (const auto & [key, value] : _in)
308 8 : defined.push_back(value);
309 4 : const auto batch_shape = neml2::utils::broadcast_batch_sizes(defined);
310 :
311 : // Make all inputs conformal
312 12 : for (auto & [key, value] : _in)
313 8 : if (value.batch_sizes() != batch_shape)
314 0 : _in[key] = value.batch_unsqueeze(0).batch_expand(batch_shape);
315 4 : }
316 :
317 : bool
318 645 : NEML2ModelExecutor::solve()
319 : {
320 : try
321 : {
322 : // Evaluate the NEML2 material model
323 645 : TIME_SECTION("NEML2 solve", 3, "Solving NEML2 material model");
324 :
325 : // NEML2 requires double precision
326 645 : auto prev_dtype = neml2::get_default_dtype();
327 645 : neml2::set_default_dtype(neml2::kFloat64);
328 :
329 645 : if (scheduler())
330 : {
331 : // We only need consistent batch sizes if we are using the dispatcher
332 4 : expandInputs();
333 4 : neml2::ValueMapLoader loader(_in, 0);
334 4 : std::tie(_out, _dout_din) = dispatcher()->run(loader);
335 4 : }
336 : else
337 641 : std::tie(_out, _dout_din) = model().value_and_dvalue(_in);
338 644 : _in.clear();
339 :
340 : // Restore the default dtype
341 644 : neml2::set_default_dtype(prev_dtype);
342 645 : }
343 1 : catch (std::exception & e)
344 : {
345 1 : _error_message = e.what();
346 1 : _error = true;
347 1 : }
348 :
349 645 : return !_error;
350 : }
351 :
352 : void
353 644 : NEML2ModelExecutor::extractOutputs()
354 : {
355 : try
356 : {
357 644 : const auto N = _batch_index_generator.getBatchIndex();
358 :
359 : // retrieve outputs
360 1328 : for (auto & [y, target] : _retrieved_outputs)
361 684 : target = _out[y].to(torch::kCPU);
362 :
363 : // retrieve parameter derivatives
364 644 : for (auto & [y, dy] : _retrieved_parameter_derivatives)
365 0 : for (auto & [p, target] : dy)
366 0 : target = neml2::jacrev(_out[y],
367 0 : model().get_parameter(p),
368 : /*retain_graph=*/true,
369 : /*create_graph=*/false,
370 : /*allow_unused=*/false)
371 0 : .to(torch::kCPU);
372 :
373 : // clear output
374 644 : _out.clear();
375 :
376 : // retrieve derivatives
377 1288 : for (auto & [y, dy] : _retrieved_derivatives)
378 1288 : for (auto & [x, target] : dy)
379 : {
380 644 : const auto & source = _dout_din[y][x];
381 644 : if (source.defined())
382 1932 : target = source.to(torch::kCPU).batch_expand({neml2::Size(N)});
383 : }
384 :
385 : // clear derivatives
386 644 : _dout_din.clear();
387 : }
388 0 : catch (std::exception & e)
389 : {
390 0 : mooseError("An error occurred while retrieving outputs from the NEML2 model. Error message:\n",
391 0 : e.what(),
392 : NEML2Utils::NEML2_help_message);
393 0 : }
394 1288 : }
395 :
396 : void
397 1121 : NEML2ModelExecutor::finalize()
398 : {
399 1121 : if (!NEML2Utils::shouldCompute(_fe_problem))
400 460 : return;
401 :
402 : // See if any rank failed
403 : processor_id_type pid;
404 661 : _communicator.maxloc(_error, pid);
405 :
406 : // Fail the next nonlinear convergence check if any rank failed
407 661 : if (_error)
408 : {
409 1 : _communicator.broadcast(_error_message, pid);
410 1 : if (_communicator.rank() == 0)
411 : {
412 2 : std::string msg = "NEML2 model execution failed on at least one processor with ID " +
413 3 : std::to_string(pid) + ". Error message:\n";
414 1 : msg += _error_message;
415 1 : if (_fe_problem.isTransient())
416 : msg += "\nTo recover, the solution will fail and then be re-attempted with a reduced time "
417 1 : "step.";
418 1 : _console << COLOR_YELLOW << msg << COLOR_DEFAULT << std::endl;
419 1 : }
420 1 : _fe_problem.setFailNextNonlinearConvergenceCheck();
421 : }
422 660 : else if (_t_step > 0)
423 644 : _output_ready = true;
424 : }
425 :
426 : void
427 141 : NEML2ModelExecutor::checkExecutionStage() const
428 : {
429 141 : if (_fe_problem.startedInitialSetup())
430 0 : mooseError("NEML2 output variables and derivatives must be retrieved during object "
431 : "construction. This is a code problem.");
432 141 : }
433 :
434 : const neml2::Tensor &
435 90 : NEML2ModelExecutor::getOutput(const neml2::VariableName & output_name) const
436 : {
437 90 : checkExecutionStage();
438 :
439 90 : if (!model().output_axis().has_variable(output_name))
440 0 : mooseError("Trying to retrieve a non-existent NEML2 output variable '", output_name, "'.");
441 :
442 90 : return _retrieved_outputs[output_name];
443 : }
444 :
445 : const neml2::Tensor &
446 51 : NEML2ModelExecutor::getOutputDerivative(const neml2::VariableName & output_name,
447 : const neml2::VariableName & input_name) const
448 : {
449 51 : checkExecutionStage();
450 :
451 51 : if (!model().output_axis().has_variable(output_name))
452 0 : mooseError("Trying to retrieve the derivative of NEML2 output variable '",
453 : output_name,
454 : "' with respect to NEML2 input variable '",
455 : input_name,
456 : "', but the NEML2 output variable does not exist.");
457 :
458 51 : if (!model().input_axis().has_variable(input_name))
459 0 : mooseError("Trying to retrieve the derivative of NEML2 output variable '",
460 : output_name,
461 : "' with respect to NEML2 input variable '",
462 : input_name,
463 : "', but the NEML2 input variable does not exist.");
464 :
465 51 : return _retrieved_derivatives[output_name][input_name];
466 : }
467 :
468 : const neml2::Tensor &
469 0 : NEML2ModelExecutor::getOutputParameterDerivative(const neml2::VariableName & output_name,
470 : const std::string & parameter_name) const
471 : {
472 0 : checkExecutionStage();
473 :
474 0 : if (!model().output_axis().has_variable(output_name))
475 0 : mooseError("Trying to retrieve the derivative of NEML2 output variable '",
476 : output_name,
477 : "' with respect to NEML2 model parameter '",
478 : parameter_name,
479 : "', but the NEML2 output variable does not exist.");
480 :
481 0 : if (model().named_parameters().count(parameter_name) != 1)
482 0 : mooseError("Trying to retrieve the derivative of NEML2 output variable '",
483 : output_name,
484 : "' with respect to NEML2 model parameter '",
485 : parameter_name,
486 : "', but the NEML2 model parameter does not exist.");
487 :
488 0 : return _retrieved_parameter_derivatives[output_name][parameter_name];
489 : }
490 :
491 : #endif
|
