From 345e450ebf4d0c4d36cf99762bafafdebd1f0dd5 Mon Sep 17 00:00:00 2001
From: Birte Kristina Friesel <birte.friesel@uos.de>
Date: Tue, 16 Jan 2024 13:45:55 +0100
Subject: Add LMT support for --export-json / ModelAttribute.to_json

---
 lib/parameters.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/lib/parameters.py b/lib/parameters.py
index 0c4401c..58af96f 100644
--- a/lib/parameters.py
+++ b/lib/parameters.py
@@ -611,7 +611,11 @@ class ModelAttribute:
         return f"ModelAttribute<{self.name}, {self.attr}, mean={mean}>"
 
     def to_json(self, **kwargs):
-        if type(self.model_function) in (df.CARTFunction, df.XGBoostFunction):
+        if type(self.model_function) in (
+            df.CARTFunction,
+            df.LMTFunction,
+            df.XGBoostFunction,
+        ):
             import sklearn.tree
 
             feature_names = list(
-- 
cgit v1.2.3