Skip to content

Commit

Permalink
use pytest parameterized tests
Browse files Browse the repository at this point in the history
  • Loading branch information
jcapriot committed Sep 26, 2024
1 parent 8fa774b commit dd14dc5
Show file tree
Hide file tree
Showing 5 changed files with 112 additions and 174 deletions.
5 changes: 3 additions & 2 deletions tests/test_Basic.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import numpy as np
import numpy.testing as npt
import pytest
import scipy.sparse as sp
from pymatsolver import Diagonal
Expand All @@ -20,5 +21,5 @@ def test_DiagonalSolver():
with pytest.raises(TypeError):
Diagonal(A, accuracy_tol=0)

assert np.linalg.norm(sol-X, np.inf) < TOL
assert np.linalg.norm(sol[:, 0]-x, np.inf) < TOL
npt.assert_allclose(sol, X, atol=TOL)
npt.assert_allclose(sol[:, 0], x, atol=TOL)
117 changes: 33 additions & 84 deletions tests/test_BicgJacobi.py
Original file line number Diff line number Diff line change
@@ -1,87 +1,36 @@
from pymatsolver import BicgJacobi
import numpy as np
import numpy.testing as npt
import scipy.sparse as sp

TOL = 1e-6


class TestBicgJacobi:

@classmethod
def setup_class(cls):

nSize = 100
A = sp.rand(nSize, nSize, 0.05, format='csr', random_state=100)
A = A + sp.spdiags(np.ones(nSize), 0, nSize, nSize)
A = A.T*A
A = A.tocsr()
np.random.seed(1)
sol = np.random.rand(nSize, 4)
rhs = A.dot(sol)

cls.A = A
cls.rhs = rhs
cls.sol = sol

def test(self):
rhs = self.rhs
Ainv = BicgJacobi(self.A, symmetric=True)
solb = Ainv*rhs
for i in range(3):
err = np.linalg.norm(
self.A*solb[:, i] - rhs[:, i]) / np.linalg.norm(rhs[:, i])
assert err < TOL
Ainv.clean()

def test_T(self):
rhs = self.rhs
Ainv = BicgJacobi(self.A, symmetric=True)
Ainv.maxIter = 2000
AinvT = Ainv.T
solb = AinvT*rhs
for i in range(3):
err = np.linalg.norm(
self.A.T*solb[:, i] - rhs[:, i]) / np.linalg.norm(rhs[:, i])
assert err < TOL
Ainv.clean()


class TestBicgJacobiComplex:

@classmethod
def setup_class(cls):
nSize = 100
A = sp.rand(nSize, nSize, 0.05, format='csr', random_state=100)
A.data = A.data + 1j*np.random.rand(A.nnz)
A = A.T.dot(A) + sp.spdiags(np.ones(nSize), 0, nSize, nSize)
A = A.tocsr()

np.random.seed(1)
sol = np.random.rand(nSize, 5) + 1j*np.random.rand(nSize, 5)
rhs = A.dot(sol)

cls.A = A
cls.rhs = rhs
cls.sol = sol

def test(self):
rhs = self.rhs
Ainv = BicgJacobi(self.A, symmetric=True)
solb = Ainv*rhs
for i in range(3):
err = np.linalg.norm(
self.A*solb[:, i] - rhs[:, i]) / np.linalg.norm(rhs[:, i])
assert err < TOL
Ainv.clean()

def test_T(self):
rhs = self.rhs
Ainv = BicgJacobi(self.A, symmetric=True)
Ainv.maxIter = 2000
AinvT = Ainv.T
solb = AinvT*rhs
for i in range(3):
err = np.linalg.norm(
self.A.T*solb[:, i] - rhs[:, i]) / np.linalg.norm(rhs[:, i])
assert err < TOL
Ainv.clean()
import pytest

TOL = 2e-6

@pytest.fixture()
def test_mat_data():
nSize = 100
A = sp.rand(nSize, nSize, 0.05, format='csr', random_state=100)
A = A + sp.spdiags(np.ones(nSize), 0, nSize, nSize)
A = A.T*A
A = A.tocsr()
np.random.seed(1)
sol = np.random.rand(nSize, 4)
rhs = A.dot(sol)
return A, sol, rhs


