From d1ae85e71a9c026713851e8a5af3c049632fc5b2 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Mon, 7 Sep 2026 00:04:30 +0000 Subject: [PATCH] refactor(import): defer matplotlib.pyplot out of `import autolens` MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `autolens/__init__.py` imports `potential_correction`, whose `__init__` imports `visualize`, which imported `matplotlib.pyplot` (and `cm`, `ticker`, `MaxNLocator`, `make_axes_locatable`, `matplotlib`) at module level: ~0.32 s on every `import autolens` for six drawing functions almost no process calls (census O3 of the ci-timing-fast-tests epic, PyAutoHeart timings/unit_import_census_2026-09.md). The imports now live inside the five functions that draw, matching what `potential_correction/mesh.py` and `iterative.py` already do. Nothing outside the module referenced `visualize.plt` et al. `scipy.interpolate` / `scipy.spatial` stay at module level (autoarray already imports scipy.spatial). Measured: `-X importtime` matplotlib.pyplot 325 ms → absent; `import autolens` 1.54 → 1.09 s with the sibling PyAutoFit change, 1.17 s with this one alone (best of 3, fresh processes). Suite 623 passed, 1 xfailed; three new tests pin the clean import, that no matplotlib name is bound on the module, and that the drawing functions are still exported. Four of the five drawing functions were exercised under Agg; the two fit-summary figures need a fit object and are covered by the import tests only. Refs PyAutoLabs/PyAutoFit#1565 Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01WYcmRnZf3cbC5Wth9bCkEg --- autolens/potential_correction/visualize.py | 28 ++++++++--- .../test_visualize_import.py | 49 +++++++++++++++++++ 2 files changed, 71 insertions(+), 6 deletions(-) create mode 100644 test_autolens/potential_correction/test_visualize_import.py diff --git a/autolens/potential_correction/visualize.py b/autolens/potential_correction/visualize.py index 895a12e29..43481c674 100644 --- a/autolens/potential_correction/visualize.py +++ b/autolens/potential_correction/visualize.py @@ -14,15 +14,15 @@ import os import numpy as np -import matplotlib as mpl -import matplotlib.cm as cm -import matplotlib.pyplot as plt -import matplotlib.ticker as ticker -from matplotlib.ticker import MaxNLocator -from mpl_toolkits.axes_grid1 import make_axes_locatable from scipy.interpolate import griddata from scipy.spatial import Voronoi, voronoi_plot_2d +# ``matplotlib`` is imported inside the drawing functions below, not here: +# importing ``matplotlib.pyplot`` at module level cost every ``import autolens`` +# ~0.32s, since ``autolens/__init__.py`` imports this sub-package (census O3). +# This matches ``potential_correction/mesh.py`` and ``iterative.py``, which +# already defer their pyplot import the same way. + def _plot_anchor_points(ax, anchor_points): if anchor_points is not None and np.shape(anchor_points) == (3, 2): @@ -45,6 +45,11 @@ def imshow_masked_data( Shows a 1D (slim) data vector on its 2D mask, with masked pixels blanked, a colorbar and optional contours. """ + from matplotlib import pyplot as plt + from matplotlib import ticker + from matplotlib.ticker import MaxNLocator + from mpl_toolkits.axes_grid1 import make_axes_locatable + data_1d = np.asarray(data_1d) mask_2d = np.asarray(mask_2d) if centralize: @@ -111,6 +116,8 @@ def show_image_irregular_interpolate( Shows values on irregular (y, x) points by linear interpolation onto a regular grid. """ + from matplotlib import pyplot as plt + points = np.asarray(points) points = points[:, ::-1] # to scipy (x, y) order @@ -141,6 +148,11 @@ def show_image_irregular( """ Shows values on irregular (y, x) points as coloured Voronoi cells. """ + import matplotlib as mpl + from matplotlib import cm + from matplotlib import pyplot as plt + from mpl_toolkits.axes_grid1 import make_axes_locatable + points = np.asarray(points) points = points[:, ::-1] # to scipy (x, y) order @@ -180,6 +192,8 @@ def show_fit_dpsi(fit, output="result.png"): The six-panel summary of a ``FitDpsiImaging``: residual data, model, residual-of-residual, normalized residual, and the dpsi and dkappa maps. """ + from matplotlib import pyplot as plt + plt.figure(figsize=(15, 10)) cmap = copy.copy(plt.get_cmap("jet")) cmap.set_bad(color="white") @@ -256,6 +270,8 @@ def show_fit_dpsi_src(fit, output="result.png", show_src_grid=True, interpolate= model, residual, normalized residual, dpsi map, dkappa map and the source reconstruction. """ + from matplotlib import pyplot as plt + fig = plt.figure(figsize=(15, 10)) cmap = copy.copy(plt.get_cmap("jet")) cmap.set_bad(color="white") diff --git a/test_autolens/potential_correction/test_visualize_import.py b/test_autolens/potential_correction/test_visualize_import.py new file mode 100644 index 000000000..561129157 --- /dev/null +++ b/test_autolens/potential_correction/test_visualize_import.py @@ -0,0 +1,49 @@ +""" +Census O3: ``import autolens`` must not pay for ``matplotlib.pyplot``. + +``autolens/__init__.py`` imports the ``potential_correction`` sub-package, +whose ``__init__`` imports ``visualize``. That module used to import +``matplotlib.pyplot`` (and ``cm``/``ticker``/``mpl_toolkits``) at module level, +so every process that touched ``autolens`` at all paid ~0.32s to build a +pyplot state machine that only its six drawing functions use. Those imports now +live inside the functions, matching ``potential_correction/mesh.py`` and +``iterative.py``. +""" +import subprocess +import sys + +from autolens.potential_correction import visualize + + +def test__importing_autolens_does_not_import_pyplot(): + result = subprocess.run( + [ + sys.executable, + "-c", + "import autolens, sys; print('matplotlib.pyplot' in sys.modules)", + ], + capture_output=True, + text=True, + check=True, + ) + + assert result.stdout.strip() == "False" + + +def test__visualize_module_binds_no_matplotlib_names_at_module_level(): + # The deferral is only real while these stay off the module: a name put back + # at the top of ``visualize.py`` would re-import pyplot on ``import autolens`` + # without failing any drawing test. + for name in ("plt", "mpl", "cm", "ticker", "MaxNLocator", "make_axes_locatable"): + assert not hasattr(visualize, name) + + +def test__drawing_functions_are_still_exported(): + for name in ( + "imshow_masked_data", + "show_image_irregular_interpolate", + "show_image_irregular", + "show_fit_dpsi", + "show_fit_dpsi_src", + ): + assert callable(getattr(visualize, name))