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 "MaternHalfIntCovariance.h"
12 : #include "LibtorchUtils.h"
13 : #include <cmath>
14 :
15 : registerMooseObject("StochasticToolsApp", MaternHalfIntCovariance);
16 :
17 : InputParameters
18 48 : MaternHalfIntCovariance::validParams()
19 : {
20 48 : InputParameters params = CovarianceFunctionBase::validParams();
21 48 : params.addClassDescription("Matern half-integer covariance function.");
22 96 : params.addRequiredParam<std::vector<Real>>("length_factor",
23 : "Length factors to use for Covariance Kernel");
24 96 : params.addRequiredParam<Real>("signal_variance",
25 : "Signal Variance ($\\sigma_f^2$) to use for kernel calculation.");
26 96 : params.addParam<Real>(
27 96 : "noise_variance", 0.0, "Noise Variance ($\\sigma_n^2$) to use for kernel calculation.");
28 96 : params.addRequiredParam<unsigned int>(
29 : "p", "Integer p to use for Matern Half Integer Covariance Kernel");
30 48 : return params;
31 0 : }
32 :
33 24 : MaternHalfIntCovariance::MaternHalfIntCovariance(const InputParameters & parameters)
34 : : CovarianceFunctionBase(parameters),
35 24 : _length_factor(addVectorRealHyperParameter(
36 : "length_factor", getParam<std::vector<Real>>("length_factor"), true)),
37 24 : _sigma_f_squared(
38 72 : addRealHyperParameter("signal_variance", getParam<Real>("signal_variance"), true)),
39 24 : _sigma_n_squared(
40 72 : addRealHyperParameter("noise_variance", getParam<Real>("noise_variance"), true)),
41 72 : _p(addRealHyperParameter("p", getParam<unsigned int>("p"), false))
42 : {
43 24 : }
44 :
45 : void
46 19524 : MaternHalfIntCovariance::computeCovarianceMatrix(torch::Tensor & K,
47 : const torch::Tensor & x,
48 : const torch::Tensor & xp,
49 : const bool is_self_covariance) const
50 : {
51 19524 : if ((unsigned)x.sizes()[1] != _length_factor.numel())
52 0 : mooseError("length_factor size does not match dimension of trainer input.");
53 :
54 19524 : maternHalfIntFunction(
55 : K, x, xp, _length_factor, _sigma_f_squared, _sigma_n_squared, _p, is_self_covariance);
56 19524 : }
57 :
58 : void
59 35524 : MaternHalfIntCovariance::maternHalfIntFunction(torch::Tensor & K,
60 : const torch::Tensor & x,
61 : const torch::Tensor & xp,
62 : const torch::Tensor & length_factor,
63 : const torch::Tensor & sigma_f_squared,
64 : const torch::Tensor & sigma_n_squared,
65 : const torch::Tensor & p,
66 : const bool is_self_covariance)
67 : {
68 71048 : const auto p_value = cast_int<unsigned int>(LibtorchUtils::toCPUContiguous(p).item<Real>());
69 : mooseAssert(x.sizes()[1] == xp.sizes()[1],
70 : "Number of parameters do not match in covariance kernel calculation");
71 :
72 : const auto l_factor = length_factor.unsqueeze(0);
73 71048 : K = torch::cdist(torch::div(x, l_factor), torch::div(xp, l_factor), 2.0);
74 35524 : const Real factor = std::sqrt(2 * p_value + 1);
75 35524 : const Real normalization = std::tgamma(p_value + 1) / std::tgamma(2 * p_value + 1);
76 :
77 35524 : auto summation = torch::zeros_like(K);
78 142096 : for (const auto tt : make_range(p_value + 1))
79 : {
80 : const Real coefficient =
81 106572 : std::tgamma(p_value + tt + 1) / (std::tgamma(tt + 1) * std::tgamma(p_value - tt + 1));
82 426288 : summation = summation + coefficient * torch::pow(2.0 * factor * K, Real(p_value - tt));
83 : }
84 :
85 71048 : K = sigma_f_squared * torch::exp(-factor * K) * normalization * summation;
86 35524 : if (is_self_covariance)
87 35548 : K.diagonal().add_(sigma_n_squared);
88 35524 : }
89 :
90 : bool
91 48000 : MaternHalfIntCovariance::computedKdhyper(torch::Tensor & dKdhp,
92 : const torch::Tensor & x,
93 : const std::string & hyper_param_name,
94 : unsigned int ind) const
95 : {
96 48000 : if (name().length() + 1 > hyper_param_name.length())
97 : return false;
98 :
99 48000 : const std::string name_without_prefix = hyper_param_name.substr(name().length() + 1);
100 :
101 48000 : if (name_without_prefix == "noise_variance")
102 : {
103 0 : const auto options = x.options().dtype(at::kDouble);
104 0 : maternHalfIntFunction(dKdhp,
105 : x,
106 : x,
107 : _length_factor,
108 0 : torch::tensor(0.0, options),
109 0 : torch::tensor(1.0, options),
110 : _p,
111 : true);
112 : return true;
113 : }
114 :
115 48000 : if (name_without_prefix == "signal_variance")
116 : {
117 16000 : const auto options = x.options().dtype(at::kDouble);
118 32000 : maternHalfIntFunction(dKdhp,
119 : x,
120 : x,
121 : _length_factor,
122 32000 : torch::tensor(1.0, options),
123 32000 : torch::tensor(0.0, options),
124 : _p,
125 : false);
126 : return true;
127 : }
128 :
129 32000 : if (name_without_prefix == "length_factor")
130 : {
131 32000 : computedKdlf(dKdhp, x, _length_factor, _sigma_f_squared, _p, ind);
132 : return true;
133 : }
134 :
135 : return false;
136 : }
137 :
138 : void
139 32000 : MaternHalfIntCovariance::computedKdlf(torch::Tensor & K,
140 : const torch::Tensor & x,
141 : const torch::Tensor & length_factor,
142 : const torch::Tensor & sigma_f_squared,
143 : const torch::Tensor & p,
144 : const int ind)
145 : {
146 64000 : const auto p_value = cast_int<unsigned int>(LibtorchUtils::toCPUContiguous(p).item<Real>());
147 :
148 : mooseAssert(ind < x.sizes()[1], "Incorrect length factor index");
149 :
150 : const auto l_factor = length_factor.unsqueeze(0);
151 96000 : const auto scaled_distance = torch::cdist(torch::div(x, l_factor), torch::div(x, l_factor), 2.0);
152 32000 : const auto nonzero_distance = scaled_distance > 0;
153 : const auto safe_scaled_distance =
154 32000 : torch::where(nonzero_distance, scaled_distance, torch::ones_like(scaled_distance));
155 32000 : const Real factor = std::sqrt(2 * p_value + 1);
156 32000 : const Real normalization = std::tgamma(p_value + 1) / std::tgamma(2 * p_value + 1);
157 :
158 32000 : auto summation = torch::zeros_like(safe_scaled_distance);
159 128000 : for (const auto tt : make_range(p_value + 1))
160 : {
161 : const Real coefficient =
162 96000 : std::tgamma(p_value + tt + 1) / (std::tgamma(tt + 1) * std::tgamma(p_value - tt + 1));
163 96000 : summation = summation +
164 288000 : coefficient * torch::pow(2.0 * factor * safe_scaled_distance, Real(p_value - tt));
165 : }
166 :
167 32000 : auto summation_derivative = torch::zeros_like(safe_scaled_distance);
168 96000 : for (const auto tt : make_range(p_value))
169 : {
170 : const Real coefficient =
171 64000 : std::tgamma(p_value + tt + 1) / (std::tgamma(tt + 1) * std::tgamma(p_value - tt + 1));
172 : summation_derivative =
173 64000 : summation_derivative +
174 128000 : coefficient * 2.0 * factor * (p_value - tt) *
175 128000 : torch::pow(2.0 * factor * safe_scaled_distance, Real(p_value - tt - 1));
176 : }
177 :
178 32000 : const auto coordinate = x.select(1, ind);
179 : const auto coordinate_distance_squared =
180 96000 : torch::pow(coordinate.unsqueeze(1) - coordinate.unsqueeze(0), 2);
181 32000 : const auto length_factor_ind = length_factor.select(0, ind);
182 : const auto dr_dlength_factor =
183 64000 : -coordinate_distance_squared / (torch::pow(length_factor_ind, 3) * safe_scaled_distance);
184 32000 : const auto dK_dlength_factor = sigma_f_squared * normalization *
185 64000 : torch::exp(-factor * safe_scaled_distance) *
186 64000 : (summation_derivative - factor * summation) * dr_dlength_factor;
187 :
188 64000 : K = torch::where(nonzero_distance, dK_dlength_factor, torch::zeros_like(dK_dlength_factor));
189 32000 : }
190 :
191 : #endif
|
