perf: defer matplotlib.pyplot via wrapper functions in plot/utils

Replace all module-level `import matplotlib.pyplot as plt` across plot
files with `subplots()` and `get_cmap()` wrappers from plot/utils that
import matplotlib lazily on first call. This prevents matplotlib from
loading during `import autoarray`, deferring ~0.3s of import cost to
first plot use.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: 2 people

autoarray/dataset/plot/imaging_plots.py

@@ -1,8 +1,7 @@
 from typing import Optional
 
-import matplotlib.pyplot as plt
 
-from autoarray.plot.utils import subplot_save, conf_subplot_figsize, tight_layout
+from autoarray.plot.utils import subplots, subplot_save, conf_subplot_figsize, tight_layout
 
 
 def subplot_imaging_dataset(
@@ -51,7 +50,7 @@ def subplot_imaging_dataset(
 
     from autoarray.plot.array import plot_array
 
-    fig, axes = plt.subplots(3, 3, figsize=conf_subplot_figsize(3, 3))
+    fig, axes = subplots(3, 3, figsize=conf_subplot_figsize(3, 3))
     axes = axes.flatten()
 
     plot_array(
@@ -172,7 +171,7 @@ def subplot_imaging_dataset_list(
     from autoarray.plot.array import plot_array
 
     n = len(dataset_list)
-    fig, axes = plt.subplots(n, 3, figsize=conf_subplot_figsize(n, 3))
+    fig, axes = subplots(n, 3, figsize=conf_subplot_figsize(n, 3))
     if n == 1:
         axes = [axes]
     for i, dataset in enumerate(dataset_list):

autoarray/dataset/plot/interferometer_plots.py

@@ -1,12 +1,11 @@
 import numpy as np
 from typing import Optional
 
-import matplotlib.pyplot as plt
 
 from autoarray.plot.array import plot_array
 from autoarray.plot.grid import plot_grid
 from autoarray.plot.yx import plot_yx
-from autoarray.plot.utils import subplot_save, hide_unused_axes, conf_subplot_figsize, tight_layout
+from autoarray.plot.utils import subplots, subplot_save, hide_unused_axes, conf_subplot_figsize, tight_layout
 from autoarray.structures.grids.irregular_2d import Grid2DIrregular
 
 
@@ -39,7 +38,7 @@ def subplot_interferometer_dataset(
     use_log10
         Apply log10 normalisation to image panels.
     """
-    fig, axes = plt.subplots(2, 3, figsize=conf_subplot_figsize(2, 3))
+    fig, axes = subplots(2, 3, figsize=conf_subplot_figsize(2, 3))
     axes = axes.flatten()
 
     plot_grid(dataset.data.in_grid, ax=axes[0], title="Visibilities", xlabel="", ylabel="")
@@ -117,7 +116,7 @@ def subplot_interferometer_dirty_images(
     use_log10
         Apply log10 normalisation.
     """
-    fig, axes = plt.subplots(1, 3, figsize=conf_subplot_figsize(1, 3))
+    fig, axes = subplots(1, 3, figsize=conf_subplot_figsize(1, 3))
 
     plot_array(
         dataset.dirty_image,

autoarray/fit/plot/fit_imaging_plots.py

@@ -1,9 +1,8 @@
 from typing import Optional
 
-import matplotlib.pyplot as plt
 
 from autoarray.plot.array import plot_array
-from autoarray.plot.utils import subplot_save, symmetric_vmin_vmax, hide_unused_axes, conf_subplot_figsize, tight_layout
+from autoarray.plot.utils import subplots, subplot_save, symmetric_vmin_vmax, hide_unused_axes, conf_subplot_figsize, tight_layout
 
 
 def subplot_fit_imaging(
@@ -43,7 +42,7 @@ def subplot_fit_imaging(
     grid, positions, lines
         Optional overlays forwarded to every panel.
     """
-    fig, axes = plt.subplots(2, 3, figsize=conf_subplot_figsize(2, 3))
+    fig, axes = subplots(2, 3, figsize=conf_subplot_figsize(2, 3))
     axes = axes.flatten()
 
     plot_array(

autoarray/fit/plot/fit_interferometer_plots.py

@@ -1,11 +1,10 @@
 import numpy as np
 from typing import Optional
 
-import matplotlib.pyplot as plt
 
 from autoarray.plot.array import plot_array
 from autoarray.plot.yx import plot_yx
-from autoarray.plot.utils import subplot_save, symmetric_vmin_vmax, hide_unused_axes, conf_subplot_figsize, tight_layout
+from autoarray.plot.utils import subplots, subplot_save, symmetric_vmin_vmax, hide_unused_axes, conf_subplot_figsize, tight_layout
 
 
 def subplot_fit_interferometer(
@@ -40,7 +39,7 @@ def subplot_fit_interferometer(
         Not used here (UV-plane residuals are scatter plots); kept for API
         consistency.
     """
-    fig, axes = plt.subplots(2, 3, figsize=conf_subplot_figsize(2, 3))
+    fig, axes = subplots(2, 3, figsize=conf_subplot_figsize(2, 3))
     axes = axes.flatten()
 
     uv = fit.dataset.uv_distances / 10**3.0
@@ -135,7 +134,7 @@ def subplot_fit_interferometer_dirty_images(
     residuals_symmetric_cmap
         Centre residual colour scale symmetrically around zero.
     """
-    fig, axes = plt.subplots(2, 3, figsize=conf_subplot_figsize(2, 3))
+    fig, axes = subplots(2, 3, figsize=conf_subplot_figsize(2, 3))
     axes = axes.flatten()
 
     plot_array(

autoarray/inversion/plot/inversion_plots.py

@@ -4,12 +4,11 @@
 from pathlib import Path
 from typing import Optional, Union
 
-import matplotlib.pyplot as plt
 from autoconf import conf
 
 from autoarray.inversion.mappers.abstract import Mapper
 from autoarray.plot.array import plot_array
-from autoarray.plot.utils import numpy_grid, numpy_lines, numpy_positions, subplot_save, hide_unused_axes, conf_subplot_figsize, tight_layout
+from autoarray.plot.utils import subplots, numpy_grid, numpy_lines, numpy_positions, subplot_save, hide_unused_axes, conf_subplot_figsize, tight_layout
 from autoarray.inversion.plot.mapper_plots import plot_mapper
 from autoarray.structures.arrays.uniform_2d import Array2D
 
@@ -53,7 +52,7 @@ def subplot_of_mapper(
     """
     mapper = inversion.cls_list_from(cls=Mapper)[mapper_index]
 
-    fig, axes = plt.subplots(3, 4, figsize=conf_subplot_figsize(3, 4))
+    fig, axes = subplots(3, 4, figsize=conf_subplot_figsize(3, 4))
     axes = axes.flatten()
 
     # panel 0: data subtracted
@@ -279,7 +278,7 @@ def subplot_mappings(
     )
     mapper.slim_indexes_for_pix_indexes(pix_indexes=pix_indexes)
 
-    fig, axes = plt.subplots(2, 2, figsize=conf_subplot_figsize(2, 2))
+    fig, axes = subplots(2, 2, figsize=conf_subplot_figsize(2, 2))
     axes = axes.flatten()
 
     # panel 0: data subtracted

autoarray/inversion/plot/mapper_plots.py

@@ -1,11 +1,10 @@
 import logging
 from typing import Optional
 
-import matplotlib.pyplot as plt
 
 from autoarray.plot.array import plot_array
 from autoarray.plot.inversion import plot_inversion_reconstruction
-from autoarray.plot.utils import numpy_grid, numpy_lines, subplot_save, conf_subplot_figsize, tight_layout
+from autoarray.plot.utils import subplots, numpy_grid, numpy_lines, subplot_save, conf_subplot_figsize, tight_layout
 
 logger = logging.getLogger(__name__)
 
@@ -114,7 +113,7 @@ def subplot_image_and_mapper(
     lines
         Lines to overlay on both panels.
     """
-    fig, axes = plt.subplots(1, 2, figsize=conf_subplot_figsize(1, 2))
+    fig, axes = subplots(1, 2, figsize=conf_subplot_figsize(1, 2))
 
     plot_array(
         image,

autoarray/plot/array.py

@@ -5,11 +5,9 @@
 import os
 from typing import List, Optional, Tuple
 
-import matplotlib.pyplot as plt
 import numpy as np
-from matplotlib.colors import LogNorm, Normalize
-
 from autoarray.plot.utils import (
+    subplots,
     apply_extent,
     apply_labels,
     conf_figsize,
@@ -29,7 +27,7 @@
 
 def plot_array(
     array,
-    ax: Optional[plt.Axes] = None,
+    ax=None,
     # --- spatial metadata -------------------------------------------------------
     extent: Optional[Tuple[float, float, float, float]] = None,
     # --- overlays ---------------------------------------------------------------
@@ -158,7 +156,7 @@ def plot_array(
     owns_figure = ax is None
     if owns_figure:
         figsize = figsize or conf_figsize("figures")
-        fig, ax = plt.subplots(1, 1, figsize=figsize)
+        fig, ax = subplots(1, 1, figsize=figsize)
     else:
         fig = ax.get_figure()
 
@@ -181,8 +179,10 @@ def plot_array(
                 vmax_log = np.nanmax(clipped)
         if not np.isfinite(vmax_log) or vmax_log <= vmin_log:
             vmax_log = vmin_log * 10.0
+        from matplotlib.colors import LogNorm
         norm = LogNorm(vmin=vmin_log, vmax=vmax_log)
     elif vmin is not None or vmax is not None:
+        from matplotlib.colors import Normalize
         norm = Normalize(vmin=vmin, vmax=vmax)
     else:
         norm = None

autoarray/plot/grid.py

@@ -6,10 +6,11 @@
 
 from typing import Iterable, List, Optional, Tuple
 
-import matplotlib.pyplot as plt
 import numpy as np
 
 from autoarray.plot.utils import (
+    subplots,
+    get_cmap,
     apply_extent,
     apply_labels,
     conf_figsize,
@@ -20,7 +21,7 @@
 
 def plot_grid(
     grid,
-    ax: Optional[plt.Axes] = None,
+    ax=None,
     # --- errors -----------------------------------------------------------------
     y_errors: Optional[np.ndarray] = None,
     x_errors: Optional[np.ndarray] = None,
@@ -108,13 +109,13 @@ def plot_grid(
     owns_figure = ax is None
     if owns_figure:
         figsize = figsize or conf_figsize("figures")
-        fig, ax = plt.subplots(1, 1, figsize=figsize)
+        fig, ax = subplots(1, 1, figsize=figsize)
     else:
         fig = ax.get_figure()
 
     # --- scatter / errorbar ----------------------------------------------------
     if color_array is not None:
-        cmap = plt.get_cmap(colormap)
+        cmap = get_cmap(colormap)
         colors = cmap((color_array - color_array.min()) / (np.ptp(color_array) or 1))
 
         if y_errors is None and x_errors is None:

autoarray/plot/inversion.py

@@ -6,17 +6,16 @@
 
 from typing import List, Optional, Tuple
 
-import matplotlib.pyplot as plt
 import numpy as np
 from matplotlib.colors import LogNorm, Normalize
 
-from autoarray.plot.utils import apply_extent, apply_labels, conf_figsize, save_figure, _conf_imshow_origin
+from autoarray.plot.utils import subplots, apply_extent, apply_labels, conf_figsize, save_figure, _conf_imshow_origin
 
 
 def plot_inversion_reconstruction(
     pixel_values: np.ndarray,
     mapper,
-    ax: Optional[plt.Axes] = None,
+    ax=None,
     # --- cosmetics --------------------------------------------------------------
     title: str = "Reconstruction",
     xlabel: str = 'x (")',
@@ -84,7 +83,7 @@ def plot_inversion_reconstruction(
     owns_figure = ax is None
     if owns_figure:
         figsize = figsize or conf_figsize("figures")
-        fig, ax = plt.subplots(1, 1, figsize=figsize)
+        fig, ax = subplots(1, 1, figsize=figsize)
     else:
         fig = ax.get_figure()