1use crate::{convert::*, error::*, layout::*, types::*};
6use ndarray::*;
7
8pub trait SVD {
10 type U;
11 type VT;
12 type Sigma;
13 fn svd(
14 &self,
15 calc_u: bool,
16 calc_vt: bool,
17 ) -> Result<(Option<Self::U>, Self::Sigma, Option<Self::VT>)>;
18}
19
20pub trait SVDInto {
22 type U;
23 type VT;
24 type Sigma;
25 fn svd_into(
26 self,
27 calc_u: bool,
28 calc_vt: bool,
29 ) -> Result<(Option<Self::U>, Self::Sigma, Option<Self::VT>)>;
30}
31
32pub trait SVDInplace {
34 type U;
35 type VT;
36 type Sigma;
37 fn svd_inplace(
38 &mut self,
39 calc_u: bool,
40 calc_vt: bool,
41 ) -> Result<(Option<Self::U>, Self::Sigma, Option<Self::VT>)>;
42}
43
44impl<A, S> SVDInto for ArrayBase<S, Ix2>
45where
46 A: Scalar + Lapack,
47 S: DataMut<Elem = A>,
48{
49 type U = Array2<A>;
50 type VT = Array2<A>;
51 type Sigma = Array1<A::Real>;
52
53 fn svd_into(
54 mut self,
55 calc_u: bool,
56 calc_vt: bool,
57 ) -> Result<(Option<Self::U>, Self::Sigma, Option<Self::VT>)> {
58 self.svd_inplace(calc_u, calc_vt)
59 }
60}
61
62impl<A, S> SVD for ArrayBase<S, Ix2>
63where
64 A: Scalar + Lapack,
65 S: Data<Elem = A>,
66{
67 type U = Array2<A>;
68 type VT = Array2<A>;
69 type Sigma = Array1<A::Real>;
70
71 fn svd(
72 &self,
73 calc_u: bool,
74 calc_vt: bool,
75 ) -> Result<(Option<Self::U>, Self::Sigma, Option<Self::VT>)> {
76 let a = self.to_owned();
77 a.svd_into(calc_u, calc_vt)
78 }
79}
80
81impl<A, S> SVDInplace for ArrayBase<S, Ix2>
82where
83 A: Scalar + Lapack,
84 S: DataMut<Elem = A>,
85{
86 type U = Array2<A>;
87 type VT = Array2<A>;
88 type Sigma = Array1<A::Real>;
89
90 fn svd_inplace(
91 &mut self,
92 calc_u: bool,
93 calc_vt: bool,
94 ) -> Result<(Option<Self::U>, Self::Sigma, Option<Self::VT>)> {
95 let l = self.layout()?;
96 let svd_res = A::svd(l, calc_u, calc_vt, self.as_allocated_mut()?)?;
97 let (n, m) = l.size();
98
99 let u = svd_res.u.map(|u| into_matrix(l.resized(n, n), u).unwrap());
100 let vt = svd_res
101 .vt
102 .map(|vt| into_matrix(l.resized(m, m), vt).unwrap());
103 let s = ArrayBase::from(svd_res.s);
104 Ok((u, s, vt))
105 }
106}