Using NumPy in the correct way to apply a function to specific points in an array

Question

cArr is an array of the form:

cArr=np.array([[0,x0,y0,z0,1],[1,x1,y1,z1,1]])

The middle three numbers of each row represent the coordinates of two points, (points 0 and 1 for reference) in 3D. The first and last values in each row are used by other functions.

cEDA is a 4D array of shape (100,100,100,4), filled with essentially random numbers.

For each point 0 and 1 there is a neighbourhood of points around each point where the coordinates differs by only one, in the 6 cardinal directions.

x0,y0,z0
x0+1,y0,z0
x0,y0+1,z0
x0,y0,z0+1
x0-1,y0,z0
x0,y0-1,z0
x0,y0,z0-1

The distance between these 7 points relating to point 0 and point 1 (so the distance between point 1 and the neighbourhood of points around point 0 and point 0 itself) are calculated and combined with the corresponding value in the -2th (100,100,100) array in cEDA and placed in the ith 100,100,100 cube in place in cEDA where i is point 0 or 1 respectively.

I am new to python from a C background and am struggling to get out of the loop mindset. I have tried precalculating distances and summing arrays but is much slower but the quickest way so far is going through each individual value and assigning it. The function is run MANY times and is a bottleneck in the code and am sure there is a better way to do it.

I conceptually understand and can use NumPy but cannot seem to apply it in this case - any ideas on how to speed this function up by using NumPy is the correct way? I've tried my best to explain the functionality of this function sorry I know it's confusing!

def cF2(cArr, cEDA):
    # the coordinates of points 0 and 1 for readability
    x0 = cArr[0,1]
    y0 = cArr[0,2]
    z0 = cArr[0,3]
    x1 = cArr[1,1]
    y1 = cArr[1,2]
    z1 = cArr[1,3]

    # for each point around point 0 and the point itself, calculate the distance between this point and point 1
    # use this value and the corresponding value in cEDA(x,y,z,-2) and place result in  cEDA(x,y,z,0)
    cEDA[x0,y0,z0,0] = cEDA[x0,y0,z0,-2]-0.4799/(np.linalg.norm([x0,y0,z0]-cArr[1,1:4]))
    cEDA[x0-1,y0,z0,0] = cEDA[x0-1,y0,z0,-2]-0.4799/(np.linalg.norm([x0-1,y0,z0]-cArr[1,1:4]))
    cEDA[x0+1,y0,z0,0] = cEDA[x0+1,y0,z0,-2]-0.4799/(np.linalg.norm([x0+1,y0,z0]-cArr[1,1:4]))
    cEDA[x0,y0-1,z0,0] = cEDA[x0,y0-1,z0,-2]-0.4799/(np.linalg.norm([x0,y0-1,z0]-cArr[1,1:4]))
    cEDA[x0,y0+1,z0,0] = cEDA[x0,y0+1,z0,-2]-0.4799/(np.linalg.norm([x0,y0+1,z0]-cArr[1,1:4]))
    cEDA[x0,y0,z0-1,0] = cEDA[x0,y0,z0-1,-2]-0.4799/(np.linalg.norm([x0,y0,z0-1]-cArr[1,1:4]))
    cEDA[x0,y0,z0+1,0] = cEDA[x0,y0,z0+1,-2]-0.4799/(np.linalg.norm([x0,y0,z0+1]-cArr[1,1:4]))

    cEDA[x1,y1,z1,1] = cEDA[x1,y1,z1,-2]+0.4799/(np.linalg.norm([x1,y1,z1]-cArr[0,1:4]))
    cEDA[x1-1,y1,z1,1] = cEDA[x1-1,y1,z1,-2]+0.4799/(np.linalg.norm([x1-1,y1,z1]-cArr[0,1:4]))
    cEDA[x1+1,y1,z1,1] = cEDA[x1+1,y1,z1,-2]+0.4799/(np.linalg.norm([x1+1,y1,z1]-cArr[0,1:4]))
    cEDA[x1,y1-1,z1,1] = cEDA[x1,y1-1,z1,-2]+0.4799/(np.linalg.norm([x1,y1-1,z1]-cArr[0,1:4]))
    cEDA[x1,y1+1,z1,1] = cEDA[x1,y1+1,z1,-2]+0.4799/(np.linalg.norm([x1,y1+1,z1]-cArr[0,1:4]))
    cEDA[x1,y1,z1-1,1] = cEDA[x1,y1,z1-1,-2]+0.4799/(np.linalg.norm([x1,y1,z1-1]-cArr[0,1:4]))
    cEDA[x1,y1,z1+1,1] = cEDA[x1,y1,z1+1,-2]+0.4799/(np.linalg.norm([x1,y1,z1+1]-cArr[0,1:4]))
    return cEDA

Answer 1

If you want to call this a lot of time, you need to convert the whole for loop in to a C-loop by numba or Cython. However there are some method to increase the speed of your function, here is your original function:

