code optimisation

BlueBrain · Feb 12, 2024 · 1439e8c · 1439e8c
1 parent 70c036c
commit 1439e8c
Showing 1 changed file with 33 additions and 34 deletions.
diff --git a/bluepyefe/cell.py b/bluepyefe/cell.py
@@ -239,52 +239,51 @@ def plot_recordings(self, protocol_name, output_dir=None, show=False):
             dirname = pathlib.Path(output_dir) / self.name
             dirname.mkdir(parents=True, exist_ok=True)
             filepath = dirname / filename
-            pdf = PdfPages(filepath)
 
-        for page in range(total_pages):
-            start_idx = page * max_plots_per_page
-            end_idx = start_idx + max_plots_per_page
-            page_recordings = recordings_sorted[start_idx:end_idx]
+            with PdfPages(filepath) as pdf:
+                for page in range(total_pages):
+                    start_idx = page * max_plots_per_page
+                    end_idx = start_idx + max_plots_per_page
+                    page_recordings = recordings_sorted[start_idx:end_idx]
 
-            n_rows = int(numpy.ceil(len(page_recordings) / n_cols)) * 2
+                    n_rows = int(numpy.ceil(len(page_recordings) / n_cols)) * 2
 
-            fig, axs = plt.subplots(
-                n_rows, n_cols,
-                figsize=[3.0 * n_cols, 2.5 * n_rows],
-                squeeze=False
-            )
+                    fig, axs = plt.subplots(
+                        n_rows, n_cols,
+                        figsize=[3.0 * n_cols, 2.5 * n_rows],
+                        squeeze=False
+                    )
 
-            for i, rec in enumerate(page_recordings):
-                col = i % n_cols
-                row = (i // n_cols) * 2
+                    for i, rec in enumerate(page_recordings):
+                        col = i % n_cols
+                        row = (i // n_cols) * 2
 
-                display_ylabel = col == 0
-                display_xlabel = (row // 2) + 1 == n_rows // 2
+                        display_ylabel = col == 0
+                        display_xlabel = (row // 2) + 1 == n_rows // 2
 
-                rec.plot(
-                    axis_current=axs[row][col],
-                    axis_voltage=axs[row + 1][col],
-                    display_xlabel=display_xlabel,
-                    display_ylabel=display_ylabel
-                )
+                        rec.plot(
+                            axis_current=axs[row][col],
+                            axis_voltage=axs[row + 1][col],
+                            display_xlabel=display_xlabel,
+                            display_ylabel=display_ylabel
+                        )
 
-            fig.suptitle(f"Cell: {self.name}, Experiment: {protocol_name}, Page: {page + 1}")
-            plt.subplots_adjust(wspace=0.53, hspace=0.7)
+                    fig.suptitle(f"Cell: {self.name}, Experiment: {protocol_name}, Page: {page + 1}")
+                    plt.subplots_adjust(wspace=0.53, hspace=0.7)
 
-            for ax in axs.flatten():
-                if not ax.lines:
-                    ax.set_visible(False)
+                    for ax in axs.flatten():
+                        if not ax.lines:
+                            ax.set_visible(False)
 
-            plt.margins(0, 0)
+                    plt.margins(0, 0)
 
-            if show:
-                fig.show()
+                    if show:
+                        plt.show()
 
-            if output_dir is not None:
-                pdf.savefig(fig)
+                    pdf.savefig(fig, dpi=80)
+                    plt.close("all")
+                    plt.clf()
 
-        if output_dir is not None:
-            pdf.close()
 
     def plot_all_recordings(self, output_dir=None, show=False):
         """Plot all the recordings of the cell.