1 #ifndef STAN_MATH_REV_FUN_VARIANCE_HPP
2 #define STAN_MATH_REV_FUN_VARIANCE_HPP
3
4 #include <stan/math/rev/meta.hpp>
5 #include <stan/math/rev/core.hpp>
6 #include <stan/math/rev/core/typedefs.hpp>
7 #include <stan/math/prim/err.hpp>
8 #include <stan/math/prim/fun/Eigen.hpp>
9 #include <stan/math/prim/fun/typedefs.hpp>
10 #include <vector>
11
12 namespace stan {
13 namespace math {
14 namespace internal {
15
calc_variance(size_t size,const var * dtrs)16 inline var calc_variance(size_t size, const var* dtrs) {
17 vari** varis = ChainableStack::instance_->memalloc_.alloc_array<vari*>(size);
18 double* partials
19 = ChainableStack::instance_->memalloc_.alloc_array<double>(size);
20
21 Eigen::Map<const vector_v> dtrs_map(dtrs, size);
22 Eigen::Map<vector_vi>(varis, size) = dtrs_map.vi();
23 vector_d dtrs_vals = dtrs_map.val();
24
25 vector_d diff = dtrs_vals.array() - dtrs_vals.mean();
26 double size_m1 = size - 1;
27 Eigen::Map<vector_d>(partials, size) = 2 * diff.array() / size_m1;
28 double variance = diff.squaredNorm() / size_m1;
29
30 return {new stored_gradient_vari(variance, size, varis, partials)};
31 }
32
33 } // namespace internal
34
35 /**
36 * Return the sample variance of the specified standard
37 * vector. Raise domain error if size is not greater than zero.
38 *
39 * @param[in] v a vector
40 * @return sample variance of specified vector
41 */
variance(const std::vector<var> & v)42 inline var variance(const std::vector<var>& v) {
43 check_nonzero_size("variance", "v", v);
44 if (v.size() == 1) {
45 return var{0.0};
46 }
47 return {internal::calc_variance(v.size(), &v[0])};
48 }
49
50 /**
51 * Return the sample variance of the specified vector, row vector,
52 * or matrix. Raise domain error if size is not greater than
53 * zero.
54 *
55 * @tparam EigMat type inheriting from `EigenBase` that has a `var`
56 * scalar type.
57 * @param[in] m input matrix
58 * @return sample variance of specified matrix
59 */
60 template <typename EigMat, require_eigen_vt<is_var, EigMat>* = nullptr>
variance(const EigMat & m)61 var variance(const EigMat& m) {
62 const auto& mat = to_ref(m);
63 check_nonzero_size("variance", "m", mat);
64 if (mat.size() == 1) {
65 return var{0.0};
66 }
67 return {internal::calc_variance(mat.size(), mat.data())};
68 }
69
70 /**
71 * Return the sample variance of the var_value matrix
72 * Raise domain error if size is not greater than zero.
73 *
74 * @tparam Mat input matrix type
75 * @param[in] x a input
76 * @return sample variance of input
77 */
78 template <typename Mat, require_var_matrix_t<Mat>* = nullptr>
variance(const Mat & x)79 inline var variance(const Mat& x) {
80 check_nonzero_size("variance", "x", x);
81
82 if (x.size() == 1) {
83 return 0.0;
84 }
85
86 double mean = x.val().mean();
87 arena_t<promote_scalar_t<double, Mat>> arena_diff(x.rows(), x.cols());
88
89 double squaredNorm = 0.0;
90 for (Eigen::Index i = 0; i < arena_diff.size(); ++i) {
91 double diff = x.val().coeff(i) - mean;
92 arena_diff.coeffRef(i) = diff;
93 squaredNorm += diff * diff;
94 }
95
96 var res = squaredNorm / (x.size() - 1);
97
98 reverse_pass_callback([x, res, arena_diff]() mutable {
99 x.adj() += (2.0 * res.adj() / (x.size() - 1)) * arena_diff;
100 });
101
102 return res;
103 }
104
105 } // namespace math
106 } // namespace stan
107 #endif
108