Deleting points above a surface

Question

Given a set of 3D coordinates and its corresponding 2nd-order polynomial surface, I want to remove all points above the surface. The code snippet below shows an example of the fitted surface and points in 3D space.

import numpy as np
import matplotlib.pyplot as plt
from sklearn.preprocessing import PolynomialFeatures
from sklearn.linear_model import LinearRegression


np.random.seed(10)
x = np.random.rand(10)
y = np.random.rand(10)
z = np.random.rand(10)
X = np.column_stack((x, y))

poly = PolynomialFeatures(degree=2)
X_poly = poly.fit_transform(X)
model = LinearRegression()
model.fit(X_poly, z)

xi, yi = np.meshgrid(np.linspace(np.min(x), np.max(x), 10),
                     np.linspace(np.min(y), np.max(y), 10))

X_grid = np.column_stack((xi.flatten(), yi.flatten()))
X_grid_poly = poly.transform(X_grid)
z_grid = model.predict(X_grid_poly)
zi = np.reshape(z_grid, xi.shape)


fig = plt.figure(figsize=(10,10))
ax = fig.add_subplot(projection='3d')
ax.scatter(x,y,z, c='r', s=20)
ax.plot_surface(xi,yi,zi)
plt.show()

How can I remove all points above the fitted surface?

Answer 1

You should compare the z-coordinates of your original points with the corresponding z-coordinates predicted by the model at the x, y coordinates. If the original z-coordinate is greater than the predicted z-coordinate, you can remove it.

import numpy as np
import matplotlib.pyplot as plt
from sklearn.preprocessing import PolynomialFeatures
from sklearn.linear_model import LinearRegression


np.random.seed(10)
x = np.random.rand(10)
y = np.random.rand(10)
z = np.random.rand(10)
X = np.column_stack((x, y))

poly = PolynomialFeatures(degree=2)
X_poly = poly.fit_transform(X)
model = LinearRegression()
model.fit(X_poly, z)

xi, yi = np.meshgrid(np.linspace(np.min(x), np.max(x), 10),
                     np.linspace(np.min(y), np.max(y), 10))

X_grid = np.column_stack((xi.flatten(), yi.flatten()))
X_grid_poly = poly.transform(X_grid)
z_grid = model.predict(X_grid_poly)
zi = np.reshape(z_grid, xi.shape)

z_pred = model.predict(poly.transform(X))
mask = z <= z_pred
x_filtered, y_filtered, z_filtered = x[mask], y[mask], z[mask]

fig = plt.figure(figsize=(10,10))
ax = fig.add_subplot(projection='3d')
ax.scatter(x_filtered, y_filtered, z_filtered, c='r', s=20)
ax.plot_surface(xi,yi,zi)
plt.show()