Fix plots in test scripts

pcpostprocess/leak_correct.py

@@ -149,7 +149,7 @@ def fit_linear_leak(current, voltage, times, ramp_start_index, ramp_end_index,
         ax4.plot(times, I_obs, label=r'$I_\mathrm{obs}$')
         ax4.plot(times, I_leak, linestyle='--', label=r'$I_\mathrm{L}$')
         ax4.plot(times, I_obs - I_leak,
-                 linestyle='--', alpha=0.5, label=r'$I_\mathrm{obs} - I_\mathrm{L}$')
+                 alpha=0.5, label=r'$I_\mathrm{obs} - I_\mathrm{L}$')
         ax4.legend(frameon=False)
 
         if not os.path.exists(output_dir):

pcpostprocess/subtraction_plots.py

@@ -53,6 +53,7 @@ def do_subtraction_plot(fig, times, sweeps, before_currents, after_currents,
 
     all_leak_params_before = []
     all_leak_params_after = []
+
     for i in range(len(sweeps)):
         before_params, _ = fit_linear_leak(before_currents, voltages, times,
                                            *ramp_bounds)
@@ -119,7 +120,6 @@ def do_subtraction_plot(fig, times, sweeps, before_currents, after_currents,
         # ax.tick_params(axis='y', rotation=90)
         # ax.yaxis.set_major_formatter(mtick.FormatStrFormatter('%.1e'))
 
-    ax = subtracted_ax
     for i, sweep in enumerate(sweeps):
         before_trace = before_currents[i, :].flatten()
         after_trace = after_currents[i, :].flatten()
@@ -130,13 +130,14 @@ def do_subtraction_plot(fig, times, sweeps, before_currents, after_currents,
 
         subtracted_currents = before_currents[i, :] - before_leak_currents[i, :] - \
             (after_currents[i, :] - after_leak_currents[i, :])
-        ax.plot(times*1e-3, subtracted_currents, label=f"sweep {sweep}", alpha=.5)
+
+        subtracted_ax.plot(times*1e-3, subtracted_currents, label=f"sweep {sweep}", alpha=.5)
 
         #  Cycle to next colour
-        ax.plot([np.nan], [np.nan], label=f"sweep {sweep}", alpha=.5)
+        subtracted_ax.plot([np.nan], [np.nan], label=f"sweep {sweep}", alpha=.5)
 
-    ax.set_ylabel(r'$I_\mathrm{obs} - I_\mathrm{L}$ (mV)')
-    ax.set_xlabel('$t$ (s)')
+    subtracted_ax.set_ylabel(r'$I_\mathrm{obs} - I_\mathrm{L}$ (mV)')
+    subtracted_ax.set_xlabel('$t$ (s)')
 
     long_protocol_ax.plot(times*1e-3, voltages, color='black')
     long_protocol_ax.set_xlabel('time (s)')

tests/test_leak_correct.py

@@ -5,6 +5,7 @@
 from syncropatch_export.trace import Trace
 
 from pcpostprocess import leak_correct
+from pcpostprocess.detect_ramp_bounds import detect_ramp_bounds
 
 
 class TestLeakCorrect(unittest.TestCase):
@@ -18,15 +19,20 @@ def setUp(self):
         if not os.path.exists(self.output_dir):
             os.makedirs(self.output_dir)
 
-        self.ramp_bounds = [1700, 2500]
         self.test_trace = Trace(test_data_dir, json_file)
 
         # get currents and QC from trace object
         self.currents = self.test_trace.get_all_traces(leakcorrect=False)
         self.currents['times'] = self.test_trace.get_times()
         self.currents['voltages'] = self.test_trace.get_voltage()
+
         self.QC = self.test_trace.get_onboard_QC_values()
 
+        # Find first times ahead of these times
+        voltage_protocol = self.test_trace.get_voltage_protocol().get_all_sections()
+        times = self.currents['times'].flatten()
+        self.ramp_bound_indices = detect_ramp_bounds(times, voltage_protocol, ramp_no=0)
+
     def test_plot_leak_fit(self):
         well = 'A01'
         sweep = 0
@@ -37,7 +43,7 @@ def test_plot_leak_fit(self):
         current = self.test_trace.get_trace_sweeps(sweeps=[sweep])[well][0, :]
 
         leak_correct.fit_linear_leak(current, voltage, times,
-                                     *self.ramp_bounds,
+                                     *self.ramp_bound_indices,
                                      output_dir=self.output_dir,
                                      save_fname=f"{well}_sweep{sweep}_leak_correction")
 
@@ -50,7 +56,8 @@ def test_get_leak_correct(self):
         times = trace.get_times()
 
         current = currents[well][sweep, :]
-        x = leak_correct.get_leak_corrected(current, voltage, times, *self.ramp_bounds)
+        x = leak_correct.get_leak_corrected(current, voltage, times,
+                                            *self.ramp_bound_indices)
         self.assertEqual(x.shape, (30784,))
 
 

tests/test_subtraction_plots.py

@@ -11,23 +11,26 @@
 
 class TestSubtractionPlots(unittest.TestCase):
     def setUp(self):
-        test_data_dir = os.path.join('tests', 'test_data', '13112023_MW2_FF',
-                                     "staircaseramp (2)_2kHz_15.01.07")
-        json_file = "staircaseramp (2)_2kHz_15.01.07.json"
+        test_data_dir_before = os.path.join('tests', 'test_data', '13112023_MW2_FF',
+                                            "staircaseramp (2)_2kHz_15.01.07")
+
+        test_data_dir_after = os.path.join('tests', 'test_data', '13112023_MW2_FF',
+                                           "staircaseramp (2)_2kHz_15.11.33")
+
+        json_file_before = "staircaseramp (2)_2kHz_15.01.07.json"
+        json_file_after = "staircaseramp (2)_2kHz_15.11.33.json"
 
         self.output_dir = os.path.join('test_output', 'test_trace_class')
 
         if not os.path.exists(self.output_dir):
             os.makedirs(self.output_dir)
 
-        self.ramp_bounds = [1700, 2500]
-
         # Use identical traces for purpose of the test
-        self.before_trace = Trace(test_data_dir, json_file)
-        self.after_trace = Trace(test_data_dir, json_file)
+        self.before_trace = Trace(test_data_dir_before, json_file_before)
+        self.after_trace = Trace(test_data_dir_after, json_file_after)
 
     def test_do_subtraction_plot(self):
-        fig = plt.figure(layout='constrained')
+        fig = plt.figure(figsize=(5, 9), layout='constrained')
         times = self.before_trace.get_times()
 
         well = 'A01'
@@ -44,6 +47,8 @@ def test_do_subtraction_plot(self):
         do_subtraction_plot(fig, times, sweeps, before_current, after_current,
                             voltages, ramp_bounds, well=well)
 
+        fig.savefig(os.path.join(self.output_dir, f"subtraction_plot_{well}"))
+
 
 if __name__ == "__main__":
     pass