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/misc/assertions.h"
31 :
32 : namespace neml2
33 : {
34 : register_NEML2_object(ImplicitUpdate);
35 :
36 : OptionSet
37 2 : ImplicitUpdate::expected_options()
38 : {
39 2 : OptionSet options = Model::expected_options();
40 2 : options.doc() =
41 2 : "Update an implicit model by solving the underlying implicit system of equations.";
42 :
43 4 : options.set<std::string>("implicit_model");
44 6 : options.set("implicit_model").doc() =
45 2 : "The implicit model defining the implicit system of equations to be solved";
46 :
47 4 : options.set<std::string>("solver");
48 4 : options.set("solver").doc() = "Solver used to solve the implicit system";
49 :
50 : // No jitting :/
51 4 : options.set<bool>("jit") = false;
52 2 : options.set("jit").suppressed() = true;
53 :
54 2 : return options;
55 0 : }
56 :
57 10 : ImplicitUpdate::ImplicitUpdate(const OptionSet & options)
58 : : Model(options),
59 10 : _model(register_model("implicit_model", /*nonlinear=*/true)),
60 30 : _solver(get_solver<NonlinearSolver>("solver"))
61 : {
62 10 : neml_assert(_model.output_axis().has_residual(),
63 : "The implicit model'",
64 10 : _model.name(),
65 : "' registered in '",
66 10 : name(),
67 : "' does not have the residual output axis.");
68 : // Take care of dependency registration:
69 : // 1. Input variables of the "implicit_model" should be *consumed* by *this* model. This has
70 : // already been taken care of by the `register_model` call.
71 : // 2. Output variables of the "implicit_model" on the "residual" subaxis should be *provided* by
72 : // *this* model.
73 33 : for (auto && [name, var] : _model.output_variables())
74 23 : clone_output_variable(*var, name.remount(STATE));
75 10 : }
76 :
77 : void
78 5 : ImplicitUpdate::diagnose() const
79 : {
80 5 : Model::diagnose();
81 5 : diagnostic_assert(_model.output_axis().nsubaxis() == 1,
82 : "The implicit model's output contains non-residual subaxis:\n",
83 5 : _model.output_axis());
84 5 : diagnostic_assert(_model.input_axis().has_state(),
85 : "The implicit model's input does not have a state subaxis:\n",
86 5 : _model.input_axis());
87 5 : diagnostic_assert(!_model.input_axis().has_residual(),
88 : "The implicit model's input cannot have a residual subaxis:\n",
89 5 : _model.input_axis());
90 15 : diagnostic_assert(
91 5 : _model.input_axis().subaxis(STATE) == _model.output_axis().subaxis(RESIDUAL),
92 : "The implicit model should have conformal trial state and residual. The input state "
93 : "subaxis is\n",
94 5 : _model.input_axis().subaxis(STATE),
95 : "\nThe output residual subaxis is\n",
96 5 : _model.output_axis().subaxis(RESIDUAL));
97 5 : }
98 :
99 : void
100 10 : ImplicitUpdate::link_output_variables()
101 : {
102 10 : Model::link_output_variables();
103 33 : for (auto && [name, var] : output_variables())
104 23 : var->ref(input_variable(name), /*ref_is_mutable=*/true);
105 10 : }
106 :
107 : void
108 191 : ImplicitUpdate::set_value(bool out, bool dout_din, bool /*d2out_din2*/)
109 : {
110 : // The trial state is used as the initial guess
111 191 : const auto sol_assember = VectorAssembler(_model.input_axis().subaxis(STATE));
112 191 : auto x0 = NonlinearSystem::Sol<false>(sol_assember.assemble_by_variable(_model.collect_input()));
113 :
114 : // Perform automatic scaling (using the trial state)
115 : // TODO: Add an interface to allow user to specify where (and when) to evaluate the Jacobian for
116 : // automatic scaling.
117 191 : _model.init_scaling(x0, _solver->verbose);
118 :
119 : // Solve for the next state
120 191 : NonlinearSolver::Result res;
121 : {
122 191 : SolvingNonlinearSystem solving;
123 191 : res = _solver->solve(_model, x0);
124 191 : _last_iterations = res.iterations;
125 191 : neml_assert(res.ret == NonlinearSolver::RetCode::SUCCESS, "Nonlinear solve failed.");
126 191 : }
127 :
128 191 : 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 189 : 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 191 : 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 191 : }
169 : } // namespace neml2
|
