clear all
use http://www.stata-press.com/data/r16/iris, clear
describe

// Step 1: split the original dataset into training and test 
// dataset which contains 80% and 20% of observations respectively
splitsample, generate(group, replace) split(0.8 0.2) show rseed(16)

// create two frames holding the the two datasets
frame put if group==1, into(training)
frame put if group==2, into(test)

// Step 2: draw a 3D plot using the training dataset
python:
from sfi import Frame
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# use the sfi Frame class to pull data from Stata training 
# frame into Python
train = Frame.connect('training')
train_X = np.array(train.get("seplen sepwid petlen petwid"))
train_y = np.array(train.get("iris"))

fig = plt.figure(1, figsize=(10, 8))
ax = Axes3D(fig, elev=-155, azim=105)
ax.scatter(train_X[:, 0], train_X[:, 1], train_X[:, 2], c=train_y, s=30)
ax.set_xlabel("Sepal length (.cm)")
ax.set_ylabel("Sepal width (.cm)")
ax.set_zlabel("Petal length (.cm)")
plt.savefig("img2.png")
end

// Step 3: train the model with the training dataset using 
//         SVC classifier
python:
from sklearn.svm import SVC
svc_clf = SVC(gamma='auto', random_state=16)
svc_clf.fit(train_X, train_y)
end

// Step 4: make predictions with the test dataset and   
//         store results back into the test frame
python:
test = Frame.connect('test')
test_X = np.array(test.get("seplen sepwid petlen petwid"))
irispr = svc_clf.predict(test_X)
test.addVarByte('irispr')
test.store('irispr', None, irispr)
end

// Step 5: see prediction results in Stata
frame change test 
label values irispr species
label variable irispr predicted
tabulate iris irispr, row
