store decision tree attributes of xv models in dataref export

1 files changed, 7 insertions, 0 deletions

diff --git a/lib/functions.py b/lib/functions.py
index 320e8ed..5358d8e 100644
--- a/lib/functions.py
+++ b/lib/functions.py
@@ -339,6 +339,13 @@ class SplitFunction(ModelFunction):
                 ret += 1
         return ret
 
+    def get_max_depth(self):
+        ret = [0]
+        for v in self.child.values():
+            if type(v) is SplitFunction:
+                ret.append(v.get_max_depth())
+        return 1 + max(ret)
+
     @classmethod
     def from_json(cls, data):
         assert data["type"] == "split"