#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Tools for sparse matrix, not present or ineficient for my usage in scipy
Created on Wed Feb 28 09:49:32 2024
@author: nchevaug
"""
from typing import cast
import numpy as np
import scipy as sp
import scipy.sparse.linalg as spla
from scipy.linalg import cho_solve, cho_factor, lu_factor, lu_solve
from scipy.linalg import lapack
from pyxfem.core.config import HAVE_UMFPACK
if HAVE_UMFPACK:
from scikits import umfpack
#pylint: disable=invalid-name
#USE_UMFPACK = True
spla.use_solver(useUmfpack = HAVE_UMFPACK)
[docs]
def factorized(A, assume_a="gen", overwrite_a=False):
"""Returns a function solve(b) using pre-factored dense LAPACK routines.
or scipy.sparse.factorized
assume_a : 'gen' (LU), 'pos' (Cholesky), or 'sym' (LDL^T)
"""
if sp.sparse.issparse(A):
if HAVE_UMFPACK:
A = sp.sparse.csc_matrix(A, dtype = float)
# fall back since umfpack.splu is not fully compatible with csc_array
lu = umfpack.splu(A)
def solve(b):
return lu.solve(b)
return solve
return spla.factorized(A)
if assume_a in ["pos", "positive"]:
# Cholesky LL^T
c, low = cho_factor(A, lower=True, overwrite_a=overwrite_a)
return lambda b: cho_solve((c, low), b)
elif assume_a in ["sym", "symmetric"]:
# Pre-factors symmetric matrix A using LAPACK dsytrf
# and returns a fast solve(b) function.
# dsytrf requires Fortran-contiguous array, and float64 (or float32 for ssytrf)
A_f = np.asfortranarray(A, dtype=np.float64)
#uplo = 1 if lower else 0
uplo =1
# 1. Factorize once (O(N^3))
(sytrf,) = lapack.get_lapack_funcs(('sytrf',), (A,))
ldl_packed, ipiv, info = sytrf(A_f, lower=uplo, overwrite_a = overwrite_a)
if info != 0:
raise np.linalg.LinAlgError(f"dsytrf failed with info={info} (singular matrix)")
# 2. Return direct solve closure (O(N^2))
def ldl_solve(b):
b_f = np.asfortranarray(b, dtype=np.float64)
(sytrs,) = lapack.get_lapack_funcs(('sytrs',), (A,))
x, solve_info = sytrs(ldl_packed, ipiv, b_f,
lower=uplo, overwrite_a = overwrite_a)
if solve_info != 0:
raise np.linalg.LinAlgError(f"dsytrs failed with info={solve_info}")
return x
return ldl_solve
else:
# General LU
lu, piv = lu_factor(A)
return lambda b: lu_solve((lu, piv), b)
[docs]
def regular_block_diag_to_csr(block_diag):
''' from an np.array of shape (nblock, sblock, sblock) representing a block diagonal
matrix A of shape (nblock*sblock) with block size sblock
constuct the csr_array representing A.
'''
fun_name = regular_block_diag_to_csr.__name__
if len(block_diag.shape) != 3:
raise ValueError('Error in '+fun_name+' input block_diag must be 3d.')
nblock = block_diag.shape[0] #number of diagonal block
xblock = block_diag.shape[1] #size of the block
yblock = block_diag.shape[2] #size of the block
lin, col, shape = block_diag_compute_index(nblock, xblock, yblock)
return sp.sparse.csr_array((block_diag.flatten(), (lin, col)), shape)
[docs]
def regular_block_diag_to_csc(block_diag):
''' from an np.array of shape (nblock, sblock, sblock) representing a block diagonal
matrix A of shape (nblock*sblock) with block size sblock
constuct the csr_array representing A.
'''
fun_name = regular_block_diag_to_csc.__name__
if len(block_diag.shape) != 3:
raise ValueError('Error in '+fun_name+' input block_diag must be 3d.')
nblock = block_diag.shape[0] #number of diagonal block
xblock = block_diag.shape[1] #size of the block
yblock = block_diag.shape[2] #size of the block
lin, col, shape = block_diag_compute_index(nblock, xblock, yblock)
return sp.sparse.csc_array((block_diag.flatten(), (lin, col)), shape)
[docs]
def block_diag_compute_index(nblock, xblock, yblock):
""" compute index of block diag matrix in order to pass to csc_array """
col = np.repeat(np.arange(yblock*nblock).reshape((-1, yblock)), xblock, axis=0).flatten()
lin = np.arange(xblock*nblock).repeat(yblock)
shape = (xblock*nblock, yblock*nblock)
return lin, col, shape
[docs]
def vblock(x0, x1):
return [x0, x1]
[docs]
def split_vblock(x):
return x[0], x[1]
[docs]
def split_mblock(A):
A00 = A[0][0]
A01 = A[0][1]
A10 = A[1][0]
A11 = A[1][1]
return A00, A01, A10, A11
[docs]
def mblock(A00,A01, A10, A11):
return [[A00, A01], [A10, A11]]
[docs]
def block_norm(x):
return np.linalg.norm(np.concatenate([x[0], x[1]]))
[docs]
def block_isub(x,y):
x[0] -= y[0]
x[1] -= y[1]
[docs]
def block_linsolve(A, b, solveA00 = None):
A00, A01, A10, A11 = split_mblock(A)
b0, b1 = split_vblock(b)
if solveA00 is None:
solveA00 = factorized(A00)
x0, x1 = linsolve_schurr(solveA00, A01, A10, A11, b0, b1)
return vblock(x0, x1)
[docs]
def linsolve_schurr(solveA11, A12, A21, A22, b1, b2):
''' Solve block linear system using shcurr complement '''
A12dense = A12 if isinstance(A12, np.ndarray) else A12.todense()
K = A22 - A21@solveA11(A12dense)
y = b2 - A21@solveA11(b1)
x2 = sp.linalg.solve(K, y, lower=False, overwrite_a=True,
overwrite_b=True, check_finite=False,
assume_a=None, transposed=False)
x1 = solveA11(b1 - A12@x2)
return x1, x2
[docs]
def rcond_estimate(A, luA = None):
''' return an estimate of the reciproque condition number
need A and an lu factorization of A.
Warning, if luA is not given, it will be computed (and destroyed ..)
inside the function
'''
if luA is None :
luA = spla.splu(A)
norm_A = spla.norm(A, 1)
class LUInverseOperator(spla.LinearOperator):
"""Opérateur linéaire représentant l'inverse d'une factorisation LU."""
def __init__(self, luA, shape, dtype):
super().__init__(dtype=dtype, shape=shape)
self.luA = luA
def _matvec(self, x):
return self.luA.solve(x)
def _rmatvec(self, x):
return self.luA.solve(x, trans='T')
A_inv_op = LUInverseOperator(luA, shape=A.shape, dtype=A.dtype)
norm_A_inv = cast(float,spla.onenormest(A_inv_op))
cond = norm_A * norm_A_inv
return 1./cond