Sync avc

commit 5827b1a6fcaee2fd244f36b835ec1ace324e49be Author: Dominik Jain <[email protected]> Date: Wed Jul 20 18:34:48 2022 +0200 SkLearnDecisionTreeVectorRegressionModel: Added method plotGraphvizPDF sensai/sklearn/sklearn_regression.py
opcode81 · Jul 20, 2022 · 28fac86 · 28fac86
1 parent 250764e
commit 28fac86
Showing 1 changed file with 14 additions and 1 deletion.
diff --git a/src/sensai/sklearn/sklearn_regression.py b/src/sensai/sklearn/sklearn_regression.py
@@ -120,4 +120,17 @@ def plot(self, predictedVarName=None, figsize=None) -> plt.Figure:
         model = self.getSkLearnModel(predictedVarName)
         fig = plt.figure(figsize=figsize)
         sklearn.tree.plot_tree(model, feature_names=self.getModelInputVariableNames())
-        return fig
+        return fig
+
+    def plotGraphvizPDF(self, dotPath, predictedVarName=None):
+        """
+        :param path: the path to a .dot file that will be created, alongside which a rendered PDF file (with added suffix ".pdf")
+            will be placed
+        :param predictedVarName: the predicted variable name for which to plot the model (if multiple; None is admissible if
+            there is only one predicted variable)
+        """
+        import graphviz
+        dot = sklearn.tree.export_graphviz(self.getSkLearnModel(predictedVarName), out_file=None,
+            feature_names=self.getModelInputVariableNames(), filled=True)
+        graphviz.Source(dot).render(dotPath)
+