From 6c19fbf76a88ff223ffa5bb8c476c6a81dace9a5 Mon Sep 17 00:00:00 2001
From: msilvafe <msilvafe@ucsd.edu>
Date: Thu, 25 May 2023 10:23:48 -0700
Subject: [PATCH] Add arg for plotting, add saving, default bgmap to not plot.

---
 sodetlib/operations/bias_steps.py | 26 +++++++++++++++++++++-----
 1 file changed, 21 insertions(+), 5 deletions(-)

diff --git a/sodetlib/operations/bias_steps.py b/sodetlib/operations/bias_steps.py
index 6562c40d..69f3f6f5 100644
--- a/sodetlib/operations/bias_steps.py
+++ b/sodetlib/operations/bias_steps.py
@@ -997,7 +997,8 @@ def plot_Rfrac(bsa, text_loc=(0.6, 0.8)):
 def take_bgmap(S, cfg, bgs=None, dc_voltage=0.3, step_voltage=0.01,
                step_duration=0.05, nsteps=20, high_current_mode=True,
                hcm_wait_time=0, analysis_kwargs=None, dacs='pos',
-               use_waveform=True, show_plots=True,g3_tag=None):
+               use_waveform=True, show_plots=True, g3_tag=None, plot_rfrac=False,
+               show_plots=False):
     """
     Function to easily create a bgmap. This will set all bias group voltages
     to 0 (since this is best for generating the bg map), and run bias-steps
@@ -1039,6 +1040,10 @@ def take_bgmap(S, cfg, bgs=None, dc_voltage=0.3, step_voltage=0.01,
         g3_tag: string, optional
             if not None, overrides the default tag "oper,bgmap" sent to the g3
             file
+        plot_rfrac : bool
+            Create rfrac plot, publish it, and save it. Default is False.
+        show_plots : bool
+            Show plot in addition to saving when running interactively. Default is False.
     """
     if bgs is None:
         bgs = cfg.dev.exp['active_bgs']
@@ -1062,8 +1067,8 @@ def take_bgmap(S, cfg, bgs=None, dc_voltage=0.3, step_voltage=0.01,
         S, cfg, bgs, step_voltage=step_voltage, step_duration=step_duration,
         nsteps=nsteps, high_current_mode=high_current_mode,
         hcm_wait_time=hcm_wait_time, run_analysis=True, dacs=dacs,
-        use_waveform=use_waveform, g3_tag=g3_tag,
-        analysis_kwargs=_analysis_kwargs
+        use_waveform=use_waveform, g3_tag=g3_tag, plot_rfrac=plot_rfrac,
+        show_plots=show_plots, analysis_kwargs=_analysis_kwargs
     )
 
     if hasattr(bsa, 'bgmap'):
@@ -1081,7 +1086,8 @@ def take_bgmap(S, cfg, bgs=None, dc_voltage=0.3, step_voltage=0.01,
 def take_bias_steps(S, cfg, bgs=None, step_voltage=0.05, step_duration=0.05,
                     nsteps=20, high_current_mode=True, hcm_wait_time=3,
                     run_analysis=True, analysis_kwargs=None, dacs='pos',
-                    use_waveform=True, channel_mask=None, g3_tag=None):
+                    use_waveform=True, channel_mask=None, g3_tag=None,
+                    plot_rfrac=True, show_plots=False):
     """
     Takes bias step data at the current DC voltage. Assumes bias lines
     are already in low-current mode (if they are in high-current this will
@@ -1132,6 +1138,10 @@ def take_bias_steps(S, cfg, bgs=None, step_voltage=0.05, step_duration=0.05,
         g3_tag: string, optional
             if not None, overrides the default tag "oper,bias_steps" sent to the 
             g3 file
+        plot_rfrac : bool
+            Create rfrac plot, publish it, and save it. Default is True.
+        show_plots : bool
+            Show plot in addition to saving when running interactively. Default is False.
     """
     if bgs is None:
         bgs = cfg.dev.exp['active_bgs']
@@ -1205,7 +1215,13 @@ def take_bias_steps(S, cfg, bgs=None, step_voltage=0.05, step_duration=0.05,
             if analysis_kwargs is None:
                 analysis_kwargs = {}
             bsa.run_analysis(save=True, **analysis_kwargs)
-            plot_Rfrac(bsa)
+            if plot_rfrac:
+                fig, ax = plot_Rfrac(bsa)
+                path = sdl.make_filename(S, 'bs_rfrac_summary.png', plot=True)
+                fig.savefig(path)
+                S.pub.register_file(path, 'bs_rfrac_summary', plot=True, format='png')
+                if not show_plots:
+                    plt.close(fig)
         except Exception:
             print(f"Bias step analysis failed with exception:")
             print(traceback.format_exc())