Allow only plotting specified graph types

Some graphs can be slow to plot and are not necessary for all users.
realrate · Sep 24, 2021 · 743d2bc · 743d2bc
1 parent 731bdac
commit 743d2bc
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/causing/graph.py b/causing/graph.py
@@ -207,11 +207,15 @@ def create_graphs(
     show_nr_indiv,
     final_var_in_percent=False,
     ids: Optional[Sequence[str]] = None,
+    graph_types=("ADE", "AME", "ATE", "IDE", "IME", "ITE"),
 ):
     """creates direct, total and mediation graph,
     for theoretical model and estimated model"""
 
     def make_graph(filename, x_weights_idmat_nodeff, y_weights_idmat_nodeff, **kwargs):
+        if filename[:3] not in graph_types:
+            return
+
         print(filename)
         xnodes = [str(var) for var in m.xvars]
         ynodes = [str(var) for var in m.yvars]