From 7e8ae10de9af40a9ce00d061e8f00bcf6b400527 Mon Sep 17 00:00:00 2001
From: Charlie Windolf <cwindolf95@gmail.com>
Date: Mon, 15 Jan 2024 13:28:57 -0500
Subject: [PATCH] Expose figsize

---
 src/dartsort/vis/scatterplots.py | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/src/dartsort/vis/scatterplots.py b/src/dartsort/vis/scatterplots.py
index 95a4a414..6147e169 100644
--- a/src/dartsort/vis/scatterplots.py
+++ b/src/dartsort/vis/scatterplots.py
@@ -31,6 +31,7 @@ def scatter_spike_features(
     limits="probe_margin",
     label_axes=True,
     random_seed=0,
+    figsize=None,
     **scatter_kw,
 ):
     """3-axis scatter plot of spike depths vs. horizontal pos, amplitude, and time
@@ -47,7 +48,10 @@ def scatter_spike_features(
         figure = plt.gcf()
     if axes is None:
         axes = figure.subplots(
-            ncols=3, sharey=True, gridspec_kw=dict(width_ratios=width_ratios)
+            ncols=3,
+            sharey=True,
+            gridspec_kw=dict(width_ratios=width_ratios),
+            figsize=figsize,
         )
 
     if hdf5_filename is not None:
@@ -263,7 +267,9 @@ def scatter_x_vs_depth(
     if show_geom and geom is not None:
         ax.scatter(*geom.T, **geom_scatter_kw)
     if limits == "probe_margin" and geom is not None:
-        ax.set_xlim([geom[:, 0].min() - probe_margin_um, geom[:, 0].max() + probe_margin_um])
+        ax.set_xlim(
+            [geom[:, 0].min() - probe_margin_um, geom[:, 0].max() + probe_margin_um]
+        )
     return ax, s1
 
 
@@ -407,5 +413,7 @@ def scatter_feature_vs_depth(
         **scatter_kw,
     )
     if limits == "probe_margin" and geom is not None:
-        ax.set_ylim([geom[:, 1].min() - probe_margin_um, geom[:, 1].max() + probe_margin_um])
+        ax.set_ylim(
+            [geom[:, 1].min() - probe_margin_um, geom[:, 1].max() + probe_margin_um]
+        )
     return ax, s