Skip to content

Commit 91d8b97

Browse files
committed
Defer the eager scipy.sparse and scipy.spatial imports (~280 ms)
`import autoarray` drops from 464.4 ms to 183.7 ms (medians of 15 runs, Python 3.13, dev extras) — a 281 ms saving. The filed task named `derivative_util.py:30`'s module-scope `from scipy.sparse import csr_matrix` as the target, on the evidence that `scipy.sparse` costs ~0.11 s of every import. Deferring it changed nothing, because `scipy.sparse` was never being imported from there: autoarray/__init__.py:80 -> inversion/mesh/mesh_geometry/delaunay.py:2 import scipy.spatial -> scipy/spatial/__init__.py:111 from ._kdtree import * -> scipy/spatial/_kdtree.py:4 from ._ckdtree import cKDTree `scipy.spatial` (134 ms) pulls `scipy.sparse` (154 ms) in transitively, so the csr_matrix import was riding on a subtree that was already paid for. Deferring `scipy.spatial` as well is what actually removes both, and is required to meet the task's own acceptance criterion. Changes: - `inversion/mesh/mesh_geometry/delaunay.py` — drop the module-scope `import scipy.spatial`. Two of its three use sites already had local imports; only `voronoi_neighbors_from` needed one added, so this finishes a deferral that had been started and left half-done. - `operators/derivative_util.py`, `operators/coarse_interp_util.py` — move `from scipy.sparse import csr_matrix` into the four functions that use it. No longer load-bearing for the import time on its own, but it keeps `scipy.sparse` off the import path independently of what `scipy.spatial` happens to pull in. Every use site is inside a function, so a plain local import suffices — no module-level cache as in `transformer.py:_load_nufftax()`, which needed one only because unpickled instances in multiprocessing workers never re-run `__init__`. Function-local imports run on every call and hit `sys.modules` after the first. Verified: `python -X importtime -c "import autoarray"` shows neither `scipy.sparse` nor `scipy.spatial`. Suites green — autoarray 1179 passed, autogalaxy 1103 passed / 1 skipped, autolens 532 passed / 1 skipped.
1 parent c330e3c commit 91d8b97

4 files changed

Lines changed: 11 additions & 5 deletions

File tree

autoarray/inversion/mesh/mesh_geometry/delaunay.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
11
import numpy as np
2-
import scipy.spatial
32
from typing import Tuple
43

54
from autonerves import cached_property
@@ -147,6 +146,7 @@ def neighbors(self) -> Neighbors:
147146
The neighbors of a Voronoi mesh are computed using the `ridge_points` attribute of the scipy `Voronoi`
148147
object, as described in the method `voronoi_neighbors_from`.
149148
"""
149+
import scipy.spatial
150150

151151
delaunay = scipy.spatial.Delaunay(self.mesh_grid_xy)
152152

autoarray/operators/coarse_interp_util.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@
1717
"""
1818

1919
import numpy as np
20-
from scipy.sparse import csr_matrix
2120

2221
from autoarray import exc
2322
from autoarray import numba_util
@@ -242,6 +241,8 @@ def coarse_interp_matrix_from(
242241
A ``scipy.sparse.csr_matrix`` of shape
243242
[n_unmasked_fine_pixels, n_unmasked_coarse_pixels].
244243
"""
244+
from scipy.sparse import csr_matrix
245+
245246
mask_itp_box = np.asarray(mask_itp_box)
246247
if np.count_nonzero(~mask_itp_box) == 0:
247248
raise exc.MeshException(

autoarray/operators/derivative_util.py

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,6 @@
2727
"""
2828

2929
import numpy as np
30-
from scipy.sparse import csr_matrix
3130

3231
from autoarray import exc
3332
from autoarray import numba_util
@@ -297,9 +296,11 @@ def derivative_1st_operators_from(mask, pixel_scale: float = 1.0):
297296
-------
298297
The (Hy, Hx) operators as ``scipy.sparse.csr_matrix``.
299298
"""
299+
from scipy.sparse import csr_matrix
300+
300301
mask, diff_types = _diff_types_of_cleaned_mask_from(mask)
301-
rows_hx, cols_hx, data_hx, rows_hy, cols_hy, data_hy = (
302-
derivative_1st_triplets_from(mask, diff_types, dpix=pixel_scale)
302+
rows_hx, cols_hx, data_hx, rows_hy, cols_hy, data_hy = derivative_1st_triplets_from(
303+
mask, diff_types, dpix=pixel_scale
303304
)
304305

305306
n_unmasked = np.count_nonzero(~mask)
@@ -443,6 +444,8 @@ def derivative_2nd_operators_from(mask, pixel_scale: float = 1.0):
443444
-------
444445
The (Hyy, Hxx) operators as ``scipy.sparse.csr_matrix``.
445446
"""
447+
from scipy.sparse import csr_matrix
448+
446449
mask, diff_types = _diff_types_of_cleaned_mask_from(mask)
447450
rows_hxx, cols_hxx, data_hxx, rows_hyy, cols_hyy, data_hyy = (
448451
derivative_2nd_triplets_from(mask, diff_types, dpix=pixel_scale)
@@ -574,6 +577,8 @@ def forward_difference_operators_from(
574577
-------
575578
The (Hy, Hx) operators as ``scipy.sparse.csr_matrix``.
576579
"""
580+
from scipy.sparse import csr_matrix
581+
577582
if max_order not in (1, 2, 3, 4):
578583
raise ValueError(f"max_order must be in 1..4, got {max_order}")
579584

Binary file not shown.

0 commit comments

Comments
 (0)