fsb_linalg.h

#pragma once
 
#include <cmath>
#include <cstddef>
 
#include "openblas/lapack.h"
#include "fsb_types.h"
 
#define FSB_MAX(a, b) (((a) > (b)) ? (a) : (b))
#define FSB_MIN(a, b) (((a) < (b)) ? (a) : (b))
 
namespace fsb {
 
enum class LinalgErrorType : uint8_t
{
    ERROR_NONE = 0,
    ERROR_INPUT = 1,
    ERROR_MEMORY = 2,
    ERROR_CONVERGE = 3,
    ERROR_QUERY = 4,
    NOT_POSITIVE_DEFINITE = 5,
    SINGULAR = 6,
    ERROR_NOT_FULL_RANK = 7, 
};
 
LinalgErrorType linalg_cholesky_decomposition(
    Span<const Real> mat, size_t dim,
    Span<Real> mat_chol);
 
LinalgErrorType linalg_cholesky_solve(
    Span<const Real> mat, size_t dim,
    Span<const Real> b_vec, size_t nrhs,
    Span<Real> work,
    Span<Real> x_vec);
 
bool linalg_is_posdef(
    Span<const Real> mat, size_t dim,
    Span<Real> work);
 
LinalgErrorType linalg_matrix_sqr_solve(
    Span<const Real> mat, size_t dim,
    Span<const Real> y_vec, size_t nrhs,
    Span<Real> work,
    Span<lapack_int> iwork,
    Span<Real> x_vec);
 
LinalgErrorType linalg_pseudoinverse(
    Span<const Real> mat, size_t rows, size_t cols,
    Span<Real> work,
    Span<Real> inv_mat);
 
LinalgErrorType linalg_leastsquares_solve(
    Span<const Real> mat, size_t rows, size_t cols,
    Span<const Real> b_vec, size_t nrhs,
    Span<Real> work,
    Span<Real> x_vec);
 
LinalgErrorType linalg_svd(
    Span<const Real> mat, size_t rows, size_t cols,
    bool u_full, bool v_full,
    Span<Real> work,
    Span<Real> unitary_u,
    Span<Real> sing_val,
    Span<Real> unitary_vt);
 
LinalgErrorType linalg_matrix_eig(
    Span<const Real> mat, size_t dim,
    Span<Real> work,
    Span<Real> val_real,
    Span<Real> val_imag,
    Span<Real> vec_real,
    Span<Real> vec_imag);
 
LinalgErrorType linalg_sym_lt_eig(
    Span<const Real> mat, size_t dim,
    Span<Real> work,
    Span<Real> val,
    Span<Real> vec);
 
template <size_t MaxMat, size_t MaxWorkLen, size_t MaxU, size_t MaxS, size_t MaxVT>
inline LinalgErrorType linalg_svd_array(
    const Array<MaxMat>& mat, size_t rows, size_t cols,
    bool u_full, bool v_full,
    Array<MaxWorkLen>& work, size_t work_len,
    Array<MaxU>& unitary_u,
    Array<MaxS>& sing_val, size_t s_dim,
    Array<MaxVT>& unitary_vt)
{
    return linalg_svd(
        mat.active(rows * cols), rows, cols,
        u_full, v_full,
        work.active(work_len),
        unitary_u.active(rows * rows),
        sing_val.active(s_dim),
        unitary_vt.active(s_dim * cols));
}
 
template <size_t MaxMat, size_t MaxWorkLen, size_t MaxEig, size_t MaxVec>
inline LinalgErrorType linalg_matrix_eig_array(
    const Array<MaxMat>& mat, size_t dim,
    Array<MaxWorkLen>& work, size_t work_len,
    Array<MaxEig>& val_real,
    Array<MaxEig>& val_imag,
    Array<MaxVec>& vec_real,
    Array<MaxVec>& vec_imag)
{
    return linalg_matrix_eig(
        mat.active(dim * dim), dim,
        work.active(work_len),
        val_real.active(dim),
        val_imag.active(dim),
        vec_real.active(dim * dim),
        vec_imag.active(dim * dim));
}
 
template <size_t MaxMat, size_t MaxWorkLen, size_t MaxEig, size_t MaxVec>
inline LinalgErrorType linalg_sym_lt_eig_array(
    const Array<MaxMat>& mat, size_t dim,
    Array<MaxWorkLen>& work, size_t work_len,
    Array<MaxEig>& val,
    Array<MaxVec>& vec)
{
    return linalg_sym_lt_eig(
        mat.active(dim * dim), dim,
        work.active(work_len),
        val.active(dim),
        vec.active(dim * dim));
}
 
template <size_t MaxMat>
inline LinalgErrorType linalg_cholesky_decomposition_array(
    const Array<MaxMat>& mat, size_t dim,
    Array<MaxMat>& mat_chol)
{
    return linalg_cholesky_decomposition(
        mat.active(dim * dim), dim,
        mat_chol.active(dim * dim));
}
 
template <size_t MaxMat, size_t MaxWorkLen>
inline bool linalg_is_posdef_array(
    const Array<MaxMat>& mat, size_t dim,
    Array<MaxWorkLen>& work, size_t work_len)
{
    return linalg_is_posdef(mat.active(dim * dim), dim, work.active(work_len));
}
 
template <size_t MaxMat, size_t MaxVec, size_t MaxWorkLen>
inline LinalgErrorType linalg_cholesky_solve_array(
    const Array<MaxMat>& mat, size_t dim,
    const Array<MaxVec>& b_vec, size_t nrhs,
    Array<MaxWorkLen>& work, size_t work_len,
    Array<MaxVec>& x_vec)
{
    return linalg_cholesky_solve(
        mat.active(dim * dim), dim,
        b_vec.active(dim * nrhs), nrhs,
        work.active(work_len),
        x_vec.active(dim * nrhs));
}
 
template <size_t MaxMat, size_t MaxVec, size_t MaxWorkLen, size_t MaxIWorkLen>
inline LinalgErrorType linalg_matrix_sqr_solve_array(
    const Array<MaxMat>& mat, size_t dim,
    const Array<MaxVec>& y_vec, size_t nrhs,
    Array<MaxWorkLen>& work, size_t work_len,
    Array<MaxIWorkLen, lapack_int>& iwork, size_t iwork_len,
    Array<MaxVec>& x_vec)
{
    return linalg_matrix_sqr_solve(
        mat.active(dim * dim), dim,
        y_vec.active(dim * nrhs), nrhs,
        work.active(work_len),
        iwork.active(iwork_len),
        x_vec.active(dim * nrhs));
}
 
template <size_t MaxMat, size_t MaxInv, size_t MaxWorkLen>
inline LinalgErrorType linalg_pseudoinverse_array(
    const Array<MaxMat>& mat, size_t rows, size_t cols,
    Array<MaxWorkLen>& work, size_t work_len,
    Array<MaxInv>& inv_mat)
{
    return linalg_pseudoinverse(
        mat.active(rows * cols), rows, cols,
        work.active(work_len),
        inv_mat.active(cols * rows));
}
 
template <size_t MaxMat, size_t MaxB, size_t MaxX, size_t MaxWorkLen>
inline LinalgErrorType linalg_leastsquares_solve_array(
    const Array<MaxMat>& mat, size_t rows, size_t cols,
    const Array<MaxB>& b_vec, size_t nrhs,
    Array<MaxWorkLen>& work, size_t work_len,
    Array<MaxX>& x_vec)
{
    return linalg_leastsquares_solve(
        mat.active(rows * cols), rows, cols,
        b_vec.active(rows * nrhs), nrhs,
        work.active(work_len),
        x_vec.active(cols * nrhs));
}
 
} // namespace fsb