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))