Sparse-Matrix Vector Multiplication

spmv.py

# MxN matrix A and N sized vector b. Ab = 

# [1, 0, 3]
# [0, 0, 0]
# [0, 2, 4]

# [1, 2, 3]

import numpy as np
cols = np.array([0, 2, 1, 2])
vals = np.array([1, 3, 2, 4])
rows = np.array([0, 2, 2, 4])
vec = np.array([1, 2, 3])

def spmv(vals, rows, cols, vec):
    results = np.empty(rows.size - 1)
    for i in range(rows.size - 1):
        # slice the corresponding column coordinates
        # and values for each row
        col_coords = cols[rows[i]:rows[i + 1]]
        data = vals[rows[i]:rows[i + 1]]
        
        # dot product with the vector
        result = 0
        for col_coord, datum in zip(col_coords, data):
            result += (vec[col_coord] * datum)
            
        results[i] = result
    
    return results

results = spmv(vals, rows, cols, vec)
print(results)
# [ 10.   0.  16.]