# -*- coding: utf-8 -*-
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at https://mozilla.org/MPL/2.0/.
# SPDX-License-Identifier: MPL-2.0
import numpy as np
from enum import Enum
from typing import Union
from collections.abc import Callable
from scipy.sparse import csr_matrix, csc_matrix
from VeraGridEngine.basic_structures import Vec, Mat
from VeraGridEngine.enumerations import SparseSolver
# list of available linear algebra frameworks
available_sparse_solvers = list()
try:
import cvxopt
from cvxoptklu import klu
spsolve = klu.linsolve
available_sparse_solvers.append(SparseSolver.KLU)
except ImportError:
pass
# print(SparseSolver.KLU.value + ' failed')
try:
from scipy.sparse.linalg import spsolve as scipy_spsolve, splu, spilu, gmres, spsolve_triangular
available_sparse_solvers.append(SparseSolver.UMFPACK) # default linsolve solver
available_sparse_solvers.append(SparseSolver.ILU)
available_sparse_solvers.append(SparseSolver.SuperLU)
available_sparse_solvers.append(SparseSolver.GMRES)
available_sparse_solvers.append(SparseSolver.UMFPACKTriangular)
except ImportError:
pass
# print(SparseSolver.BLAS_LAPACK.value + ' failed')
try:
from pypardiso import spsolve as pardiso_spsolve
available_sparse_solvers.append(SparseSolver.Pardiso) # pypardiso
except ImportError:
pass
# print(SparseSolver.Pardiso.value + ' failed')
preferred_type = SparseSolver.SuperLU
if preferred_type not in available_sparse_solvers:
if len(available_sparse_solvers) > 0:
preferred_type = available_sparse_solvers[0]
# print('Falling back to', preferred_type)
else:
raise Exception('No linear algebra solver!!!! VeraGrid cannot work without one.')
# print('Using', preferred_type)
[docs]
def get_sparse_type(solver_type: SparseSolver = preferred_type):
"""
GEt sparse matrix type matching the selected sparse linear systems solver
:param solver_type:
:return: sparse matrix type
"""
if solver_type in [SparseSolver.Pardiso, SparseSolver.GMRES]:
return csr_matrix
elif solver_type in [SparseSolver.KLU, SparseSolver.SuperLU, SparseSolver.ILU, SparseSolver.UMFPACK]:
return csc_matrix
else:
raise Exception('Unknown solver' + str(solver_type))
[docs]
def super_lu_linsolver(A: csc_matrix, b: Union[Vec, Mat]) -> Union[Vec, Mat]:
"""
SuperLU wrapper function for linear system solve A x = b
:param A: System matrix
:param b: right hand side
:return: solution
"""
return splu(A).solve(b)
[docs]
def ilu_linsolver(A: csc_matrix, b: Union[Vec, Mat]) -> Union[Vec, Mat]:
"""
ILU wrapper function for linear system solve A x = b
:param A: System matrix
:param b: right hand side
:return: solution
"""
return spilu(A).solve(b)
[docs]
def klu_linsolve(A: csc_matrix, b: Union[Vec, Mat]) -> Union[Vec, Mat]:
"""
KLU wrapper function for linear system solve A x = b
:param A: System matrix
:param b: right hand side
:return: solution
"""
A2 = A.tocoo()
A_cvxopt = cvxopt.spmatrix(A2.data, A2.row, A2.col, A2.shape, 'd')
x = cvxopt.matrix(b)
klu.linsolve(A_cvxopt, x)
return np.array(x)[:, 0]
[docs]
def gmres_linsolve(A: csc_matrix, b: Union[Vec, Mat]) -> Union[Vec, Mat]:
"""
:param A:
:param b:
:return:
"""
x, info = gmres(A, b)
return x
[docs]
def get_linear_solver(solver_type: SparseSolver = preferred_type) -> Callable[[csc_matrix, Union[Vec, Mat]], Union[Vec, Mat]]:
"""
Privide the chosen linear solver_type function pointer to
solver_type linear systems of the type A x = b, with x = f(A,b)
:param solver_type: SparseSolver option
:return: function pointer f(A, b)
"""
if solver_type in available_sparse_solvers:
if solver_type == SparseSolver.UMFPACK:
return scipy_spsolve
elif solver_type == SparseSolver.UMFPACKTriangular:
return spsolve_triangular
elif solver_type == SparseSolver.KLU:
return klu_linsolve
elif solver_type == SparseSolver.SuperLU:
return super_lu_linsolver
elif solver_type == SparseSolver.Pardiso:
return pardiso_spsolve
elif solver_type == SparseSolver.ILU:
return ilu_linsolver
elif solver_type == SparseSolver.GMRES:
return gmres_linsolve
else:
raise Exception('Unrecognized LU solver')
else:
return scipy_spsolve