How to realize select specific elements in array while keeping other indices general in python?-CodePudding

For a rank-4 array, M[n0,n1,n2,n3]

I would like to do M[m0,m1,:,:]*(some operation) for all m0,m1

or M[:,m1,:,m3]*(some operation) for all m1,m3 with some operation

I can do

import numpy as np
import time
n0=2
n1=2
n2=2
n3=2
M = np.zeros((n0,n1,n2,n3))
M2 = np.zeros((n0,n1,n2,n3))
M3 = np.zeros((n0,n1,n2,n3))

i = 0
for m0 in range(M.shape[0]):
    for m1 in range(M.shape[1]):  
        for m2 in range(M.shape[2]):  
            for m3 in range(M.shape[3]):  
                M[m0,m1,m2,m3] = i
                i = i   1

input = 0


if input == 0:
    for m0 in range(M.shape[0]):
        for m1 in range(M.shape[1]):   
            M2[m0,m1,:,:] =  M[m0,m1,:,:]*(m0 m1)

            
elif input == 1:
    for m1 in range(M.shape[1]):
        for m3 in range(M.shape[3]):  
            M2[:,m1,:,m3] = M[:,m1,:,m3]*(m1-m3)

for m0 in range(M2.shape[0]):
    for m1 in range(M2.shape[1]):  
        for m2 in range(M2.shape[2]):  
            for m3 in range(M2.shape[3]):  
               # M2[m0,m1,m2,m3] = i
                print(m0, m1, m2, m3, 'M', M[m0,m1,m2,m3], 'M2', M2[m0,m1,m2,m3])

(using : to skip a few loops seems faster than explicitly looping over all indices m0 - m3. That is the motivation for this question: utilize :. m0 m1 and m1-m3 are somehow randomly picked up)

sample output is

0 0 0 0 M 0.0 M2 0.0
0 0 0 1 M 1.0 M2 0.0
0 0 1 0 M 2.0 M2 0.0
0 0 1 1 M 3.0 M2 0.0
0 1 0 0 M 4.0 M2 4.0
0 1 0 1 M 5.0 M2 5.0
0 1 1 0 M 6.0 M2 6.0
0 1 1 1 M 7.0 M2 7.0
1 0 0 0 M 8.0 M2 8.0
1 0 0 1 M 9.0 M2 9.0
1 0 1 0 M 10.0 M2 10.0
1 0 1 1 M 11.0 M2 11.0
1 1 0 0 M 12.0 M2 24.0
1 1 0 1 M 13.0 M2 26.0
1 1 1 0 M 14.0 M2 28.0
1 1 1 1 M 15.0 M2 30.0

My question is, is there any simple approach to realize the above code by inputting the position of the target array elements, e.g., 0,1,M[m0,m1,:.:] to realize the upper part (input = 0) of the code; 1,3 M[:,m1,:,m3]for the lower part (input = 1) of the above code? And include other cases, e.g., 0,2; 0,3;1,2;2,3. Essentially, change the positions of labels of indices and :.

I can let python print code that includes all cases, but I hope there is something simpler

CodePudding user response：

I don't know what complex operations you will need, but for now, I think you are not familiar with numpy calculation. To this end, I rewrite the whole answer and complete your three requirements in a simpler and faster way:

# Yours
n0 = n1 = n2 = n3 = 2
M = np.zeros((n0, n1, n2, n3))
M2 = np.zeros((n0, n1, n2, n3))
M3 = np.zeros((n0, n1, n2, n3))

i = 0
for m0 in range(M.shape[0]):
    for m1 in range(M.shape[1]):  
        for m2 in range(M.shape[2]):  
            for m3 in range(M.shape[3]):  
                M[m0,m1,m2,m3] = i
                i = i   1

for m0 in range(M.shape[0]):
    for m1 in range(M.shape[1]):   
        M2[m0,m1,:,:] =  M[m0,m1,:,:]*(m0 m1)


for m1 in range(M.shape[1]):
    for m3 in range(M.shape[3]):  
        M3[:,m1,:,m3] = M[:,m1,:,m3]*(m1-m3)


# Mine
N = np.arange(n0 * n1 * n2 * n3, dtype=float).reshape(n0, n1, n2, n3)
m0, m1, m2, m3 = np.indices(N.shape, sparse=True)
N2 = N * (m0   m1)
N3 = N * (m1 - m3)

You can use the equal sign and the method ndarray.all to determine whether two arrays are exactly equal:

print((M == N).all())
print((M2 == N2).all())
print((M3 == N3).all())

You can try it yourself, they're all True. If you want to know the principle, I can try to write it, but I don't have much time now.

I also want to remind you that when operating numpy arrays, you should avoid using loops as much as possible.