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 : #ifdef MOOSE_LIBTORCH_ENABLED
10 :
11 : #include "LMC.h"
12 : #include "MooseRandom.h"
13 : #include "MathUtils.h"
14 :
15 : registerMooseObject("StochasticToolsApp", LMC);
16 :
17 : InputParameters
18 48 : LMC::validParams()
19 : {
20 48 : InputParameters params = CovarianceFunctionBase::validParams();
21 48 : params.addClassDescription("Covariance function for multioutput Gaussian Processes based on the "
22 : "Linear Model of Coregionalization (LMC).");
23 96 : params.addParam<unsigned int>(
24 96 : "num_latent_funcs", 1., "The number of latent functions for the expansion of the outputs.");
25 48 : params.makeParamRequired<unsigned int>("num_outputs");
26 48 : params.makeParamRequired<std::vector<UserObjectName>>("covariance_functions");
27 48 : return params;
28 0 : }
29 :
30 24 : LMC::LMC(const InputParameters & parameters)
31 : : CovarianceFunctionBase(parameters),
32 48 : _num_expansion_terms(getParam<unsigned int>("num_latent_funcs"))
33 : {
34 : // We use a random number generator to obtain the initial guess for the
35 : // hyperparams
36 : MooseRandom generator_latent;
37 : generator_latent.seed(0, 1980);
38 :
39 : // First add and initialize the a A coefficients in the (aa^T+lambda*I) matrix
40 48 : for (const auto exp_i : make_range(_num_expansion_terms))
41 : {
42 24 : const std::string a_coeff_name = "acoeff_" + std::to_string(exp_i);
43 24 : std::vector<Real> acoeff_values(_num_outputs);
44 72 : for (const auto out_i : make_range(_num_outputs))
45 48 : acoeff_values[out_i] = 3.0 * generator_latent.rand(0) + 1.0;
46 24 : auto & acoeff_vector = addVectorRealHyperParameter(a_coeff_name, acoeff_values, true);
47 24 : _a_coeffs.push_back(&acoeff_vector);
48 24 : }
49 :
50 : // Then add and initialize the lambda coefficients in the (aa^T+lambda*I) matrix
51 48 : for (const auto exp_i : make_range(_num_expansion_terms))
52 : {
53 24 : const std::string lambda_name = "lambda_" + std::to_string(exp_i);
54 24 : std::vector<Real> lambda_values(_num_outputs);
55 72 : for (const auto out_i : make_range(_num_outputs))
56 48 : lambda_values[out_i] = 3.0 * generator_latent.rand(0) + 1.0;
57 24 : auto & lambda_vector = addVectorRealHyperParameter(lambda_name, lambda_values, true);
58 24 : _lambdas.push_back(&lambda_vector);
59 24 : }
60 24 : }
61 :
62 : void
63 8166 : LMC::computeCovarianceMatrix(torch::Tensor & K,
64 : const torch::Tensor & x,
65 : const torch::Tensor & xp,
66 : const bool is_self_covariance) const
67 : {
68 8166 : const auto options = x.options().dtype(at::kDouble);
69 : // Create temporary vectors for constructing the covariance matrix
70 8166 : torch::Tensor K_params = torch::zeros({x.sizes()[0], xp.sizes()[0]}, options);
71 8166 : torch::Tensor B = torch::zeros({_num_outputs, _num_outputs}, options);
72 16332 : K = torch::zeros({x.sizes()[0] * _num_outputs, xp.sizes()[0] * _num_outputs}, options);
73 : torch::Tensor K_working;
74 :
75 : // For every expansion term we add the contribution to the covariance matrix
76 16332 : for (const auto exp_i : make_range(_num_expansion_terms))
77 : {
78 8166 : _covariance_functions[exp_i]->computeCovarianceMatrix(K_params, x, xp, is_self_covariance);
79 8166 : computeBMatrix(B, exp_i);
80 0 : K_working = torch::kron(B, K_params);
81 8166 : K += K_working;
82 : }
83 8166 : }
84 :
85 : bool
86 56000 : LMC::computedKdhyper(torch::Tensor & dKdhp,
87 : const torch::Tensor & x,
88 : const std::string & hyper_param_name,
89 : unsigned int ind) const
90 : {
91 : // Early return in the paramter name is longer than the expected [name] prefix.
92 : // We prefix the parameter names with the name of the covariance function.
93 56000 : if (name().length() + 1 > hyper_param_name.length())
94 : return false;
95 :
96 : // Strip the prefix from the given parameter name
97 56000 : const std::string name_without_prefix = hyper_param_name.substr(name().length() + 1);
98 :
99 : // Check if the parameter is tunable
100 56000 : if (_tunable_hp.find(hyper_param_name) != _tunable_hp.end())
101 : {
102 32000 : const std::string acoeff_prefix = "acoeff_";
103 32000 : const std::string lambda_prefix = "lambda_";
104 :
105 : // Allocate storage for the factors of the total gradient matrix
106 32000 : const auto options = x.options().dtype(at::kDouble);
107 32000 : torch::Tensor dBdhp = torch::zeros({_num_outputs, _num_outputs}, options);
108 32000 : torch::Tensor K_params = torch::zeros({x.sizes()[0], x.sizes()[0]}, options);
109 :
110 32000 : if (name_without_prefix.find(acoeff_prefix) != std::string::npos)
111 : {
112 : // Automatically grab the expansion index
113 16000 : const int number = std::stoi(name_without_prefix.substr(acoeff_prefix.length()));
114 16000 : computeAGradient(dBdhp, number, ind);
115 16000 : _covariance_functions[number]->computeCovarianceMatrix(K_params, x, x, true);
116 : }
117 16000 : else if (name_without_prefix.find(lambda_prefix) != std::string::npos)
118 : {
119 : // Automatically grab the expansion index
120 16000 : const int number = std::stoi(name_without_prefix.substr(lambda_prefix.length()));
121 16000 : computeLambdaGradient(dBdhp, number, ind);
122 16000 : _covariance_functions[number]->computeCovarianceMatrix(K_params, x, x, true);
123 : }
124 32000 : dKdhp = torch::kron(dBdhp, K_params);
125 : return true;
126 : }
127 : else
128 : {
129 : // Allocate storage for the matrix factors
130 24000 : const auto options = x.options().dtype(at::kDouble);
131 24000 : torch::Tensor B_tmp = torch::zeros({_num_outputs, _num_outputs}, options);
132 24000 : torch::Tensor B = torch::zeros({_num_outputs, _num_outputs}, options);
133 24000 : torch::Tensor dKdhp_sub = torch::zeros({x.sizes()[0], x.sizes()[0]}, options);
134 :
135 : // First, check the dependent covariances
136 : bool found = false;
137 48000 : for (const auto dependent_covar : _covariance_functions)
138 24000 : if (!found)
139 24000 : found = dependent_covar->computedKdhyper(dKdhp_sub, x, hyper_param_name, ind);
140 :
141 24000 : if (!found)
142 0 : mooseError("Hyperparameter ", hyper_param_name, "not found!");
143 :
144 : // Then we compute the output covariance
145 48000 : for (const auto exp_i : make_range(_num_expansion_terms))
146 : {
147 24000 : computeBMatrix(B_tmp, exp_i);
148 24000 : B += B_tmp;
149 : }
150 :
151 24000 : dKdhp = torch::kron(B, dKdhp_sub);
152 :
153 : return true;
154 : }
155 :
156 : return false;
157 : }
158 :
159 : void
160 32166 : LMC::computeBMatrix(torch::Tensor & Bmat, const unsigned int exp_i) const
161 : {
162 32166 : const auto & a_coeffs = *_a_coeffs[exp_i];
163 32166 : const auto & lambda_coeffs = *_lambdas[exp_i];
164 64332 : Bmat = torch::outer(a_coeffs, a_coeffs) + torch::diag(lambda_coeffs);
165 32166 : }
166 :
167 : void
168 16000 : LMC::computeAGradient(torch::Tensor & grad,
169 : const unsigned int exp_i,
170 : const unsigned int index) const
171 : {
172 16000 : const auto & a_coeffs = *_a_coeffs[exp_i];
173 : mooseAssert(cast_int<int64_t>(index) < a_coeffs.numel(), "Incorrect LMC coefficient index.");
174 16000 : auto basis = torch::zeros_like(a_coeffs);
175 : const auto index_tensor =
176 16000 : torch::tensor({cast_int<int64_t>(index)},
177 32000 : torch::TensorOptions().dtype(torch::kLong).device(a_coeffs.device()));
178 16000 : basis.index_fill_(0, index_tensor, 1.0);
179 48000 : grad = torch::outer(basis, a_coeffs) + torch::outer(a_coeffs, basis);
180 16000 : }
181 :
182 : void
183 16000 : LMC::computeLambdaGradient(torch::Tensor & grad,
184 : const unsigned int exp_i,
185 : const unsigned int index) const
186 : {
187 : mooseAssert(index < _num_outputs, "Incorrect LMC lambda index.");
188 16000 : auto basis = torch::zeros_like(*_lambdas[exp_i]);
189 : const auto index_tensor =
190 16000 : torch::tensor({cast_int<int64_t>(index)},
191 32000 : torch::TensorOptions().dtype(torch::kLong).device(_lambdas[exp_i]->device()));
192 16000 : basis.index_fill_(0, index_tensor, 1.0);
193 16000 : grad = torch::diag(basis);
194 16000 : }
195 :
196 : #endif
|
