brahmap.math.cg

A replacement of scipy.sparse.linalg.cg where np.linalg.norm is replaced with brahmap.math.parallel_norm when the parameter parallel is set True. Also all the matrices and vectors are assumed to be real.

Parameters:

Name	Type	Description	Default
A	LinearOperator	description	required
b	ndarray	description	required
x0	ndarray	description, by default None	None
atol	float	description, by default 1.0e-12	1e-12
maxiter	int	description, by default 100	100
M	LinearOperator	description, by default None	None
callback	Callable	description, by default None	None
parallel	bool	description, by default False	False

Returns:

Type	Description
_type_	description

Source code in brahmap/math/linalg.py

def cg(
    A: LinearOperator,
    b: np.ndarray,
    x0: np.ndarray = None,
    atol: float = 1.0e-12,
    maxiter: int = 100,
    M: LinearOperator = None,
    callback: Callable = None,
    parallel: bool = False,
):
    """A replacement of `scipy.sparse.linalg.cg` where `np.linalg.norm` is
    replaced with `brahmap.math.parallel_norm` when the parameter `parallel`
    is set `True`. Also all the matrices and vectors are assumed to be real.

    Parameters
    ----------
    A : LinearOperator
        _description_
    b : np.ndarray
        _description_
    x0 : np.ndarray, optional
        _description_, by default None
    atol : float, optional
        _description_, by default 1.0e-12
    maxiter : int, optional
        _description_, by default 100
    M : LinearOperator, optional
        _description_, by default None
    callback : Callable, optional
        _description_, by default None
    parallel : bool, optional
        _description_, by default False

    Returns
    -------
    _type_
        _description_
    """
    temp_tuple = scipy.sparse.linalg._isolve.utils.make_system(
        A,
        M,
        x0,
        b,
    )

    # Starting from SciPy 1.16.0, `make_system` returns 4 objects instead of 5.
    # Even in earlier versions, only the first 4 objects were needed for our
    # use. The following unpacking ensures compatibility across all versions.
    # This logic can be simplified once support for versions below 1.16.0 is
    # dropped.
    A = temp_tuple[0]
    M = temp_tuple[1]
    x = temp_tuple[2]
    b = temp_tuple[3]

    if parallel:
        norm_function: Callable = parallel_norm
    else:
        def norm_function(x): return np.sqrt(x.dot(x))

    b_norm = norm_function(b)

    if b_norm == 0:
        return b, 0

    # r = b - A@x if x has any non-zero element, otherwise r = b
    r = b - A * x if x.any() else b.copy()

    # Dummy initialization
    rho_prev, p = None, None

    norm_residual = 1.0

    for iteration in range(maxiter):
        if norm_residual < atol:
            return x, 0

        z = M * r
        rho_cur = np.dot(r, z)
        if iteration > 0:
            beta = rho_cur / rho_prev
            p *= beta
            p += z
        else:
            p = z.copy()

        q = A * p
        alpha = rho_cur / np.dot(p, q)
        x += alpha * p
        r -= alpha * q
        rho_prev = rho_cur

        norm_residual = norm_function(r) / b_norm

        if callback:
            callback(x, r, norm_residual)

    else:
        return x, maxiter