diff --git a/src/cmdtools/analysis/pcca.py b/src/cmdtools/analysis/pcca.py index 5f9afdd..19fe5b5 100644 --- a/src/cmdtools/analysis/pcca.py +++ b/src/cmdtools/analysis/pcca.py @@ -128,8 +128,7 @@ def krylovschur(A, n, massmatrix=None, onseperation="continue"): from petsc4py import PETSc from slepc4py import SLEPc - M = PETSc.Mat().create() - M.createDense(list(np.shape(A)), array=A) + M = petsc_matrix(A) E = SLEPc.EPS().create() E.setOperators(M) E.setDimensions(nev=n) @@ -137,3 +136,17 @@ def krylovschur(A, n, massmatrix=None, onseperation="continue"): E.solve() X = np.column_stack([x.array for x in E.getInvariantSubspace()]) return X[:, :n] + + +def petsc_matrix(A): + from scipy import sparse + from petsc4py import PETSc + from slepc4py import SLEPc + + M = PETSc.Mat() + if sparse.isspmatrix_csr(A): + nrows = np.size(A, 0) + M.createAIJWithArrays(nrows, (A.indptr, A.indices, A.data)) + else: + M.createDense(list(np.shape(A)), array=A) + return M diff --git a/tests/test_pcca.py b/tests/test_pcca.py index d844b1f..f74fd36 100644 --- a/tests/test_pcca.py +++ b/tests/test_pcca.py @@ -51,22 +51,48 @@ def test_krylovschur(n=30, m=5, N=100): assert R == m -def test_bench_scipyschur(benchmark, n=500, m=6): +N_BENCHMARK = 1000 +M_BENCHMARK = 10 + + +def test_bench_scipyschur(benchmark, n=N_BENCHMARK, m=M_BENCHMARK): T = utils.randompropagator(n) massmatrix = np.diag(np.ones(n)) solver = pcca.ScipySchur() benchmark(solver.solve, T, m, massmatrix) -def test_bench_scipyqz(benchmark, n=500, m=6): +def test_bench_scipyqz(benchmark, n=N_BENCHMARK, m=M_BENCHMARK): T = utils.randompropagator(n) solver = pcca.ScipySchur() benchmark(solver.solve, T, m) -def test_bench_krylovschur(benchmark, n=500, m=6): +def test_bench_krylovschur(benchmark, n=N_BENCHMARK, m=M_BENCHMARK): if not pcca.HAS_SLEPC: return T = utils.randompropagator(n) solver = pcca.KrylovSchur(onseperation="continue") benchmark(solver.solve, T, m) + + +def test_bench_krylovschursparse_dense(benchmark, n=N_BENCHMARK, m=M_BENCHMARK): + if not pcca.HAS_SLEPC: + return + from scipy import sparse + T = utils.randompropagator(n) + T = sparse.csr_matrix(T) + solver = pcca.KrylovSchur(onseperation="continue") + benchmark(solver.solve, T, m) + + +def test_bench_krylovschursparse_sparse(benchmark, n=N_BENCHMARK, m=M_BENCHMARK, p=0.01): + if not pcca.HAS_SLEPC: + return + from scipy import sparse + T = sparse.random(n, n, p) + T = np.array(np.identity(n) + T) + T = utils.rowstochastic(T) + T = sparse.csr_matrix(T) + solver = pcca.KrylovSchur(onseperation="continue") + benchmark(solver.solve, T, m)