Optimize function slicing numpy arrays

Question

I have the following function, which takes a numpy array of floats and an integer as its arguments. Each row in the array 'counts' is the result of some experiment, and I want to randomly draw a list of the experiments and add them up, then repeat this process to create lots of samples groups.

def my_function(counts,nSamples):
    ''' Create multiple randomly drawn (with replacement)
        samples from the raw data '''
    nSat,nRegions = counts.shape
    sampleData = np.zeros((nSamples,nRegions))
    for i in range(nSamples):
        rc = np.random.randint(0,nSat,size=nSat)
        sampleData[i] = counts[rc].sum(axis=0)
    return sampleData

This function seems quite slow, typically counts has around 100,000 rows (and 4 columns) and nSamples is around 2000. I have tried using numba and implicit for loops to try and speed up this code with no success. What are some other methods to try and increase the speed?

I have run cProfile on the function and got the following output.

8005 function calls in 60.208 seconds

Ordered by: standard name

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)

    1    0.000    0.000   60.208   60.208 <string>:1(<module>)

 2000    0.010    0.000   13.306    0.007 _methods.py:31(_sum)

    1   40.950   40.950   60.208   60.208 optimize_bootstrap.py:25(bootstrap)

    1    0.000    0.000    0.000    0.000 {method 'disable' of '_lsprof.Profiler' objects}

 2000    5.938    0.003    5.938    0.003 {method 'randint' of 'mtrand.RandomState' objects}

 2000   13.296    0.007   13.296    0.007 {method 'reduce' of 'numpy.ufunc' objects}

 2000    0.015    0.000   13.321    0.007 {method 'sum' of 'numpy.ndarray' objects}

    1    0.000    0.000    0.000    0.000 {numpy.core.multiarray.zeros}

    1    0.000    0.000    0.000    0.000 {range}

Answer 1

Are you sure that

rc = np.random.randint(0,nSat,size=nSat)

is what you want, instead of size=someconstant ? Otherwise you're summing over all the rows with many repeats.

---

edit does it help to replace the slicing altogether with a matrix product:

rcvec=np.zeros(nSat,np.int) for i in rc: rcvec[i]+=1 sampleData[i] = rcvec.dot(counts)

(maybe there is a function in numpy that can give you rcvec faster)

Answer 2

Simply generate all indices in one go with a 2D size for np.random.randint , use those to index into counts array and then sum along the first axis, just like you were doing with the loopy one.

Thus, one vectorized way and as such faster one, would be like so -

RC = np.random.randint(0,nSat,size=(nSat, nSamples))
sampleData_out = counts[RC].sum(axis=0)