def cF2(cArr, cEDA):
    # the coordinates of points 0 and 1 for readability
    x0 = cArr[0,1]
    y0 = cArr[0,2]
    z0 = cArr[0,3]
    x1 = cArr[1,1]
    y1 = cArr[1,2]
    z1 = cArr[1,3]

    # for each point around point 0 and the point itself, calculate the distance between this point and point 1
    # use this value and the corresponding value in cEDA(x,y,z,-2) and place result in  cEDA(x,y,z,0)
    cEDA[x0,y0,z0,0] = cEDA[x0,y0,z0,-2]-0.4799/(np.linalg.norm([x0,y0,z0]-cArr[1,1:4]))
    cEDA[x0-1,y0,z0,0] = cEDA[x0-1,y0,z0,-2]-0.4799/(np.linalg.norm([x0-1,y0,z0]-cArr[1,1:4]))
    cEDA[x0+1,y0,z0,0] = cEDA[x0+1,y0,z0,-2]-0.4799/(np.linalg.norm([x0+1,y0,z0]-cArr[1,1:4]))
    cEDA[x0,y0-1,z0,0] = cEDA[x0,y0-1,z0,-2]-0.4799/(np.linalg.norm([x0,y0-1,z0]-cArr[1,1:4]))
    cEDA[x0,y0+1,z0,0] = cEDA[x0,y0+1,z0,-2]-0.4799/(np.linalg.norm([x0,y0+1,z0]-cArr[1,1:4]))
    cEDA[x0,y0,z0-1,0] = cEDA[x0,y0,z0-1,-2]-0.4799/(np.linalg.norm([x0,y0,z0-1]-cArr[1,1:4]))
    cEDA[x0,y0,z0+1,0] = cEDA[x0,y0,z0+1,-2]-0.4799/(np.linalg.norm([x0,y0,z0+1]-cArr[1,1:4]))

    cEDA[x1,y1,z1,1] = cEDA[x1,y1,z1,-2]+0.4799/(np.linalg.norm([x1,y1,z1]-cArr[0,1:4]))
    cEDA[x1-1,y1,z1,1] = cEDA[x1-1,y1,z1,-2]+0.4799/(np.linalg.norm([x1-1,y1,z1]-cArr[0,1:4]))
    cEDA[x1+1,y1,z1,1] = cEDA[x1+1,y1,z1,-2]+0.4799/(np.linalg.norm([x1+1,y1,z1]-cArr[0,1:4]))
    cEDA[x1,y1-1,z1,1] = cEDA[x1,y1-1,z1,-2]+0.4799/(np.linalg.norm([x1,y1-1,z1]-cArr[0,1:4]))
    cEDA[x1,y1+1,z1,1] = cEDA[x1,y1+1,z1,-2]+0.4799/(np.linalg.norm([x1,y1+1,z1]-cArr[0,1:4]))
    cEDA[x1,y1,z1-1,1] = cEDA[x1,y1,z1-1,-2]+0.4799/(np.linalg.norm([x1,y1,z1-1]-cArr[0,1:4]))
    cEDA[x1,y1,z1+1,1] = cEDA[x1,y1,z1+1,-2]+0.4799/(np.linalg.norm([x1,y1,z1+1]-cArr[0,1:4]))
    return cEDA

np.random.seed(42)
cArr = np.random.randint(0, 100, (2, 5))
cEDA = np.random.rand(100, 100, 100, 4)
r1 = cF2(cArr, cEDA)

Here is the optimized function:

tmp = np.eye(3)
offsets = np.concatenate((tmp, -tmp, [[0, 0, 0]]), 0).astype(int)[:, None, :]

def fast_cF2(cArr, cEDA):
    cArr = cArr[:, 1:4]
    t = cArr[None,:,:] + offsets
    X0, Y0, Z0 = t[:, 0, :].T
    X1, Y1, Z1 = t[:, 1, :].T

    d1 = 0.4799/np.linalg.norm(t[:, 0, :] - cArr[1], axis=1)
    d0 = 0.4799/np.linalg.norm(t[:, 1, :] - cArr[0], axis=1)

    cEDA[X0, Y0, Z0, 0] = cEDA[X0, Y0, Z0, -2] - d1
    cEDA[X1, Y1, Z1, 1] = cEDA[X1, Y1, Z1, -2] + d0
    return cEDA

np.random.seed(42)
cArr = np.random.randint(0, 100, (2, 5))
cEDA = np.random.rand(100, 100, 100, 4)
r2 = fast_cF2(cArr, cEDA)
print np.allclose(r1, r2)

the timeit result:

%timeit cF2(cArr, cEDA)
%timeit fast_cF2(cArr, cEDA)

output:

1000 loops, best of 3: 320 µs per loop
10000 loops, best of 3: 91.6 µs per loop

Answer 2

I think that scipy.ndimage.generic_filter would do it...

def neighborly_function(in_arr1d):
    # put your equation below - will have to figure out the details
    # below is just an example
    # has to return a scalar
    return sum(in_arr1d)

import scipy.ndimage as nd

huge_arr3d = [....]

# Since you want the cardinal directions you have to 'footprint'
footprint = np.array([[[False, False, False],
    [False,  True, False],
    [False, False, False]],

   [[False,  True, False],
    [ True, False,  True],
    [False,  True, False]],

   [[False, False, False],
    [False,  True, False],
    [False, False, False]]])

out = nd.generic_filter(huge_arr3d, neighborly_function, footprint=footprint)

Check out the ndimage docs at http://docs.scipy.org/doc/scipy/reference/tutorial/ndimage.html and generic_filter at http://docs.scipy.org/doc/scipy/reference/tutorial/ndimage.html#generic-filter-functions