How to fit a plane to a 3D dataset in Python

Question

I have a set of x, y and z points and am trying to fit a plane to this three-dimensional data so that z=f(x,y) can be calculated for any x and y.

I am hoping to get an equation for the plane and plot the graph in a Jupyter notebook for visualization.

This is the (working) code I've been using to plot my data:

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

x = np.arange(-12, 1)
y = np.arange(-40,-25) 
Z = array([[402., 398., 395., 391., 387., 383., 379., 375., 371., 367., 363.,358., 354.],
       [421., 417., 413., 409., 406., 402., 398., 393., 389., 385., 381.,
        376., 372.],
       [440., 436., 432., 429., 425., 421., 416., 412., 408., 404., 399.,
        395., 391.],
       [460., 456., 452., 448., 444., 440., 436., 432., 427., 423., 419.,
        414., 410.],
       [480., 477., 473., 469., 465., 460., 456., 452., 447., 443., 438.,
        434., 429.],
       [501., 498., 494., 490., 485., 481., 477., 472., 468., 463., 459.,
        454., 449.],
       [523., 519., 515., 511., 507., 502., 498., 494., 489., 484., 480.,
        475., 470.],
       [545., 541., 537., 533., 529., 524., 520., 515., 511., 506., 501.,
        496., 492.],
       [568., 564., 560., 556., 551., 547., 542., 538., 533., 528., 523.,
        518., 513.],
       [592., 588., 583., 579., 575., 570., 565., 561., 556., 551., 546.,
        541., 536.],
       [616., 612., 607., 603., 598., 594., 589., 584., 579., 575., 569.,
        564., 559.],
       [640., 636., 632., 627., 623., 618., 613., 609., 604., 599., 593.,
        588., 583.],
       [666., 662., 657., 653., 648., 643., 638., 633., 628., 623., 618.,
        613., 607.],
       [692., 688., 683., 679., 674., 669., 664., 659., 654., 649., 643.,
        638., 632.],
       [ nan, 714., 710., 705., 700., 695., 690., 685., 680., 675., 669.,
        664., 658.]])

X, Y = np.meshgrid(x, y)

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

print (X.shape, Y.shape, Z.shape)

ax.plot_surface(X, Y, Z)
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')

plt.show()

I have tried implementing these solutions:

https://gist.github.com/amroamroamro/1db8d69b4b65e8bc66a6

http://inversionlabs.com/2016/03/21/best-fit-surfaces-for-3-dimensional-data.html

However, since my x and y arrays don't have the same length, I get this error message:

ValueError: all the input array dimensions except for the concatenation axis must match exactly

Answer 1

The data you shared seemed to work for me during plotting. Your X, Y, Z are all having the same size. There is one nan value in your Z array. You can remove that point while estimating equation of plane.

You want to fit your data to a plan in 3D. Thus, it is a linear regression problem. You can use multivariate regression from scikit-learn package to estimate the coefficient of the equation of plane.

Equation of plane is given by the following:

Z = a1 * X + a2 * Y + c

You can flatten your data as follows and use scikit-learn's linear_model to fit a plane to the data. Please refer below:

# your data is stored as X, Y, Z
print(X.shape, Y.shape, Z.shape)

x1, y1, z1 = X.flatten(), Y.flatten(), Z.flatten()

X_data = np.array([x1, y1]).reshape((-1, 2))
Y_data = z1

from sklearn import linear_model

reg = linear_model.LinearRegression().fit(X_data, Y_data)

print("coefficients of equation of plane, (a1, a2): ", reg.coef_)

print("value of intercept, c:", reg.intercept_)

The above code will fit a plane to the given data which is linear.

To fit a second degree surface, read further.

You will have Second Degree Surface equation for the following form:

Z = a1*X + a2*Y + a3*X*Y + a4*X*X + a5*Y*Y + c

To fit this curve using linear regression, you will have to modify the above code in the following manner:

# your data is stored as X, Y, Z
print(X.shape, Y.shape, Z.shape)

x1, y1, z1 = X.flatten(), Y.flatten(), Z.flatten()
x1y1, x1x1, y1y1 = x1*y1, x1*x1, y1*y1

X_data = np.array([x1, y1, x1y1, x1x1, y1y1]).T  # X_data shape: n, 5
Y_data = z1

from sklearn import linear_model

reg = linear_model.LinearRegression().fit(X_data, Y_data)

print("coefficients of equation of plane, (a1, a2, a3, a4, a5): ", reg.coef_)

print("value of intercept, c:", reg.intercept_)