FOLFunction: Implement to_dot

2 files changed, 11 insertions, 0 deletions

diff --git a/lib/functions.py b/lib/functions.py
index 74b5509..dfa5b94 100644
--- a/lib/functions.py
+++ b/lib/functions.py
@@ -834,6 +834,16 @@ class FOLFunction(ModelFunction):
     def get_complexity_score(self):
         return len(self.model_args)
 
+    def to_dot(self, pydot, graph, feature_names, parent=None):
+        model_function = self.model_function
+        for i, arg in enumerate(self.model_args):
+            model_function = model_function.replace(
+                f"regression_arg({i})", f"{arg:.2f}"
+            )
+        graph.add_node(
+            pydot.Node(str(id(self)), label=model_function, shape="rectangle")
+        )
+
     def to_json(self, **kwargs):
         ret = super().to_json(**kwargs)
         ret.update(
diff --git a/lib/parameters.py b/lib/parameters.py
index bc0d2a1..e6d5561 100644
--- a/lib/parameters.py
+++ b/lib/parameters.py
@@ -681,6 +681,7 @@ class ModelAttribute:
             df.SplitFunction,
             df.StaticFunction,
             df.AnalyticFunction,
+            df.FOLFunction,
         ):
             import pydot