Search code examples
arrayspython-3.xregressionplane

Fitting a plane to a 2D array

I have a topological image that I am attempting to perform a plane subtraction on using Python. The image is a 256x256 2-D array of float32 values between 0 and 1.

What I wish to do is to use linear regression to fit a plane to this data and subsequently subtract this plane from the original values.

I am unsure how to go about achieving this.

I am new to the Python language and appreciate any help.

Solution

  • At first you need to represent your data in the proper way.

    You have two arguments X1 and X2, which define the coordinates of your topological image, and one target value Y, which defines the height of each point. For regression analysis you need to expand the list of arguments, by adding X0, which is always equal to one.

    Then you need to unroll the parameters and the target into matrices [m*m x 3] and [m*m x 1] respectively. You want to find vector theta, which will describe the desired plane. For this purpose you can use the Normal Equation:

    enter image description here

    To demonstrate the approach I generated some topological surface. You can see the surface, the surface with the fitted plane and the surface after subtraction on the picture:

    regression plane

    Here is the code:

    from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import numpy as np

m = 256 #size of the matrix

X1, X2 = np.mgrid[:m, :m]

fig = plt.figure()
ax = fig.add_subplot(3,1,1, projection='3d')
jet = plt.get_cmap('jet')

#generation of the surface
F = 3        
i = np.minimum(X1, m-X1-1)
j = np.minimum(X2, m-X2-1)
H = np.exp(-.5*(np.power(i, 2)  +  np.power(j, 2)   )/(F*F))
Y = np.real(  np.fft.ifft2   (H  *  np.fft.fft2(  np.random.randn(m, m))))
a = 0.0005; b = 0.0002; #parameters of the tilted plane
Y = Y + (a*X1 + b*X2); #adding the plane
Y = (Y - np.min(Y)) / (np.max(Y) - np.min(Y)) #data scaling

#plot the initial topological surface
ax.plot_surface(X1,X2,Y, rstride = 1, cstride = 1, cmap = jet, linewidth = 0)


#Regression
X = np.hstack(   ( np.reshape(X1, (m*m, 1)) , np.reshape(X2, (m*m, 1)) ) )
X = np.hstack(   ( np.ones((m*m, 1)) , X ))
YY = np.reshape(Y, (m*m, 1))

theta = np.dot(np.dot( np.linalg.pinv(np.dot(X.transpose(), X)), X.transpose()), YY)

plane = np.reshape(np.dot(X, theta), (m, m));

ax = fig.add_subplot(3,1,2, projection='3d')
ax.plot_surface(X1,X2,plane)
ax.plot_surface(X1,X2,Y, rstride = 1, cstride = 1, cmap = jet, linewidth = 0)


#Subtraction
Y_sub = Y - plane
ax = fig.add_subplot(3,1,3, projection='3d')
ax.plot_surface(X1,X2,Y_sub, rstride = 1, cstride = 1, cmap = jet, linewidth = 0)

plt.show()