Numpy select checkers pattern to new array - numpy

Question

I have a square 2D matrix with odd number of rows and columns. For example
11 12 13 14 15
21 22 23 24 25
31 32 33 34 35
41 42 43 44 45
51 52 53 54 55
I need to rotate it 45 degrees clockwise and select the maximal square matrix. In this case:
13 24 35
22 33 44
31 42 53
I can do this in two cycles:
new_arr = np.zeros(((orig_range+1)//2, (orig_range+1)//2)
for new_h in range((orig_range+1)//2):
for new_w in range((orig_range+1)//2):
old_h = new_h + new_w
old_w = ((orig_range+1)//2) - new_h + new_w
new_arr[new_h, new_w] = orig_arr[old_h, old_w]
But this approach is very slow. Rotation in cv2 is reasonably fast, but the "pixels" don't align well. Forward rotation with sqrt(2) scaling followed by backward rotation with sqrt(2) scaling results in a altered colors of pixels in the center region of the image due to the rounding errors.
What is the efficient way to rotate such a matrix?

Answer 1

You can use the Numba's JIT to drastically speed up the operation, especially by running it in parallel and natively. Moreover, note that the array do not need to be filled with zeros. Here is an untested example:
import numba as nb
#nb.njit(parallel=True)
def compute(orig_range, orig_arr):
new_arr = np.empty(((orig_range+1)//2, (orig_range+1)//2)
for new_h in nb.prange((orig_range+1)//2):
for new_w in range((orig_range+1)//2):
old_h = new_h + new_w
old_w = ((orig_range+1)//2) - new_h + new_w
new_arr[new_h, new_w] = orig_arr[old_h, old_w]
return new_arr
You can specify the type of the inputs to compile the function ahead of time and so avoid the first call to be significantly slower. This operation should be very fast on array fitting in cache. For big arrays, one can use tiling and a different read/write ordering to speed the operation up a bit more.