dataref export: include XGB hyper-parameters

2 files changed, 15 insertions, 1 deletions

diff --git a/lib/functions.py b/lib/functions.py
index 1b8612e..82e516e 100644
--- a/lib/functions.py
+++ b/lib/functions.py
@@ -731,6 +731,17 @@ class XGBoostFunction(SKLearnRegressionFunction):
     def get_complexity_score(self):
         return self.get_number_of_nodes()
 
+    def to_dref(self):
+        return {
+            "hyper/n estimators": self.regressor.n_estimators,
+            "hyper/max depth": self.regressor.max_depth,
+            "hyper/subsample": self.regressor.subsample,
+            "hyper/eta": self.regressor.learning_rate,
+            "hyper/gamma": self.regressor.gamma,
+            "hyper/alpha": self.regressor.reg_alpha,
+            "hyper/lambda": self.regressor.reg_lambda,
+        }
+
 
 # first-order linear function (no feature interaction)
 class FOLFunction(ModelFunction):
diff --git a/lib/parameters.py b/lib/parameters.py
index 0f23ef8..6183ccc 100644
--- a/lib/parameters.py
+++ b/lib/parameters.py
@@ -692,6 +692,9 @@ class ModelAttribute:
             ret["decision tree/inner nodes"] = 0
             ret["decision tree/max depth"] = 0
 
+        if type(self.model_function) == df.XGBoostFunction:
+            ret.update(self.model_function.to_dref())
+
         return ret
 
     def to_dot(self):
@@ -1056,7 +1059,7 @@ class ModelAttribute:
                 n_estimators=int(os.getenv("DFATOOL_XGB_N_ESTIMATORS", "100")),
                 max_depth=int(os.getenv("DFATOOL_XGB_MAX_DEPTH", "10")),
                 subsample=float(os.getenv("DFATOOL_XGB_SUBSAMPLE", "0.7")),
-                eta=float(os.getenv("DFATOOL_XGB_ETA", "0.3")),
+                learning_rate=float(os.getenv("DFATOOL_XGB_ETA", "0.3")),
                 gamma=float(os.getenv("DFATOOL_XGB_GAMMA", "0.01")),
                 reg_alpha=float(os.getenv("DFATOOL_XGB_REG_ALPHA", "0.0006")),
                 reg_lambda=float(os.getenv("DFATOOL_XGB_REG_LAMBDA", "1")),