Notebooks · lib/sparse

Sparse matrices in compressed sparse row form

A matrix with few nonzeros per row is stored as three arrays, the compressed sparse row form that NPB CG, SciPy's csr_matrix and most iterative solvers use. vals[k] is the k-th nonzero, row by row and left to right within a row, and cols[k] is its column. rowptr[i] is where row i starts in vals, so row i is vals[rowptr[i]], ..., vals[rowptr[i + 1] - 1], and rowptr[n] is the number of nonzeros.

rowptr and cols are Int arrays, so an index is never a rounded double.

csr_matvec(rowptr, cols, vals, x, y): y = A x, written into y. Each row's sum runs left to right, the order NPB CG's Fortran uses, so the result is the reference's to the bit. Rows are independent, so the loop is a parallel for: each pass writes only y[i], which the compiler checks.
csr_from_dense(M): the CSR arrays of a dense matrix, dropping zeros.
csr_to_dense(rowptr, cols, vals, ncols): the dense matrix back, the check member the tests compare against @.

Used with from sparse import csr_matvec.

Look inside this program

Atlas
The program’s map: every function and which calls which.
There’s no recorded run of this program yet.

Shift+Enter in a cell runs it and moves to the next. Running a cell runs every cell above it too, as one program.

[1]lines 25–81

def csr_matvec(rowptr: [Int], cols: [Int], vals: Arr, x: Arr, y: Arr) {
    rows: Int = len(rowptr) - 1
    parallel for i to rows {
        s = 0.0
        for k from rowptr[i] to rowptr[i + 1] {
            s += vals[k] * x[cols[k]]
        }
        y[i] = s
    }
}

def csr_from_dense(M) {
    n: Int = int(shape(M)[0])
    m: Int = int(shape(M)[1])
    nnz: Int = int(sum(M != 0.0))
    rowptr = int_zeros(n + 1)
    cols = int_zeros(nnz)
    vals = zeros(nnz)
    k: Int = 0
    for i to n {
        rowptr[i] = k
        for j to m {
            if M[i, j] != 0.0 {
                cols[k] = j
                vals[k] = M[i, j]
                k += 1
            }
        }
    }
    rowptr[n] = k
    return rowptr, cols, vals
}

def csr_to_dense(rowptr: [Int], cols: [Int], vals: Arr, ncols) {
    n = len(rowptr) - 1
    M = zeros(n, ncols)
    for i to n {
        for k from rowptr[i] to rowptr[i + 1] {
            M[i, cols[k]] += vals[k]
        }
    }
    return M
}

test "csr_matvec is A @ x for a small matrix with an empty row" {
    A = [[4.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 0.0], [2.0, -1.0, 0.0, 3.0]]
    rowptr, cols, vals = csr_from_dense(A)
    assert_true(rowptr[3] == 5)
    x = [1.0, 2.0, 3.0, 4.0]
    y = zeros(3)
    csr_matvec(rowptr, cols, vals, x, y)
    assert_allclose(y, A @ x, 0.0)
    assert_allclose(csr_to_dense(rowptr, cols, vals, 4), A, 0.0)
}

# SPDX-License-Identifier: MIT OR Apache-2.0
# SPDX-FileCopyrightText: 2026 The Viv authors

def csr_matvec(rowptr: [Int], cols: [Int], vals: Arr, x: Arr, y: Arr) {
    rows: Int = len(rowptr) - 1
    parallel for i to rows {
        s = 0.0
        for k from rowptr[i] to rowptr[i + 1] {
            s += vals[k] * x[cols[k]]
        }
        y[i] = s
    }
}

def csr_from_dense(M) {
    n: Int = int(shape(M)[0])
    m: Int = int(shape(M)[1])
    nnz: Int = int(sum(M != 0.0))
    rowptr = int_zeros(n + 1)
    cols = int_zeros(nnz)
    vals = zeros(nnz)
    k: Int = 0
    for i to n {
        rowptr[i] = k
        for j to m {
            if M[i, j] != 0.0 {
                cols[k] = j
                vals[k] = M[i, j]
                k += 1
            }
        }
    }
    rowptr[n] = k
    return rowptr, cols, vals
}

def csr_to_dense(rowptr: [Int], cols: [Int], vals: Arr, ncols) {
    n = len(rowptr) - 1
    M = zeros(n, ncols)
    for i to n {
        for k from rowptr[i] to rowptr[i + 1] {
            M[i, cols[k]] += vals[k]
        }
    }
    return M
}

test "csr_matvec is A @ x for a small matrix with an empty row" {
    A = [[4.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 0.0], [2.0, -1.0, 0.0, 3.0]]
    rowptr, cols, vals = csr_from_dense(A)
    assert_true(rowptr[3] == 5)
    x = [1.0, 2.0, 3.0, 4.0]
    y = zeros(3)
    csr_matvec(rowptr, cols, vals, x, y)
    assert_allclose(y, A @ x, 0.0)
    assert_allclose(csr_to_dense(rowptr, cols, vals, 4), A, 0.0)
}

# SPDX-License-Identifier: MIT OR Apache-2.0
# SPDX-FileCopyrightText: 2026 The Viv authors