Numpy split cube into cubes

Question

There is a function np.split() which can split an array along 1 axis. I was wondering if there was a multi axis version where you can split along axes (0,1,2) for e

Answer 1

In addition to my extra question to @unutbu's answer I think I have got the reverse to work (in case you want to split a cube into cubes, apply a function to each one and then combine them back).

import numpy as np
import pdb
np.set_printoptions(precision=3,linewidth=300)

class Cubify():
    def __init__(self,oldshape,newshape):
        self.newshape = np.array(newshape)
        self.oldshape = np.array(oldshape)
        self.repeats = (oldshape / newshape).astype(int)
        self.tmpshape = np.column_stack([self.repeats, newshape]).ravel()
        order = np.arange(len(self.tmpshape))
        self.order = np.concatenate([order[::2], order[1::2]])
        self.reverseOrder = self.order.copy()
        self.reverseOrder = np.arange(len(self.tmpshape)).reshape(2, -1).ravel(order='F')
        self.reverseReshape = np.concatenate([self.repeats,self.newshape])

    def cubify(self,arr):
        # newshape must divide oldshape evenly or else ValueError will be raised
        return arr.reshape(self.tmpshape).transpose(self.order).reshape(-1, *self.newshape)

    def uncubify(self,arr):
        return arr.reshape(self.reverseReshape).transpose(self.reverseOrder).reshape(self.oldshape)

if __name__ == "__main__":
    N = 9
    x = np.arange(N**3).reshape(N,N,N)
    oldshape = x.shape
    newshape = np.array([3,3,3])
    cuber = Cubify(oldshape,newshape)
    out = cuber.cubify(x)
    back = cuber.uncubify(out)