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