benchmark.py

import numpy as onp
import mindspore.scipy as msp
import mindspore.numpy as mnp
from mindspore import context, nn, ms_function
from mindspore.ops import functional as F
from mindspore.common import Tensor
# from mindspore.scipy.sparse.linalg import IterativeGmres
from mindspore.scipy.utils import _to_tensor

onp.random.seed(0)


def rotate_vectors(H, i, cs, sn):
    x1 = H[i]
    y1 = H[i + 1]
    x2 = cs * x1 - sn * y1
    y2 = sn * x1 + cs * y1
    H[i] = x2
    H[i + 1] = y2
    return H


class GivensRotation(nn.Cell):
    """ do the Givens Rotation"""

    def __init__(self):
        super(GivensRotation, self).__init__()

    def construct(self, H_row, givens, k):
        i = 0

        while i < k:
            H_row = rotate_vectors(H_row, i, givens[i, 0], givens[i, 1])
            i = i + 1

        print(H_row)

        t = -H_row[k] / H_row[k + 1]
        givens[k, 0] = t
        givens[k, 1] = 1 / t

        R_row = rotate_vectors(H_row, k, givens[k, 0], givens[k, 1])
        return R_row, givens

    # def rotate_vectors(self, H, i, cs, sn):
    #     x1 = H[i]
    #     y1 = H[i + 1]
    #     x2 = cs * x1 - sn * y1
    #     y2 = sn * x1 + cs * y1
    #     H[i] = x2
    #     H[i + 1] = y2
    #     return H


class IterativeGmres(nn.Cell):
    """
    Implements a iterative GMRES. While building the ``restart``-dimensional
    Krylov subspace iteratively using Givens Rotation method, the algorithm
    constructs a Triangular matrix R which could be more easily solved.
    """

    def __init__(self):
        super(IterativeGmres, self).__init__()
        self.givens_rotation = GivensRotation()

    def construct(self, R, givens, restart):
        k = 0
        while k < restart:
            R[k, :], givens = self.givens_rotation(R[k, :], givens, k)
            print(givens[0, 1])
            k += 1

        return R


if __name__ == '__main__':
    preconditioner = 'identity'
    n = 5
    restart = 5

    # dtype = onp.float64
    # A = create_full_rank_matrix((n, n), dtype)
    # b = onp.random.rand(n).astype(dtype)
    # x0 = onp.zeros_like(b).astype(dtype)

    R = onp.random.random((restart, restart + 1))
    givens = onp.zeros((restart, 2), dtype=R.dtype)

    context.set_context(mode=context.PYNATIVE_MODE)
    native_x = IterativeGmres()(Tensor(R), Tensor(givens), restart)

    context.set_context(mode=context.GRAPH_MODE)
    graph_x = IterativeGmres()(Tensor(R), Tensor(givens), restart)

    onp.testing.assert_almost_equal(native_x.asnumpy(), graph_x.asnumpy(), decimal=7)