I have the following function in pure python:
import numpy as np
def subtractPython(a, b):
xAxisCount = a.shape[0]
yAxisCount = a.shape[1]
shape = (xAxisCount, yAxisCount, xAxisCount)
results = np.zeros(shape)
for index in range(len(b)):
subtracted = (a - b[index])
results[:, :, index] = subtracted
return results
I tried to cythonize it this way:
import numpy as np
cimport numpy as np
DTYPE = np.int
ctypedef np.int_t DTYPE_t
def subtractPython(np.ndarray[DTYPE_t, ndim=2] a, np.ndarray[DTYPE_t, ndim=2] b):
cdef int xAxisCount = a.shape[0]
cdef int yAxisCount = a.shape[1]
cdef np.ndarray[DTYPE_t, ndim=3] results = np.zeros([xAxisCount, yAxisCount, xAxisCount], dtype=DTYPE)
cdef int lenB = len(b)
cdef np.ndarray[DTYPE_t, ndim=2] subtracted
for index in range(lenB):
subtracted = (a - b[index])
results[:, :, index] = subtracted
return results
However, Im not seeing any speedup. Is there something I'm missing or this process can't be sped up?
EDIT -> I've realized that I'm not actually cythonizing the subtraction algorithm in the above code. I've managed to cythonize it, but it has the exact same runtime as a - b[:, None], so I guess this is the maximum speed of this operation.
This is basically a - b[:, None] -> has same runtime
%%cython
import numpy as np
cimport numpy as np
DTYPE = np.int
ctypedef np.int_t DTYPE_t
cimport cython
@cython.boundscheck(False) # turn off bounds-checking for entire function
@cython.wraparound(False) # turn off negative index wrapping for entire function
def subtract(np.ndarray[DTYPE_t, ndim=2] a, np.ndarray[DTYPE_t, ndim=2] b):
cdef np.ndarray[DTYPE_t, ndim=3] result = np.zeros([b.shape[0], a.shape[0], a.shape[1]], dtype=DTYPE)
cdef int lenB = b.shape[0]
cdef int lenA = a.shape[0]
cdef int lenColB = b.shape[1]
cdef int rowA, rowB, column
for rowB in range(lenB):
for rowA in range(lenA):
for column in range(lenColB):
result[rowB, rowA, column] = a[rowA, column] - b[rowB, column]
return result
results = a[:,:,np.newaxis]-bgets the same output.index- it's possible Cython can't guess it). However, this might just be how long the calculation takes.