1"""benchmarks for the scipy.sparse.linalg._expm_multiply module"""
2import math
3
4import numpy as np
5from .common import Benchmark, safe_import
6
7with safe_import():
8    import scipy.linalg
9    from scipy.sparse.linalg import expm_multiply
10
11
12def random_sparse_csr(m, n, nnz_per_row):
13    # Copied from the scipy.sparse benchmark.
14    rows = np.arange(m).repeat(nnz_per_row)
15    cols = np.random.randint(0, n, size=nnz_per_row*m)
16    vals = np.random.random_sample(m*nnz_per_row)
17    M = scipy.sparse.coo_matrix((vals, (rows, cols)), (m, n), dtype=float)
18    return M.tocsr()
19
20
21def random_sparse_csc(m, n, nnz_per_row):
22    # Copied from the scipy.sparse benchmark.
23    rows = np.arange(m).repeat(nnz_per_row)
24    cols = np.random.randint(0, n, size=nnz_per_row*m)
25    vals = np.random.random_sample(m*nnz_per_row)
26    M = scipy.sparse.coo_matrix((vals, (rows, cols)), (m, n), dtype=float)
27    # Use csc instead of csr, because sparse LU decomposition
28    # raises a warning when I use csr.
29    return M.tocsc()
30
31
32class ExpmMultiply(Benchmark):
33    def setup(self):
34        self.n = 2000
35        self.i = 100
36        self.j = 200
37        nnz_per_row = 25
38        self.A = random_sparse_csr(self.n, self.n, nnz_per_row)
39
40    def time_expm_multiply(self):
41        # computing only column', j, 'of expm of the sparse matrix
42        v = np.zeros(self.n, dtype=float)
43        v[self.j] = 1
44        A_expm_col_j = expm_multiply(self.A, v)
45        A_expm_col_j[self.i]
46
47
48class Expm(Benchmark):
49    params = [
50        [30, 100, 300],
51        ['sparse', 'dense']
52    ]
53    param_names = ['n', 'format']
54
55    def setup(self, n, format):
56        np.random.seed(1234)
57
58        # Let the number of nonzero entries per row
59        # scale like the log of the order of the matrix.
60        nnz_per_row = int(math.ceil(math.log(n)))
61
62        # time the sampling of a random sparse matrix
63        self.A_sparse = random_sparse_csc(n, n, nnz_per_row)
64
65        # first format conversion
66        self.A_dense = self.A_sparse.toarray()
67
68    def time_expm(self, n, format):
69        if format == 'sparse':
70            scipy.linalg.expm(self.A_sparse)
71        elif format == 'dense':
72            scipy.linalg.expm(self.A_dense)
73