zenu_matrix/operation/
var.rs

1use crate::{
2    device::Device,
3    dim::DimDyn,
4    matrix::{Matrix, Owned, Repr},
5    num::Num,
6};
7
8// use crate::{
9//     dim::DimDyn, matrix::ToViewMatrix, matrix_impl::Matrix, memory_impl::OwnedMem, num::Num,
10// };
11//
12// use super::mean::Mean;
13//
14// pub trait Variance<T: Num> {
15//     fn variance(&self, axis: Option<usize>, keep_dim: bool) -> Matrix<OwnedMem<T>, DimDyn>;
16// }
17//
18impl<T: Num, R: Repr<Item = T>, D: Device> Matrix<R, DimDyn, D> {
19    #[must_use]
20    pub fn variance(&self, axis: Option<usize>, keep_dim: bool) -> Matrix<Owned<T>, DimDyn, D> {
21        let mean = self.mean(axis, true);
22        let diff = self.to_ref() - mean;
23        let diff = diff.to_ref() * diff.to_ref();
24        diff.mean(axis, keep_dim)
25    }
26}
27
28#[cfg(test)]
29mod variance {
30    use crate::{device::Device, dim::DimDyn, matrix::Matrix};
31
32    fn variance_1d<D: Device>() {
33        let x = vec![1.0, 2.0, 3.0, 4.0];
34        let x = Matrix::<_, DimDyn, D>::from_vec(x, [4]);
35        let ans = x.variance(None, false);
36        assert!((ans - 1.25).asum() < 1e-6);
37    }
38    #[test]
39    fn variance_1d_cpu() {
40        variance_1d::<crate::device::cpu::Cpu>();
41    }
42    #[cfg(feature = "nvidia")]
43    #[test]
44    fn variance_1d_cuda() {
45        variance_1d::<crate::device::nvidia::Nvidia>();
46    }
47
48    fn variance_1d_<D: Device>() {
49        let x = vec![1.0, 2.0];
50        let x = Matrix::<_, DimDyn, D>::from_vec(x, [2]);
51        let ans = x.variance(None, false);
52        assert!((ans - 0.25).asum() < 1e-6);
53    }
54    #[test]
55    fn variance_1d_cpu_() {
56        variance_1d_::<crate::device::cpu::Cpu>();
57    }
58    #[cfg(feature = "nvidia")]
59    #[test]
60    fn variance_1d_cuda_() {
61        variance_1d_::<crate::device::nvidia::Nvidia>();
62    }
63
64    fn variance_2d<D: Device>() {
65        let x = vec![1.0, 2.0, 3.0, 4.0];
66        let x = Matrix::<_, DimDyn, D>::from_vec(x, [2, 2]);
67        let ans = x.variance(None, false);
68        assert!((ans - 1.25).asum() < 1e-6);
69    }
70    #[test]
71    fn variance_2d_cpu() {
72        variance_2d::<crate::device::cpu::Cpu>();
73    }
74    #[cfg(feature = "nvidia")]
75    #[test]
76    fn variance_2d_cuda() {
77        variance_2d::<crate::device::nvidia::Nvidia>();
78    }
79
80    fn variance_2d_0<D: Device>() {
81        let x = vec![1.0, 2.0, 3.0, 4.0];
82        let x = Matrix::<_, DimDyn, D>::from_vec(x, [2, 2]);
83        let ans = x.variance(Some(0), false);
84        assert!((ans - 1.0).asum() < 1e-6);
85    }
86    #[test]
87    fn variance_2d_0_cpu() {
88        variance_2d_0::<crate::device::cpu::Cpu>();
89    }
90    #[cfg(feature = "nvidia")]
91    #[test]
92    fn variance_2d_0_cuda() {
93        variance_2d_0::<crate::device::nvidia::Nvidia>();
94    }
95
96    fn variance_2d_1<D: Device>() {
97        let x = vec![1.0, 2.0, 3.0, 4.0];
98        let x = Matrix::<_, DimDyn, D>::from_vec(x, [2, 2]);
99        let ans = x.variance(Some(1), false);
100        assert!((ans - 0.25).asum() < 1e-6);
101    }
102    #[test]
103    fn variance_2d_1_cpu() {
104        variance_2d_1::<crate::device::cpu::Cpu>();
105    }
106    #[cfg(feature = "nvidia")]
107    #[test]
108    fn variance_2d_1_cuda() {
109        variance_2d_1::<crate::device::nvidia::Nvidia>();
110    }
111}