From d6c99e9e568e32f58020e11444a7820233664702 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ricardo=20Montan=CC=83ana?= <rmontanana@gmail.com>
Date: Wed, 30 Mar 2022 11:58:16 +0200
Subject: [PATCH] Add graphviz representation of the tree

---
 stree/Splitter.py         | 24 ++++++++++++++++++++++++
 stree/Strees.py           | 14 ++++++++++++++
 stree/tests/Stree_test.py | 17 +++++++++++++++++
 3 files changed, 55 insertions(+)

diff --git a/stree/Splitter.py b/stree/Splitter.py
index 2829503..e9ac169 100644
--- a/stree/Splitter.py
+++ b/stree/Splitter.py
@@ -145,6 +145,30 @@ class Snode:
             except IndexError:
                 self._class = None
 
+    def graph(self):
+        """
+        Return a string representing the node in graphviz format
+        """
+        output = ""
+        count_values = np.unique(self._y, return_counts=True)
+        if self.is_leaf():
+            output += (
+                f'N{id(self)} [shape=box style=filled label="'
+                f"class={self._class} belief={self._belief: .3f} "
+                f"impurity={self._impurity:.3f} "
+                f'classes/samples={count_values}"];\n'
+            )
+        else:
+            output += (
+                f'N{id(self)} [label="#features={len(self._features)} '
+                f'classes/samples={count_values}"];\n'
+            )
+            output += f'N{id(self)} -> N{id(self.get_up())} [label="Up"];\n'
+            output += (
+                f'N{id(self)} -> N{id(self.get_down())} [label="Down"];\n'
+            )
+        return output
+
     def __str__(self) -> str:
         count_values = np.unique(self._y, return_counts=True)
         if self.is_leaf():
diff --git a/stree/Strees.py b/stree/Strees.py
index ccb1044..97dc181 100644
--- a/stree/Strees.py
+++ b/stree/Strees.py
@@ -476,6 +476,20 @@ class Stree(BaseEstimator, ClassifierMixin):
             tree = None
         return Siterator(tree)
 
+    def graph(self) -> str:
+        """Graphviz code representing the tree
+
+        Returns
+        -------
+        str
+            graphviz code
+        """
+        output = "digraph STree {\n"
+        for node in self:
+            output += node.graph()
+        output += "}\n"
+        return output
+
     def __str__(self) -> str:
         """String representation of the tree
 
diff --git a/stree/tests/Stree_test.py b/stree/tests/Stree_test.py
index 3813de4..b46e6ff 100644
--- a/stree/tests/Stree_test.py
+++ b/stree/tests/Stree_test.py
@@ -666,3 +666,20 @@ class Stree_test(unittest.TestCase):
     def test_version(self):
         clf = Stree()
         self.assertEqual(__version__, clf.version())
+
+    def test_graph(self):
+        X, y = load_wine(return_X_y=True)
+        clf = Stree(random_state=self._random_state)
+        clf.fit(X, y)
+        expected_head = "digraph STree {\n"
+        expected_tail = (
+            ' [shape=box style=filled label="class=1 belief= '
+            '1.000 impurity=0.000 classes/samples=(array([1]), array([1]))"]'
+            ";\n}\n"
+        )
+        computed = clf.graph()
+        computed_head = computed[: len(expected_head)]
+        num = -len(expected_tail)
+        computed_tail = computed[num:]
+        self.assertEqual(computed_head, expected_head)
+        self.assertEqual(computed_tail, expected_tail)