mdarray_linalg_lapack/lu/
context.rs

1//! LU Decomposition:
2//!     P * A = L * U
3//! where:
4//!     - A is m × n (input matrix)
5//!     - P is m × m (permutation matrix, represented by pivot vector)
6//!     - L is m × min(m,n) (lower triangular matrix with unit diagonal)
7//!     - U is min(m,n) × n (upper triangular matrix)
8//! This decomposition is used to solve linear systems, compute matrix determinants, and matrix inversion.
9//! The function `getrf` (LAPACK) computes the LU factorization of a general m-by-n matrix A using partial pivoting.
10//! The matrix L is lower triangular with unit diagonal, and U is upper triangular.
11use mdarray::{Array, Dim, Layout, Shape, Slice};
12use mdarray_linalg::{
13    lu::{InvError, LU},
14    utils::{into_i32, ipiv_to_perm_mat, transpose_in_place},
15};
16use num_complex::ComplexFloat;
17
18use super::{
19    scalar::{LapackScalar, Workspace},
20    simple::{getrf, getri, potrf},
21};
22use crate::Lapack;
23
24impl<T, D0: Dim, D1: Dim> LU<T, D0, D1> for Lapack
25where
26    T: ComplexFloat + Default + LapackScalar + Workspace,
27    T::Real: Into<T>,
28{
29    fn lu_write<L: Layout, Ll: Layout, Lu: Layout, Lp: Layout>(
30        &self,
31        a: &mut Slice<T, (D0, D1), L>,
32        l: &mut Slice<T, (D0, D0), Ll>,
33        u: &mut Slice<T, (D0, D1), Lu>,
34        p: &mut Slice<T, (D0, D0), Lp>,
35    ) {
36        let ash = *a.shape();
37        let m = ash.dim(0);
38
39        let ipiv = getrf(a, l, u);
40
41        let p_matrix = ipiv_to_perm_mat::<T, D0, D1>(&ipiv, m);
42
43        for i in 0..m {
44            for j in 0..m {
45                p[[i, j]] = p_matrix[[i, j]];
46            }
47        }
48    }
49
50    fn lu<L: Layout>(
51        &self,
52        a: &mut Slice<T, (D0, D1), L>,
53    ) -> (
54        Array<T, (D0, D0)>,
55        Array<T, (D0, D1)>,
56        Array<T, (D0, D0)>,
57    ) {
58        let ash = *a.shape();
59        let (m, n) = (ash.dim(0), ash.dim(1));
60
61        let min_mn = m.min(n);
62
63        let l_shape = <(D0, D0) as Shape>::from_dims(&[m, min_mn]);
64        let u_shape = <(D0, D1) as Shape>::from_dims(&[min_mn, n]);
65
66        let mut l = Array::from_elem(l_shape, T::default());
67        let mut u = Array::from_elem(u_shape, T::default());
68
69        let ipiv = getrf::<T, D0, D1, _, _, _>(a, &mut l, &mut u);
70
71        let p_matrix = ipiv_to_perm_mat::<T, D0, D0>(&ipiv, m);
72
73        (l, u, p_matrix)
74    }
75
76    fn inv_write<L: Layout>(&self, a: &mut Slice<T, (D0, D1), L>) -> Result<(), InvError> {
77        let ash = *a.shape();
78        let (m, n) = (ash.dim(0), ash.dim(1));
79
80        if m != n {
81            return Err(InvError::NotSquare {
82                rows: into_i32(m),
83                cols: into_i32(n),
84            });
85        }
86
87        let min_mn = m.min(n);
88
89        let l_shape = <(D0, D0) as Shape>::from_dims(&[m, min_mn]);
90        let u_shape = <(D0, D1) as Shape>::from_dims(&[min_mn, n]);
91
92        let mut l = Array::from_elem(l_shape, T::default());
93        let mut u = Array::from_elem(u_shape, T::default());
94        let mut ipiv = getrf::<T, D0, D1, _, _, _>(a, &mut l, &mut u);
95
96        match getri::<T, D0, D1, _>(a, &mut ipiv) {
97            0 => Ok(()),
98            i if i > 0 => Err(InvError::Singular { pivot: i }),
99            i => Err(InvError::BackendError(i)),
100        }
101    }
102
103    fn inv<L: Layout>(
104        &self,
105        a: &mut Slice<T, (D0, D1), L>,
106    ) -> Result<Array<T, (D0, D1)>, InvError> {
107        let ash = *a.shape();
108        let (m, n) = (ash.dim(0), ash.dim(1));
109
110        if m != n {
111            return Err(InvError::NotSquare {
112                rows: into_i32(m),
113                cols: into_i32(n),
114            });
115        }
116
117        let mut a_inv = Array::<T, (D0, D1)>::zeros(ash);
118
119        // let mut a_inv_mut = a_inv.view_mut(.., ..);
120
121        for i in 0..n {
122            for j in 0..m {
123                a_inv[[i, j]] = a[[i, j]];
124            }
125        }
126
127        let min_mn = m.min(n);
128
129        let l_shape = <(D0, D0) as Shape>::from_dims(&[m, min_mn]);
130        let u_shape = <(D0, D1) as Shape>::from_dims(&[min_mn, n]);
131
132        let mut l = Array::from_elem(l_shape, T::default());
133        let mut u = Array::from_elem(u_shape, T::default());
134        let mut ipiv = getrf::<T, D0, D1, _, _, _>(&mut a_inv, &mut l, &mut u);
135
136        match getri::<T, D0, D1, L>(a, &mut ipiv) {
137            0 => Ok(a.to_tensor()),
138            i if i > 0 => Err(InvError::Singular { pivot: i }),
139            i => Err(InvError::BackendError(i)),
140        }
141    }
142
143    fn det<L: Layout>(&self, a: &mut Slice<T, (D0, D1), L>) -> T {
144        let ash = *a.shape();
145        let (m, n) = (ash.dim(0), ash.dim(1));
146        assert_eq!(m, n, "determinant is only defined for square matrices");
147
148        let l_shape = <(D0, D0) as Shape>::from_dims(&[n, n]);
149        let u_shape = <(D0, D1) as Shape>::from_dims(&[n, n]);
150        let mut l = Array::from_elem(l_shape, T::default());
151        let mut u = Array::from_elem(u_shape, T::default());
152
153        let ipiv = getrf::<T, D0, D1, _, _, _>(a, &mut l, &mut u);
154
155        let mut det = T::one();
156        for i in 0..n {
157            det = det * u[[i, i]];
158        }
159
160        let mut sign = T::one();
161        for (i, &pivot) in ipiv.iter().enumerate() {
162            if (i as i32) != (pivot - 1) {
163                sign = sign * (-T::one());
164            }
165        }
166        det * sign
167    }
168
169    /// Computes the Cholesky decomposition, returning a lower-triangular matrix
170    fn cholesky<L: Layout>(
171        &self,
172        a: &mut Slice<T, (D0, D1), L>,
173    ) -> Result<Array<T, (D0, D1)>, InvError> {
174        let ash = *a.shape();
175        let (m, n) = (ash.dim(0), ash.dim(1));
176        assert_eq!(m, n, "Matrix must be square for Cholesky decomposition");
177
178        let mut l = Array::<T, (D0, D1)>::zeros(ash);
179
180        match potrf::<T, D0, D1, _>(a, 'L') {
181            0 => {
182                for i in 0..m {
183                    for j in 0..n {
184                        if i >= j {
185                            l[[i, j]] = a[[j, i]];
186                        } else {
187                            l[[i, j]] = T::zero();
188                        }
189                    }
190                }
191                Ok(l)
192            }
193            i if i > 0 => Err(InvError::NotPositiveDefinite { lpm: i }),
194            i => Err(InvError::BackendError(i)),
195        }
196    }
197
198    /// Computes the Cholesky decomposition in-place, overwriting the input matrix
199    fn cholesky_write<L: Layout>(&self, a: &mut Slice<T, (D0, D1), L>) -> Result<(), InvError> {
200        let ash = *a.shape();
201        let (m, n) = (ash.dim(0), ash.dim(1));
202        assert_eq!(m, n, "Matrix must be square for Cholesky decomposition");
203
204        match potrf::<T, D0, D1, _>(a, 'L') {
205            0 => {
206                transpose_in_place(a);
207                Ok(())
208            }
209            i if i > 0 => Err(InvError::NotPositiveDefinite { lpm: i }),
210            i => Err(InvError::BackendError(i)),
211        }
212    }
213}
mdarray_linalg_lapack/lu/context.rs

mdarray_linalg_lapack/lu/
context.rs
