In this demo we will look at the cost of a sparse matrix-vector multipy ($A * v$) for different sparse matrix formats.
from pyamg import gallery
import numpy as np
from time import time
import matplotlib.pyplot as plt
%matplotlib inline
Set up the tests to do ntrials
of $w \leftarrow A * v$ at different sizes in nlist
with nnzperrow
non-zeros per row on average
types = ['csr', 'coo', 'csc', 'lil']
k = np.arange(1, 10)
nv = np.array([1e2, 1e3])
nlist = np.kron(nv, k)
ntrials = 10
nnzperrow = 50
times = {a: np.zeros((len(nlist), ntrials)) for a in types}
A = {a: [] for a in types}
nnz = []
for j, n in enumerate(nlist):
print("n = %d" % n)
tmp = time()
A['csr'] = gallery.sprand(n, n, float(nnzperrow) / n, format='csr')
tmp = time() - tmp
print(" setup: %g" % tmp)
v = np.random.rand(n,)
w = 0 * v.copy()
for tp in types:
if tp is 'csr':
nnz.append(A[tp].nnz)
if tp is not 'csr':
A[tp] = getattr(A['csr'], 'to' + tp)()
for i in range(ntrials):
tmp = time()
w = A[tp] * v
tmp = time() - tmp
times[tp][j, i] = tmp
print(" mat-vec %s: %g" % (tp, tmp))
nnz = np.array(nnz)
plt.figure(figsize=(12,12))
plt.hold(True)
for tp in types:
plt.loglog(nlist, times[tp].min(axis=1), label=tp, lw=3)
plt.loglog(nlist, nnz / nnz[0] * times[types[0]][0, 0], label='nnz')
plt.legend()
plt.show()