Skip to content

Commit 062a345

Browse files
termoshtttermoshtt
committed
Merge Cholesky_ into Lapack
1 parent d42fc3e commit 062a345

File tree

2 files changed

+69
-149
lines changed

2 files changed

+69
-149
lines changed

‎lax/src/cholesky.rs

Lines changed: 27 additions & 148 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,15 @@ use super::*;
22
use crate::{error::*, layout::*};
33
use cauchy::*;
44

5+
/// Compute Cholesky decomposition according to [UPLO]
6+
///
7+
/// LAPACK correspondance
8+
/// ----------------------
9+
///
10+
/// | f32 | f64 | c32 | c64 |
11+
/// |:-------|:-------|:-------|:-------|
12+
/// | spotrf | dpotrf | cpotrf | zpotrf |
13+
///
514
pub trait CholeskyImpl: Scalar {
615
fn cholesky(l: MatrixLayout, uplo: UPLO, a: &mut [Self]) -> Result<()>;
716
}
@@ -32,6 +41,15 @@ impl_cholesky_!(c32, lapack_sys::cpotrf_);
3241
impl_cholesky_!(f64, lapack_sys::dpotrf_);
3342
impl_cholesky_!(f32, lapack_sys::spotrf_);
3443

44+
/// Compute inverse matrix using Cholesky factroization result
45+
///
46+
/// LAPACK correspondance
47+
/// ----------------------
48+
///
49+
/// | f32 | f64 | c32 | c64 |
50+
/// |:-------|:-------|:-------|:-------|
51+
/// | spotri | dpotri | cpotri | zpotri |
52+
///
3553
pub trait InvCholeskyImpl: Scalar {
3654
fn inv_cholesky(l: MatrixLayout, uplo: UPLO, a: &mut [Self]) -> Result<()>;
3755
}
@@ -62,6 +80,15 @@ impl_inv_cholesky!(c32, lapack_sys::cpotri_);
6280
impl_inv_cholesky!(f64, lapack_sys::dpotri_);
6381
impl_inv_cholesky!(f32, lapack_sys::spotri_);
6482

