NumPy Scipy optimization I have an optimization issue and I'm not sure if I can improve the overall speed of my function.

The function draw_w is my actual implementation that gives me the right update for my vector w. You can see that the code is running for 1000 iterations in 9s.
I had a first implementation (draw_w_fw) that was quiet good because, well-vectorized, that is an order of magnitude quicker by 0.8s. Unfortunately, this implementation is wrong because the cache[0] is not updated correctly.

I would like to know if someone knows a way to speed up draw_w from what I've wrote.

import scipy.sparse as sps
import numpy as np
import timeit


def draw_w(X, w, cache):

    X = X.tocsr()

    for i in xrange(w.shape[0]):
        x_li = X.getrow(i)
        Y = w[i] * x_li
        Y.data -= np.take(cache[0], Y.indices) 
        h = x_li.multiply(-Y)

        w_mean = h.sum()
        w_sigma_sqr = (x_li.multiply(x_li)).sum() 

        w_sigma_sqr = 1.0 / w_sigma_sqr
        w_mean = - w_sigma_sqr * w_mean

        # update w:
        w_old = np.copy(w[i])

        if np.isinf(w[i]):
            w[i] = 0
        elif np.isnan(w[i]):
            w[i] = 0
        else:
            w[i] = w_mean

        # update error:
        cache[0] -= (w_old - w[i]) * x_li

    #print 'draw_w', w
    #True result w [ 2.34125626  2.37726446  4.00792293  3.71059779  4.00792293  0.11100713
    # -0.28899287 -0.04393113  0.21429929]

####################


def draw_w_fw(X, w, cache):

    x_li = X.tocsr()
    nnz_per_row = np.diff(x_li.indptr)
    Y = sps.csr_matrix((x_li.data * np.repeat(w, nnz_per_row), x_li.indices, x_li.indptr), shape=x_li.shape)
    Y.data -= np.take(cache[0], Y.indices)   #not good because cache[0] is updated...
    h = x_li.multiply(-Y)

    w_mean = np.asarray(h.sum(axis=1).transpose())[0]
    w_sigma_sqr = np.asarray(( x_li.multiply(x_li) ).sum(axis=1).transpose())[0]

    w_sigma_sqr = 1.0 / w_sigma_sqr
    w_mean = - w_sigma_sqr * w_mean

    # update w:
    w_old = np.copy(w)

    w[~np.isnan(w) &amp; ~np.isinf(w)] = w_mean[~np.isnan(w) &amp; ~np.isinf(w)]

    w[np.isinf(w)] = 0.0
    w[np.isnan(w)] = 0.0

    # update error:
    cache[0] -= (w_old - w) * x_li

    #print 'draw_w_fw', w

##########################


def test():

    data = [ 1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.]
    cols = [ 0,0,1,1,2,2,3,3,4,4,5,5,6,6,7,7,8,8,9,9,10,10,11,11,12,12,13,13,14,14]
    rows = [0,5,1,5,2,5,3,5,4,5,0,6,1,6,2,6,3,6,4,6,1,7,3,7,0,8,2,8,4,8]

    X = sps.coo_matrix((data, (rows, cols)), shape = (9,15))
    X = X.tocsr()

    w = np.array([ 0.16243454, -0.06117564, -0.05281718, -0.10729686,  0.08654076, -0.23015387,  0.17448118, -0.07612069,  0.03190391])

    cache = np.zeros((2, 15))
    cache[0] =  np.asarray([-5.06771933, -5.29132951, -4.28297104, -1.33745073, -2.14361311, -0.66308429,
     -0.88669446, -2.878336,  -4.93281569, -4.73897806, -1.13729633, -5.18341755,
     -0.80566155, -5.02091327, -4.88155533])
    #draw_w(X, w, cache)
    draw_w_fw(X, w, cache)

if __name__ == '__main__':
    print(timeit.timeit("test()", setup="from __main__ import test", number=1000))
    #draw_w:    9.26 for 1000 iterations
    #draw_w_fw: 0.80 for 1000 iterations
Before streamlining the code, it almost always pays off to streamline the math. The h you calculate is the dot product of two vectors, X[i] and Y[i] = cache[0] - w[i]*X[i], so we can rewrite it as

h = X[i].cache[0] - w[i]*X[i].X[i]


When you go on to calculate w_mean, you are actually negating h and then dividing it by a number that turns out to be X[i].X[i], so you can calculate w_mean as

w_mean = w[i] - X[i].cache[0] / X[i].X[i]


Setting aside the np.isnan and np.isinf checks you are doing, you then update cache[0] with a value where w[i] cancels out, and you are basically doing:

cache[0] -= X[i] * (X[i].cache[0] / X[i].X[i])


Putting it all together into a single function, I got this very compact code:

from numpy.lib.stride_tricks import as_strided

def draw_w_bis(X, w, cache):
    X = X.tocsr()
    x_rows_sqr = np.add.reduceat(X.data*X.data, X.indptr[:-1])
    rows, cols = X.shape
    row_start_stop = as_strided(X.indptr, shape=(rows, 2),
                                strides=2*X.indptr.strides)
    for row, (start, stop) in enumerate(row_start_stop):
        data = X.data[start:stop]
        cols = X.indices[start:stop]
        delta = np.dot(data, cache[0, cols]) / x_rows_sqr[row]
        w[row] -= delta
        cache[0, cols] -= delta * data


If you are not comfortable with calling as_strided you could replace the definition of row_start_stop with:

row_start_stop = np.column_stack((X.indptr[:-1], X.indptr[1:]))


On my system your test code runs the 1000 iterations over an order of magnitude faster (6 sec vs. 0.4 sec) which seems to be what you were after. There still is a for loop in there, so if you are going to run this on large arrays, it is still going to be slow, and you may want to consider using things like Cython.