Add solve and svn to dispatched function

qutip/core/data/__init__.py

@@ -24,6 +24,7 @@
 from .reshape import *
 from .tidyup import *
 from .trace import *
+from .solve import *
 # For operations with mulitple related versions, we just import the module.
 from . import norm, permute
 

qutip/core/data/eigen.py

@@ -1,6 +1,7 @@
 import numpy as np
 import scipy.linalg
 import scipy.sparse as sp
+import scipy.sparse.linalg
 from itertools import combinations
 
 from .dense import Dense, from_csr
@@ -10,6 +11,7 @@
 
 __all__ = [
     'eigs', 'eigs_csr', 'eigs_dense',
+    'svd', 'svd_csr', 'svd_dense',
 ]
 
 
@@ -276,6 +278,8 @@ def eigs_dense(data, isherm=None, vecs=True, sort='low', eigvals=0):
 
 
 from .dispatch import Dispatcher as _Dispatcher
+import inspect as _inspect
+
 
 # We use eigs_dense as the signature source, since in this case it has the
 # complete signature that we allow, so we don't need to manually set it.
@@ -323,4 +327,127 @@ def eigs_dense(data, isherm=None, vecs=True, sort='low', eigvals=0):
     (Dense, eigs_dense),
 ], _defer=True)
 
+
+def svd_csr(data, vecs=True, k=6, **kw):
+    """
+    Singular Value Decomposition:
+
+    ``data = U @ S @ Vh``
+
+    Where ``S`` is diagonal.
+
+    Parameters
+    ----------
+    data : Data
+        Input matrix
+    vecs : bool, optional (True)
+        Whether the singular vectors (``U``, ``Vh``) should be returned.
+    k : int, optional (6)
+        Number of state to compute, default is ``6`` to match scipy's default.
+    **kw : dict
+        Options to pass to ``scipy.sparse.linalg.svds``.
+
+    Returns
+    -------
+    U : Dense
+        Left singular vectors as columns. Only returned if ``vecs == True``.
+        shape = (data.shape[0], k)
+    S : np.ndarray
+        The ``k``'s largest singular values.
+    Vh : Dense
+        Right singular vectors as rows. Only returned if ``vecs == True``.
+        shape = (k, data.shape[1])
+
+    .. note::
+        svds cannot compute all states at once. While it could find the
+        largest and smallest in 2 calls, it has issues converging with when
+        solving for the smallest (finding the 5 smallest in a 50x50 matrix
+        can fail with default options). It should be used when not all states
+        are needed.
+    """
+    out = scipy.sparse.linalg.svds(
+        data.as_scipy(), k, return_singular_vectors=vecs, **kw
+    )
+    if vecs:
+        u, s, vh = out
+        return Dense(u, copy=False), s, Dense(vh, copy=False)
+    return out
+
+
+def svd_dense(data, vecs=True, **kw):
+    """
+    Singular Value Decomposition:
+
+    ``data = U @ S @ Vh``
+
+    Where ``S`` is diagonal.
+
+    Parameters
+    ----------
+    data : Data
+        Input matrix
+    vecs : bool, optional (True)
+        Whether the singular vectors (``U``, ``Vh``) should be returned.
+    **kw : dict
+        Options to pass to ``scipy.linalg.svd``.
+
+    Returns
+    -------
+    U : Dense
+        Left singular vectors as columns. Only returned if ``vecs == True``.
+    S : np.ndarray
+        Singular values.
+    Vh : Dense
+        Right singular vectors as rows. Only returned if ``vecs == True``.
+    """
+    out = scipy.linalg.svd(
+        data.to_array(), compute_uv=vecs, **kw
+    )
+    if vecs:
+        u, s, vh = out
+        return Dense(u, copy=False), s, Dense(vh, copy=False)
+    return out
+
+
+svd = _Dispatcher(
+    _inspect.Signature([
+        _inspect.Parameter('data', _inspect.Parameter.POSITIONAL_OR_KEYWORD),
+        _inspect.Parameter('vecs', _inspect.Parameter.POSITIONAL_OR_KEYWORD),
+    ]),
+    name='svd',
+    module=__name__,
+    inputs=('data',),
+    out=False)
+svd.__doc__ =\
+    """
+    Singular Value Decomposition:
+
+    ``data = U @ S @ Vh``
+
+    Where ``S`` is diagonal.
+
+    Parameters
+    ----------
+    data : Data
+        Input matrix
+    vecs : bool, optional (True)
+        Whether the singular vectors (``U``, ``Vh``) should be returned.
+
+    Returns
+    -------
+    U : Dense
+        Left singular vectors as columns. Only returned if ``vecs == True``.
+    S : np.ndarray
+        Singular values.
+    Vh : Dense
+        Right singular vectors as rows. Only returned if ``vecs == True``.
+    """
+# Dense implementation return all states, but sparse implementation compute
+# only a few states. So only the dense version is registed.
+svd.add_specialisations([
+    (Dense, svd_dense),
+], _defer=True)
+
+
 del _Dispatcher
+del _inspect

qutip/core/data/solve.py

