Numpy: slicing a volume using a matrix

Question

I have a 3D numpy volume and a 2D numpy matrix:

foo = np.random.rand(20,20,10)
amin = np.argmin(foo, axis=2)

i would like to use amin variable to slice the volume in the same way np.min would do:

grid = np.indices(min.shape)
idcs = np.stack([grid[0], grid[1], min])

fmin = foo[idcs[0], idcs[1], idcs[2]]

problem is that i can't use np.min because i also need the amin neighbors for interpolation reasons, something that i would obtain doing:

pre  = foo[idcs[0], idcs[1], np.clip(idcs[2]-1, 0, 9)]
post = foo[idcs[0], idcs[1], np.clip(idcs[2]+1, 0, 9)]

Is there a more pythonic (nupyic) way to do this without creating an np.grid? something like:

foo[:,:,amin-1:amin+1]

that actually works (i would care about margin handling with an early-padding)

Answer 1

You could use np.ogrid instead of np.indices to save memory. np.ogrid returns an "open" meshgrid:

In [24]: np.ogrid[:5,:5]
Out[24]: 
[array([[0],
        [1],
        [2],
        [3],
        [4]]), array([[0, 1, 2, 3, 4]])]

ogrid returns component arrays which can be used as indices in the same way as one would use np.indices. NumPy will automatically broadcast the values in the open mesh when they are used as indices:

In [49]: (np.indices((5,5)) == np.broadcast_arrays(*np.ogrid[:5, :5])).all()
Out[49]: True

---

import numpy as np
h, w, d = 20, 20, 10
foo = np.random.rand(h, w, d)
amin = np.argmin(foo, axis=2)
X, Y = np.ogrid[:h, :w]
amins = np.stack([np.clip(amin+i, 0, d-1) for i in [-1, 0, 1]])
fmins = foo[X, Y, amins]

It's better to store fmin, pre and post in one array, fmins, since some NumPy/Scipy operations (like argmin or griddata) may need the values in one array. If, later, you need to operate on the 3 components individually, you can always access them using fmins[i] or define

pre, fmin, post = fmins