2

I have a 3D numpy array of voxels, i.e. where the indices of each point correspond to its position in 3D space.

There are a number of calculations I wish to perform that involve knowing whether the points satisfy various geometric conditions. This means converting the indices into vectors by multiplying by the basis, and then calculating dot and cross products, norms etc.. I am looking for a reasonably fast way to do this, as my efforts so far seem very slow.

If I have an arbitrary basis a, b, c:

basis = np.array([[a1, a2, a3],
                  [b1, b2, b3],
                  [c1, c2, c3]])

where a1, a2, a3 are the x, y and z components of a, and likewise for b and c. I can calculate the Cartesian coordinate p = (x, y, z) of each voxel in the following way:

for i in range(vox.shape[0]):
    for j in range(vox.shape[1]):
        for k in range(vox.shape[2]):
            p = np.dot(basis, np.array([i, j, k]))

where "vox" is the 3D array of voxels. If, for example, I wish to calculate the dot product of each of these vectors with a single other (Cartesian) vector (e.g. q = np.array([qx, qy, qz])) and store that index if the result meets a given condition (here greater than 0.0), I could do something like this (within the same loop as above):

            if np.dot(p, q) > 0.0:
                desired_vox_indices.append([i, j, k])

The problem is that this is very slow. Can I do this in a more pythonic way, or by using more numpy tools? I realise I am not even accessing the values of the vox array at this stage either.

EDIT: attempt at cross product based on answer by Divakar

# Get q_x (matrix multiplication version of cross product)
q_x = np.array([[0.0, -q[2], q[1]],
                [q[2], 0.0, -q[0]],
                [-q[1], q[0], 0.0]])

# transpose (as per cross product definition) and matrix multiply with basis
u = np.matmul(q_x.T, basis)

# Get open range arrays
m,n,r = vox.shape
I,J,K = np.ogrid[:m,:n,:r]

# writing everything explicitly, since I am unsure how ogrid objects behave
pxq = np.empty(3)
pxq[0] = u[0,0]*I + u[0,1]*J + u[0,2]*K
pxq[1] = u[1,0]*I + u[1,1]*J + u[1,2]*K
pxq[2] = u[2,0]*I + u[2,1]*J + u[2,2]*K

which might be the same as just:

pxq = np.dot(u, np.array([I, J, K]))

but I am not sure...

Idle_92
  • 69
  • 7

2 Answers2

2

We will build up the summations using range arrays for scaling without generating all of the indices at once. We will do it in three steps corresponding to the three nested loops. The idea is very similar to the one explored in this answer to - Python vectorizing nested for loops. This would be memory efficient and hence hopefully good on performance too. Then, we will compare the summations against the threshold 0.0 and use np.argwhere to get the corresponding indices. Thus, we would have a solution, like so -

# Get q scaled version
s = q.dot(basis)

# Get open range arrays and scale and sum-reduce s array
m,n,r = vox.shape
I,J,K = np.ogrid[:m,:n,:r]
sums = s[0]*I + s[1]*J + s[2]*K

# Finally compare the sums against threshold amd get corresponding indices
out = np.argwhere(sums > 0.0)

Timings on a large vox array -

# Setup
In [371]: np.random.seed(0)
     ...: basis = np.random.randn(3,3)
     ...: vox = np.random.randn(100,100,100)
     ...: q = np.random.randn(3)

# Original soln
In [372]: %%timeit
     ...: desired_vox_indices = []
     ...: for i in range(vox.shape[0]):
     ...:     for j in range(vox.shape[1]):
     ...:         for k in range(vox.shape[2]):
     ...:             p = np.dot(basis, np.array([i, j, k]))            
     ...:             if np.dot(p, q) > 0.0:
     ...:                 desired_vox_indices.append([i, j, k])
1 loop, best of 3: 2.13 s per loop

# @jdehesa's soln
In [373]: %%timeit
     ...: ind = np.indices(vox.shape).reshape(3, -1)
     ...: p = basis.dot(ind)
     ...: d = q.dot(p)
     ...: desired_vox_indices = ind[:, d > 0.0].T
10 loops, best of 3: 35.9 ms per loop

# From this post
In [374]: %%timeit
     ...: s = q.dot(basis)
     ...: m,n,r = vox.shape
     ...: I,J,K = np.ogrid[:m,:n,:r]
     ...: sums = s[0]*I + s[1]*J + s[2]*K
     ...: out = np.argwhere(sums > 0.0)
100 loops, best of 3: 7.56 ms per loop
Divakar
  • 218,885
  • 19
  • 262
  • 358
  • Thanks for the timing information, that's really helpful. For your solution, `q.dot(basis)` is q as a row vector multiplied by the basis, if I understand the documentation correctly, which means you obtain the desired inner product when you dot it with (i, j, k), which you have done explicitly with `sums`. This is to perform a single matrix multiplication rather than one for each voxel, I assume. If I were to extend this to calculate a cross product, I would have to change a few things. That attempt to follow in an edit of my question (let me know if that isn't the right place to put it). – Idle_92 Feb 19 '19 at 12:13
  • @Idle_92 Yes, please go ahead and edit the question with those additional details. – Divakar Feb 19 '19 at 12:21
  • I suppose there are a number of ways to have done this, but I tried to combine all the matrix/vector operations at the start, as you did with the dot product solution – Idle_92 Feb 19 '19 at 13:47
1

You can do that in a vectorized manner like this:

# Array of indices for your voxel data
ind = np.indices(vox.shape).reshape(3, -1)
# Multiply the basis times each coordinate
p = basis @ ind
# Dot product of each result with vector
d = q @ p
# Select coordinates where criteria is met
desired_vox_indices = ind[:, d > 0.0].T
jdehesa
  • 58,456
  • 7
  • 77
  • 121
  • In `d = q @ p`, `q` seems to be not defined. – fountainhead Feb 18 '19 at 18:00
  • @fountainhead I'm assuming `q` is an input like `vox` and `basis`, since it is not specified in the question either. – jdehesa Feb 18 '19 at 18:03
  • Is it supposed to be `reshape(3, 1)` and not `reshape(3, -1)`? – ap21 Feb 18 '19 at 18:40
  • 1
    @ap21 No, it's `reshape(3, -1)` so it reshapes it into a matrix with three rows (X, Y and Z) and as many columns as voxels. – jdehesa Feb 19 '19 at 09:43
  • This method is very clear and understandable, thank you. I take it this would naturally extend to cross products and norms with the appropriate function/operation, e.g. replace `d = q @ p` with `c = np.cross(p, q)`. Perhaps you could point me to where I can read about how this method works quickly, since it will still need to perform that operation once per voxel (which in my systems is in the 10e7 range). – Idle_92 Feb 19 '19 at 14:18
  • 1
    @Idle_92 Yes, you can adapt the code to your needs, it should work as long as NumPy supports the vectorized operation you want to perform. In the case of cross product, for example, you just need to note that `p` here has shape `3xN` (`N` being the number of voxels), so maybe you could do for example `c = np.cross(p.T, q)` (and you would get an array with shape `Nx3` in this case). What do you mean by "about how this method works quickly"? Is there something you don't understand in the answer? – jdehesa Feb 19 '19 at 14:23