calculate feature importance

pkg/ml/ensemble/iforest/forest.go

@@ -105,3 +105,15 @@ func (f *IsolationForest) Predict(samples *mat.Dense) []int {
 	}
 	return predictions
 }
+
+func (f *IsolationForest) FeatureImportance(sample []float64) []float64 {
+	o := make([]float64, len(sample))
+
+	for _, tree := range f.Forest {
+		for i, c := range tree.FeatureImportance(sample) {
+			o[i] += float64(c) / float64(f.NumTrees)
+		}
+	}
+
+	return o
+}

pkg/ml/ensemble/iforest/tree.go

@@ -11,3 +11,21 @@ type TreeNode struct {
 func (node *TreeNode) IsLeaf() bool {
 	return node.Left == nil && node.Right == nil
 }
+
+func (node *TreeNode) trackUsedFeatures(sample []float64, featureUsed []float64) []float64 {
+	if node.IsLeaf() {
+		return featureUsed
+	}
+
+	featureUsed[node.SplitIndex]++
+
+	if sample[node.SplitIndex] < node.SplitValue {
+		return node.Left.trackUsedFeatures(sample, featureUsed)
+	} else {
+		return node.Right.trackUsedFeatures(sample, featureUsed)
+	}
+}
+
+func (node *TreeNode) FeatureImportance(sample []float64) []float64 {
+	return node.trackUsedFeatures(sample, make([]float64, len(sample)))
+}