Skip to content

Commit ef90578

Browse files
Jammy2211claude
andcommitted
perf(interferometer): cache curvature_matrix / data_vector on sparse and mapping inversions (#581)
On the NumPy path one figure_of_merit of InversionInterferometerSparse / InversionInterferometerSparseNumba built F twice and D four times, because data_vector, curvature_matrix and curvature_matrix_diag were plain properties read by curvature_reg_matrix, reconstruction and fast_chi_squared (43-50 % of an alma RAL CPU call, autolens_profiling #326). They are now autonerves cached_property, as the imaging inversions' already are; the numba curvature_matrix_diag override and the dense InversionInterferometerMapping data_vector / curvature_matrix follow. An inversion is built per fit, so no invalidation is needed; the preloads.curvature_matrix short-circuit is unchanged. Tests: evaluation-count tests (descriptor-preserving counters plus the underlying operator / numba kernel) assert F == 1 and D == 1 per likelihood for Sparse, SparseNumba (serial + parallel kernel) and Mapping; red on the unfixed source (2 / 2 each). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1 parent 4383ea8 commit ef90578

5 files changed

Lines changed: 220 additions & 6 deletions

File tree

‎autoarray/inversion/inversion/interferometer/mapping.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
import numpy as np
22
from typing import Dict, List, Union
33

4+
from autonerves import cached_property
5+
46
from autoarray.dataset.interferometer.dataset import Interferometer
57
from autoarray.inversion.inversion.dataset_interface import DatasetInterface
68
from autoarray.inversion.inversion.interferometer.abstract import (
@@ -50,7 +52,7 @@ def __init__(
5052
dataset=dataset, linear_obj_list=linear_obj_list, settings=settings, xp=xp
5153
)
5254

53-
@property
55+
@cached_property
5456
def data_vector(self) -> np.ndarray:
5557
"""
5658
The `data_vector` is a 1D vector whose values are solved for by the simultaneous linear equations constructed
@@ -71,7 +73,7 @@ def data_vector(self) -> np.ndarray:
7173
noise_map=self.noise_map,
7274
)
7375

74-
@property
76+
@cached_property
7577
def curvature_matrix(self) -> np.ndarray:
7678
"""
7779
The `curvature_matrix` is a 2D matrix which uses the mappings between the data and the linear objects to

‎autoarray/inversion/inversion/interferometer/sparse.py‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
import numpy as np
22
from typing import Dict, List, Optional, Union
33

4+
from autonerves import cached_property
5+
46
from autoarray import exc
57
from autoarray.dataset.interferometer.dataset import Interferometer
68
from autoarray.inversion.inversion.dataset_interface import DatasetInterface
@@ -67,7 +69,7 @@ def __init__(
6769
preloads=preloads,
6870
)
6971

70-
@property
72+
@cached_property
7173
def data_vector(self) -> np.ndarray:
7274
"""
7375
The `data_vector` is a 1D vector whose values are solved for by the simultaneous linear equations constructed
@@ -129,7 +131,7 @@ def _sparse_triplets_curvature_from(self, mapper: Mapper):
129131
xp=self._xp,
130132
)
131133

132-
@property
134+
@cached_property
133135
def curvature_matrix(self) -> np.ndarray:
134136
"""
135137
The `curvature_matrix` is a 2D matrix which uses the mappings between the data and the linear objects to
@@ -187,7 +189,7 @@ def curvature_matrix(self) -> np.ndarray:
187189

188190
return curvature_matrix
189191

190-
@property
192+
@cached_property
191193
def curvature_matrix_diag(self) -> np.ndarray:
192194
"""
193195
The `curvature_matrix` is a 2D matrix which uses the mappings between the data and the linear objects to

‎autoarray/inversion/inversion/interferometer_numba/sparse.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -181,7 +181,7 @@ def kernel_index_arrays(self) -> dict:
181181

182182
return inputs
183183

184-
@property
184+
@cached_property
185185
def curvature_matrix_diag(self) -> np.ndarray:
186186
"""
187187
`F = Aᵀ W~ A` for the inversion's single mapper, from the `direct_conv` numba

‎test_autoarray/inversion/inversion/interferometer/test_interferometer.py‎

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1044,3 +1044,135 @@ def test__interferometer_sparse_operator__numpy_inversion_matches_jax_inversion(
10441044
).max()
10451045

10461046
assert difference_sparse <= max(10.0 * difference_dense, 1.0e-14)
1047+
1048+
1049+
def _count_evaluations(monkeypatch, cls, name, counts):
1050+
"""
1051+
Wrap the body of the `cls.name` property with a counter while keeping its descriptor
1052+
type, so a `cached_property` stays cached (and a plain `property` stays uncached) and
1053+
the count is the number of times the body actually runs.
1054+
"""
1055+
import functools
1056+
1057+
descriptor = cls.__dict__[name]
1058+
func = descriptor.fget if isinstance(descriptor, property) else descriptor.func
1059+
1060+
@functools.wraps(func)
1061+
def counted(self):
1062+
counts[name] += 1
1063+
return func(self)
1064+
1065+
monkeypatch.setattr(cls, name, type(descriptor)(counted))
1066+
1067+
1068+
def _log_evidence_terms_from(inversion):
1069+
"""
1070+
The inversion terms of `FitInterferometer.log_evidence`, which is what a figure of merit
1071+
evaluates: the fast chi-squared (reads F, D and the reconstruction), the regularization
1072+
term and both log-determinants (the cached `curvature_reg_matrix` reads F).
1073+
"""
1074+
return (
1075+
float(inversion.fast_chi_squared)
1076+
+ float(inversion.regularization_term)
1077+
+ float(inversion.log_det_curvature_reg_matrix_term)
1078+
- float(inversion.log_det_regularization_matrix_term)
1079+
)
1080+
1081+
1082+
def test__interferometer_sparse_operator__curvature_matrix_and_data_vector_evaluated_once_per_likelihood(
1083+
monkeypatch,
1084+
):
1085+
"""
1086+
`fast_chi_squared`, `curvature_reg_matrix` and `reconstruction` all read `curvature_matrix`
1087+
(F) and `data_vector` (D). Each must be built once per inversion, not once per reader:
1088+
on the NumPy path nothing merges the repeated builds, which were ~45 % of an alma call
1089+
(autolens_profiling #326).
1090+
"""
1091+
mask, _, _, dataset_sparse = _sparse_parity_setup()
1092+
1093+
mapper = _mapper_from(
1094+
mask=mask, pixels=9, shape=(3, 3), regularization=aa.reg.Constant(coefficient=1.0)
1095+
)
1096+
1097+
# Built directly (not via `aa.Inversion`) so the NumPy FFT class runs even when numba is
1098+
# installed and the factory would route to `InversionInterferometerSparseNumba`.
1099+
def inversion_from():
1100+
return aa.InversionInterferometerSparse(
1101+
dataset=dataset_sparse, linear_obj_list=[mapper], xp=np
1102+
)
1103+
1104+
reference = _log_evidence_terms_from(inversion_from())
1105+
1106+
counts = {
1107+
"curvature_matrix_diag": 0,
1108+
"data_vector": 0,
1109+
"curvature_matrix_diag_from": 0,
1110+
}
1111+
1112+
_count_evaluations(
1113+
monkeypatch, aa.InversionInterferometerSparse, "curvature_matrix_diag", counts
1114+
)
1115+
_count_evaluations(
1116+
monkeypatch, aa.InversionInterferometerSparse, "data_vector", counts
1117+
)
1118+
1119+
from autoarray.inversion.inversion.interferometer.inversion_interferometer_util import (
1120+
InterferometerSparseOperator,
1121+
)
1122+
1123+
curvature_matrix_diag_from = InterferometerSparseOperator.curvature_matrix_diag_from
1124+
1125+
def counted_curvature_matrix_diag_from(self, *args, **kwargs):
1126+
counts["curvature_matrix_diag_from"] += 1
1127+
return curvature_matrix_diag_from(self, *args, **kwargs)
1128+
1129+
monkeypatch.setattr(
1130+
InterferometerSparseOperator,
1131+
"curvature_matrix_diag_from",
1132+
counted_curvature_matrix_diag_from,
1133+
)
1134+
1135+
inversion = inversion_from()
1136+
1137+
log_evidence_terms = _log_evidence_terms_from(inversion)
1138+
1139+
assert counts == {
1140+
"curvature_matrix_diag": 1,
1141+
"data_vector": 1,
1142+
"curvature_matrix_diag_from": 1,
1143+
}
1144+
assert log_evidence_terms == reference
1145+
1146+
1147+
def test__interferometer_mapping__curvature_matrix_and_data_vector_evaluated_once_per_likelihood(
1148+
monkeypatch,
1149+
):
1150+
"""
1151+
The dense mapping inversion's `curvature_matrix` and `data_vector` are cached per inversion
1152+
too, as the imaging mapping inversion's are.
1153+
"""
1154+
mask, _, dataset, _ = _sparse_parity_setup()
1155+
1156+
mapper = _mapper_from(
1157+
mask=mask, pixels=9, shape=(3, 3), regularization=aa.reg.Constant(coefficient=1.0)
1158+
)
1159+
1160+
reference = _log_evidence_terms_from(
1161+
aa.Inversion(dataset=dataset, linear_obj_list=[mapper])
1162+
)
1163+
1164+
counts = {"curvature_matrix": 0, "data_vector": 0}
1165+
1166+
for name in counts:
1167+
_count_evaluations(
1168+
monkeypatch, aa.InversionInterferometerMapping, name, counts
1169+
)
1170+
1171+
inversion = aa.Inversion(dataset=dataset, linear_obj_list=[mapper])
1172+
1173+
assert isinstance(inversion, aa.InversionInterferometerMapping)
1174+
1175+
log_evidence_terms = _log_evidence_terms_from(inversion)
1176+
1177+
assert counts == {"curvature_matrix": 1, "data_vector": 1}
1178+
assert log_evidence_terms == reference

‎test_autoarray/inversion/inversion/interferometer_numba/test_interferometer_numba.py‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -512,3 +512,81 @@ def test__existing_sparse_path_is_unchanged_when_the_gate_rejects_the_geometry()
512512
rtol=1.0e-12,
513513
atol=0.0,
514514
)
515+
516+
517+
@pytest.mark.parametrize("parallel", [False, True])
518+
def test__numba_inversion__curvature_matrix_and_data_vector_evaluated_once_per_likelihood(
519+
monkeypatch, parallel
520+
):
521+
"""
522+
`fast_chi_squared`, `curvature_reg_matrix` and `reconstruction` all read `curvature_matrix`
523+
(F, via the numba `curvature_matrix_diag`) and `data_vector` (D). The `direct_conv` kernel
524+
must run once per inversion, not once per reader (autolens_profiling #326: F twice and
525+
D four times on one figure of merit).
526+
"""
527+
import functools
528+
529+
from autoarray.inversion.inversion.interferometer_numba import sparse as numba_sparse
530+
531+
monkeypatch.setattr(numba_sparse, "_numba_parallel", lambda: parallel)
532+
533+
mask, _, dataset_sparse = _dataset_from()
534+
535+
mapper = _delaunay_mapper_from(mask=mask)
536+
537+
def log_evidence_terms_from(inversion):
538+
return (
539+
float(inversion.fast_chi_squared)
540+
+ float(inversion.regularization_term)
541+
+ float(inversion.log_det_curvature_reg_matrix_term)
542+
- float(inversion.log_det_regularization_matrix_term)
543+
)
544+
545+
reference = log_evidence_terms_from(
546+
InversionInterferometerSparseNumba(
547+
dataset=dataset_sparse, linear_obj_list=[mapper], xp=np
548+
)
549+
)
550+
551+
counts = {"kernel": 0, "data_vector": 0}
552+
553+
def counted(kernel):
554+
@functools.wraps(kernel)
555+
def wrapper(*args, **kwargs):
556+
counts["kernel"] += 1
557+
return kernel(*args, **kwargs)
558+
559+
return wrapper
560+
561+
monkeypatch.setattr(
562+
numba_util, "curvature_direct_conv", counted(numba_util.curvature_direct_conv)
563+
)
564+
parallel_kernel = numba_util.direct_conv_parallel_kernel
565+
monkeypatch.setattr(
566+
numba_util,
567+
"direct_conv_parallel_kernel",
568+
lambda: counted(parallel_kernel()),
569+
)
570+
571+
descriptor = InversionInterferometerSparse.__dict__["data_vector"]
572+
func = descriptor.fget if isinstance(descriptor, property) else descriptor.func
573+
574+
@functools.wraps(func)
575+
def counted_data_vector(self):
576+
counts["data_vector"] += 1
577+
return func(self)
578+
579+
monkeypatch.setattr(
580+
InversionInterferometerSparse,
581+
"data_vector",
582+
type(descriptor)(counted_data_vector),
583+
)
584+
585+
inversion = InversionInterferometerSparseNumba(
586+
dataset=dataset_sparse, linear_obj_list=[mapper], xp=np
587+
)
588+
589+
log_evidence_terms = log_evidence_terms_from(inversion)
590+
591+
assert counts == {"kernel": 1, "data_vector": 1}
592+
assert log_evidence_terms == reference

0 commit comments

Comments
 (0)