1

I'm trying to optimize the following code, potentially by rewriting it in Cython: it simply takes a low dimensional but relatively long numpy arrays, looks into of its columns for 0 values, and marks those as -1 in an array. The code is:

import numpy as np

def get_data():
    data = np.array([[1,5,1]] * 5000 + [[1,0,5]] * 5000 + [[0,0,0]] * 5000)
    return data

def get_cols(K):
    cols = np.array([2] * K)
    return cols

def test_nonzero(data):
    K = len(data)
    result = np.array([1] * K)
    # Index into columns of data
    cols = get_cols(K)
    # Mark zero points with -1
    idx = np.nonzero(data[np.arange(K), cols] == 0)[0]
    result[idx] = -1

import time
t_start = time.time()
data = get_data()
for n in range(5000):
    test_nonzero(data)
t_end = time.time()
print (t_end - t_start)

data is the data. cols is the array of columns of data to look for non-zero values (for simplicity, I made it all the same column). The goal is to compute a numpy array, result, which has a 1 value for each row where the column of interest is non-zero, and -1 for the rows where the corresponding columns of interest have a zero.

Running this function 5000 times on a not-so-large array of 15,000 rows by 3 columns takes about 20 seconds. Is there a way this can be sped up? It appears that most of the work goes into finding the nonzero elements and retrieving them with indices (the call to nonzero and subsequent use of its index.) Can this be optimized or is this the best that can be done? How could a Cython implementation gain speed on this?

Improve this question
1
  • The nonzero is a good try (not sure if it helps much or at all though). If you are desperate and know cols are valid, you could try making a linear index. If K is constant in the loop, you could not redo the np.arange every time too... Commented Apr 16, 2013 at 7:46

1 Answer 1

Reset to default
2 
cols = np.array([2] * K)

That's going to be really slow. That's create a very large python list and then converts it into a numpy array. Instead, do something like:

cols = np.ones(K, int)*2

That'll be way faster

result = np.array([1] * K)

Here you should do:

result = np.ones(K, int)

That will produce the numpy array directly.

idx = np.nonzero(data[np.arange(K), cols] == 0)[0]
result[idx] = -1

The cols is an array, but you can just pass a 2. Furthermore, using nonzero adds an extra step.

idx = data[np.arange(K), 2] == 0
result[idx] = -1

Should have the same effect.

Improve this answer
Sign up to request clarification or add additional context in comments.

2 Comments

thanks but cols is not part of the computation, it's just an example. and as I mentioned, cols is not always just one value like 2 -- I just made it that for simplicity. it's usually a vector will different columns. so I don't think these suggestions speed it up
@user248237dfsf, ok, I missed the part about cols not being real. However, numpy.array([2] * K) is really really slow, and its why the code as it stands is slow. Without seeing your real code I can't really guess why it would be slow.

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.