@@ -0,0 +1,208 @@
+from qutip.core.data import CSR, Data, csr, Dense
+import qutip.core.data as _data
+import scipy.sparse.linalg as splinalg
+import numpy as np
+from qutip.settings import settings
+if settings.has_mkl:
+    from qutip._mkl.spsolve import mkl_spsolve
+else:
+    mkl_spsolve = None
+
+
+def _splu(A, B, **kwargs):
+    lu = splinalg.splu(A, **kwargs)
+    return lu.solve(B)
+
+
+def solve_csr_dense(matrix: CSR, target: Dense, method=None,
+                    options: dict={}) -> Dense:
+    """
+    Solve ``Ax=b`` for ``x``.
+
+    Parameters:
+    -----------
+
+    matrix : CSR
+        The matrix ``A``.
+
+    target : Data
+        The matrix or vector ``b``.
+
+    method : str {"spsolve", "splu", "mkl_spsolve", etc.}, default="spsolve"
+        The function to use to solve the system. Any function from
+        scipy.sparse.linalg which solve the equation Ax=b can be used.
+        `splu` from the same and `mkl_spsolve` are also valid choice.
+
+    options : dict
+        Keywork options to pass to the solver. Refer to the documenentation in
+        scipy.sparse.linalg of the used method for a list of supported keyword.
+        The keyword "csc" can be set to ``True`` to convert the sparse matrix
+        before passing it to the solver.
+
+    .. note::
+        Options for ``mkl_spsolve`` are presently only found in the source
+        code.
+
+    Returns:
+    --------
+    x : Dense
+        Solution to the system Ax = b.
+    """
+    if matrix.shape[0] != matrix.shape[1]:
+        raise ValueError("can only solve using square matrix")
+    if matrix.shape[1] != target.shape[0]:
+        raise ValueError("target does not match the system")
+
+    b = target.as_ndarray()
+
+    method = method or "spsolve"
+
+    if method == "splu":
+        solver = _splu
+    elif hasattr(splinalg, method):
+        solver = getattr(splinalg, method)
+    elif method == "mkl_spsolve" and mkl_spsolve is None:
+        raise ValueError("mkl is not available")
+    elif method == "mkl_spsolve":
+        solver = mkl_spsolve
+    else:
+        raise ValueError(f"Unknown sparse solver {method}.")
+
+    options = options.copy()
+    M = matrix.as_scipy()
+    if options.pop("csc", False):
+        M = M.tocsc()
+
+    out = solver(M, b, **options)
+
+    if isinstance(out, tuple) and len(out) == 2:
+        # iterative method return a success flag
+        out, check = out
+        if check == 0:
+            # Successful
+            pass
+        elif check > 0:
+            raise RunTimeError(
+                f"scipy.sparse.linalg.{method} error: Tolerance was not"
+                " reached. Error code: {check}"
+            )
+
+        elif check < 0:
+            raise RunTimeError(
+                f"scipy.sparse.linalg.{method} error: Bad input. "
+                "Error code: {check}"
+            )
+    elif isinstance(out, tuple) and len(out) > 2:
+        # least sqare method return residual, flag, etc.
+        out, *_ = out
+    return Dense(out, copy=False)
+
+
+def solve_dense(matrix: Dense, target: Data, method=None,
+                options: dict={}) -> Dense:
+    """
+    Solve ``Ax=b`` for ``x``.
+
+    Parameters:
+    -----------
+
+    matrix : Dense
+        The matrix ``A``.
+
+    target : Data
+        The matrix or vector ``b``.
+
+    method : str {"solve", "lstsq"}, default="solve"
+        The function from numpy.linalg to use to solve the system.
+
+    options : dict
+        Options to pass to the solver. "lstsq" use "rcond" while, "solve" do
+        not use any.
+
+    Returns:
+    --------
+    x : Dense
+        Solution to the system Ax = b.
+    """
+    if matrix.shape[0] != matrix.shape[1]:
+        raise ValueError("can only solve using square matrix")
+    if matrix.shape[1] != target.shape[0]:
+        raise ValueError("target does not match the system")
+
+    if isinstance(target, Dense):
+        b = target.as_ndarray()
+    else:
+        b = target.to_array()
+
+    if method in ["solve", None]:
+        out = np.linalg.solve(matrix.as_ndarray(), b)
+    elif method == "lstsq":
+        out, *_ = np.linalg.lstsq(
+            matrix.as_ndarray(),
+            b,
+            rcond=options.get("rcond", None)
+        )
+    else:
+        raise ValueError(f"Unknown solver {method},"
+                         " 'solve' and 'lstsq' are supported.")
+    return Dense(out, copy=False)
+
+
+from .dispatch import Dispatcher as _Dispatcher
+import inspect as _inspect
+
+
+solve = _Dispatcher(
+    _inspect.Signature([
+        _inspect.Parameter('matrix', _inspect.Parameter.POSITIONAL_OR_KEYWORD),
+        _inspect.Parameter('target', _inspect.Parameter.POSITIONAL_OR_KEYWORD),
+        _inspect.Parameter('method', _inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                           default=None),
+        _inspect.Parameter('options', _inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                           default={}),
+    ]),
+    name='solve',
+    module=__name__,
+    inputs=('matrix', 'target'),
+    out=True,
+)
+solve.__doc__ = """
+    Solve ``Ax=b`` for ``x``.
+
+    Parameters:
+    -----------
+    matrix : Data
+    The matrix ``A``.
+
+    target : Data
+    The matrix or vector ``b``.
+
+    method : str
+        The function to use to solve the system. Function which solve the
+        equation Ax=b from scipy.sparse.linalg (CSR ``matrix``) or
+        numpy.linalg (Dense ``matrix``) can be used.
+        Sparse cases also accept `splu` and `mkl_spsolve`.
+
+    options : dict
+        Keywork options to pass to the solver. Refer to the documenentation
+        for the chosen method in scipy.sparse.linalg or numpy.linalg.
+        The keyword "csc" can be set to ``True`` to convert the sparse matrix
+        in sparse cases.
+
+    .. note::
+        Options for ``mkl_spsolve`` are presently only found in the source
+        code.
+
+    Returns:
+    --------
+    x : Data
+        Solution to the system Ax = b.
+"""
+solve.add_specialisations([
+    (CSR, Dense, Dense, solve_csr_dense),
+    (Dense, Dense, Dense, solve_dense),
+], _defer=True)
+
+
+del _Dispatcher
+del _inspect