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/LinearCombination.h"
26 : #include "neml2/tensors/Scalar.h"
27 : #include "neml2/tensors/Vec.h"
28 : #include "neml2/tensors/R2.h"
29 : #include "neml2/tensors/SR2.h"
30 : #include "neml2/tensors/SSR4.h"
31 : #include "neml2/misc/assertions.h"
32 :
33 : namespace neml2
34 : {
35 : template <typename T>
36 : OptionSet
37 9 : LinearCombination<T>::expected_options()
38 : {
39 : // This is the only way of getting tensor type in a static method like this...
40 : // Trim 6 chars to remove 'neml2::'
41 9 : auto tensor_type = utils::demangle(typeid(T).name()).substr(7);
42 :
43 9 : OptionSet options = Model::expected_options();
44 9 : options.doc() =
45 : "Calculate linear combination of multiple " + tensor_type +
46 : " tensors as \\f$ u = c_i v_i + s \\f$ (Einstein summation assumed), where \\f$ c_i "
47 : "\\f$ are the coefficients, and \\f$ v_i \\f$ are the variables to be summed. \\f$ s \\f$ is "
48 : "a constant offset.";
49 :
50 18 : options.set<bool>("define_second_derivatives") = true;
51 :
52 9 : options.set<std::vector<VariableName>>("from_var");
53 27 : options.set("from_var").doc() = tensor_type + " tensors to be summed";
54 :
55 18 : options.set_output("to_var");
56 9 : options.set("to_var").doc() = "The sum";
57 :
58 27 : options.set_parameter<std::vector<TensorName<Scalar>>>("coefficients") = {
59 : TensorName<Scalar>("1")};
60 18 : options.set("coefficients").doc() =
61 : "Weights associated with each variable. This option takes a list of weights, one for each "
62 : "coefficient. When the length of this list is 1, the same weight applies to all "
63 : "coefficients.";
64 :
65 36 : options.set_parameter<TensorName<Scalar>>("constant_coefficient") = {TensorName<Scalar>("0")};
66 18 : options.set("constant_coefficient").doc() =
67 : "The constant coefficient added to the final summation";
68 :
69 18 : options.set<bool>("constant_coefficient_as_parameter") = false;
70 9 : options.set("constant_coefficient_as_parameter").doc() =
71 : "By default, the constant_coefficient are declared as buffers. Set this option to true to "
72 : "declare "
73 : "them as (trainable) parameters.";
74 :
75 27 : options.set<std::vector<bool>>("coefficient_as_parameter") = {false};
76 9 : options.set("coefficient_as_parameter").doc() =
77 : "By default, the coefficients are declared as buffers. Set this option to true to declare "
78 : "them as (trainable) parameters. This option takes a list of booleans, one for each "
79 : "coefficient. When the length of this list is 1, the boolean applies to all coefficients.";
80 :
81 18 : return options;
82 45 : }
83 :
84 : template <typename T>
85 25 : LinearCombination<T>::LinearCombination(const OptionSet & options)
86 : : Model(options),
87 25 : _to(declare_output_variable<T>("to_var"))
88 : {
89 101 : for (const auto & fv : options.get<std::vector<VariableName>>("from_var"))
90 51 : _from.push_back(&declare_input_variable<T>(fv));
91 :
92 25 : auto coef_as_param = options.get<std::vector<bool>>("coefficient_as_parameter");
93 25 : neml_assert(coef_as_param.size() == 1 || coef_as_param.size() == _from.size(),
94 : "Expected 1 or ",
95 25 : _from.size(),
96 : " entries in coefficient_as_parameter, got ",
97 25 : coef_as_param.size(),
98 : ".");
99 :
100 : // Expand the list of booleans to match the number of coefficients
101 25 : if (coef_as_param.size() == 1)
102 19 : coef_as_param = std::vector<bool>(_from.size(), coef_as_param[0]);
103 :
104 25 : const auto coef_refs = options.get<std::vector<TensorName<Scalar>>>("coefficients");
105 25 : neml_assert(coef_refs.size() == 1 || coef_refs.size() == _from.size(),
106 : "Expected 1 or ",
107 25 : _from.size(),
108 : " coefficients, got ",
109 25 : coef_refs.size(),
110 : ".");
111 :
112 : // Declare parameters or buffers
113 25 : _coefs.resize(_from.size());
114 76 : for (std::size_t i = 0; i < _from.size(); i++)
115 : {
116 51 : const auto & coef_ref = coef_refs.size() == 1 ? coef_refs[0] : coef_refs[i];
117 51 : if (coef_as_param[i])
118 24 : _coefs[i] =
119 12 : &declare_parameter<Scalar>("c_" + std::to_string(i), coef_ref, /*allow_nonlinear=*/true);
120 : else
121 39 : _coefs[i] = &declare_buffer<Scalar>("c_" + std::to_string(i), coef_ref);
122 : }
123 :
124 50 : if (options.user_specified("constant_coefficient"))
125 : {
126 4 : auto s_as_param = options.get<bool>("constant_coefficient_as_parameter");
127 4 : if (s_as_param)
128 10 : _s = &declare_parameter<Scalar>("s", "constant_coefficient", /*allow_nonlinear=*/true);
129 : else
130 10 : _s = &declare_buffer<Scalar>("s", "constant_coefficient");
131 : }
132 25 : }
133 :
134 : template <typename T>
135 : void
136 44 : LinearCombination<T>::set_value(bool out, bool dout_din, bool d2out_din2)
137 : {
138 44 : if (out)
139 : {
140 41 : auto value = _s ? (*_s) + (*_coefs[0]) * (*_from[0]) : (*_coefs[0]) * (*_from[0]);
141 82 : for (std::size_t i = 1; i < _from.size(); i++)
142 41 : value = value + (*_coefs[i]) * (*_from[i]);
143 41 : _to = value;
144 41 : }
145 :
146 44 : if (dout_din)
147 : {
148 25 : const auto I = T::identity_map(_from[0]->options());
149 75 : for (std::size_t i = 0; i < _from.size(); i++)
150 : {
151 50 : if (_from[i]->is_dependent())
152 48 : _to.d(*_from[i]) = (*_coefs[i]) * I;
153 :
154 50 : if (const auto * const pi = nl_param("c_" + std::to_string(i)))
155 7 : _to.d(*pi) = (*_from[i]);
156 : }
157 75 : if (const auto * const s = nl_param("s"))
158 2 : _to.d(*s) = neml2::Scalar::full(1.0, _from[0]->options());
159 25 : }
160 :
161 : if (d2out_din2)
162 : {
163 : // zero
164 : }
165 44 : }
166 :
167 : #define REGISTER(T) \
168 : using T##LinearCombination = LinearCombination<T>; \
169 : register_NEML2_object(T##LinearCombination); \
170 : template class LinearCombination<T>
171 : REGISTER(Scalar);
172 : REGISTER(Vec);
173 : REGISTER(SR2);
174 : REGISTER(R2);
175 : } // namespace neml2
|
