diff --git a/pcpostprocess/leak_correct.py b/pcpostprocess/leak_correct.py index 77af47a2..c1fc76d1 100644 --- a/pcpostprocess/leak_correct.py +++ b/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): diff --git a/pcpostprocess/subtraction_plots.py b/pcpostprocess/subtraction_plots.py index 23410760..ae16a2d0 100644 --- a/pcpostprocess/subtraction_plots.py +++ b/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)') diff --git a/tests/test_leak_correct.py b/tests/test_leak_correct.py index 2fa717c5..3d831278 100755 --- a/tests/test_leak_correct.py +++ b/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,)) diff --git a/tests/test_subtraction_plots.py b/tests/test_subtraction_plots.py index 85bf47b6..14da5a0a 100755 --- a/tests/test_subtraction_plots.py +++ b/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