From 5adc39c8de1fbda6e1ba9abcaaaf17516eb046ae Mon Sep 17 00:00:00 2001
From: Birte Kristina Friesel <birte.friesel@uos.de>
Date: Wed, 10 Jan 2024 09:41:50 +0100
Subject: --show-model=param: add CART support

---
 lib/cli.py | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

(limited to 'lib')

diff --git a/lib/cli.py b/lib/cli.py
index 314c6dd..a1655c6 100644
--- a/lib/cli.py
+++ b/lib/cli.py
@@ -94,6 +94,26 @@ def print_staticinfo(prefix, info):
     print(f"{prefix}: {info.value}")
 
 
+def print_cartinfo(prefix, info, feature_names):
+    _print_cartinfo(prefix, info.to_json(feature_names=feature_names), feature_names)
+
+
+def _print_cartinfo(prefix, model, feature_names):
+    if model["type"] == "static":
+        print(f"""{prefix}: {model["value"]}""")
+    else:
+        _print_cartinfo(
+            f"""{prefix} {model["paramName"]}<{model["paramDecisionValue"]} """,
+            model["left"],
+            feature_names,
+        )
+        _print_cartinfo(
+            f"""{prefix} {model["paramName"]}≥{model["paramDecisionValue"]} """,
+            model["right"],
+            feature_names,
+        )
+
+
 def print_splitinfo(param_names, info, prefix=""):
     if type(info) is SplitFunction:
         for k, v in info.child.items():
-- 
cgit v1.2.3