Re: Stata/Python integration part 5: Three-dimensional surface plots of marginal predictions

Hi Luis and Shawn,
Here is a code that works for me:

webuse nhanes2, clear
logistic highbp c.age##c.weight
quietly margins, at(age=(20(5)80) weight=(40(5)180)) ///
saving(predictions, replace)
use predictions, clear
rename _at1 age
rename _at2 weight
rename _margin pr_highbp
save predictions, replace

// Create the three-dimensional surface plot with Python
python:
# Import the necessary Python packages
import numpy as np
import pandas as pd
import matplotlib
import matplotlib.pyplot as plt
matplotlib.use('TkAgg')

# Read (import) the Stata dataset "predictions.dta"
# into a pandas data frame named "data"
data = pd.read_stata("predictions.dta")

# Define a 3-D graph named "ax"
ax = plt.axes(projection='3d')

# Render the graph
ax.plot_trisurf(data['age'], data['weight'], data['pr_highbp'],
cmap=plt.cm.Spectral_r)


# Specify the axis ticks
ax.set_xticks(np.arange(20, 90, step=10))
ax.set_yticks(np.arange(40, 200, step=40))
ax.set_zticks(np.arange( 0, 1.2, step=0.2))

# Specify the title and axis titles
ax.set_title("Probability of Hypertension by Age and Weight")
ax.set_xlabel("Age (years)")
ax.set_ylabel("Weight (kg)")
ax.zaxis.set_rotate_label(False)
ax.set_zlabel("Probability of Hypertension", rotation=90)

# Specify the view angle of the graph
ax.view_init(elev=30, azim=240)

plt.show()

# Save the graph
#plt.savefig("Margins3d.png", dpi=1200)
end