mindspore2022/tests/st/scipy_st/sparse/test_linalg.py

459 lines
19 KiB
Python

# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""st for scipy.sparse.linalg."""
import pytest
import numpy as onp
import scipy as osp
import scipy.sparse.linalg
import mindspore.ops as ops
import mindspore.nn as nn
import mindspore.scipy as msp
from mindspore import context
from mindspore.common import Tensor
from tests.st.scipy_st.utils import create_sym_pos_matrix, create_full_rank_matrix, to_tensor
def _fetch_preconditioner(preconditioner, A):
"""
Returns one of various preconditioning matrices depending on the identifier
`preconditioner' and the input matrix A whose inverse it supposedly
approximates.
"""
if preconditioner == 'identity':
M = onp.eye(A.shape[0], dtype=A.dtype)
elif preconditioner == 'random':
random_metrix = create_sym_pos_matrix(A.shape, A.dtype)
M = onp.linalg.inv(random_metrix)
elif preconditioner == 'exact':
M = onp.linalg.inv(A)
else:
M = None
return M
@pytest.mark.level0
@pytest.mark.platform_x86_gpu_training
@pytest.mark.platform_x86_cpu
@pytest.mark.env_onecard
@pytest.mark.parametrize('tensor_type, dtype, tol', [('Tensor', onp.float32, 1e-5), ('Tensor', onp.float64, 1e-12),
('CSRTensor', onp.float32, 1e-5)])
@pytest.mark.parametrize('shape', [(7, 7)])
@pytest.mark.parametrize('preconditioner', [None, 'identity', 'exact', 'random'])
@pytest.mark.parametrize('maxiter', [3, None])
def test_cg_against_scipy(tensor_type, dtype, tol, shape, preconditioner, maxiter):
"""
Feature: ALL TO ALL
Description: test cases for cg using function way in pynative/graph mode
Expectation: the result match scipy
"""
onp.random.seed(0)
a = create_sym_pos_matrix(shape, dtype)
b = onp.random.random(shape[:1]).astype(dtype)
m = _fetch_preconditioner(preconditioner, a)
osp_res = scipy.sparse.linalg.cg(a, b, M=m, maxiter=maxiter, atol=tol, tol=tol)
a = to_tensor((a, tensor_type))
b = Tensor(b)
m = to_tensor((m, tensor_type)) if m is not None else m
# using PYNATIVE MODE
context.set_context(mode=context.PYNATIVE_MODE)
msp_res_dyn = msp.sparse.linalg.cg(a, b, M=m, maxiter=maxiter, atol=tol, tol=tol)
# using GRAPH MODE
context.set_context(mode=context.GRAPH_MODE)
msp_res_sta = msp.sparse.linalg.cg(a, b, M=m, maxiter=maxiter, atol=tol, tol=tol)
kw = {"atol": tol, "rtol": tol}
onp.testing.assert_allclose(osp_res[0], msp_res_dyn[0].asnumpy(), **kw)
onp.testing.assert_allclose(osp_res[0], msp_res_sta[0].asnumpy(), **kw)
assert osp_res[1] == msp_res_dyn[1].asnumpy().item()
assert osp_res[1] == msp_res_sta[1].asnumpy().item()
@pytest.mark.level0
@pytest.mark.platform_x86_gpu_training
@pytest.mark.platform_x86_cpu
@pytest.mark.env_onecard
@pytest.mark.parametrize('dtype', [onp.float32, onp.float64])
@pytest.mark.parametrize('shape', [(2, 2)])
def test_cg_against_numpy(dtype, shape):
"""
Feature: ALL TO ALL
Description: test cases for cg
Expectation: the result match numpy
"""
onp.random.seed(0)
a = create_sym_pos_matrix(shape, dtype)
b = onp.random.random(shape[:1]).astype(dtype)
expected = onp.linalg.solve(a, b)
# using PYNATIVE MODE
context.set_context(mode=context.PYNATIVE_MODE)
actual_dyn, _ = msp.sparse.linalg.cg(Tensor(a), Tensor(b))
# using GRAPH MODE
context.set_context(mode=context.GRAPH_MODE)
actual_sta, _ = msp.sparse.linalg.cg(Tensor(a), Tensor(b))
kw = {"atol": 1e-5, "rtol": 1e-5}
onp.testing.assert_allclose(expected, actual_dyn.asnumpy(), **kw)
onp.testing.assert_allclose(expected, actual_sta.asnumpy(), **kw)
@pytest.mark.level0
@pytest.mark.platform_x86_gpu_training
@pytest.mark.platform_x86_cpu
@pytest.mark.env_onecard
@pytest.mark.parametrize('tensor_type, dtype, tol', [('Tensor', onp.float32, 1e-5), ('Tensor', onp.float64, 1e-12),
('CSRTensor', onp.float32, 1e-5)])
@pytest.mark.parametrize('shape', [(7, 7)])
@pytest.mark.parametrize('preconditioner', [None, 'identity', 'exact', 'random'])
@pytest.mark.parametrize('maxiter', [3, None])
def test_cg_against_scipy_graph(tensor_type, dtype, tol, shape, preconditioner, maxiter):
"""
Feature: ALL TO ALL
Description: test cases for cg within Cell object in pynative/graph mode
Expectation: the result match scipy
"""
class Net(nn.Cell):
def construct(self, a, b, m, maxiter, tol):
return msp.sparse.linalg.cg(a, b, M=m, maxiter=maxiter, atol=tol, tol=tol)
onp.random.seed(0)
a = create_sym_pos_matrix(shape, dtype)
b = onp.random.random(shape[:1]).astype(dtype)
m = _fetch_preconditioner(preconditioner, a)
osp_res = scipy.sparse.linalg.cg(a, b, M=m, maxiter=maxiter, atol=tol, tol=tol)
a = to_tensor((a, tensor_type))
b = Tensor(b)
m = to_tensor((m, tensor_type)) if m is not None else m
# using PYNATIVE MODE
context.set_context(mode=context.PYNATIVE_MODE)
msp_res_dyn = Net()(a, b, m, maxiter, tol)
# using GRAPH MODE
context.set_context(mode=context.GRAPH_MODE)
msp_res_sta = Net()(a, b, m, maxiter, tol)
kw = {"atol": tol, "rtol": tol}
onp.testing.assert_allclose(osp_res[0], msp_res_dyn[0].asnumpy(), **kw)
onp.testing.assert_allclose(osp_res[0], msp_res_sta[0].asnumpy(), **kw)
assert osp_res[1] == msp_res_dyn[1].asnumpy().item()
assert osp_res[1] == msp_res_sta[1].asnumpy().item()
@pytest.mark.level0
@pytest.mark.platform_x86_cpu
@pytest.mark.platform_x86_gpu_training
@pytest.mark.env_onecard
@pytest.mark.parametrize('tensor_type, dtype, tol', [('Tensor', onp.float32, 1e-5), ('Tensor', onp.float64, 1e-8),
('CSRTensor', onp.float32, 1e-5)])
@pytest.mark.parametrize('a, b, grad_a, grad_b', [
([[1.96822833, 0.82204467, 1.03749232, 0.88915326, 0.44986806, 1.11167143],
[0.82204467, 2.25216591, 1.40235719, 0.70838919, 0.81377919, 1.06000368],
[1.03749232, 1.40235719, 2.90618746, 0.7126087, 0.81029544, 1.28673025],
[0.88915326, 0.70838919, 0.7126087, 2.17515263, 0.40443765, 1.02082996],
[0.44986806, 0.81377919, 0.81029544, 0.40443765, 1.60570668, 0.62292701],
[1.11167143, 1.06000368, 1.28673025, 1.02082996, 0.62292701, 2.30795277]],
[0.79363745, 0.58000418, 0.1622986, 0.70075235, 0.96455108, 0.50000836],
[[-0.07867674, -0.01521201, 0.06394698, -0.03854052, -0.13523701, 0.01326866],
[-0.03508505, -0.00678363, 0.02851647, -0.01718673, -0.06030749, 0.00591702],
[-0.00586019, -0.00113306, 0.00476305, -0.00287067, -0.01007304, 0.00098831],
[-0.07704304, -0.01489613, 0.06261914, -0.03774023, -0.13242886, 0.01299314],
[-0.14497008, -0.02802971, 0.11782896, -0.07101491, -0.24918826, 0.02444888],
[-0.01868565, -0.00361284, 0.01518735, -0.00915334, -0.03211867, 0.00315129]],
[0.22853142, 0.10191113, 0.01702201, 0.22378603, 0.42109291, 0.054276]),
([[1.85910724, 0.73233206, 0.65960803, 1.03821349, 0.55277616],
[0.73233206, 1.69548841, 0.59992146, 1.01518264, 0.50824059],
[0.65960803, 0.59992146, 1.98169091, 1.45565213, 0.47901749],
[1.03821349, 1.01518264, 1.45565213, 3.3133049, 0.75598147],
[0.55277616, 0.50824059, 0.47901749, 0.75598147, 1.46831254]],
[0.59674531, 0.226012, 0.10694568, 0.22030621, 0.34982629],
[[-0.07498642, 0.00167461, 0.01353184, 0.01008293, -0.03770084],
[-0.09940184, 0.00221986, 0.01793778, 0.01336592, -0.04997616],
[-0.09572781, 0.00213781, 0.01727477, 0.01287189, -0.04812897],
[0.03135044, -0.00070012, -0.00565741, -0.00421549, 0.01576203],
[-0.14053766, 0.00313851, 0.02536103, 0.01889718, -0.07065797]],
[0.23398106, 0.31016481, 0.29870068, -0.09782316, 0.43852141]),
])
def test_cg_grad(tensor_type, dtype, tol, a, b, grad_a, grad_b):
"""
Feature: ALL TO ALL
Description: test cases for grad implementation of cg in graph mode
Expectation: the result match expectation
"""
context.set_context(mode=context.GRAPH_MODE)
a = to_tensor((a, tensor_type), dtype)
b = Tensor(onp.array(b, dtype=dtype))
expect_grad_a = onp.array(grad_a, dtype=dtype)
expect_grad_b = onp.array(grad_b, dtype=dtype)
kw = {"atol": tol, "rtol": tol}
# Function
grad_net = ops.GradOperation(get_all=True)(msp.sparse.linalg.cg)
grad_a, grad_b = grad_net(a, b)[:2]
onp.testing.assert_allclose(expect_grad_a, grad_a.asnumpy(), **kw)
onp.testing.assert_allclose(expect_grad_b, grad_b.asnumpy(), **kw)
# Cell
class Net(nn.Cell):
def __init__(self):
super(Net, self).__init__()
self.sum = ops.ReduceSum()
self.cg = msp.sparse.linalg.cg
def construct(self, a, b):
x, _ = self.cg(a, b)
return self.sum(x)
grad_net = ops.GradOperation(get_all=True)(Net())
grad_a, grad_b = grad_net(a, b)[:2]
onp.testing.assert_allclose(expect_grad_a, grad_a.asnumpy(), **kw)
onp.testing.assert_allclose(expect_grad_b, grad_b.asnumpy(), **kw)
@pytest.mark.level0
@pytest.mark.platform_x86_cpu
@pytest.mark.platform_x86_gpu_training
@pytest.mark.env_onecard
@pytest.mark.parametrize('tensor_type, dtype, tol', [('Tensor', onp.float32, 1e-5), ('Tensor', onp.float64, 1e-8)])
@pytest.mark.parametrize('a, b, grad_a, grad_b', [
([[1.96822833, 0.82204467, 1.03749232, 0.88915326, 0.44986806, 1.11167143],
[0.82204467, 2.25216591, 1.40235719, 0.70838919, 0.81377919, 1.06000368],
[1.03749232, 1.40235719, 2.90618746, 0.7126087, 0.81029544, 1.28673025],
[0.88915326, 0.70838919, 0.7126087, 2.17515263, 0.40443765, 1.02082996],
[0.44986806, 0.81377919, 0.81029544, 0.40443765, 1.60570668, 0.62292701],
[1.11167143, 1.06000368, 1.28673025, 1.02082996, 0.62292701, 2.30795277]],
[0.79363745, 0.58000418, 0.1622986, 0.70075235, 0.96455108, 0.50000836],
[[-0.07867674, -0.01521201, 0.06394698, -0.03854052, -0.13523701, 0.01326866],
[-0.03508505, -0.00678363, 0.02851647, -0.01718673, -0.06030749, 0.00591702],
[-0.00586019, -0.00113306, 0.00476305, -0.00287067, -0.01007304, 0.00098831],
[-0.07704304, -0.01489613, 0.06261914, -0.03774023, -0.13242886, 0.01299314],
[-0.14497008, -0.02802971, 0.11782896, -0.07101491, -0.24918826, 0.02444888],
[-0.01868565, -0.00361284, 0.01518735, -0.00915334, -0.03211867, 0.00315129]],
[0.22853142, 0.10191113, 0.01702201, 0.22378603, 0.42109291, 0.054276]),
([[1.85910724, 0.73233206, 0.65960803, 1.03821349, 0.55277616],
[0.73233206, 1.69548841, 0.59992146, 1.01518264, 0.50824059],
[0.65960803, 0.59992146, 1.98169091, 1.45565213, 0.47901749],
[1.03821349, 1.01518264, 1.45565213, 3.3133049, 0.75598147],
[0.55277616, 0.50824059, 0.47901749, 0.75598147, 1.46831254]],
[0.59674531, 0.226012, 0.10694568, 0.22030621, 0.34982629],
[[-0.07498642, 0.00167461, 0.01353184, 0.01008293, -0.03770084],
[-0.09940184, 0.00221986, 0.01793778, 0.01336592, -0.04997616],
[-0.09572781, 0.00213781, 0.01727477, 0.01287189, -0.04812897],
[0.03135044, -0.00070012, -0.00565741, -0.00421549, 0.01576203],
[-0.14053766, 0.00313851, 0.02536103, 0.01889718, -0.07065797]],
[0.23398106, 0.31016481, 0.29870068, -0.09782316, 0.43852141]),
])
def test_cg_grad_pynative(tensor_type, dtype, tol, a, b, grad_a, grad_b):
"""
Feature: ALL TO ALL
Description: test cases for grad implementation of cg in pynative mode
Expectation: the result match expectation
"""
context.set_context(mode=context.PYNATIVE_MODE)
a = to_tensor((a, tensor_type), dtype)
b = Tensor(onp.array(b, dtype=dtype))
expect_grad_a = onp.array(grad_a, dtype=dtype)
expect_grad_b = onp.array(grad_b, dtype=dtype)
kw = {"atol": tol, "rtol": tol}
# Function
grad_net = ops.GradOperation(get_all=True)(msp.sparse.linalg.cg)
grad_a, grad_b = grad_net(a, b)[:2]
onp.testing.assert_allclose(expect_grad_a, grad_a.asnumpy(), **kw)
onp.testing.assert_allclose(expect_grad_b, grad_b.asnumpy(), **kw)
# Cell
class Net(nn.Cell):
def __init__(self):
super(Net, self).__init__()
self.sum = ops.ReduceSum()
self.cg = msp.sparse.linalg.cg
def construct(self, a, b):
x, _ = self.cg(a, b)
return self.sum(x)
grad_net = ops.GradOperation(get_all=True)(Net())
grad_a, grad_b = grad_net(a, b)[:2]
onp.testing.assert_allclose(expect_grad_a, grad_a.asnumpy(), **kw)
onp.testing.assert_allclose(expect_grad_b, grad_b.asnumpy(), **kw)
@pytest.mark.level0
@pytest.mark.platform_x86_cpu
@pytest.mark.platform_x86_gpu_training
@pytest.mark.env_onecard
@pytest.mark.parametrize('n', [3, 5, 7])
@pytest.mark.parametrize('dtype,tol', [(onp.float64, 7), (onp.float32, 3)])
@pytest.mark.parametrize('preconditioner', [None, 'identity', 'exact', 'random'])
def test_gmres_incremental_against_scipy(n, tol, dtype, preconditioner):
"""
Feature: ALL TO ALL
Description: test cases for [N x N] X [N X 1]
Expectation: the result match scipy
"""
onp.random.seed(0)
context.set_context(mode=context.PYNATIVE_MODE)
A = create_full_rank_matrix((n, n), dtype)
b = onp.random.rand(n).astype(dtype)
x0 = onp.zeros_like(b).astype(dtype)
M = _fetch_preconditioner(preconditioner, A)
scipy_x, _ = osp.sparse.linalg.gmres(A, b, x0, tol=1e-07, atol=0, M=M)
A = Tensor(A)
b = Tensor(b)
x0 = Tensor(x0)
if M is not None:
M = Tensor(M)
gmres_x, _ = msp.sparse.linalg.gmres(A, b, x0, tol=1e-07, atol=0, solve_method='incremental', M=M)
onp.testing.assert_almost_equal(scipy_x, gmres_x.asnumpy(), decimal=tol)
@pytest.mark.level0
@pytest.mark.platform_x86_cpu
@pytest.mark.platform_x86_gpu_training
@pytest.mark.env_onecard
@pytest.mark.parametrize('n', [3, 5, 7])
@pytest.mark.parametrize('dtype, tol', [(onp.float64, 7), (onp.float32, 3)])
@pytest.mark.parametrize('preconditioner', [None, 'identity', 'exact', 'random'])
def test_gmres_incremental_against_scipy_graph(n, tol, dtype, preconditioner):
"""
Feature: ALL TO ALL
Description: test cases for [N x N] X [N X 1]
Expectation: the result match scipy
"""
onp.random.seed(0)
context.set_context(mode=context.GRAPH_MODE)
A = create_full_rank_matrix((n, n), dtype)
b = onp.random.rand(n).astype(dtype)
x0 = onp.zeros_like(b).astype(dtype)
M = _fetch_preconditioner(preconditioner, A)
scipy_x, _ = osp.sparse.linalg.gmres(A, b, x0, tol=1e-07, atol=0, M=M)
A = Tensor(A)
b = Tensor(b)
x0 = Tensor(x0)
if M is not None:
M = Tensor(M)
gmres_x, _ = msp.sparse.linalg.gmres(A, b, x0, tol=1e-07, atol=0, solve_method='incremental', M=M)
onp.testing.assert_almost_equal(scipy_x, gmres_x.asnumpy(), decimal=tol)
@pytest.mark.level0
@pytest.mark.platform_x86_gpu_training
@pytest.mark.platform_x86_cpu
@pytest.mark.env_onecard
@pytest.mark.parametrize('n', [4, 5, 6])
@pytest.mark.parametrize('dtype, tol', [(onp.float64, 7), (onp.float32, 3)])
@pytest.mark.parametrize('preconditioner', [None, 'identity', 'exact', 'random'])
@pytest.mark.parametrize('maxiter', [1, 2])
def test_pynative_batched_gmres_against_scipy(n, dtype, tol, preconditioner, maxiter):
"""
Feature: ALL TO ALL
Description: test cases for gmres
Expectation: the result match scipy
"""
onp.random.seed(0)
context.set_context(mode=context.PYNATIVE_MODE)
shape = (n, n)
a = create_full_rank_matrix(shape, dtype)
b = onp.random.rand(n).astype(dtype=dtype)
M = _fetch_preconditioner(preconditioner, a)
tensor_a = Tensor(a)
tensor_b = Tensor(b)
M = Tensor(M) if M is not None else M
osp_x, _ = osp.sparse.linalg.gmres(a, b, maxiter=maxiter, atol=1e-6)
msp_x, _ = msp.sparse.linalg.gmres(tensor_a, tensor_b, maxiter=maxiter, M=M, atol=1e-6,
solve_method='batched')
onp.testing.assert_almost_equal(msp_x.asnumpy(), osp_x, decimal=tol)
@pytest.mark.level0
@pytest.mark.platform_x86_gpu_training
@pytest.mark.platform_x86_cpu
@pytest.mark.env_onecard
@pytest.mark.parametrize('n', [5, 6])
@pytest.mark.parametrize('dtype, tol', [(onp.float64, 7), (onp.float32, 3)])
@pytest.mark.parametrize('preconditioner', [None, 'identity', 'exact', 'random'])
@pytest.mark.parametrize('maxiter', [1, 2])
def test_graph_batched_gmres_against_scipy(n, dtype, tol, preconditioner, maxiter):
"""
Feature: ALL TO ALL
Description: test cases for gmres
Expectation: the result match scipy
"""
onp.random.seed(0)
context.set_context(mode=context.GRAPH_MODE)
shape = (n, n)
a = create_full_rank_matrix(shape, dtype)
b = onp.random.rand(n).astype(dtype=dtype)
tensor_a = Tensor(a)
tensor_b = Tensor(b)
M = _fetch_preconditioner(preconditioner, a)
M = Tensor(M) if M is not None else M
osp_x, _ = osp.sparse.linalg.gmres(a, b, maxiter=maxiter, atol=0.0)
msp_x, _ = msp.sparse.linalg.gmres(tensor_a, tensor_b, maxiter=maxiter, M=M, atol=0.0, solve_method='batched')
onp.testing.assert_almost_equal(msp_x.asnumpy(), osp_x, decimal=tol)
@pytest.mark.level0
@pytest.mark.platform_x86_gpu_training
@pytest.mark.platform_x86_cpu
@pytest.mark.env_onecard
@pytest.mark.parametrize('dtype_tol', [(onp.float64, 1e-10)])
@pytest.mark.parametrize('shape', [(4, 4), (7, 7)])
@pytest.mark.parametrize('preconditioner', [None, 'identity', 'exact', 'random'])
@pytest.mark.parametrize('maxiter', [1, 3])
def test_bicgstab_against_scipy(dtype_tol, shape, preconditioner, maxiter):
"""
Feature: ALL TO ALL
Description: test cases for bicgstab
Expectation: the result match scipy
"""
onp.random.seed(0)
dtype, tol = dtype_tol
A = create_full_rank_matrix(shape, dtype)
b = onp.random.random(shape[:1]).astype(dtype)
M = _fetch_preconditioner(preconditioner, A)
osp_res = scipy.sparse.linalg.bicgstab(A, b, M=M, maxiter=maxiter, atol=tol, tol=tol)[0]
A = Tensor(A)
b = Tensor(b)
M = Tensor(M) if M is not None else M
# using PYNATIVE MODE
context.set_context(mode=context.PYNATIVE_MODE)
msp_res_dyn = msp.sparse.linalg.bicgstab(A, b, M=M, maxiter=maxiter, atol=tol, tol=tol)[0]
# using GRAPH MODE
context.set_context(mode=context.GRAPH_MODE)
msp_res_sta = msp.sparse.linalg.bicgstab(A, b, M=M, maxiter=maxiter, atol=tol, tol=tol)[0]
kw = {"atol": tol, "rtol": tol}
onp.testing.assert_allclose(osp_res, msp_res_dyn.asnumpy(), **kw)
onp.testing.assert_allclose(osp_res, msp_res_sta.asnumpy(), **kw)