Numpy array assignment along axis and index

I have a 3D volume and I modify slices along different axis.

for idx in range(len(self.volume)): 
    for axe in range(self.volume.ndim): # = range(3)
        slice_ = np.take(self.volume, idx, axis = axe)
        ''' Do something '''

(np.take is equivalent of writing self.volume[idx], self.volume[:, idx] and self.volume[:, :, idx])

Finally, I want to assign a new slice in my volume along the axis :

    if axe == 0:
        self.volume[idx] = new_slice
    elif axe == 1:
        self.volume[:,idx] = new_slice
    else:
        self.volume[:,:,idx] = new_slice

This is where I need some help. I can’t figure out a cleaner way of doing this assignment. (I would like something as clean as np.take())

I have tried np.insert, np.insert_along_axis, np.put, np.put_along_axis… but I am clearly missing something out.

Any ideas ? :)

Have a great day

Answer

There may be a more elegant solution but the following should work:

s = [slice(None)]*self.volume.ndim
s[axe] = slice(idx,idx+1)
self.volume[tuple(s)] = np.expand_dims(new_slice, axe)

Alternatively, you can try:

self.volume  = np.swapaxes(self.volume, 0, axe)
self.volume[idx] = new_slice
self.volume  = np.swapaxes(self.volume, 0, axe)