83+
/// Solve linear equation using Cholesky factroization result
84+
///
85+
/// LAPACK correspondance
86+
/// ----------------------
87+
///
88+
/// | f32 | f64 | c32 | c64 |
89+
/// |:-------|:-------|:-------|:-------|
90+
/// | spotrs | dpotrs | cpotrs | zpotrs |
91+
///
6592
pub trait SolveCholeskyImpl: Scalar {
6693
fn solve_cholesky(l: MatrixLayout, uplo: UPLO, a: &[Self], b: &mut [Self]) -> Result<()>;
6794
}
@@ -111,151 +138,3 @@ impl_solve_cholesky!(c64, lapack_sys::zpotrs_);
111138
impl_solve_cholesky!(c32, lapack_sys::cpotrs_);
112139
impl_solve_cholesky!(f64, lapack_sys::dpotrs_);
113140
impl_solve_cholesky!(f32, lapack_sys::spotrs_);
114-
115-
#[cfg_attr(doc, katexit::katexit)]
116-
/// Solve symmetric/hermite positive-definite linear equations using Cholesky decomposition
117-
///
118-
/// For a given positive definite matrix $A$,
119-
/// Cholesky decomposition is described as $A = U^T U$ or $A = LL^T$ where
120-
///
121-
/// - $L$ is lower matrix
122-
/// - $U$ is upper matrix
123-
///
124-
/// This is designed as two step computation according to LAPACK API
125-
///
126-
/// 1. Factorize input matrix $A$ into $L$ or $U$
127-
/// 2. Solve linear equation $Ax = b$ or compute inverse matrix $A^{-1}$
128-
/// using $U$ or $L$.
129-
pub trait Cholesky_: Sized {
130-
/// Compute Cholesky decomposition $A = U^T U$ or $A = L L^T$ according to [UPLO]
131-
///
132-
/// LAPACK correspondance
133-
/// ----------------------
134-
///
135-
/// | f32 | f64 | c32 | c64 |
136-
/// |:-------|:-------|:-------|:-------|
137-
/// | spotrf | dpotrf | cpotrf | zpotrf |
138-
///
139-
fn cholesky(l: MatrixLayout, uplo: UPLO, a: &mut [Self]) -> Result<()>;
140-
141-
/// Compute inverse matrix $A^{-1}$ using $U$ or $L$
142-
///
143-
/// LAPACK correspondance
144-
/// ----------------------
145-
///
146-
/// | f32 | f64 | c32 | c64 |
147-
/// |:-------|:-------|:-------|:-------|
148-
/// | spotri | dpotri | cpotri | zpotri |
149-
///
150-
fn inv_cholesky(l: MatrixLayout, uplo: UPLO, a: &mut [Self]) -> Result<()>;
151-
152-
/// Solve linear equation $Ax = b$ using $U$ or $L$
153-
///
154-
/// LAPACK correspondance
155-
/// ----------------------
156-
///
157-
/// | f32 | f64 | c32 | c64 |
158-
/// |:-------|:-------|:-------|:-------|
159-
/// | spotrs | dpotrs | cpotrs | zpotrs |
160-
///
161-
fn solve_cholesky(l: MatrixLayout, uplo: UPLO, a: &[Self], b: &mut [Self]) -> Result<()>;
162-
}
163-
164-
macro_rules! impl_cholesky {
165-
($scalar:ty, $trf:path, $tri:path, $trs:path) => {
166-
impl Cholesky_ for $scalar {
167-
fn cholesky(l: MatrixLayout, uplo: UPLO, a: &mut [Self]) -> Result<()> {
168-
let (n, _) = l.size();
169-
if matches!(l, MatrixLayout::C { .. }) {
170-
square_transpose(l, a);
171-
}
172-
let mut info = 0;
173-
unsafe {
174-
$trf(uplo.as_ptr(), &n, AsPtr::as_mut_ptr(a), &n, &mut info);
175-
}
176-
info.as_lapack_result()?;
177-
if matches!(l, MatrixLayout::C { .. }) {
178-
square_transpose(l, a);
179-
}
180-
Ok(())
181-
}
182-
183-
fn inv_cholesky(l: MatrixLayout, uplo: UPLO, a: &mut [Self]) -> Result<()> {
184-
let (n, _) = l.size();
185-
if matches!(l, MatrixLayout::C { .. }) {
186-
square_transpose(l, a);
187-
}
188-
let mut info = 0;
189-
unsafe {
190-
$tri(uplo.as_ptr(), &n, AsPtr::as_mut_ptr(a), &l.lda(), &mut info);
191-
}
192-
info.as_lapack_result()?;
193-
if matches!(l, MatrixLayout::C { .. }) {
194-
square_transpose(l, a);
195-
}
196-
Ok(())
197-
}
198-
199-
fn solve_cholesky(
200-
l: MatrixLayout,
201-
mut uplo: UPLO,
202-
a: &[Self],
203-
b: &mut [Self],
204-
) -> Result<()> {
205-
let (n, _) = l.size();
206-
let nrhs = 1;
207-
let mut info = 0;
208-
if matches!(l, MatrixLayout::C { .. }) {
209-
uplo = uplo.t();
210-
for val in b.iter_mut() {
211-
*val = val.conj();
212-
}
213-
}
214-
unsafe {
215-
$trs(
216-
uplo.as_ptr(),
217-
&n,
218-
&nrhs,
219-
AsPtr::as_ptr(a),
220-
&l.lda(),
221-
AsPtr::as_mut_ptr(b),
222-
&n,
223-
&mut info,
224-
);
225-
}
226-
info.as_lapack_result()?;
227-
if matches!(l, MatrixLayout::C { .. }) {
228-
for val in b.iter_mut() {
229-
*val = val.conj();
230-
}
231-
}
232-
Ok(())
233-
}
234-
}
235-
};
236-
} // end macro_rules
237-
238-
impl_cholesky!(
239-
f64,
240-
lapack_sys::dpotrf_,
241-
lapack_sys::dpotri_,
242-
lapack_sys::dpotrs_
243-
);
244-
impl_cholesky!(
245-
f32,
246-
lapack_sys::spotrf_,
247-
lapack_sys::spotri_,
248-
lapack_sys::spotrs_
249-
);
250-
impl_cholesky!(
251-
c64,
252-
lapack_sys::zpotrf_,
253-
lapack_sys::zpotri_,
254-
lapack_sys::zpotrs_
255-
);
256-
impl_cholesky!(
257-
c32,
258-
lapack_sys::cpotrf_,
259-
lapack_sys::cpotri_,
260-
lapack_sys::cpotrs_
261-
);

