Create toy dataset for regression

create_toy.py

from numpy import *
from pylab import *
import numpy
from sklearn.mixture import GMM
from mpl_toolkits.mplot3d import axes3d

numpy.random.seed(0)

X = arange(0, 2*numpy.pi, 0.01)
N = len(X)
Y = sin(X)**3 + numpy.random.randn(*X.shape)*0.1
num_clusters = 3

data = numpy.hstack((X[:, None], Y[:, None]))
g = GMM(num_clusters, random_state=0).fit(data)
cl = g.predict_proba(data)


figure()
training = []
testing = []
for n in range(N):
    c = (cl[n, 0], cl[n, 1], cl[n, 2])
    plot(X[n], Y[n], "o", color=c)
    if not(cl[n, 0] > cl[n, 1] and cl[n, 0] > cl[n, 2]):
        training.append([X[n], Y[n]])
    else:
        testing.append([X[n], Y[n]])

training = numpy.array(training)
training_out = (training[:, 0] * training[:, 1])**2
testing  = numpy.array(testing)
testing_out = (testing[:, 0] * testing[:, 1])**2

figure()
plot(training[:, 0], training[:, 1], "ro", label="training")
plot(testing[:, 0], testing[:, 1], "bo", label="testing")
legend()

fig = figure()
ax = axes3d.Axes3D(fig)
ax.plot3D(training[:, 0], training[:, 1], training_out, "ro")
ax.plot3D(testing[:, 0], testing[:, 1], testing_out, "bo")

show()

# Use training + training_out for training
# and testing + testing_out for validation