From 2b3e4bec4faf95fab469de9375b9c0932ae204b4 Mon Sep 17 00:00:00 2001 From: conrad Date: Sat, 3 Jul 2021 01:26:28 +1000 Subject: [PATCH] initial implementation of pivoted chol --- include/armadillo_bits/auxlib_bones.hpp | 3 + include/armadillo_bits/auxlib_meat.hpp | 55 +++++++++++++++ include/armadillo_bits/def_lapack.hpp | 24 ++++++- include/armadillo_bits/fn_chol.hpp | 75 +++++++++++++++++++++ include/armadillo_bits/translate_lapack.hpp | 22 ++++++ src/wrapper1.cpp | 22 ++++++ src/wrapper2.cpp | 22 ++++++ 7 files changed, 222 insertions(+), 1 deletion(-) diff --git a/include/armadillo_bits/auxlib_bones.hpp b/include/armadillo_bits/auxlib_bones.hpp index 3e0ba34a..4c9ee42a 100644 --- a/include/armadillo_bits/auxlib_bones.hpp +++ b/include/armadillo_bits/auxlib_bones.hpp @@ -174,6 +174,9 @@ class auxlib template inline static bool chol_band_common(Mat& X, const uword KD, const uword layout); + template + inline static bool chol_pivot(Mat& X, Mat& P, const uword layout); + // // hessenberg decomposition diff --git a/include/armadillo_bits/auxlib_meat.hpp b/include/armadillo_bits/auxlib_meat.hpp index 4ae40473..a7d2edce 100644 --- a/include/armadillo_bits/auxlib_meat.hpp +++ b/include/armadillo_bits/auxlib_meat.hpp @@ -2636,6 +2636,61 @@ auxlib::chol_band_common(Mat& X, const uword KD, const uword layout) } + +template +inline +bool +auxlib::chol_pivot(Mat& X, Mat& P, const uword layout) + { + arma_extra_debug_sigprint(); + + #if defined(ARMA_USE_LAPACK) + { + typedef typename get_pod_type::result T; + + arma_debug_assert_blas_size(X); + + char uplo = (layout == 0) ? 'U' : 'L'; + blas_int n = blas_int(X.n_rows); + blas_int rank = 0; + T tol = T(-1); + blas_int info = 0; + + podarray ipiv( X.n_rows); + podarray work(2*X.n_rows); + + ipiv.zeros(); + + arma_extra_debug_print("lapack::pstrf()"); + lapack::pstrf(&uplo, &n, X.memptr(), &n, ipiv.memptr(), &rank, &tol, work.memptr(), &info); + + if(info != 0) { return false; } + + X = (layout == 0) ? trimatu(X) : trimatl(X); // trimatu() and trimatl() return the same type + + P.set_size(X.n_rows, 1); + + for(uword i=0; i < X.n_rows; ++i) + { + P[i] = uword(ipiv[i] - 1); // take into account that Fortran counts from 1 + } + + return true; + } + #else + { + arma_ignore(X); + arma_ignore(P); + arma_ignore(layout); + + arma_stop_logic_error("chol(): use of LAPACK must be enabled"); + return false; + } + #endif + } + + + // // hessenberg decomposition template diff --git a/include/armadillo_bits/def_lapack.hpp b/include/armadillo_bits/def_lapack.hpp index 9d341126..192caef0 100644 --- a/include/armadillo_bits/def_lapack.hpp +++ b/include/armadillo_bits/def_lapack.hpp @@ -265,6 +265,11 @@ #define arma_cgehrd cgehrd #define arma_zgehrd zgehrd + #define arma_spstrf spstrf + #define arma_dpstrf dpstrf + #define arma_cpstrf cpstrf + #define arma_zpstrf zpstrf + #else #define arma_sgetrf SGETRF @@ -499,6 +504,11 @@ #define arma_cgehrd CGEHRD #define arma_zgehrd ZGEHRD + #define arma_spstrf SPSTRF + #define arma_dpstrf DPSTRF + #define arma_cpstrf CPSTRF + #define arma_zpstrf ZPSTRF + #endif @@ -838,6 +848,12 @@ extern "C" void arma_fortran(arma_cgehrd)(const blas_int* n, const blas_int* ilo, const blas_int* ihi, blas_cxf* a, const blas_int* lda, blas_cxf* tao, blas_cxf* work, const blas_int* lwork, blas_int* info) ARMA_NOEXCEPT; void arma_fortran(arma_zgehrd)(const blas_int* n, const blas_int* ilo, const blas_int* ihi, blas_cxd* a, const blas_int* lda, blas_cxd* tao, blas_cxd* work, const blas_int* lwork, blas_int* info) ARMA_NOEXCEPT; + // pivoted cholesky + void arma_fortran(arma_spstrf)(const char* uplo, const blas_int* n, float* a, const blas_int* lda, blas_int* piv, blas_int* rank, const float* tol, float* work, blas_int* info, blas_len uplo_len) ARMA_NOEXCEPT; + void arma_fortran(arma_dpstrf)(const char* uplo, const blas_int* n, double* a, const blas_int* lda, blas_int* piv, blas_int* rank, const double* tol, double* work, blas_int* info, blas_len uplo_len) ARMA_NOEXCEPT; + void arma_fortran(arma_cpstrf)(const char* uplo, const blas_int* n, blas_cxf* a, const blas_int* lda, blas_int* piv, blas_int* rank, const float* tol, float* work, blas_int* info, blas_len uplo_len) ARMA_NOEXCEPT; + void arma_fortran(arma_zpstrf)(const char* uplo, const blas_int* n, blas_cxd* a, const blas_int* lda, blas_int* piv, blas_int* rank, const double* tol, double* work, blas_int* info, blas_len uplo_len) ARMA_NOEXCEPT; + #else // prototypes without hidden arguments @@ -1159,7 +1175,13 @@ extern "C" void arma_fortran(arma_dgehrd)(const blas_int* n, const blas_int* ilo, const blas_int* ihi, double* a, const blas_int* lda, double* tao, double* work, const blas_int* lwork, blas_int* info) ARMA_NOEXCEPT; void arma_fortran(arma_cgehrd)(const blas_int* n, const blas_int* ilo, const blas_int* ihi, blas_cxf* a, const blas_int* lda, blas_cxf* tao, blas_cxf* work, const blas_int* lwork, blas_int* info) ARMA_NOEXCEPT; void arma_fortran(arma_zgehrd)(const blas_int* n, const blas_int* ilo, const blas_int* ihi, blas_cxd* a, const blas_int* lda, blas_cxd* tao, blas_cxd* work, const blas_int* lwork, blas_int* info) ARMA_NOEXCEPT; - + + // pivoted cholesky + void arma_fortran(arma_spstrf)(const char* uplo, const blas_int* n, float* a, const blas_int* lda, blas_int* piv, blas_int* rank, const float* tol, float* work, blas_int* info) ARMA_NOEXCEPT; + void arma_fortran(arma_dpstrf)(const char* uplo, const blas_int* n, double* a, const blas_int* lda, blas_int* piv, blas_int* rank, const double* tol, double* work, blas_int* info) ARMA_NOEXCEPT; + void arma_fortran(arma_cpstrf)(const char* uplo, const blas_int* n, blas_cxf* a, const blas_int* lda, blas_int* piv, blas_int* rank, const float* tol, float* work, blas_int* info) ARMA_NOEXCEPT; + void arma_fortran(arma_zpstrf)(const char* uplo, const blas_int* n, blas_cxd* a, const blas_int* lda, blas_int* piv, blas_int* rank, const double* tol, double* work, blas_int* info) ARMA_NOEXCEPT; + #endif } diff --git a/include/armadillo_bits/fn_chol.hpp b/include/armadillo_bits/fn_chol.hpp index 0fa6d397..07460639 100644 --- a/include/armadillo_bits/fn_chol.hpp +++ b/include/armadillo_bits/fn_chol.hpp @@ -69,4 +69,79 @@ chol +template +inline +typename enable_if2< is_supported_blas_type::value, bool >::result +chol + ( + Mat& out, + Mat& P, + const Base& X, + const char* layout = "upper", + const char* P_mode = "matrix" + ) + { + arma_extra_debug_sigprint(); + + typedef typename T1::elem_type eT; + + const char sig_layout = (layout != nullptr) ? layout[0] : char(0); + const char sig_P_mode = (P_mode != nullptr) ? P_mode[0] : char(0); + + arma_debug_check( ((sig_layout != 'u') && (sig_layout != 'l')), "chol(): argument 'layout' must be \"upper\" or \"lower\"" ); + arma_debug_check( ((sig_P_mode != 'm') && (sig_P_mode != 'v')), "chol(): argument 'P_mode' must be \"vector\" or \"matrix\"" ); + + out = X.get_ref(); + + arma_debug_check( (out.is_square() == false), "chol(): given matrix must be square sized" ); + + if(out.is_empty()) + { + P.reset(); + return true; + } + + if((arma_config::debug) && (auxlib::rudimentary_sym_check(out) == false)) + { + if(is_cx::no ) { arma_debug_warn_level(1, "chol(): given matrix is not symmetric"); } + if(is_cx::yes) { arma_debug_warn_level(1, "chol(): given matrix is not hermitian"); } + } + + bool status = false; + + if(sig_P_mode == 'v') + { + status = auxlib::chol_pivot(out, P, ((sig_layout == 'u') ? 0 : 1)); + } + else + if(sig_P_mode == 'm') + { + Mat P_vec; + + status = auxlib::chol_pivot(out, P_vec, ((sig_layout == 'u') ? 0 : 1)); + + if(status) + { + // construct P + + const uword N = P_vec.n_rows; + + P.zeros(N,N); + + for(uword i=0; i < N; ++i) { P.at(P_vec[i], i) = uword(1); } + } + } + + if(status == false) + { + out.soft_reset(); + P.soft_reset(); + arma_debug_warn_level(3, "chol(): decomposition failed"); + } + + return status; + } + + + //! @} diff --git a/include/armadillo_bits/translate_lapack.hpp b/include/armadillo_bits/translate_lapack.hpp index 031f03a7..a31179a8 100644 --- a/include/armadillo_bits/translate_lapack.hpp +++ b/include/armadillo_bits/translate_lapack.hpp @@ -1329,6 +1329,28 @@ namespace lapack } + + template + inline + void + pstrf(const char* uplo, const blas_int* n, eT* a, const blas_int* lda, blas_int* piv, blas_int* rank, const typename get_pod_type::result* tol, const typename get_pod_type::result* work, blas_int* info) + { + arma_type_check(( is_supported_blas_type::value == false )); + + #if defined(ARMA_USE_FORTRAN_HIDDEN_ARGS) + if( is_float::value) { typedef float pod_T; typedef float T; arma_fortran(arma_spstrf)(uplo, n, (T*)a, lda, piv, rank, (const pod_T*)tol, (pod_T*)work, info, 1); } + else if( is_double::value) { typedef double pod_T; typedef double T; arma_fortran(arma_dpstrf)(uplo, n, (T*)a, lda, piv, rank, (const pod_T*)tol, (pod_T*)work, info, 1); } + else if( is_cx_float::value) { typedef float pod_T; typedef blas_cxf T; arma_fortran(arma_cpstrf)(uplo, n, (T*)a, lda, piv, rank, (const pod_T*)tol, (pod_T*)work, info, 1); } + else if(is_cx_double::value) { typedef double pod_T; typedef blas_cxd T; arma_fortran(arma_zpstrf)(uplo, n, (T*)a, lda, piv, rank, (const pod_T*)tol, (pod_T*)work, info, 1); } + #else + if( is_float::value) { typedef float pod_T; typedef float T; arma_fortran(arma_spstrf)(uplo, n, (T*)a, lda, piv, rank, (const pod_T*)tol, (pod_T*)work, info); } + else if( is_double::value) { typedef double pod_T; typedef double T; arma_fortran(arma_dpstrf)(uplo, n, (T*)a, lda, piv, rank, (const pod_T*)tol, (pod_T*)work, info); } + else if( is_cx_float::value) { typedef float pod_T; typedef blas_cxf T; arma_fortran(arma_cpstrf)(uplo, n, (T*)a, lda, piv, rank, (const pod_T*)tol, (pod_T*)work, info); } + else if(is_cx_double::value) { typedef double pod_T; typedef blas_cxd T; arma_fortran(arma_zpstrf)(uplo, n, (T*)a, lda, piv, rank, (const pod_T*)tol, (pod_T*)work, info); } + #endif + } + + } diff --git a/src/wrapper1.cpp b/src/wrapper1.cpp index 923ef030..315574e8 100644 --- a/src/wrapper1.cpp +++ b/src/wrapper1.cpp @@ -1196,6 +1196,28 @@ extern "C" arma_fortran_sans_prefix(arma_zgehrd)(n, ilo, ihi, a, lda, tao, work, lwork, info); } + + + void arma_fortran_with_prefix(arma_spstrf)(const char* uplo, const blas_int* n, float* a, const blas_int* lda, blas_int* piv, blas_int* rank, const float* tol, float* work, blas_int* info) + { + arma_fortran_sans_prefix(arma_spstrf)(uplo, n, a, lda, piv, rank, tol, work, info); + } + + void arma_fortran_with_prefix(arma_dpstrf)(const char* uplo, const blas_int* n, double* a, const blas_int* lda, blas_int* piv, blas_int* rank, const double* tol, double* work, blas_int* info) + { + arma_fortran_sans_prefix(arma_dpstrf)(uplo, n, a, lda, piv, rank, tol, work, info); + } + + void arma_fortran_with_prefix(arma_cpstrf)(const char* uplo, const blas_int* n, blas_cxf* a, const blas_int* lda, blas_int* piv, blas_int* rank, const float* tol, float* work, blas_int* info) + { + arma_fortran_sans_prefix(arma_cpstrf)(uplo, n, a, lda, piv, rank, tol, work, info); + } + + void arma_fortran_with_prefix(arma_zpstrf)(const char* uplo, const blas_int* n, blas_cxd* a, const blas_int* lda, blas_int* piv, blas_int* rank, const double* tol, double* work, blas_int* info) + { + arma_fortran_sans_prefix(arma_zpstrf)(uplo, n, a, lda, piv, rank, tol, work, info); + } + #endif diff --git a/src/wrapper2.cpp b/src/wrapper2.cpp index 453df1eb..0b544182 100644 --- a/src/wrapper2.cpp +++ b/src/wrapper2.cpp @@ -1149,6 +1149,28 @@ extern "C" arma_fortran_sans_prefix(arma_zgehrd)(n, ilo, ihi, a, lda, tao, work, lwork, info); } + + + void arma_fortran_with_prefix(arma_spstrf)(const char* uplo, const blas_int* n, float* a, const blas_int* lda, blas_int* piv, blas_int* rank, const float* tol, float* work, blas_int* info, blas_len uplo_len) + { + arma_fortran_sans_prefix(arma_spstrf)(uplo, n, a, lda, piv, rank, tol, work, info, uplo_len); + } + + void arma_fortran_with_prefix(arma_dpstrf)(const char* uplo, const blas_int* n, double* a, const blas_int* lda, blas_int* piv, blas_int* rank, const double* tol, double* work, blas_int* info, blas_len uplo_len) + { + arma_fortran_sans_prefix(arma_dpstrf)(uplo, n, a, lda, piv, rank, tol, work, info, uplo_len); + } + + void arma_fortran_with_prefix(arma_cpstrf)(const char* uplo, const blas_int* n, blas_cxf* a, const blas_int* lda, blas_int* piv, blas_int* rank, const float* tol, float* work, blas_int* info, blas_len uplo_len) + { + arma_fortran_sans_prefix(arma_cpstrf)(uplo, n, a, lda, piv, rank, tol, work, info, uplo_len); + } + + void arma_fortran_with_prefix(arma_zpstrf)(const char* uplo, const blas_int* n, blas_cxd* a, const blas_int* lda, blas_int* piv, blas_int* rank, const double* tol, double* work, blas_int* info, blas_len uplo_len) + { + arma_fortran_sans_prefix(arma_zpstrf)(uplo, n, a, lda, piv, rank, tol, work, info, uplo_len); + } + #endif