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