import numpy as np
n = 5
wR = np.random.choice(a=[0, 1, 2], size=(n, n), p=[0.5, 0.25,0.25])
identity_3d = np.zeros((n, n, n))
idx = np.arange(n)
identity_3d[:, idx, idx] = 1
I,J = np.nonzero(wR==0)
identity_3d[I,:,J]=0
identity_3d[I,J,:]=0
identity_3d