From 8382877dae4e0014fce0808cba9a83156ec50e53 Mon Sep 17 00:00:00 2001
From: xadupre <xadupre@microsoft.com>
Date: Fri, 31 Mar 2023 19:41:50 +0200
Subject: [PATCH 1/2] improve tree plotting

---
 _unittests/ut_plotting/test_text_plot.py | 12 ++++++++++++
 onnx_array_api/plotting/text_plot.py     |  6 +++---
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/_unittests/ut_plotting/test_text_plot.py b/_unittests/ut_plotting/test_text_plot.py
index 4ebdcc8..8f677aa 100644
--- a/_unittests/ut_plotting/test_text_plot.py
+++ b/_unittests/ut_plotting/test_text_plot.py
@@ -65,6 +65,18 @@ def test_onnx_text_plot_tree_cls(self):
         self.assertIn("         T y=", res)
         self.assertIn("n_classes=3", res)
 
+    def test_onnx_text_plot_tree_cls_2(self):
+        iris = load_iris()
+        X_train, y_train = iris.data.astype(numpy.float32), iris.target
+        clr = DecisionTreeClassifier()
+        clr.fit(X_train, y_train)
+        model_def = to_onnx(
+            clr, X_train.astype(numpy.float32), options={"zipmap": False}
+        )
+        res = onnx_text_plot_tree(model_def.graph.node[0])
+        self.assertIn("n_classes=3", res)
+        print(res)
+
     @ignore_warnings((UserWarning, FutureWarning))
     def test_onnx_simple_text_plot_kmeans(self):
         x = numpy.random.randn(10, 3)
diff --git a/onnx_array_api/plotting/text_plot.py b/onnx_array_api/plotting/text_plot.py
index 669b1ab..5dd84ca 100644
--- a/onnx_array_api/plotting/text_plot.py
+++ b/onnx_array_api/plotting/text_plot.py
@@ -114,9 +114,9 @@ def process_tree(atts, treeid):
         for i in range(len(short[f"{prefix}_treeids"])):
             idn = short[f"{prefix}_nodeids"][i]
             node = nodes[idn]
-            node.target_nodeids = idn
-            node.target_ids = short[f"{prefix}_ids"][i]
-            node.target_weights = short[f"{prefix}_weights"][i]
+            node.append_target(
+                id=short[f"{prefix}_ids"][i], weight=short[f"{prefix}_weights"][i]
+            )
 
         def iterate(nodes, node, depth=0, true_false=""):
             node.depth = depth

From c28aeb234bbaa7c2f80c815934cddb8fcaba862d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?xavier=20dupr=C3=A9?= <xavier.dupre@gmail.com>
Date: Sat, 1 Apr 2023 14:03:29 +0200
Subject: [PATCH 2/2] fix onnx_text_plot_tree

---
 .../data/onnx_text_plot_tree_cls_2.onnx       | Bin 0 -> 1410 bytes
 _unittests/ut_plotting/test_text_plot.py      |  41 +++++++++++++----
 onnx_array_api/plotting/text_plot.py          |  42 +++++++++++++-----
 3 files changed, 64 insertions(+), 19 deletions(-)
 create mode 100644 _unittests/ut_plotting/data/onnx_text_plot_tree_cls_2.onnx

