Doctorado-ML/stree
1
1
Fork 0
mirror of https://github.com/Doctorado-ML/STree.git synced 2025-08-18 00:46:02 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add optional title to tree graph

Browse Source
This commit is contained in:
Ricardo Montañana Gómez
2022-04-04 20:22:47 +02:00
parent cf63863e64
commit 08f8ac018b
3 changed files with 41 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
stree/Splitter.py
View File

@@ -154,19 +154,17 @@ class Snode:
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'
f"class={self._class} impurity={self._impurity:.3f} "
f'classes={count_values[0]} samples={count_values[1]}"];\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'
f"classes={count_values[0]} samples={count_values[1]} "
f'({sum(count_values[1])})"];\n'
)
output += f"N{id(self)} -> N{id(self.get_up())};\n"
output += f"N{id(self)} -> N{id(self.get_down())};\n"
return output
def __str__(self) -> str:

7
stree/Strees.py
View File

@@ -476,7 +476,7 @@ class Stree(BaseEstimator, ClassifierMixin):
tree = None
return Siterator(tree)
def graph(self) -> str:
def graph(self, title="") -> str:
"""Graphviz code representing the tree
Returns
@@ -484,7 +484,10 @@ class Stree(BaseEstimator, ClassifierMixin):
str
graphviz code
"""
output = "digraph STree {\n"
output = (
"digraph STree {\nlabel=<STree "
f"{title}>\nfontsize=30\nfontcolor=blue\nlabelloc=t\n"
)
for node in self:
output += node.graph()
output += "}\n"

37
stree/tests/Stree_test.py
View File

@@ -688,17 +688,40 @@ class Stree_test(unittest.TestCase):
"""Check graphviz representation of the tree."""
X, y = load_wine(return_X_y=True)
clf = Stree(random_state=self._random_state)
self.assertEqual(clf.graph(), "digraph STree {\n}\n")
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"
expected_head = (
"digraph STree {\nlabel=<STree >\nfontsize=30\n"
"fontcolor=blue\nlabelloc=t\n"
)
expected_tail = (
' [shape=box style=filled label="class=1 impurity=0.000 '
'classes=[1] samples=[1]"];\n}\n'
)
self.assertEqual(clf.graph(), expected_head + "}\n")
clf.fit(X, y)
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)
def test_graph_title(self):
X, y = load_wine(return_X_y=True)
clf = Stree(random_state=self._random_state)
expected_head = (
"digraph STree {\nlabel=<STree Sample title>\nfontsize=30\n"
"fontcolor=blue\nlabelloc=t\n"
)
expected_tail = (
' [shape=box style=filled label="class=1 impurity=0.000 '
'classes=[1] samples=[1]"];\n}\n'
)
self.assertEqual(clf.graph("Sample title"), expected_head + "}\n")
clf.fit(X, y)
computed = clf.graph("Sample title")
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)