direct_solver#

Solve \(a @ x = b\) for \(x\). This function-form API is a wrapper around the stateful DirectSolver object APIs and is meant for single use (the user needs to perform just one sparse direct solve, for example), in which case there is no possibility of amortizing preparatory costs.

Detailed information on what’s happening within this function can be obtained by passing in a logging.Logger object to DirectSolverOptions or by setting the appropriate options in the root logger object, which is used by default:

>>> import logging
>>> logging.basicConfig(
...     level=logging.INFO,
...     format="%(asctime)s %(levelname)-8s %(message)s",
...     datefmt="%m-%d %H:%M:%S",
... )

A user can select the desired logging level and, in general, take advantage of all of the functionality offered by the Python logging module.

Parameters:

a – The sparse operand (or sequence of operands) representing the left-hand side (LHS) of the system of equations. The LHS operand may be a (sequence of) scipy.sparse.csr_matrix, scipy.sparse.csr_array, cupyx.scipy.sparse.csr_matrix, or torch.sparse_csr_tensor(). That is, the LHS is a sparse matrix or tensor in Compressed Sparse Row (CSR) format from one of the supported packages: SciPy, CuPy, PyTorch. Refer to the semantics section for details.
b – The ndarray/tensor or (sequence of ndarray/tensors) representing the dense right-hand side (RHS) of the system of equations. The RHS operand may be a (sequence of) numpy.ndarray, cupy.ndarray, and torch.Tensor. Refer to the semantics section for details.
options – Specify options for the direct sparse solver as a DirectSolverOptions object. Alternatively, a dict containing the parameters for the DirectSolverOptions constructor can also be provided. If not specified, the value will be set to the default-constructed DirectSolverOptions object.
execution – Specify execution space options for the direct solver as a ExecutionCUDA or ExecutionHybrid object. Alternatively, a string (‘cuda’ or ‘hybrid’), or a dict with the ‘name’ key set to ‘cuda’ or ‘hybrid’ and optional parameters relevant to the given execution space. The default execution space is ‘cuda’ and the corresponding ExecutionCUDA object will be default-constructed.
stream – Provide the CUDA stream to use for executing the operation. Acceptable inputs include cudaStream_t (as Python int), cupy.cuda.Stream, and torch.cuda.Stream. If a stream is not provided, the current stream from the operand package will be used.

Returns:

The result of the specified sparse direct solve, which has the same shape, remains on the same device, and belongs to the same package as the RHS b. If b is a sequence, the result x is also a sequence of ndarray/tensor, each of which has the same shape as the corresponding tensor in b.

Semantics:

The sparse direct solver solves \(a @ x = b\) for x given the left-hand side (LHS) a and the right-hand side (RHS) b.

In the simplest version with no batching, a is sparse (square) matrix of size n in Compressed Sparse Row (CSR) format, and b is a dense vector or matrix. A matrix (2D ndarray or tensor) in the RHS is treated as multiple column vectors corresponding to multiple solution vectors (i.e. x is the same shape as the RHS) to solve for.

Important

Currently, only column-major (Fortran) layout is supported for the RHS.
Batching can be specified either explicitly or implicitly.
- An explicitly-batched LHS is provided as a Python sequence of sparse CSR matrices \([a_0, a_1, ..., a_n]\). Likewise an explicitly-batched RHS is provided as a sequence of vectors or matrices \([b_0, b_1, ..., b_n]\). The solver will solve all \(n\) systems \(a_i @ x_i = b_i\) for the solution sequence \(x_i\). Each sample in explicit batching can be of a different size, with the only constraint being that a given LHS size is consistent with that of its corresponding RHS.
- An implicitly-batched LHS is provided as a higher-dimensional \(N \geq 3D\) sparse tensor in CSR format, where the leading \(N - 2\) dimensions are the batch dimensions, and the last two dimensions correspond to that of the \(n \times n\) sparse system. Currently, only PyTorch supports higher-dimensional sparse CSR tensors. Likewise, an implicitly-batched RHS is provided as a \(N \geq 3D\) dense ndarray/tensor, where the leading \(N - 2\) dimensions are the batch dimensions, and the last two dimensions correspond to the \(n \times 1\) vector or \(n \times m\) matrix for each sample. The solver solves \(a_i @ x_i = b_i\) for each sample \(i\) in the batch, and the solution x has the same shape as the RHS b.
- Each sample \(a_i\) and \(b_i\) in the (explicitly- or implicitly-specified) batch are essentially treated as, and subject to, the same rules as the case with no batching.
- The LHS and RHS batch specification is independent: for example, the LHS can be explicitly-batched while the RHS is implicitly-batched (or vice-versa). The same batch specification can be used for both as well.
- The solution x always has the same form as the RHS b. It is a sequence of matrices or vectors if b is explicitly-batched, or a higher-dimensional ndarray/tensor if b is implicitly-batched.

Examples

>>> import cupy as cp
>>> import cupyx.scipy.sparse as sp
>>> import nvmath

Create a sparse float32 ndarray in CSR format on the CPU for the LHS.

>>> n = 16
>>> a = sp.random(n, n, density=0.5, format="csr", dtype="float32")

Ensure that the randomly-generated LHS is not singular.

>>> a += sp.diags([2.0] * n, format="csr", dtype="float32")

The RHS can be a vector or matrix. Here we create a random matrix with 4 columns (indicating 4 vectors to be solved for) in column-major format.

>>> b = cp.random.rand(4, n, dtype="float32").T

Solve a @ x = b for x.

>>> x = nvmath.sparse.advanced.direct_solver(a, b)

Batching can be specified, explicitly or implicitly, following the semantics described above. Here we explicitly batch the LHS since CuPy doesn’t support 3D CSR, while we implicitly batch the RHS.

Create an explicit batch of two CSR matrices:

>>> batch = 2
>>> a = [a] * batch
>>> a[1] *= 10.0

Create a 3D ndarray, with each sample in the batch having column-major layout.

>>> b = cp.random.rand(batch, 4, n, dtype="float32").transpose(0, 2, 1)

Solve the batched system a @ x = b for x, where x has the same shape as b.

>>> x = nvmath.sparse.advanced.direct_solver(a, b)

Options can be provided to the sparse direct solver using DirectSolverOptions, and the execution space can be specified using the execution option. Refer to DirectSolver and the GitHub link below for examples.

Notes

This function is a convenience wrapper around DirectSolver and is specifically meant for single use.

Further examples can be found in the nvmath/examples/sparse/advanced/direct_solver directory.