diff options
| author | Daniel Friesel <daniel.friesel@uos.de> | 2022-03-17 09:15:37 +0100 |
|---|---|---|
| committer | Daniel Friesel <daniel.friesel@uos.de> | 2022-03-17 09:15:37 +0100 |
| commit | 2c14cc60d5f10bfdd8670a7409d3f83f2531c230 (patch) | |
| tree | 269b10975d01760e9d3c258b1e2a649b0fb7a4a1 | |
| parent | 49aced3b7b26c605e5be5f1382df49f05a39a452 (diff) | |
LMT: Fix max_depth and number_of_leaves getters
| -rw-r--r-- | lib/functions.py | 6 |
1 files changed, 4 insertions, 2 deletions
diff --git a/lib/functions.py b/lib/functions.py index 5aa773b..bfe0b00 100644 --- a/lib/functions.py +++ b/lib/functions.py @@ -538,9 +538,11 @@ class LMTFunction(SKLearnRegressionFunction): def get_number_of_nodes(self): return self.regressor.node_count + def get_number_of_leaves(self): + return len(self.regressor._leaves.keys()) + def get_max_depth(self): - # FIXME this returns the configured maximum depth, not the actual tree depth. - return self.regressor.max_depth + return max(map(len, self.regressor._leaves.keys())) + 1 class XGBoostFunction(SKLearnRegressionFunction): |