Idiomatic way to flatten a 6D array to 2D in numpy, while preserving the Hamiltonian structure

user_3264977 0
Viewed 73

Though there are similar questions floating around about flattening numpy arrays, there are none I can find that deal with preserving some kind of algebraic structure.

I have a numpy array H, where H.shape is (M, N, S, M, N, S). A triple (m, n, s) labels a particular basis vector for the matrix, so in order to pass this to np.eigh or similar i need to rearrange my 6D tensor into a 2D array H_flat like so:

for m in range(M):
  for m2 in range(M):
    for n in range(N):
      for n2 in range(N):
        for s in range(S):
          for s2 in range(S):
            H_flat[s + S*n + S*N*m, s2 + S*n2 + S*N*m2] = H[m,n,s, m2,n2,s2]

While this works, it is certainly not canonical numpy style and is very slow. My question is therefore

How can I massage np.reshape (and similar) to do this?

numpy python
1 Answers

Looks like this is simply:

out = H.reshape(M*N*S,-1)

Check:

M,N,S = 2,3,5

H = np.arange(M*N*S*M*N*S).reshape(M,N,S,M,N,S)

H_flat = np.zeros((M*N*S,M*N*S))
for m in range(M):
  for m2 in range(M):
    for n in range(N):
      for n2 in range(N):
        for s in range(S):
          for s2 in range(S):
            H_flat[s + S*n + S*N*m, s2 + S*n2 + S*N*m2] = H[m,n,s, m2,n2,s2]

(H_flat == H.reshape(M*N*S,-1)).all()
# True
user_42384080
Related