I would like to know if there is a way to move all elements of a numpy array without iterating over each entry. The shift I desire is to relabel the indices by a fixed XOR operation, in the following form:
import numpy as np
N = 2
x = np.arange(2**(2 * N)).reshape(2**N, 2**N)
z = np.zeros((2**N, 2**N))
k = 1
for i in range(2**N):
for j in range(2**N):
z[i][j] = x[i ^ k][j ^ k]
The problem I have is that latter I wish to take huge values of N
, which becomes a bottleneck if we wish to iterate over each entry. Any advice on how to perform the move in a single shot will be very much appreciated.
There is simply no reason to use a loop here at all. You are not shifting anything, and the problem is completely separable across the axes:
x = np.arange(2**(2 * N)).reshape(2**N, 2**N)
z = x[np.arange(2**N)[:, None] ^ k, np.arange(2**N) ^ k]
But what is x[a, b]
for some a, b
to begin with? From the original definition, you can see that it's a * 2**N + b
. You can therefore plug in a = np.arange(2**N)[:, None] ^ k
and b = np.arange(2**N) ^ k
to avoid having to generate anything but the final result:
idx = np.arange(2**N) ^ k
z = (idx[:, None] << N) + idx
It's generally faster to substitute << N
for * 2**N
.
Allocation
None of the solutions shown here pre-allocate z
. However, if you were to do that, as in the original question, you should be careful about the type.
z = np.zeros((2**N, 2**N))
This creates an array of floats by default, which is likely not what you want.
z = np.zeros((2**N, 2**N), dtype=int)
Adding an explicit dtype
make the array into integers. However, the simplest method available to you is likely
z = np.zeros_like(x)
Since you are planning on filling in every single element, you don't need to waste time filling it with zeros first. np.empty
presents a better option in this case:
z = np.empty((2**N, 2**N), dtype=int)
OR
z = np.empty_like(x)