Using scipy sparse matrices to solve system of equations

Question

This is a follow up to How to set up and solve simultaneous equations in python but I feel deserves its own reputation points for any answer.

For a fixed integer n, I have a set of 2(n-1) simultaneous equations as follows.

M(p) = 1+((n-p-1)/n)*M(n-1) + (2/n)*N(p-1) + ((p-1)/n)*M(p-1)

N(p) = 1+((n-p-1)/n)*M(n-1) + (p/n)*N(p-1)

M(1) = 1+((n-2)/n)*M(n-1) + (2/n)*N(0)

N(0) = 1+((n-1)/n)*M(n-1)

M(p) is defined for 1 <= p <= n-1. N(p) is defined for 0 <= p <= n-2. Notice also that p is just a constant integer in every equation so the whole system is linear.

Some very nice answers were given for how to set up a system of equations in python. However, the system is sparse and I would like to solve it for large n. How can I use scipy's sparse matrix representation and http://docs.scipy.org/doc/scipy/reference/sparse.linalg.html for example instead?

Answer 1

I wouldn't normally keep beating a dead horse, but it happens that my non-vectorized approach to solving your other question, has some merit when things get big. Because I was actually filling the coefficient matrix one item at a time, it is very easy to translate into COO sparse matrix format, which can efficiently be transformed to CSC and solved. The following does it:

import scipy.sparse

def sps_solve(n) :
    # Solution vector is [N[0], N[1], ..., N[n - 2], M[1], M[2], ..., M[n - 1]]
    n_pos = lambda p : p
    m_pos = lambda p : p + n - 2
    data = []
    row = []
    col = []
    # p = 0
    # n * N[0] + (1 - n) * M[n-1] = n
    row += [n_pos(0), n_pos(0)]
    col += [n_pos(0), m_pos(n - 1)]
    data += [n, 1 - n]
    for p in xrange(1, n - 1) :
        # n * M[p] + (1 + p - n) * M[n - 1] - 2 * N[p - 1] +
        #  (1 - p) * M[p - 1] = n
        row += [m_pos(p)] * (4 if p > 1 else 3)
        col += ([m_pos(p), m_pos(n - 1), n_pos(p - 1)] +
                ([m_pos(p - 1)] if p > 1 else []))
        data += [n, 1 + p - n , -2] + ([1 - p] if p > 1 else [])
        # n * N[p] + (1 + p -n) * M[n - 1] - p * N[p - 1] = n
        row += [n_pos(p)] * 3
        col += [n_pos(p), m_pos(n - 1), n_pos(p - 1)]
        data += [n, 1 + p - n, -p]
    if n > 2 :
        # p = n - 1
        # n * M[n - 1] - 2 * N[n - 2] + (2 - n) * M[n - 2] = n
        row += [m_pos(n-1)] * 3
        col += [m_pos(n - 1), n_pos(n - 2), m_pos(n - 2)]
        data += [n, -2, 2 - n]
    else :
        # p = 1 
        # n * M[1] - 2 * N[0] = n
        row += [m_pos(n - 1)] * 2
        col += [m_pos(n - 1), n_pos(n - 2)]
        data += [n, -2]
    coeff_mat = scipy.sparse.coo_matrix((data, (row, col))).tocsc()
    return scipy.sparse.linalg.spsolve(coeff_mat,
                                       np.ones(2 * (n - 1)) * n)

It is of course much more verbose than building it from vectorized blocks, as TheodorosZelleke does, but an interesting thing happens when you time both approaches:

First, and this is (very) nice, time is scaling linearly in both solutions, as one would expect from using the sparse approach. But the solution I gave in this answer is always faster, more so for larger ns. Just for the fun of it, I also timed TheodorosZelleke's dense approach from the other question, which gives this nice graph showing the different scaling of both types of solutions, and how very early, somewhere around n = 75, the solution here should be your choice:

I don't know enough about scipy.sparse to really figure out why the differences between the two sparse approaches, although I suspect heavily of the use of LIL format sparse matrices. There may be some very marginal performance gain, although a lot of compactness in the code, by turning TheodorosZelleke's answer into COO format. But that is left as an exercise for the OP!

Answer 2

This is a solution using scipy.sparse. Unfortunately the problem is not stated here. So in order to comprehend this solution, future visitors have to first look up the problem under the link provided in the question.

Solution using scipy.sparse:

from scipy.sparse import spdiags, lil_matrix, vstack, hstack
from scipy.sparse.linalg import spsolve
import numpy as np


def solve(n):
    nrange = np.arange(n)
    diag = np.ones(n-1)

    # upper left block
    n_to_M = spdiags(-2. * diag, 0, n-1, n-1)

    # lower left block
    n_to_N = spdiags([n * diag, -nrange[-1:0:-1]], [0, 1], n-1, n-1)

    # upper right block
    m_to_M = lil_matrix(n_to_N)
    m_to_M[1:, 0] = -nrange[1:-1].reshape((n-2, 1))

    # lower right block
    m_to_N = lil_matrix((n-1, n-1))
    m_to_N[:, 0] = -nrange[1:].reshape((n-1, 1))

    # build A, combine all blocks
    coeff_mat = hstack(
                       (vstack((n_to_M, n_to_N)),
                        vstack((m_to_M, m_to_N))))

    # const vector, right side of eq.
    const = n * np.ones((2 * (n-1),1))

    return spsolve(coeff_mat.tocsr(), const).reshape((-1,1))