CODEWITHSUNDEEP
pythonoptimizationnumpyscipycython
user248237
user248237

Reputation:

optimizing indexing and retrieval of elements in numpy arrays in Python?

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?

Upvotes: 1

Views: 900

Answers (1)

Winston Ewert
Winston Ewert

Reputation: 45059

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.

Upvotes: 2

Related Questions