diff options
Diffstat (limited to 'lib')
| -rw-r--r-- | lib/functions.py | 10 | ||||
| -rw-r--r-- | lib/parameters.py | 1 |
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 |