Best-fit Surfaces for 3-Dimensional Data

The following code generates best-fit planes for 3-dimensional data using linear regression techniques (1st-order and 2nd-order polynomials). Although I recently developed this code to analyze data for the Bridger-Teton Avalanche Center, below I generate a random dataset using a Gaussian function. The result approximates a slightly warped half-cyclinder surface.

I also continue with the techniques demonstrated in the previous post, showing the resulting 3-dimensional plots in both Matplotlib and Plotly. Although the Matplotlylib project (Matplotlib + Plotly = Matplotlylib) has tools to convert Matplotlib plots into Plotly interactive plots (using 'tls.mpl_to_plotly' or 'py.iplot_mpl'), these functions fail to successfully convert the fitted planes in my 3-dimensional Matplotlib plots. It is possible that this is not a supported plot type for Matplotlylib.

Instead, it is possible to reproduce the same plot by writing the code directly for a Plotly figure. For both the 1st-order and 2nd-order planes, I show the Matplotlib and Plotly figures, with separate code for each. They are very similar, except of course with the ability to interactively zoom and rotate the Plotly figures! With the Plotly figure, the best-fit plane by default is also color-coded for z-values with a corresponding scale bar. The ability to interactively zoom-in and rotate the Plotly figures is very powerful, and adds a huge amount of value to any 3-dimensional plot displayed online or in a Jupyter Notebook.

In [17]:
%matplotlib inline
In [18]:
import numpy as np
import scipy.linalg
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt

# Import plotly package
import plotly
import plotly.graph_objs as go

# Check ploltly version
plotly.__version__

# To communicate with Plotly's server, sign in with credentials file
import plotly.plotly as py
In [19]:
# Create data with x and y random over [-2, 2], and z a Gaussian function of x and y.
np.random.seed(12345)
x = 2 * (np.random.random(500) - 0.5)
y = 2 * (np.random.random(500) - 0.5)

def f(x, y):
    return np.exp(-(x + y ** 2))

z = f(x, y)

data = np.c_[x,y,z]

# regular grid covering the domain of the data
mn = np.min(data, axis=0)
mx = np.max(data, axis=0)
X,Y = np.meshgrid(np.linspace(mn[0], mx[0], 20), np.linspace(mn[1], mx[1], 20))
XX = X.flatten()
YY = Y.flatten()
    
# best-fit linear plane (1st-order)
A = np.c_[data[:,0], data[:,1], np.ones(data.shape[0])]
C,_,_,_ = scipy.linalg.lstsq(A, data[:,2])    # coefficients
    
# evaluate it on grid
Z = C[0]*X + C[1]*Y + C[2]
    
# or expressed using matrix/vector product
#Z = np.dot(np.c_[XX, YY, np.ones(XX.shape)], C).reshape(X.shape)
In [20]:
# plot points and fitted surface using Matplotlib
fig1 =  plt.figure(figsize=(10, 10))
ax = fig1.gca(projection='3d')
ax.plot_surface(X, Y, Z, rstride=1, cstride=1, alpha=0.2)
ax.scatter(data[:,0], data[:,1], data[:,2], c='r', s=50)
plt.xlabel('X')
plt.ylabel('Y')
ax.set_zlabel('Z')
ax.axis('equal')
ax.axis('tight')

# plot points and fitted surface using Plotly
trace1 = go.Scatter3d(
    x=data[:,0],
    y=data[:,1],
    z=data[:,2],
    mode='markers',
    marker=dict(size=4, color='red', line=dict(color='black', width=0.5), opacity=0.8)
)

trace2 = go.Surface(z=Z, x=X, y=Y, colorscale='RdBu', opacity=0.999)

# Package the trace dictionary into a data object
data_test1 = go.Data([trace1, trace2])

# Dictionary of style options for all axes
axis = dict(
    showbackground=True, # show axis background
    backgroundcolor="rgb(204, 204, 204)", # set background color to grey
    gridcolor="rgb(255, 255, 255)",       # set grid line color
    zerolinecolor="rgb(255, 255, 255)",   # set zero grid line color
)

# Make a layout object
layout = go.Layout(
    title='1st-order (linear) plane', # set plot title
    scene=go.Scene(  # axes are part of a 'scene' in 3d plots
        xaxis=go.XAxis(axis), # set x-axis style
        yaxis=go.YAxis(axis), # set y-axis style
        zaxis=go.ZAxis(axis)),  # set z-axis style
)

# Make a figure object
fig = go.Figure(data=data_test1, layout=layout)

# Send to Plotly and show in notebook
py.iplot(fig, filename='test1')
Out[20]:
In [21]:
# best-fit quadratic curve (2nd-order)
A = np.c_[np.ones(data.shape[0]), data[:,:2], np.prod(data[:,:2], axis=1), data[:,:2]**2]
C,_,_,_ = scipy.linalg.lstsq(A, data[:,2])
    
# evaluate it on a grid
Z = np.dot(np.c_[np.ones(XX.shape), XX, YY, XX*YY, XX**2, YY**2], C).reshape(X.shape)
In [23]:
# plot points and fitted surface using Matplotlib
fig = plt.figure(figsize=(10, 10))
ax = fig.gca(projection='3d')
ax.plot_surface(X, Y, Z, rstride=1, cstride=1, alpha=0.2)
ax.scatter(data[:,0], data[:,1], data[:,2], c='r', s=50)
plt.xlabel('X')
plt.ylabel('Y')
ax.set_zlabel('Z')
ax.axis('equal')
ax.axis('tight')

# plot points and fitted surface using Plotly
trace3 = go.Surface(
    z=Z,  
    x=X, 
    y=Y,
    colorscale='RdBu',
    opacity=0.999
)

# Package the trace dictionary into a data object
data_test2 = go.Data([trace1, trace3])

# Make a layout object
layout = go.Layout(
    title='2nd-order (quadratic) surface', # set plot title
    scene=go.Scene(  # axes are part of a 'scene' in 3d plots
        xaxis=go.XAxis(axis), # set x-axis style
        yaxis=go.YAxis(axis), # set y-axis style
        zaxis=go.ZAxis(axis)),  # set z-axis style
)

# Make a figure object
fig = go.Figure(data=data_test2, layout=layout)

# Send to Plotly and show in notebook
py.iplot(fig, filename='test2')
Out[23]: