Search code examples
pythonnumpygridtransformcoordinate

Applying a Fast Coordinate Transformation in Python


I have a simple 2x2 transformation matrix, s, which encodes some liner transformation of coordinates such that X' = sX.

I have generated a set of uniformley distributed coordinates on a grid using the np.meshgrid() function and at the moment I traverse each coordinate and apply the transformation at a coordinate by coordinate level. Unfortunately, this very slow for large arrays. Are there any fast ways of doing this? Thanks!

import numpy as np

image_dimension = 1024
image_index = np.arange(0,image_dimension,1) 
xx, yy = np.meshgrid(image_index,image_index)

# Pre-calculated Transformation Matrix.
s = np.array([[ -2.45963439e+04,  -2.54997726e-01], [  3.55680731e-02, -2.48005486e+04]])

xx_f = xx.flatten()
yy_f = yy.flatten()

for x_t in range(0, image_dimension*image_dimension):

    # Get the current (x,y) coordinate.
    x_y_in = np.matrix([[xx_f[x_t]],[yy_f[x_t]]])

    # Perform the transformation with x.
    optout =  s * x_y_in

    # Store the new coordinate.
    xx_f[x_t] = np.array(optout)[0][0]
    yy_f[x_t] = np.array(optout)[1][0]

# Reshape Output
xx_t = xx_f.reshape((image_dimension, image_dimension))
yy_t = yy_f.reshape((image_dimension, image_dimension))

Solution

  • Loops are slow in Python. It is better to use vectorization. In a nutshell, the idea is to let numpy do the loops in C, which is much faster.

    You can express your problem as matrix multiplications X' = sX, where you put all the points in X and transform them all with just one call to numpy's dot product:

    xy = np.vstack([xx.ravel(), yy.ravel()])
    xy_t = np.dot(s, xy)
    xx_t, yy_t = xy_t.reshape((2, image_dimension, image_dimension))