Explainer api for local classifiers

hiclass/LocalClassifierPerParentNode.py

@@ -225,14 +225,18 @@ def _fit_digraph(self, local_mode: bool = False, use_joblib: bool = False):
         nodes = self._get_parents()
         self._fit_node_classifier(nodes, local_mode, use_joblib)
 
-    def _get_predict_proba(model, X):
-        classifiers = [
-            model.hierarchy_.nodes[node]["classifier"]
-            for node in model.hierarchy_.nodes
-            if "classifier" in model.hierarchy_.nodes[node]
+    def get_predict_proba(self, X):
+        classifier_nodes = [
+            node
+            for node in self.hierarchy_.nodes
+            if "classifier" in self.hierarchy_.nodes[node]
         ]
 
-        # This will give list of target labels in the same order as predict_proba probabilities
-        # target_labels = [clf.classes_ for clf in classifiers]
-        predict_proba_per_classifier = [clf.predict_proba(X) for clf in classifiers]
-        return predict_proba_per_classifier
+        predict_proba_dict = {}
+        for node in classifier_nodes:
+            pred_probabilities = self.hierarchy_.nodes[node][
+                "classifier"
+            ].predict_proba(X)
+            predict_proba_dict[node] = pred_probabilities
+
+        return predict_proba_dict

tests/test_LocalClassifierPerParentNode.py

@@ -238,19 +238,20 @@ def test_explainer_tree_traversal(explainer_data):
     x_train, x_test, y_train = explainer_data
 
     lcppn.fit(x_train, y_train)
-
-    explainer = Explainer(lcppn, data=x_train, mode="tree")
-    shap_dict = explainer.explain_traversed_nodes(x_test)
-    print(shap_dict)
-
-    for key, val in shap_dict.items():
-        # Assert on shapes of shap values, must match (target_classes, num_samples, num_features)
-        model = lcppn.hierarchy_.nodes[key]["classifier"]
-        assert shap_dict[key].shape == (
-            len(model.classes_),
-            x_test.shape[0],
-            x_test.shape[1],
-        )
+    print(lcppn.get_predict_proba(x_test))
+
+    # explainer = Explainer(lcppn, data=x_train, mode="tree")
+    # shap_dict = explainer.explain_traversed_nodes(x_test)
+    # print(shap_dict)
+    #
+    # for key, val in shap_dict.items():
+    #     # Assert on shapes of shap values, must match (target_classes, num_samples, num_features)
+    #     model = lcppn.hierarchy_.nodes[key]["classifier"]
+    #     assert shap_dict[key].shape == (
+    #         len(model.classes_),
+    #         x_test.shape[0],
+    #         x_test.shape[1],
+    #     )
 
 
 # TODO: Add new test cases with hierarchies without root nodes
@@ -312,3 +313,17 @@ def test_explainer_tree_no_root(explainer_data_no_root):
             x_test.shape[0],
             x_test.shape[1],
         )
+
+
+def test_predict_proba(explainer_data):
+    rfc = RandomForestClassifier()
+    lcppn = LocalClassifierPerParentNode(
+        local_classifier=rfc, replace_classifiers=False
+    )
+
+    x_train, x_test, y_train = explainer_data
+
+    lcppn.fit(x_train, y_train)
+
+    pred_proba_dict = lcppn.get_predict_proba(x_test)
+    assert pred_proba_dict is not None