def testRecursiveSetPrune(self):
numExamples = 1000
X, y = data.make_regression(numExamples)
y = Standardiser().normaliseArray(y)
numTrain = numpy.round(numExamples * 0.66)
trainX = X[0:numTrain, :]
trainY = y[0:numTrain]
testX = X[numTrain:, :]
testY = y[numTrain:]
learner = DecisionTreeLearner()
learner.learnModel(trainX, trainY)
rootId = (0,)
learner.tree.getVertex(rootId).setTestInds(numpy.arange(testX.shape[0]))
learner.recursiveSetPrune(testX, testY, rootId)
for vertexId in learner.tree.getAllVertexIds():
tempY = testY[learner.tree.getVertex(vertexId).getTestInds()]
predY = numpy.ones(tempY.shape[0])*learner.tree.getVertex(vertexId).getValue()
error = numpy.sum((tempY-predY)**2)
self.assertAlmostEquals(error, learner.tree.getVertex(vertexId).getTestError())
#Check leaf indices form all indices
inds = numpy.array([])
for vertexId in learner.tree.leaves():
inds = numpy.union1d(inds, learner.tree.getVertex(vertexId).getTestInds())
nptst.assert_array_equal(inds, numpy.arange(testY.shape[0]))
