Add tree attribute export for XGBoost

2 files changed, 32 insertions, 10 deletions

diff --git a/lib/functions.py b/lib/functions.py
index 14893ad..43ee58a 100644
--- a/lib/functions.py
+++ b/lib/functions.py
@@ -471,19 +471,37 @@ class CARTFunction(SKLearnRegressionFunction):
 
 class XGBoostFunction(SKLearnRegressionFunction):
     def get_number_of_nodes(self):
+        import json
+
+        self.regressor.get_booster().dump_model(
+            "/tmp/xgb.json", dump_format="json", with_stats=True
+        )
+        with open("/tmp/xgb.json", "r") as f:
+            data = json.load(f)
+
+        return sum(map(self._get_number_of_nodes, data))
+
+    def _get_number_of_nodes(self, data):
         ret = 1
-        for v in self.child.values():
-            if type(v) is SplitFunction:
-                ret += v.get_number_of_nodes()
-            else:
-                ret += 1
+        for child in data.get("children", list()):
+            ret += self._get_number_of_nodes(child)
         return ret
 
     def get_max_depth(self):
+        import json
+
+        self.regressor.get_booster().dump_model(
+            "/tmp/xgb.json", dump_format="json", with_stats=True
+        )
+        with open("/tmp/xgb.json", "r") as f:
+            data = json.load(f)
+
+        return max(map(self._get_max_depth, data))
+
+    def _get_max_depth(self, data):
         ret = [0]
-        for v in self.child.values():
-            if type(v) is SplitFunction:
-                ret.append(v.get_max_depth())
+        for child in data.get("children", list()):
+            ret.append(self._get_max_depth(child))
         return 1 + max(ret)
 
 
diff --git a/lib/parameters.py b/lib/parameters.py
index 9350a06..4e98f54 100644
--- a/lib/parameters.py
+++ b/lib/parameters.py
@@ -666,7 +666,11 @@ class ModelAttribute:
     def to_dref(self, unit=None):
         ret = {"mean": (self.mean, unit), "median": (self.median, unit)}
 
-        if type(self.model_function) in (df.SplitFunction, df.CARTFunction):
+        if type(self.model_function) in (
+            df.SplitFunction,
+            df.CARTFunction,
+            df.XGBoostFunction,
+        ):
             ret["decision tree/nodes"] = self.model_function.get_number_of_nodes()
             ret["decision tree/max depth"] = self.model_function.get_max_depth()
 
@@ -961,7 +965,7 @@ class ModelAttribute:
                 self.model_function = df.StaticFunction(np.mean(data))
                 return
             xgb.fit(fit_parameters, np.reshape(data, (-1, 1)))
-            self.model_function = df.SKLearnRegressionFunction(
+            self.model_function = df.XGBoostFunction(
                 np.mean(data), xgb, category_to_index, ignore_index
             )
             output_filename = os.getenv("DFATOOL_XGB_DUMP_MODEL", None)