How to construct the following matrix elegantly in numpy?

Question

Suppose I have a 5 dimensional matrix v and now I want a new matrix D fulfilling

D[a, b, n, m, d] = v[a, b, n, n, d]-v[a, b, m, m, d].

How do I elegantly do this in numpy?

Answer 1

How do you want to change the dimensionality? You can reshape it like this

import numpy as np

a, b, n, d = 2, 3, 4, 5
v = np.zeros((a, b, n, n, d))
D = v.reshape((a, b, n*n, d))

Answer 2

I found einsum can do this:

D = np.einsum('abiic->abic', v)[..., None, :] - np.einsum('abiic->abic', v)[:, :, None, ...]