def testLearnModel(self):
maxDepth = 2
treeRankForest = TreeRankForest(self.leafRanklearner)
treeRankForest.setMaxDepth(maxDepth)
treeRankForest.learnModel(self.X, self.y)
forest = treeRankForest.getForest()
self.assertEquals(len(forest), treeRankForest.getNumTrees())
for treeRank in forest:
tree = treeRank.getTree()
self.assertTrue(tree.depth() <= maxDepth)
def testPredict(self):
maxDepth = 2
treeRankForest = TreeRankForest(self.leafRanklearner)
treeRankForest.setMaxDepth(maxDepth)
treeRankForest.learnModel(self.X, self.y)
scores = treeRankForest.predict(self.X)
scores2 = numpy.zeros(self.X.shape[0])
forest = treeRankForest.getForest()
for i in range(len(forest)):
scores2 += forest[i].predict(self.X)
scores2 /= treeRankForest.getNumTrees()
self.assertTrue((scores==scores2).all())