diff --git a/_unittests/ut_plotting/data/onnx_text_plot_tree_cls_2.onnx b/_unittests/ut_plotting/data/onnx_text_plot_tree_cls_2.onnx
new file mode 100644
index 0000000000000000000000000000000000000000..6efc78ecd49992e3840da68810251f19a0b54f4d
GIT binary patch
literal 1410
zcmb_cOKaOe5Z<*NTOL1R>riM&DEO56V4B$U5is63aR?zewDb^qF_N>2m&j7k%59JN
z8>LWu>d)x0^pIQcrSz9{XC+%UsUd-C*wJb<ul;6c48t_!mAmVE-p7sXPGhIlcvj0<
zqJe^KvP<@U0Q|`VJvY}}%jT|Gx(<BXvV|)`!R5xIw*lupPx7(t@&nhBQaFMKwH5OE
zA>@asH4qn)5k^_+p?pen^<yYSUOfMTqbI6qjUi<<84L+D6&O-k(>g<{AT4<c)j6V^
zh|1`Jgq}XpSyCUkxEzJDO~_40F+!HuJ!FYWvP--qvSjzKzgl+FnbK^13rvldE}_ES
z3!k@Q_%ZcYU?Pxm;8?EYGfc5EMV3=&A{u$M3VK&j_-TC5^aWI6Q*MQmKtqUzX_it8
z)gq2Nl+O)EP>u{oA`GmM&*9RC`d<tWe*{mWNMoUHr*}DcZ@CjbP*I0hQoSa8me=Yx
z9`!4QLhJ~^I)Uz+{r<u6;I#MpBy-x^KU$J3{$1Qhns@bAN$$#5dDEIxYSdd$it<FB
z^jVr^ba&)bU*uI^6jXIfswY%8;PK{G_wEF=dI{Ai7UDRIM-`A!G%UBF7HDWO+WZ1I
zP>d{fIb{3xc9YE-<#(s;IA2)1+5EZKL5}l1{&bvQb7ghUo6XNmS$tcR5%pmCA4B=t
zmm${Mla3#ZYnlF8zc;1nJ%|TgW9*unXZ`;9_8}h%DSQv7S~AU^zN_K{ECLwR%mRUw
jk+G`~PecQ<M$SwDvA1A(IxbnHyJc7iSfSkvkc#;iW0pA4

literal 0
HcmV?d00001

diff --git a/_unittests/ut_plotting/test_text_plot.py b/_unittests/ut_plotting/test_text_plot.py
index 8f677aa..c505e7d 100644
--- a/_unittests/ut_plotting/test_text_plot.py
+++ b/_unittests/ut_plotting/test_text_plot.py
@@ -52,7 +52,7 @@ def test_onnx_text_plot_tree_reg(self):
         onx = to_onnx(clr, X)
         res = onnx_text_plot_tree(onx.graph.node[0])
         self.assertIn("treeid=0", res)
-        self.assertIn("         T y=", res)
+        self.assertIn("         +f", res)
 
     def test_onnx_text_plot_tree_cls(self):
         iris = load_iris()
@@ -62,20 +62,43 @@ def test_onnx_text_plot_tree_cls(self):
         onx = to_onnx(clr, X)
         res = onnx_text_plot_tree(onx.graph.node[0])
         self.assertIn("treeid=0", res)
-        self.assertIn("         T y=", res)
+        self.assertIn("         +f 0:", res)
         self.assertIn("n_classes=3", res)
 
     def test_onnx_text_plot_tree_cls_2(self):
-        iris = load_iris()
-        X_train, y_train = iris.data.astype(numpy.float32), iris.target
-        clr = DecisionTreeClassifier()
-        clr.fit(X_train, y_train)
-        model_def = to_onnx(
-            clr, X_train.astype(numpy.float32), options={"zipmap": False}
+        this = os.path.join(
+            os.path.dirname(__file__), "data", "onnx_text_plot_tree_cls_2.onnx"
         )
+        with open(this, "rb") as f:
+            model_def = load(f)
         res = onnx_text_plot_tree(model_def.graph.node[0])
         self.assertIn("n_classes=3", res)
-        print(res)
+        expected = textwrap.dedent(
+            """
+            n_classes=3
+            n_trees=1
+            ----
+            treeid=0
+            n X2 <= 2.4499998
+               -n X3 <= 1.75
+                  -n X2 <= 4.85
+                     -f 0:0 1:0 2:1
+                     +n X0 <= 5.95
+                        -f 0:0 1:0 2:1
+                        +f 0:0 1:1 2:0
+                  +n X2 <= 4.95
+                     -n X3 <= 1.55
+                        -n X0 <= 6.95
+                           -f 0:0 1:0 2:1
+                           +f 0:0 1:1 2:0
+                        +f 0:0 1:0 2:1
+                     +n X3 <= 1.65
+                        -f 0:0 1:0 2:1
+                        +f 0:0 1:1 2:0
+               +f 0:1 1:0 2:0
+        """
+        ).strip(" \n\r")
+        self.assertEqual(expected, res.strip(" \n\r"))
 
     @ignore_warnings((UserWarning, FutureWarning))
     def test_onnx_simple_text_plot_kmeans(self):
diff --git a/onnx_array_api/plotting/text_plot.py b/onnx_array_api/plotting/text_plot.py
index 5dd84ca..dfb9be0 100644
--- a/onnx_array_api/plotting/text_plot.py
+++ b/onnx_array_api/plotting/text_plot.py
@@ -24,6 +24,14 @@ def _rule(r):
     raise ValueError(f"Unexpected rule {r!r}.")
 
 
+def _number2str(i):
+    if isinstance(i, int):
+        return str(i)
+    if int(i) == i:
+        return str(int(i))
+    return f"{i:1.2f}"
+
+
 def onnx_text_plot_tree(node):
     """
     Gives a textual representation of a tree ensemble.
@@ -61,18 +69,32 @@ def __init__(self, i, atts):
                     setattr(self, k, v[i])
             self.depth = 0
             self.true_false = ""
+            self.targets = []
+
+        def append_target(self, tid, weight):
+            self.targets.append(dict(target_id=tid, weight=weight))
 
         def process_node(self):
             "node to string"
             if self.nodes_modes == "LEAF":
-                text = "%s y=%r f=%r i=%r" % (
-                    self.true_false,
-                    self.target_weights,
-                    self.target_ids,
-                    self.target_nodeids,
-                )
+                if len(self.targets) == 0:
+                    text = f"{self.true_false}f"
+                elif len(self.targets) == 1:
+                    t = self.targets[0]
+                    text = (
+                        f"{self.true_false}f "
+                        f"{t['target_id']}:{_number2str(t['weight'])}"
+                    )
+                else:
+                    ts = " ".join(
+                        map(
+                            lambda t: f"{t['target_id']}:{_number2str(t['weight'])}",
+                            self.targets,
+                        )
+                    )
+                    text = f"{self.true_false}f {ts}"
             else:
-                text = "%s X%d %s %r" % (
+                text = "%sn X%d %s %r" % (
                     self.true_false,
                     self.nodes_featureids,
                     _rule(self.nodes_modes),
@@ -115,7 +137,7 @@ def process_tree(atts, treeid):
             idn = short[f"{prefix}_nodeids"][i]
             node = nodes[idn]
             node.append_target(
-                id=short[f"{prefix}_ids"][i], weight=short[f"{prefix}_weights"][i]
+                tid=short[f"{prefix}_ids"][i], weight=short[f"{prefix}_weights"][i]
             )
 
         def iterate(nodes, node, depth=0, true_false=""):
@@ -127,14 +149,14 @@ def iterate(nodes, node, depth=0, true_false=""):
                     nodes,
                     nodes[node.nodes_falsenodeids],
                     depth=depth + 1,
-                    true_false="F",
+                    true_false="-",
                 ):
                     yield n
                 for n in iterate(
                     nodes,
                     nodes[node.nodes_truenodeids],
                     depth=depth + 1,
-                    true_false="T",
+                    true_false="+",
                 ):
                     yield n