‎lax/src/lib.rs

Lines changed: 42 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -122,7 +122,7 @@ pub type Pivot = Vec<i32>;
122122

123123
#[cfg_attr(doc, katexit::katexit)]
124124
/// Trait for primitive types which implements LAPACK subroutines
125-
pub trait Lapack: OperatorNorm_ + Cholesky_ + Triangular_ + Tridiagonal_ + Rcond_ {
125+
pub trait Lapack: OperatorNorm_ + Triangular_ + Tridiagonal_ + Rcond_ {
126126
/// Compute right eigenvalue and eigenvectors for a general matrix
127127
fn eig(
128128
calc_v: bool,
@@ -238,6 +238,27 @@ pub trait Lapack: OperatorNorm_ + Cholesky_ + Triangular_ + Tridiagonal_ + Rcond
238238

239239
/// Solve symmetric/Hermitian linear equation $Ax = b$ using factroized result
240240
fn solveh(l: MatrixLayout, uplo: UPLO, a: &[Self], ipiv: &Pivot, b: &mut [Self]) -> Result<()>;
241+
242+
/// Solve symmetric/hermite positive-definite linear equations using Cholesky decomposition
243+
///
244+
/// For a given positive definite matrix $A$,
245+
/// Cholesky decomposition is described as $A = U^T U$ or $A = LL^T$ where
246+
///
247+
/// - $L$ is lower matrix
248+
/// - $U$ is upper matrix
249+
///
250+
/// This is designed as two step computation according to LAPACK API
251+
///
252+
/// 1. Factorize input matrix $A$ into $L$ or $U$
253+
/// 2. Solve linear equation $Ax = b$ or compute inverse matrix $A^{-1}$
254+
/// using $U$ or $L$.
255+
fn cholesky(l: MatrixLayout, uplo: UPLO, a: &mut [Self]) -> Result<()>;
256+
257+
/// Compute inverse matrix $A^{-1}$ using $U$ or $L$
258+
fn inv_cholesky(l: MatrixLayout, uplo: UPLO, a: &mut [Self]) -> Result<()>;
259+
260+
/// Solve linear equation $Ax = b$ using $U$ or $L$
261+
fn solve_cholesky(l: MatrixLayout, uplo: UPLO, a: &[Self], b: &mut [Self]) -> Result<()>;
241262
}
242263

243264
macro_rules! impl_lapack {
@@ -379,6 +400,26 @@ macro_rules! impl_lapack {
379400
use solveh::*;
380401
SolvehImpl::solveh(l, uplo, a, ipiv, b)
381402
}
403+
404+
fn cholesky(l: MatrixLayout, uplo: UPLO, a: &mut [Self]) -> Result<()> {
405+
use cholesky::*;
406+
CholeskyImpl::cholesky(l, uplo, a)
407+
}
408+
409+
fn inv_cholesky(l: MatrixLayout, uplo: UPLO, a: &mut [Self]) -> Result<()> {
410+
use cholesky::*;
411+
InvCholeskyImpl::inv_cholesky(l, uplo, a)
412+
}
413+
414+
fn solve_cholesky(
415+
l: MatrixLayout,
416+
uplo: UPLO,
417+
a: &[Self],
418+
b: &mut [Self],
419+
) -> Result<()> {
420+
use cholesky::*;
421+
SolveCholeskyImpl::solve_cholesky(l, uplo, a, b)
422+
}
382423
}
383424
};
384425
}

0 commit comments

Comments
 (0)