forked from abacusmodeling/abacus-develop
-
Notifications
You must be signed in to change notification settings - Fork 247
Expand file tree
/
Copy pathtest_hsolver.py
More file actions
74 lines (61 loc) · 2.4 KB
/
Copy pathtest_hsolver.py
File metadata and controls
74 lines (61 loc) · 2.4 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
from __future__ import annotations
import pytest
from pyabacus import hsolver
import numpy as np
import scipy
def diag_pyabacus(h_sparse, nband, method):
dav = {
'dav_subspace': hsolver.dav_subspace,
'davidson': hsolver.davidson
}
cg = {
'cg': hsolver.cg
}
def mm_op(x):
return h_sparse.dot(x)
nbasis = h_sparse.shape[0]
v0 = np.random.rand(nbasis, nband)
diag_elem = h_sparse.diagonal()
diag_elem = np.where(np.abs(diag_elem) < 1e-8, 1e-8, diag_elem)
precond = 1.0 / np.abs(diag_elem)
if method in dav:
algo = dav[method]
args = (mm_op, v0, nbasis, nband, precond, 8, 1e-12, 5000)
elif method in cg:
algo = cg[method]
args = (mm_op, v0, nbasis, nband, precond, 1e-12, 5000)
else:
raise ValueError(f"Method {method} not available")
e, _ = algo(*args)
return e
def diag_eigsh(h_sparse, nband):
e, _ = scipy.sparse.linalg.eigsh(h_sparse, k=nband, which='SA', maxiter=5000, tol=1e-12)
return e
# Check if CG is available
_cg_available = hsolver.cg_available()
@pytest.mark.parametrize("method", [
('dav_subspace'),
('davidson'),
pytest.param('cg', marks=pytest.mark.skipif(not _cg_available, reason="CG requires ATen support"))
])
def test_random_matrix_diag(method):
np.random.seed(12)
n = 500
h_sparse = np.random.rand(n,n)
h_sparse = h_sparse + h_sparse.conj().T + np.diag(np.random.random(n))*10
e_pyabacus = diag_pyabacus(h_sparse, 8, method)
e_scipy = diag_eigsh(h_sparse, 8)
np.testing.assert_allclose(e_pyabacus, e_scipy, atol=1e-8)
@pytest.mark.parametrize("file_name, nband, atol, method", [
('./test_diag/Si2.mat', 16, 1e-8, 'dav_subspace'),
('./test_diag/Si2.mat', 16, 1e-8, 'davidson'),
pytest.param('./test_diag/Si2.mat', 16, 1e-8, 'cg', marks=pytest.mark.skipif(not _cg_available, reason="CG requires ATen support")),
('./test_diag/Na5.mat', 16, 1e-8, 'dav_subspace'),
('./test_diag/Na5.mat', 16, 1e-8, 'davidson'),
pytest.param('./test_diag/Na5.mat', 16, 1e-8, 'cg', marks=pytest.mark.skipif(not _cg_available, reason="CG requires ATen support")),
])
def test_diag(file_name, nband, atol, method):
h_sparse = scipy.io.loadmat(file_name)['Problem']['A'][0, 0]
e_pyabacus = diag_pyabacus(h_sparse, nband, method)
e_scipy = diag_eigsh(h_sparse, nband)
np.testing.assert_allclose(e_pyabacus, e_scipy, atol=atol)