Python 3D Scatter Plot: Customize Color for Points Above and Below a Plane

This Python code demonstrates how to customize the color of points in a 3D scatter plot, differentiating between those above and below a regression plane. It uses the 'statsmodels' library for linear regression and 'matplotlib' for plotting.

from cProfile import label
import statsmodels.api as sm
import statsmodels.formula.api as smf
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from mpl_toolkits import mplot3d

# Import necessary libraries for plotting and data handling
data = pd.read_csv('cobb_douglas.csv')  # Load data from a CSV file
model = smf.ols('lny ~ lnk + lnl', data=data)  # Define the regression model
results = model.fit()  # Fit the regression model
print(results.params)  # Print regression coefficients

# Generate data for the regression plane
xx = np.linspace(data.lnk.min(), data.lnk.max(), 100)
yy = np.linspace(data.lnl.min(), data.lnl.max(), 100)
XX, YY = np.meshgrid(xx, yy)
ZZ = results.params[0] + XX * results.params[1] + YY * results.params[2]

# Create the 3D plot
fig = plt.figure()
ax = plt.axes(projection='3d')
ax.plot_surface(XX, YY, ZZ, rstride=8, cstride=8, alpha=0.4, cmap='pink_r', label='lny = 0.233054*lnk + 0.807278*lnl - 0.177310')

# Plot points, differentiating color based on position relative to the plane
ax.scatter(data.lnk, data.lnl, data.lny, c='red', label='Points below the plane')  # Red for points below the plane
ax.scatter(data.lnk, data.lnl, data.lny, c='green', label='Points above the plane')  # Green for points above the plane

ax.set_xlabel('lnk')
ax.set_ylabel('lnl')
ax.set_zlabel('lny')
ax.set_title("Figure 1.1 lny's regression to lnk and lnl", fontdict=None, loc="center", pad=None)
plt.show()

This code demonstrates a common technique in data visualization: using color to highlight different aspects of data. By changing the color of points based on their position relative to the regression plane, we can visually understand the relationship between the variables and the model's predictions.