How to compose permutations in numpy efficiently?

The idea here is very much like any other resampler: for each destination location, find the corresponding source location. Let's start with a simplified case. Say we have a 1-D array and corresponding indices:

arr1 = np.arange(6)
a1 = np.array([3, 4])
b1 = np.array([2, 3])

It's much easier to visualize what happens when you make the assignments:

arr[a1] = arr[b1] # [0, 1, 2, 2, 3, 5]
arr[a1] = arr[b1] # [0, 1, 2, 2, 2, 5]

For a destination index that is the entire array (e.g., np.arange(arr.size)), the source index is just the application of the assignment to the destination.

There are a couple of ways of generalizing this to a 3D array. One way would be to make a (*arr.shapes, arr.ndim) array containing a meshgrid of all the indices. Another is simply to convert a and b into raveled (linear) indices into the raveled version of arr. I recommend going with the latter.

a_r = np.ravel_multi_index(a, arr.shape)
b_r = np.ravel_multi_index(b, arr.shape)

You can construct a source and destination index pair of the same size as arr.ravel(), and set the sources and destinations in it:

dest = np.arange(arr.size)
src = np.arange(arr.size)
src[a] = src[b]

I wrote out the last assignment "in full", although the first time you can just do src[a] = b. You can iterate the last assignment as many times as you want. In your particular example, do it a second time:

src[a] = src[b]

If you want, you can trim off the elements that are not modified by the assignment:

mask = dest != src
dest = dest[mask]
src = src[mask]

Finally, you can unravel the index back to the original shape:

c = np.unravel_index(dest, arr.shape)
d = np.unravel_index(src, arr.shape)

If you want to write a function that accepts an arbitrary number of input indices, it might look something like this:

def compose(*args, shape=None, sparse=True):
    if shape is None:
        get_max = lambda tup: np.array([arr.max() for arr in tup])
        for a, b in args:
            if shape is None:
                shape = np.maximum(get_max(a), get_max(b))
            else:
                shape = np.maximum(shape, get_max(a))
                shape = np.maximum(shape, get_max(b))
        shape = tuple(shape + 1)

    size = np.prod(shape)
    dest = np.arange(size)
    src = np.arange(size)
    for a, b in args:
        a = np.ravel_multi_index(a, shape)
        b = np.ravel_multi_index(b, shape)
        src[a] = src[b]

    if sparse:
        mask = dest != src
        dest = dest[mask]
        src = src[mask]
    return np.unravel_index(dest, shape), np.unravel_index(src, shape)