updates so plot_graphic_file_flow will run with beast call

BEAST-Fitting · Dec 17, 2020 · ecf8733 · ecf8733
1 parent 59311e1
commit ecf8733
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 3 deletions.
diff --git a/beast/main.py b/beast/main.py
@@ -10,6 +10,7 @@
     plot_cmd_with_fits,
     plot_completeness,
     plot_filters,
+    plot_graphic_file_flow,
     plot_graphic_model,
     plot_indiv_pdfs,
     plot_noisemodel,
@@ -43,6 +44,7 @@ def main():
         plot_cmd_with_fits,
         plot_completeness,
         plot_filters,
+        plot_graphic_file_flow,
         plot_graphic_model,
         plot_indiv_pdfs,
         plot_noisemodel,

diff --git a/beast/plotting/plot_graphic_file_flow.py b/beast/plotting/plot_graphic_file_flow.py
@@ -9,13 +9,28 @@
 ]
 
 
-def plot_graphic_file_flow():
+def plot_graphic_file_flow(plot_type="both", savefig="png"):
     """
     Function to call the source density and subgrid functions
+
+    Parameters
+    ----------
+    plot_type : string
+        Choose either 'subgrid', 'sd', or 'both' (case insensitive)
+
+    savefig : str (default='png')
+        set to the file extension of desired file to save image of model
     """
 
-    plot_graphic_file_flow_sd(n_sd=2, n_sub=3)
-    plot_graphic_file_flow_subgrid(n_sg=4)
+    if plot_type.lower() == "sd":
+        plot_graphic_file_flow_sd(n_sd=2, n_sub=3, savefig=savefig)
+    elif plot_type.lower() == "subgrid":
+        plot_graphic_file_flow_subgrid(n_sg=4, savefig=savefig)
+    elif plot_type.lower() == "both":
+        plot_graphic_file_flow_sd(n_sd=2, n_sub=3, savefig=savefig)
+        plot_graphic_file_flow_subgrid(n_sg=4, savefig=savefig)
+    else:
+        raise ValueError("plot_type must be 'subgrid', 'sd', or 'both'")
 
 
 def plot_graphic_file_flow_sd(n_sd=1, n_sub=3, savefig="png"):