export number of dtree nodes to dataref

2 files changed, 13 insertions, 0 deletions

diff --git a/lib/functions.py b/lib/functions.py
index 0a488dc..698d68c 100644
--- a/lib/functions.py
+++ b/lib/functions.py
@@ -324,6 +324,15 @@ class SplitFunction(ModelFunction):
         ret.update(update)
         return ret
 
+    def get_number_of_nodes(self):
+        ret = 1
+        for v in self.child.values():
+            if type(v) is SplitFunction:
+                ret += v.get_number_of_nodes()
+            else:
+                ret += 1
+        return ret
+
     @classmethod
     def from_json(cls, data):
         assert data["type"] == "split"
diff --git a/lib/parameters.py b/lib/parameters.py
index 5ebf25c..40d0ba5 100644
--- a/lib/parameters.py
+++ b/lib/parameters.py
@@ -572,6 +572,10 @@ class ModelAttribute:
 
     def to_dref(self, unit=None):
         ret = {"mean": (self.mean, unit), "median": (self.median, unit)}
+
+        if type(self.model_function) is df.SplitFunction:
+            ret["decision tree nodes"] = self.model_function.get_number_of_nodes()
+
         return ret
 
     def webconf_function_map(self):