From f81e1b5c49d3f5bba92303d81e8598d3da4f77c8 Mon Sep 17 00:00:00 2001 From: Joren Hammudoglu Date: Wed, 11 Dec 2024 22:55:31 +0100 Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20`sparse.linalg`:=20complete=20`=5Fd?= =?UTF-8?q?solve`=20(#302)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../sparse/linalg/_dsolve/_superlu.pyi | 78 +++++++++++-- .../sparse/linalg/_dsolve/linsolve.pyi | 104 +++++++++++++----- scipy-stubs/sparse/linalg/_svdp.pyi | 30 ----- 3 files changed, 147 insertions(+), 65 deletions(-) delete mode 100644 scipy-stubs/sparse/linalg/_svdp.pyi diff --git a/scipy-stubs/sparse/linalg/_dsolve/_superlu.pyi b/scipy-stubs/sparse/linalg/_dsolve/_superlu.pyi index 98b8da19..3bf1989c 100644 --- a/scipy-stubs/sparse/linalg/_dsolve/_superlu.pyi +++ b/scipy-stubs/sparse/linalg/_dsolve/_superlu.pyi @@ -1,21 +1,79 @@ -from typing import Any, final +from collections.abc import Callable, Mapping +from typing import Any, Literal, TypeAlias, final, overload import numpy as np +import optype as op import optype.numpy as onp -from scipy._typing import Untyped -from scipy.sparse import csc_matrix +from scipy.sparse import csc_matrix, csr_matrix + +_Int1D: TypeAlias = onp.Array1D[np.int32] +_Float1D: TypeAlias = onp.Array1D[np.float64] +_Float2D: TypeAlias = onp.Array2D[np.float64] +_Complex1D: TypeAlias = onp.Array1D[np.complex128] +_Complex2D: TypeAlias = onp.Array2D[np.complex128] +_Inexact2D: TypeAlias = onp.Array2D[np.float32 | np.float64 | np.complex64 | np.complex128] + +### @final class SuperLU: - L: csc_matrix - U: csc_matrix + shape: tuple[int, int] nnz: int perm_r: onp.Array1D[np.intp] perm_c: onp.Array1D[np.intp] - shape: tuple[int, ...] + L: csc_matrix + U: csc_matrix + + @overload + def solve(self, /, rhs: onp.Array1D[np.integer[Any] | np.floating[Any]]) -> _Float1D: ... + @overload + def solve(self, /, rhs: onp.Array1D[np.complexfloating[Any, Any]]) -> _Complex1D: ... + @overload + def solve(self, /, rhs: onp.Array2D[np.integer[Any] | np.floating[Any]]) -> _Float2D: ... + @overload + def solve(self, /, rhs: onp.Array2D[np.complexfloating[Any, Any]]) -> _Complex2D: ... + @overload + def solve(self, /, rhs: onp.ArrayND[np.integer[Any] | np.floating[Any]]) -> _Float1D | _Float2D: ... + @overload + def solve(self, /, rhs: onp.ArrayND[np.complexfloating[Any, Any]]) -> _Complex1D | _Complex2D: ... + @overload + def solve(self, /, rhs: onp.ArrayND[np.number[Any]]) -> _Float1D | _Complex1D | _Float2D | _Complex2D: ... + +def gssv( + N: op.CanIndex, + nnz: op.CanIndex, + nzvals: _Inexact2D, + colind: _Int1D, + rowptr: _Int1D, + B: _Inexact2D, + csc: onp.ToBool = 0, + options: Mapping[str, object] = ..., +) -> tuple[csc_matrix | csr_matrix, int]: ... - def solve(self, /, rhs: onp.ArrayND[np.number[Any]]) -> onp.ArrayND[np.number[Any]]: ... +# +def gstrf( + N: op.CanIndex, + nnz: op.CanIndex, + nzvals: _Inexact2D, + colind: _Int1D, + rowptr: _Int1D, + csc_construct_func: type[csc_matrix] | Callable[..., csc_matrix], + ilu: onp.ToBool = 0, + options: Mapping[str, object] = ..., +) -> SuperLU: ... -def gssv(*args: Untyped, **kwargs: Untyped) -> Untyped: ... -def gstrf(*args: Untyped, **kwargs: Untyped) -> Untyped: ... -def gstrs(*args: Untyped, **kwargs: Untyped) -> Untyped: ... +# +def gstrs( + trans: Literal["N", "T"], + L_n: op.CanIndex, + L_nnz: op.CanIndex, + L_nzvals: _Inexact2D, + L_rowind: _Int1D, + L_colptr: _Int1D, + U_n: op.CanIndex, + U_nnz: op.CanIndex, + U_nzvals: _Inexact2D, + U_rowind: _Int1D, + U_colptr: _Int1D, + B: _Inexact2D, +) -> tuple[_Float1D | _Complex1D | _Float2D | _Complex2D, int]: ... diff --git a/scipy-stubs/sparse/linalg/_dsolve/linsolve.pyi b/scipy-stubs/sparse/linalg/_dsolve/linsolve.pyi index 957efa1b..0d63834e 100644 --- a/scipy-stubs/sparse/linalg/_dsolve/linsolve.pyi +++ b/scipy-stubs/sparse/linalg/_dsolve/linsolve.pyi @@ -1,36 +1,90 @@ -from scipy._typing import Untyped +from collections.abc import Mapping +from typing import Any, Literal, Protocol, TypeAlias, TypeVar, overload, type_check_only + +import numpy as np +import optype.numpy as onp +from scipy.sparse._base import _spbase +from ._superlu import SuperLU __all__ = ["MatrixRankWarning", "factorized", "spilu", "splu", "spsolve", "spsolve_triangular", "use_solver"] +_SparseT = TypeVar("_SparseT", bound=_spbase) + +_PermcSpec: TypeAlias = Literal["COLAMD", "NATURAL", "MMD_ATA", "MMD_AT_PLUS_A"] +_Float1D: TypeAlias = onp.Array1D[np.float64] +_Float2D: TypeAlias = onp.Array2D[np.float64] +_Complex1D: TypeAlias = onp.Array1D[np.complex128] +_Complex2D: TypeAlias = onp.Array2D[np.complex128] + +@type_check_only +class _Solve(Protocol): + @overload + def __call__(self, b: onp.Array1D[np.integer[Any] | np.floating[Any]], /) -> _Float1D: ... + @overload + def __call__(self, b: onp.Array1D[np.complexfloating[Any, Any]], /) -> _Complex1D: ... + @overload + def __call__(self, b: onp.Array2D[np.integer[Any] | np.floating[Any]], /) -> _Float2D: ... + @overload + def __call__(self, b: onp.Array2D[np.complexfloating[Any, Any]], /) -> _Complex2D: ... + @overload + def __call__(self, b: onp.ArrayND[np.integer[Any] | np.floating[Any]], /) -> _Float1D | _Float2D: ... + @overload + def __call__(self, b: onp.ArrayND[np.complexfloating[Any, Any]], /) -> _Complex1D | _Complex2D: ... + @overload + def __call__(self, b: onp.ArrayND[np.number[Any]], /) -> _Float1D | _Complex1D | _Float2D | _Complex2D: ... + +### + class MatrixRankWarning(UserWarning): ... def use_solver(*, useUmfpack: bool = ..., assumeSortedIndices: bool = ...) -> None: ... -def spsolve(A: Untyped, b: Untyped, permc_spec: Untyped | None = None, use_umfpack: bool = True) -> Untyped: ... -def splu( - A: Untyped, - permc_spec: Untyped | None = None, - diag_pivot_thresh: Untyped | None = None, - relax: Untyped | None = None, - panel_size: Untyped | None = None, - options: Untyped | None = {}, -) -> Untyped: ... -def spilu( - A: Untyped, - drop_tol: Untyped | None = None, - fill_factor: Untyped | None = None, - drop_rule: Untyped | None = None, - permc_spec: Untyped | None = None, - diag_pivot_thresh: Untyped | None = None, - relax: Untyped | None = None, - panel_size: Untyped | None = None, - options: Untyped | None = None, -) -> Untyped: ... -def factorized(A: Untyped) -> Untyped: ... +def factorized(A: _spbase | onp.ToComplex2D) -> _Solve: ... + +# +@overload +def spsolve( + A: _spbase | onp.ToComplex2D, + b: _SparseT, + permc_spec: _PermcSpec | None = None, + use_umfpack: bool = True, +) -> _SparseT: ... +@overload +def spsolve( + A: _spbase | onp.ToComplex2D, + b: onp.ToComplex2D | onp.ToComplex1D, + permc_spec: _PermcSpec | None = None, + use_umfpack: bool = True, +) -> _Float1D | _Complex1D | _Float2D | _Complex2D: ... + +# def spsolve_triangular( - A: Untyped, - b: Untyped, + A: _spbase | onp.ToComplex2D, + b: _spbase | onp.ToComplex2D | onp.ToComplex1D, lower: bool = True, overwrite_A: bool = False, overwrite_b: bool = False, unit_diagonal: bool = False, -) -> Untyped: ... +) -> _Float1D | _Complex1D | _Float2D | _Complex2D: ... + +# +def splu( + A: _spbase | onp.ToComplex2D, + permc_spec: _PermcSpec | None = None, + diag_pivot_thresh: onp.ToFloat | None = None, + relax: int | None = None, + panel_size: int | None = None, + options: Mapping[str, object] | None = {}, +) -> SuperLU: ... + +# +def spilu( + A: _spbase | onp.ToComplex2D, + drop_tol: onp.ToFloat | None = None, + fill_factor: onp.ToFloat | None = None, + drop_rule: str | None = None, + permc_spec: _PermcSpec | None = None, + diag_pivot_thresh: onp.ToFloat | None = None, + relax: int | None = None, + panel_size: int | None = None, + options: Mapping[str, object] | None = None, +) -> SuperLU: ... diff --git a/scipy-stubs/sparse/linalg/_svdp.pyi b/scipy-stubs/sparse/linalg/_svdp.pyi deleted file mode 100644 index 79d069a7..00000000 --- a/scipy-stubs/sparse/linalg/_svdp.pyi +++ /dev/null @@ -1,30 +0,0 @@ -from typing import Literal - -import optype.numpy as onp -from scipy._typing import Seed, UntypedArray -from scipy.sparse import sparray, spmatrix -from ._interface import LinearOperator - -__all__ = ["_svdp"] - -def _svdp( - A: onp.ToComplex2D | sparray | spmatrix | LinearOperator, - k: int, - which: Literal["LM", "SM"] = "LM", - irl_mode: bool = True, - kmax: int | None = None, - compute_u: bool = True, - compute_v: bool = True, - v0: onp.ToComplexND | None = None, - full_output: bool = False, - tol: float = 0, - delta: float | None = None, - eta: float | None = None, - anorm: float = 0, - cgs: bool = False, - elr: bool = True, - min_relgap: float | None = 0.002, - shifts: int | None = None, - maxiter: int | None = None, - random_state: Seed | None = None, -) -> UntypedArray | tuple[UntypedArray, ...]: ...