@pytest.mark.parametrize('transpose', [True, False])
@pytest.mark.parametrize('dtype', [np.float64, np.complex128])
def test_solve(test_mat_data, dtype, transpose):
A, rhs, sol = test_mat_data
A = A.astype(dtype)
rhs = rhs.astype(dtype)
sol = sol.astype(dtype)
if transpose:
A = A.T
Ainv = BicgJacobi(A, symmetric=True).T
else:
Ainv = BicgJacobi(A, symmetric=True)
Ainv.maxIter = 2000
solb = Ainv * rhs
npt.assert_allclose(rhs, A @ solb, atol=TOL)
15 changes: 8 additions & 7 deletions tests/test_Pardiso.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import pymatsolver
import numpy as np
import numpy.testing as npt
import pytest
import scipy.sparse as sp
import os
Expand Down Expand Up @@ -36,29 +37,29 @@ def test_solve(test_mat_data, dtype, transpose):
else:
Ainv = pymatsolver.Pardiso(A)
for i in range(rhs.shape[1]):
np.testing.assert_allclose(Ainv * rhs[:, i], sol[:, i], atol=TOL)
np.testing.assert_allclose(Ainv * rhs, sol, atol=TOL)
npt.assert_allclose(Ainv * rhs[:, i], sol[:, i], atol=TOL)
npt.assert_allclose(Ainv * rhs, sol, atol=TOL)

def test_symmetric_solve(test_mat_data):
A, rhs, sol = test_mat_data
Ainv = pymatsolver.Pardiso(A, is_symmetric=True)
for i in range(rhs.shape[1]):
np.testing.assert_allclose(Ainv * rhs[:, i], sol[:, i], atol=TOL)
np.testing.assert_allclose(Ainv * rhs, sol, atol=TOL)
npt.assert_allclose(Ainv * rhs[:, i], sol[:, i], atol=TOL)
npt.assert_allclose(Ainv * rhs, sol, atol=TOL)


def test_refactor(test_mat_data):
A, rhs, sol = test_mat_data
Ainv = pymatsolver.Pardiso(A, is_symmetric=True)
np.testing.assert_allclose(Ainv * rhs, sol, atol=TOL)
npt.assert_allclose(Ainv * rhs, sol, atol=TOL)

# scale rows and columns
D = sp.diags(np.random.rand(A.shape[0]) + 1.0)
A2 = D.T @ A @ D

rhs2 = A2 @ sol
Ainv.factor(A2)
np.testing.assert_allclose(Ainv * rhs2, sol, atol=TOL)
npt.assert_allclose(Ainv * rhs2, sol, atol=TOL)

def test_n_threads(test_mat_data):
A, rhs, sol = test_mat_data
Expand Down Expand Up @@ -120,4 +121,4 @@ def test_pardiso_fdem():

sol = Ainv * rhs

np.testing.assert_allclose(A @ sol, rhs, atol=TOL)
npt.assert_allclose(A @ sol, rhs, atol=TOL)
109 changes: 52 additions & 57 deletions tests/test_Scipy.py
Original file line number Diff line number Diff line change
@@ -1,77 +1,72 @@
from pymatsolver import Solver, Diagonal, SolverCG, SolverLU
import scipy.sparse as sp
import numpy as np
import numpy.testing as npt
import pytest


TOLD = 1e-10
TOLI = 1e-3
numRHS = 5

np.random.seed(77)


