Repository: Rambatino/CHAID

Issue:
model_predictions fails with categorical dependant variables
If the dependent variable is categorical, where categories are strings, the method model_predictions fails. The problem is that the the pred array is initialized as:

    pred = np.zeros(self.data_size)

and that enforces predictions to be numerical. In order to solve that, the model_predictions could be rewritten to something like the following:

    pred = [None] * self.data_size
    for node in self:
        if node.is_terminal:
            max_val = max(node.members, key=node.members.get)
            for i in node.indices:
                pred[i] = max_val
    return pred

Best regards

Discussion:

Patch:
diff --git a/CHAID/tree.py b/CHAID/tree.py
index 17e192a..86508e2 100644
--- a/CHAID/tree.py
+++ b/CHAID/tree.py
@@ -221,7 +221,7 @@ class Tree(object):
 
     def print_tree(self):
         """ prints the tree out """
-        self.to_tree().show()
+        self.to_tree().show(line_type='ascii')
 
     def node_predictions(self):
         """ Determines which rows fall into which node """
@@ -264,7 +264,7 @@ class Tree(object):
         """
         if isinstance(self.observed, ContinuousColumn):
             return ValueError("Cannot make model predictions on a continuous scale")
-        pred = np.zeros(self.data_size)
+        pred = np.zeros(self.data_size).astype('object')
         for node in self:
             if node.is_terminal:
                 pred[node.indices] = max(node.members, key=node.members.get)
@@ -275,4 +275,5 @@ class Tree(object):
         Calculates the fraction of risk associated
         with the model predictions
         """
-        return 1 - float((self.model_predictions() == self.observed.arr).sum()) / self.data_size
+        sub_observed = np.array([self.observed.metadata[i] for i in self.observed.arr])
+        return 1 - float((self.model_predictions() == sub_observed).sum()) / self.data_size


Test patch:
diff --git a/tests/test_tree.py b/tests/test_tree.py
index fcb5b60..776ab6c 100644
--- a/tests/test_tree.py
+++ b/tests/test_tree.py
@@ -570,3 +570,41 @@ class TestContinuousDependentVariable(TestCase):
         tree = CHAID.Tree.from_numpy(self.ndarr, self.normal_arr, alpha_merge=0.999, max_depth=5, min_child_node_size=11, dep_variable_type='continuous', weights=self.wt)
         assert round(tree.tree_store[0].p, 4) == 0.3681
         assert len(tree.tree_store) == 5
+
+
+class TestStringCategoricalDependentVariableForModelPrediction(TestCase):
+    """ Test to make sure we can handle string categorical dependent varaibles """
+    def setUp(self):
+        """
+        Setup data for test case
+        """
+        self.region = np.array([
+            1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
+            1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 3, 3, 2, 3, 2, 2, 2,
+            3, 2, 4, 4, 2, 4, 4, 4, 2, 2, 2, 2, 3, 2, 3, 2, 3, 2, 2, 2])
+        self.age = np.array([
+            3, 4, 4, 3, 2, 4, 2, 3, 3, 2, 2, 3, 4, 3, 4, 2, 2, 3, 2, 3,
+            2, 4, 4, 3, 2, 3, 1, 2, 4, 4, 3, 4, 4, 3, 2, 4, 2, 3, 3, 2,
+            2, 3, 4, 3, 4, 2, 2, 3, 2, 3, 2, 4, 4, 3, 2, 3, 1, 2, 4, 4])
+        self.gender = np.array([
+            1, 2, 1, 2, 2, 2, 2, 2, 1, 2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 2,
+            2, 2, 2, 1, 2, 1, 2, 1, 2, 2, 1, 2, 1, 2, 2, 2, 2, 2, 1, 2,
+            2, 2, 2, 2, 1, 1, 1, 1, 1, 2, 2, 2, 2, 1, 2, 1, 2, 1, 2, 2])
+        self.lover = np.array(['lover'] * 25 + ['non-lover'] * 35)
+        self.tree = CHAID.Tree.from_numpy(
+            np.vstack([self.region, self.age, self.gender]).transpose(),
+            self.lover,
+            alpha_merge=0.05
+        )
+
+    def test_string_dependent_categorical_variable_for_model_prediction(self):
+        assert (self.tree.model_predictions() ==  np.array(['lover'] * 30 + ['non-lover'] * 30)).all()
+
+    def test_risk_still_works(self):
+        int_lover = np.array([1] * 25 + [0] * 35)
+        other_tree = CHAID.Tree.from_numpy(
+            np.vstack([self.region, self.age, self.gender]).transpose(),
+            int_lover,
+            alpha_merge=0.05
+        )
+        assert self.tree.risk() == other_tree.risk()