dataref export: add number of inner decision tree nodes

1 files changed, 6 insertions, 8 deletions

diff --git a/lib/parameters.py b/lib/parameters.py
index 341b463..a3c5ab1 100644
--- a/lib/parameters.py
+++ b/lib/parameters.py
@@ -654,19 +654,17 @@ class ModelAttribute:
             df.LMTFunction,
         ):
             ret["decision tree/nodes"] = self.model_function.get_number_of_nodes()
+            ret["decision tree/leaves"] = self.model_function.get_number_of_leaves()
+            ret["decision tree/inner nodes"] = (
+                ret["decision tree/nodes"] - ret["decision tree/leaves"]
+            )
             ret["decision tree/max depth"] = self.model_function.get_max_depth()
         elif type(self.model_function) in (df.StaticFunction, df.AnalyticFunction):
             ret["decision tree/nodes"] = 1
+            ret["decision tree/leaves"] = 1
+            ret["decision tree/inner nodes"] = 0
             ret["decision tree/max depth"] = 0
 
-        if type(self.model_function) in (
-            df.SplitFunction,
-            df.CARTFunction,
-            df.LMTFunction,
-            df.XGBoostFunction,
-        ):
-            ret["decision tree/leaves"] = self.model_function.get_number_of_leaves()
-
         return ret
 
     def to_dot(self):