def dotest(MYSOLVER, multi=False, A=None, **solverOpts):
if A is None:
nx, ny, nz = 10, 10, 10
n = nx * ny * nz
Gz = sp.kron(
sp.eye(nx),
sp.kron(
sp.eye(ny),
sp.diags([-1, 1], [-1, 0], shape=(nz+1, nz))
)
@pytest.fixture()
def a_matrix():
nx, ny, nz = 10, 10, 10
n = nx * ny * nz
Gz = sp.kron(
sp.eye(nx),
sp.kron(
sp.eye(ny),
sp.diags([-1, 1], [-1, 0], shape=(nz+1, nz))
)
Gy = sp.kron(
sp.eye(nx),
sp.kron(
sp.diags([-1, 1], [-1, 0], shape=(ny+1, ny)),
sp.eye(nz),
)
)
Gy = sp.kron(
sp.eye(nx),
sp.kron(
sp.diags([-1, 1], [-1, 0], shape=(ny+1, ny)),
sp.eye(nz),
)
Gx = sp.kron(
sp.diags([-1, 1], [-1, 0], shape=(nx+1, nx)),
sp.kron(
sp.eye(ny),
sp.eye(nz),
)
)
Gx = sp.kron(
sp.diags([-1, 1], [-1, 0], shape=(nx+1, nx)),
sp.kron(
sp.eye(ny),
sp.eye(nz),
)
A = Gx.T @ Gx + Gy.T @ Gy + Gz.T @ Gz
else:
n = A.shape[0]
)
A = Gx.T @ Gx + Gy.T @ Gy + Gz.T @ Gz
return A


Ainv = MYSOLVER(A, **solverOpts)
if multi:
e = np.ones(n)
@pytest.mark.parametrize('n_rhs', [1, 5])
@pytest.mark.parametrize('solver', [Solver, SolverLU, SolverCG])
def test_solver(a_matrix, n_rhs, solver):
if solver is SolverCG:
tol = TOLI
else:
e = np.ones((n, numRHS))
rhs = A * e
tol = TOLD

n = a_matrix.shape[0]
b = np.linspace(0.9, 1.1, n)
if n_rhs > 1:
b = np.repeat(b[:, None], n_rhs, axis=-1)
rhs = a_matrix @ b

Ainv = solver(a_matrix)
x = Ainv * rhs
Ainv.clean()
return np.linalg.norm(e-x, np.inf)

npt.assert_allclose(x, b, atol=tol)

@pytest.mark.parametrize(
["solver", "multi"],
[
pytest.param(Solver, False),
pytest.param(Solver, True),
pytest.param(SolverLU, False),
pytest.param(SolverLU, True),
]
)
def test_direct(solver, multi):
assert dotest(solver, multi) < TOLD
@pytest.mark.parametrize('n_rhs', [1, 5])
def test_diag_solver(n_rhs):
n = 10
A = sp.diags(np.linspace(2, 3, n))
b = np.linspace(0.9, 1.1, n)
if n_rhs > 1:
b = np.repeat(b[:, None], n_rhs, axis=-1)
rhs = A @ b

Ainv = Diagonal(A)
x = Ainv * rhs

@pytest.mark.parametrize(
["solver", "multi", "A"],
[
pytest.param(Diagonal, False, sp.diags(np.random.rand(10)+1.0)),
pytest.param(Diagonal, True, sp.diags(np.random.rand(10)+1.0)),
pytest.param(SolverCG, False, None),
pytest.param(SolverCG, True, None),
]
)
def test_iterative(solver, multi, A):
assert dotest(solver, multi, A) < TOLI
npt.assert_allclose(x, b, atol=TOLD)
40 changes: 16 additions & 24 deletions tests/test_Triangle.py
Original file line number Diff line number Diff line change
@@ -1,31 +1,23 @@
import numpy as np
import numpy.testing as npt
import scipy.sparse as sp
import pymatsolver
import pytest

TOL = 1e-12

@pytest.mark.parametrize("solver", [pymatsolver.Forward, pymatsolver.Backward])
def test_solve(solver):
n = 50
nrhs = 20
A = sp.rand(n, n, 0.4) + sp.identity(n)
sol = np.ones((n, nrhs))
if solver is pymatsolver.Backward:
A = sp.triu(A)
else:
A = sp.tril(A)
rhs = A @ sol

class TestTriangle:

@classmethod
def setup_class(cls):
n = 50
nrhs = 20
cls.A = sp.rand(n, n, 0.4) + sp.identity(n)
cls.sol = np.ones((n, nrhs))
cls.rhsU = sp.triu(cls.A) * cls.sol
cls.rhsL = sp.tril(cls.A) * cls.sol

def test_directLower(self):
ALinv = pymatsolver.Forward(sp.tril(self.A))
X = ALinv * self.rhsL
x = ALinv * self.rhsL[:, 0]
assert np.linalg.norm(self.sol-X, np.inf) < TOL
assert np.linalg.norm(self.sol[:, 0]-x, np.inf) < TOL

def test_directLower_1(self):
AUinv = pymatsolver.Backward(sp.triu(self.A))
X = AUinv * self.rhsU
x = AUinv * self.rhsU[:, 0]
assert np.linalg.norm(self.sol-X, np.inf) < TOL
assert np.linalg.norm(self.sol[:, 0]-x, np.inf) < TOL
Ainv = solver(A)
npt.assert_allclose(Ainv * rhs, sol, atol=TOL)
npt.assert_allclose(Ainv * rhs[:, 0], sol[:, 0], atol=TOL)

0 comments on commit dd14dc5

Please sign in to comment.