Line data Source code
1 : // Copyright 2024, UChicago Argonne, LLC
2 : // All Rights Reserved
3 : // Software Name: NEML2 -- the New Engineering material Model Library, version 2
4 : // By: Argonne National Laboratory
5 : // OPEN SOURCE LICENSE (MIT)
6 : //
7 : // Permission is hereby granted, free of charge, to any person obtaining a copy
8 : // of this software and associated documentation files (the "Software"), to deal
9 : // in the Software without restriction, including without limitation the rights
10 : // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
11 : // copies of the Software, and to permit persons to whom the Software is
12 : // furnished to do so, subject to the following conditions:
13 : //
14 : // The above copyright notice and this permission notice shall be included in
15 : // all copies or substantial portions of the Software.
16 : //
17 : // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
18 : // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
19 : // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
20 : // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
21 : // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
22 : // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
23 : // THE SOFTWARE.
24 :
25 : #include "neml2/models/ImplicitUpdate.h"
26 : #include "neml2/models/Assembler.h"
27 : #include "neml2/solvers/NonlinearSolver.h"
28 : #include "neml2/tensors/functions/linalg/lu_factor.h"
29 : #include "neml2/tensors/functions/linalg/lu_solve.h"
30 : #include "neml2/base/guards.h"
31 : #include "neml2/misc/assertions.h"
32 :
33 : namespace neml2
34 : {
35 : register_NEML2_object(ImplicitUpdate);
36 :
37 : OptionSet
38 2 : ImplicitUpdate::expected_options()
39 : {
40 2 : OptionSet options = Model::expected_options();
41 2 : options.doc() =
42 2 : "Update an implicit model by solving the underlying implicit system of equations.";
43 :
44 4 : options.set<std::string>("implicit_model");
45 6 : options.set("implicit_model").doc() =
46 2 : "The implicit model defining the implicit system of equations to be solved";
47 :
48 4 : options.set<std::string>("solver");
49 4 : options.set("solver").doc() = "Solver used to solve the implicit system";
50 :
51 : // No jitting :/
52 4 : options.set<bool>("jit") = false;
53 2 : options.set("jit").suppressed() = true;
54 :
55 2 : return options;
56 0 : }
57 :
58 9 : ImplicitUpdate::ImplicitUpdate(const OptionSet & options)
59 : : Model(options),
60 9 : _model(register_model("implicit_model", /*nonlinear=*/true)),
61 27 : _solver(get_solver<NonlinearSolver>("solver"))
62 : {
63 9 : neml_assert(_model.output_axis().has_residual(),
64 : "The implicit model'",
65 9 : _model.name(),
66 : "' registered in '",
67 9 : name(),
68 : "' does not have the residual output axis.");
69 : // Take care of dependency registration:
70 : // 1. Input variables of the "implicit_model" should be *consumed* by *this* model. This has
71 : // already been taken care of by the `register_model` call.
72 : // 2. Output variables of the "implicit_model" on the "residual" subaxis should be *provided* by
73 : // *this* model.
74 29 : for (auto && [name, var] : _model.output_variables())
75 20 : clone_output_variable(*var, name.remount(STATE));
76 9 : }
77 :
78 : void
79 5 : ImplicitUpdate::diagnose() const
80 : {
81 5 : Model::diagnose();
82 5 : diagnostic_assert(_model.output_axis().nsubaxis() == 1,
83 : "The implicit model's output contains non-residual subaxis:\n",
84 5 : _model.output_axis());
85 5 : diagnostic_assert(_model.input_axis().has_state(),
86 : "The implicit model's input does not have a state subaxis:\n",
87 5 : _model.input_axis());
88 5 : diagnostic_assert(!_model.input_axis().has_residual(),
89 : "The implicit model's input cannot have a residual subaxis:\n",
90 5 : _model.input_axis());
91 15 : diagnostic_assert(
92 5 : _model.input_axis().subaxis(STATE) == _model.output_axis().subaxis(RESIDUAL),
93 : "The implicit model should have conformal trial state and residual. The input state "
94 : "subaxis is\n",
95 5 : _model.input_axis().subaxis(STATE),
96 : "\nThe output residual subaxis is\n",
97 5 : _model.output_axis().subaxis(RESIDUAL));
98 5 : }
99 :
100 : void
101 9 : ImplicitUpdate::link_output_variables()
102 : {
103 9 : Model::link_output_variables();
104 29 : for (auto && [name, var] : output_variables())
105 20 : var->ref(input_variable(name), /*ref_is_mutable=*/true);
106 9 : }
107 :
108 : void
109 189 : ImplicitUpdate::set_value(bool out, bool dout_din, bool /*d2out_din2*/)
110 : {
111 : // The trial state is used as the initial guess
112 189 : const auto sol_assember = VectorAssembler(_model.input_axis().subaxis(STATE));
113 189 : auto x0 = NonlinearSystem::Sol<false>(sol_assember.assemble_by_variable(_model.collect_input()));
114 :
115 : // Perform automatic scaling (using the trial state)
116 : // TODO: Add an interface to allow user to specify where (and when) to evaluate the Jacobian for
117 : // automatic scaling.
118 189 : _model.init_scaling(x0, _solver->verbose);
119 :
120 : // Solve for the next state
121 189 : NonlinearSolver::Result res;
122 : {
123 189 : SolvingNonlinearSystem solving;
124 189 : res = _solver->solve(_model, x0);
125 189 : neml_assert(res.ret == NonlinearSolver::RetCode::SUCCESS, "Nonlinear solve failed.");
126 189 : }
127 :
128 189 : if (out)
129 : {
130 : // You may be tempted to assign the solution, i.e., res.solution, to the output variables. But
131 : // we don't have to. Think about it: The output variables share the same name as those input
132 : // variables on the state subaxis, and since we don't duplicate storage for variables with the
133 : // same name, they are essentially the same variable with FType::INPUT | FType::OUTPUT. During
134 : // the nonlinear solve, we have to iteratively update the guess (i.e., the input variables on
135 : // the state subaxis) until convergece. Once the nonlinear system has converged, the input
136 : // variables on the state subaxis _must_ contain the solution. Therefore, the output variables
137 : // _must_ also contain the solution upon convergence.
138 :
139 : // All that being said, if the result has AD graph, we need to propagate the graph to the output
140 187 : if (res.solution.requires_grad())
141 26 : assign_output(sol_assember.split_by_variable(res.solution));
142 : }
143 :
144 : // Use the implicit function theorem (IFT) to calculate the other derivatives
145 189 : if (dout_din)
146 : {
147 : // IFT requires the Jacobian evaluated at the solution:
148 3 : _model.forward_maybe_jit(false, true, false);
149 3 : const auto jac_assembler = MatrixAssembler(_model.output_axis(), _model.input_axis());
150 3 : const auto J = jac_assembler.assemble_by_variable(_model.collect_output_derivatives());
151 3 : const auto derivs = jac_assembler.split_by_subaxis(J).at(RESIDUAL);
152 3 : const auto dr_ds = derivs.at(STATE);
153 :
154 : // Factorize the Jacobian once and for all
155 3 : const auto [LU, pivot] = linalg::lu_factor(dr_ds);
156 :
157 : // The actual IFT:
158 15 : for (const auto & [subaxis, deriv] : derivs)
159 : {
160 12 : if (subaxis == STATE)
161 3 : continue;
162 : const auto ift_assembler =
163 9 : MatrixAssembler(output_axis(), _model.input_axis().subaxis(subaxis));
164 9 : assign_output_derivatives(
165 18 : ift_assembler.split_by_variable(-linalg::lu_solve(LU, pivot, deriv)));
166 : }
167 3 : }
168 189 : }
169 : } // namespace neml2
|
