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 "ExponentialCovariance.h"
12 : #include <cmath>
13 :
14 : registerMooseObject("StochasticToolsApp", ExponentialCovariance);
15 :
16 : InputParameters
17 32 : ExponentialCovariance::validParams()
18 : {
19 32 : InputParameters params = CovarianceFunctionBase::validParams();
20 32 : params.addClassDescription("Exponential covariance function.");
21 64 : params.addRequiredParam<std::vector<Real>>("length_factor",
22 : "Length factors to use for Covariance Kernel");
23 64 : params.addRequiredParam<Real>("signal_variance",
24 : "Signal Variance ($\\sigma_f^2$) to use for kernel calculation.");
25 64 : params.addParam<Real>(
26 64 : "noise_variance", 0.0, "Noise Variance ($\\sigma_n^2$) to use for kernel calculation.");
27 64 : params.addRequiredParam<Real>("gamma", "Gamma to use for Exponential Covariance Kernel");
28 32 : return params;
29 0 : }
30 :
31 16 : ExponentialCovariance::ExponentialCovariance(const InputParameters & parameters)
32 : : CovarianceFunctionBase(parameters),
33 16 : _length_factor(addVectorRealHyperParameter(
34 : "length_factor", getParam<std::vector<Real>>("length_factor"), true)),
35 16 : _sigma_f_squared(
36 48 : addRealHyperParameter("signal_variance", getParam<Real>("signal_variance"), true)),
37 16 : _sigma_n_squared(
38 48 : addRealHyperParameter("noise_variance", getParam<Real>("noise_variance"), true)),
39 48 : _gamma(addRealHyperParameter("gamma", getParam<Real>("gamma"), false))
40 : {
41 16 : }
42 :
43 : void
44 10316 : ExponentialCovariance::computeCovarianceMatrix(torch::Tensor & K,
45 : const torch::Tensor & x,
46 : const torch::Tensor & xp,
47 : const bool is_self_covariance) const
48 : {
49 10316 : if ((unsigned)x.sizes()[1] != _length_factor.numel())
50 0 : mooseError("length_factor size does not match dimension of trainer input.");
51 :
52 10316 : ExponentialFunction(
53 : K, x, xp, _length_factor, _sigma_f_squared, _sigma_n_squared, _gamma, is_self_covariance);
54 10316 : }
55 :
56 : void
57 18316 : ExponentialCovariance::ExponentialFunction(torch::Tensor & K,
58 : const torch::Tensor & x,
59 : const torch::Tensor & xp,
60 : const torch::Tensor & length_factor,
61 : const torch::Tensor & sigma_f_squared,
62 : const torch::Tensor & sigma_n_squared,
63 : const torch::Tensor & gamma,
64 : const bool is_self_covariance)
65 : {
66 : mooseAssert(x.sizes()[1] == xp.sizes()[1],
67 : "Number of parameters do not match in covariance kernel calculation");
68 :
69 : const auto l_factor = length_factor.unsqueeze(0);
70 54948 : K = torch::cdist(torch::div(x, l_factor), torch::div(xp, l_factor), 2.0);
71 18316 : K = sigma_f_squared * torch::exp(-torch::pow(K, gamma));
72 18316 : if (is_self_covariance)
73 18332 : K.diagonal().add_(sigma_n_squared);
74 18316 : }
75 :
76 : bool
77 24000 : ExponentialCovariance::computedKdhyper(torch::Tensor & dKdhp,
78 : const torch::Tensor & x,
79 : const std::string & hyper_param_name,
80 : unsigned int ind) const
81 : {
82 24000 : if (name().length() + 1 > hyper_param_name.length())
83 : return false;
84 :
85 24000 : const std::string name_without_prefix = hyper_param_name.substr(name().length() + 1);
86 :
87 24000 : if (name_without_prefix == "noise_variance")
88 : {
89 0 : const auto options = x.options().dtype(at::kDouble);
90 0 : ExponentialFunction(dKdhp,
91 : x,
92 : x,
93 : _length_factor,
94 0 : torch::tensor(0.0, options),
95 0 : torch::tensor(1.0, options),
96 : _gamma,
97 : true);
98 : return true;
99 : }
100 :
101 24000 : if (name_without_prefix == "signal_variance")
102 : {
103 8000 : const auto options = x.options().dtype(at::kDouble);
104 16000 : ExponentialFunction(dKdhp,
105 : x,
106 : x,
107 : _length_factor,
108 16000 : torch::tensor(1.0, options),
109 16000 : torch::tensor(0.0, options),
110 : _gamma,
111 : false);
112 : return true;
113 : }
114 :
115 16000 : if (name_without_prefix == "length_factor")
116 : {
117 16000 : computedKdlf(dKdhp, x, _length_factor, _sigma_f_squared, _gamma, ind);
118 : return true;
119 : }
120 :
121 : return false;
122 : }
123 :
124 : void
125 16000 : ExponentialCovariance::computedKdlf(torch::Tensor & K,
126 : const torch::Tensor & x,
127 : const torch::Tensor & length_factor,
128 : const torch::Tensor & sigma_f_squared,
129 : const torch::Tensor & gamma,
130 : const int ind)
131 : {
132 : mooseAssert(ind < x.sizes()[1], "Incorrect length factor index");
133 :
134 : const auto l_factor = length_factor.unsqueeze(0);
135 48000 : const auto scaled_distance = torch::cdist(torch::div(x, l_factor), torch::div(x, l_factor), 2.0);
136 16000 : const auto nonzero_distance = scaled_distance > 0;
137 : const auto safe_scaled_distance =
138 16000 : torch::where(nonzero_distance, scaled_distance, torch::ones_like(scaled_distance));
139 16000 : const auto coordinate = x.select(1, ind);
140 : const auto coordinate_distance_squared =
141 48000 : torch::pow(coordinate.unsqueeze(1) - coordinate.unsqueeze(0), 2);
142 16000 : const auto length_factor_ind = length_factor.select(0, ind);
143 :
144 16000 : const auto dK_dlength_factor = coordinate_distance_squared / torch::pow(length_factor_ind, 3) *
145 32000 : gamma * torch::pow(safe_scaled_distance, gamma - 2.0) *
146 : sigma_f_squared *
147 16000 : torch::exp(-torch::pow(safe_scaled_distance, gamma));
148 :
149 32000 : K = torch::where(nonzero_distance, dK_dlength_factor, torch::zeros_like(dK_dlength_factor));
150 16000 : }
151 :
152 : #endif
|