Skip to content

Commit b5ef2e8

Browse files
Jammy2211claude
andcommitted
feat: route func-list-only interferometer inversions through the sparse operator (#575)
- factory: an interferometer dataset with a sparse operator now selects InversionInterferometerSparse for every linear object list, including MGE-only (no mapper); imaging routing unchanged (commented why). - DatasetInterface gains sparse_dirty_image; the sparse data vector uses it when set, else the operator's cached dirty image. - TransformerDFT.image_from / image_direct_from take xp so the adjoint DFT traces under jax.jit. - Tests: func-list-only sparse vs dense parity (1e-10), sparse_dirty_image override with a cached-image control, JAX-vs-NumPy func-list-only case, factory routing, DFT image_from under jit. Refs #575 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1 parent 0f870c3 commit b5ef2e8

8 files changed

Lines changed: 288 additions & 19 deletions

File tree

‎autoarray/inversion/inversion/dataset_interface.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ def __init__(
88
transformer=None,
99
sparse_operator=None,
1010
noise_covariance_matrix=None,
11+
sparse_dirty_image=None,
1112
):
1213
"""
1314
Generic class which acts as an interface between a dataset and an inversion.
@@ -49,6 +50,13 @@ def __init__(
4950
noise_covariance_matrix
5051
A noise-map covariance matrix representing the covariance between noise in every `data` value, which
5152
can be used via a bespoke fit to account for correlated noise in the data.
53+
sparse_dirty_image
54+
The noise-weighted dirty image `Re(Fᴴ W d)` of this interface's `data`, used by the sparse (w-tilde)
55+
interferometer inversion to form its data vector. The `sparse_operator` caches the dirty image of the
56+
visibilities it was built from, so this is only needed when `data` differs from them (e.g. when the
57+
visibilities of ordinary light profiles have been subtracted). If `None`, the operator's cached dirty
58+
image is used. This is distinct from `Interferometer.dirty_image`, the unweighted dirty image of the
59+
data used for visualization.
5260
"""
5361
self.data = data
5462
self.noise_map = noise_map
@@ -57,6 +65,7 @@ def __init__(
5765
self.transformer = transformer
5866
self.sparse_operator = sparse_operator
5967
self.noise_covariance_matrix = noise_covariance_matrix
68+
self.sparse_dirty_image = sparse_dirty_image
6069

6170
@property
6271
def mask(self):

‎autoarray/inversion/inversion/factory.py‎

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,10 @@ def inversion_imaging_from(
125125
An `Inversion` whose type is determined by the input `dataset` and `settings`.
126126
"""
127127

128+
# Unlike `inversion_interferometer_from`, imaging keeps func-list-only inversions (e.g. an MGE with no
129+
# pixelization) on the mapping formalism. A func-list-only imaging mapping inversion is cheap (one PSF
130+
# convolution per basis function), and the imaging sparse data vector's handling of profile-subtracted
131+
# data is tracked separately, so the interferometer routing change is deliberately not mirrored here.
128132
use_sparse_operator = True
129133

130134
if all(
@@ -199,15 +203,12 @@ def inversion_interferometer_from(
199203
-------
200204
An `Inversion` whose type is determined by the input `dataset` and `settings`.
201205
"""
202-
use_sparse_operator = True
203-
204-
if all(
205-
isinstance(linear_obj, AbstractLinearObjFuncList)
206-
for linear_obj in linear_obj_list
207-
):
208-
use_sparse_operator = False
209-
210-
if dataset.sparse_operator is not None and use_sparse_operator:
206+
# A sparse operator selects the w-tilde formalism for every linear object list, including func-list-only
207+
# ones (e.g. an MGE with no pixelization). For interferometer data this avoids the dense
208+
# `transform_mapping_matrix` of shape (N_vis, S), which dominates run time and memory at large N_vis, and
209+
# `InversionInterferometerSparse` handles zero mappers (the func-func curvature blocks and the data vector
210+
# need no mapper). `_use_interferometer_numba` returns False whenever a func-list is present.
211+
if dataset.sparse_operator is not None:
211212

212213
if _use_interferometer_numba(
213214
linear_obj_list=linear_obj_list,

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

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -88,10 +88,18 @@ def data_vector(self) -> np.ndarray:
8888
which is exactly the entry the mapping (dense) formalism computes via the linear
8989
function's transformed mapping matrix. Linear function lists therefore require no
9090
separate branch here.
91+
92+
The cached dirty image is that of the visibilities the `sparse_operator` was built from. When
93+
the inversion's data differs from them (e.g. a `DatasetInterface` whose data has the visibilities
94+
of ordinary light profiles subtracted), the dataset's `sparse_dirty_image` of that data is used
95+
instead, so that `D` is consistent with the data the chi-squared is computed from.
9196
"""
92-
return self._xp.dot(
93-
self.mapping_matrix.T, self.dataset.sparse_operator.dirty_image
94-
)
97+
dirty_image = getattr(self.dataset, "sparse_dirty_image", None)
98+
99+
if dirty_image is None:
100+
dirty_image = self.dataset.sparse_operator.dirty_image
101+
102+
return self._xp.dot(self.mapping_matrix.T, dirty_image)
95103

96104
def _sparse_triplets_curvature_from(self, mapper: Mapper):
97105
"""

‎autoarray/operators/transformer.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -199,9 +199,12 @@ def image_from(self, visibilities: Visibilities, xp=np) -> Array2D:
199199
mask as this transformer's `real_space_mask`.
200200
"""
201201
image_slim = transformer_util.image_direct_from(
202-
visibilities=visibilities.in_array,
202+
visibilities=xp.stack(
203+
(xp.real(visibilities.array), xp.imag(visibilities.array)), axis=-1
204+
),
203205
grid_radians=self.grid.array,
204206
uv_wavelengths=self.uv_wavelengths,
207+
xp=xp,
205208
)
206209

207210
image_native = array_2d_util.array_2d_native_from(

‎autoarray/operators/transformer_util.py‎

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,10 @@ def visibilities_from(
5050

5151

5252
def image_direct_from(
53-
visibilities: np.ndarray, grid_radians: np.ndarray, uv_wavelengths: np.ndarray
53+
visibilities: np.ndarray,
54+
grid_radians: np.ndarray,
55+
uv_wavelengths: np.ndarray,
56+
xp=np,
5457
) -> np.ndarray:
5558
"""
5659
Reconstruct a real-valued sky image from complex interferometric visibilities
@@ -69,6 +72,8 @@ def image_direct_from(
6972
7073
uv_wavelengths
7174
The (u, v) spatial frequencies in units of wavelengths for each baseline.
75+
xp
76+
The array module (`numpy` or `jax.numpy`) the transform is computed with.
7277
7378
Returns
7479
-------
@@ -78,15 +83,15 @@ def image_direct_from(
7883
# Compute the phase term for each (pixel, visibility) pair
7984
phase = (
8085
2.0
81-
* np.pi
86+
* xp.pi
8287
* (
83-
np.outer(grid_radians[:, 1], uv_wavelengths[:, 0])
84-
+ np.outer(grid_radians[:, 0], uv_wavelengths[:, 1])
88+
xp.outer(grid_radians[:, 1], uv_wavelengths[:, 0])
89+
+ xp.outer(grid_radians[:, 0], uv_wavelengths[:, 1])
8590
)
8691
)
8792

88-
real_part = np.dot(np.cos(phase), visibilities[:, 0])
89-
imag_part = np.dot(np.sin(phase), visibilities[:, 1])
93+
real_part = xp.dot(xp.cos(phase), visibilities[:, 0])
94+
imag_part = xp.dot(xp.sin(phase), visibilities[:, 1])
9095

9196
image_1d = real_part - imag_part
9297

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

Lines changed: 178 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -563,6 +563,133 @@ def test__interferometer_sparse_operator__func_list_and_mapper__identical_to_map
563563
)
564564

565565

566+
def test__interferometer_sparse_operator__func_list_only__identical_to_mapping():
567+
"""
568+
A linear function list with no mapper (e.g. an MGE with no pixelization) is routed to the sparse
569+
path whenever the dataset has a sparse operator, where its curvature matrix is the func-func block
570+
`Bᵀ W~ B` alone and its data vector is `Bᵀ d~`, reproducing the dense (mapping formalism) inversion.
571+
"""
572+
mask, rng, dataset, dataset_sparse = _sparse_parity_setup()
573+
574+
linear_obj = aa.m.MockLinearObjFuncList(
575+
parameters=3,
576+
mapping_matrix=rng.normal(size=(mask.pixels_in_mask, 3)),
577+
)
578+
579+
inversion_sparse = aa.Inversion(
580+
dataset=dataset_sparse, linear_obj_list=[linear_obj]
581+
)
582+
inversion_mapping = aa.Inversion(dataset=dataset, linear_obj_list=[linear_obj])
583+
584+
assert type(inversion_sparse) is aa.InversionInterferometerSparse
585+
assert isinstance(inversion_mapping, aa.InversionInterferometerMapping)
586+
587+
# The reconstruction comes out of a linear solve of this small, poorly conditioned system and is
588+
# compared (to its looser tolerance) by `_assert_sparse_matches_mapping` below.
589+
for name in ("curvature_matrix", "data_vector"):
590+
reference = np.asarray(getattr(inversion_mapping, name))
591+
592+
np.testing.assert_allclose(
593+
np.asarray(getattr(inversion_sparse, name)),
594+
reference,
595+
rtol=1.0e-10,
596+
atol=1.0e-10 * np.abs(reference).max(),
597+
err_msg=name,
598+
)
599+
600+
_assert_sparse_matches_mapping(
601+
dataset=dataset,
602+
dataset_sparse=dataset_sparse,
603+
linear_obj_list=[linear_obj],
604+
)
605+
606+
607+
def test__interferometer_sparse_operator__sparse_dirty_image_override__used_by_data_vector():
608+
"""
609+
The sparse operator caches the dirty image of the visibilities it was built from. When an inversion
610+
fits different visibilities (e.g. with the visibilities of ordinary light profiles subtracted), the
611+
`DatasetInterface` supplies their dirty image via `sparse_dirty_image`, which the data vector must use
612+
so that the sparse inversion reproduces the dense inversion of the subtracted visibilities.
613+
"""
614+
mask, rng, dataset, dataset_sparse = _sparse_parity_setup()
615+
616+
linear_obj = aa.m.MockLinearObjFuncList(
617+
parameters=2,
618+
mapping_matrix=rng.normal(size=(mask.pixels_in_mask, 2)),
619+
)
620+
621+
mapper = _mapper_from(
622+
mask=mask,
623+
pixels=9,
624+
shape=(3, 3),
625+
regularization=aa.reg.Constant(coefficient=1.0),
626+
)
627+
628+
subtracted = aa.Visibilities(
629+
visibilities=dataset.data.array
630+
- rng.normal(size=dataset.data.shape)
631+
- 1j * rng.normal(size=dataset.data.shape)
632+
)
633+
634+
sparse_dirty_image = dataset.transformer.image_from(
635+
visibilities=aa.Visibilities(
636+
visibilities=subtracted.array.real * dataset.noise_map.array.real**-2.0
637+
+ 1j * subtracted.array.imag * dataset.noise_map.array.imag**-2.0
638+
)
639+
).array
640+
641+
dataset_interface_mapping = aa.DatasetInterface(
642+
data=subtracted,
643+
noise_map=dataset.noise_map,
644+
grids=dataset.grids,
645+
transformer=dataset.transformer,
646+
)
647+
648+
for linear_obj_list in ([linear_obj], [linear_obj, mapper]):
649+
inversion_mapping = aa.Inversion(
650+
dataset=dataset_interface_mapping, linear_obj_list=linear_obj_list
651+
)
652+
653+
inversion_sparse = aa.Inversion(
654+
dataset=aa.DatasetInterface(
655+
data=subtracted,
656+
noise_map=dataset.noise_map,
657+
grids=dataset.grids,
658+
transformer=dataset.transformer,
659+
sparse_operator=dataset_sparse.sparse_operator,
660+
sparse_dirty_image=sparse_dirty_image,
661+
),
662+
linear_obj_list=linear_obj_list,
663+
)
664+
665+
assert isinstance(inversion_sparse, aa.InversionInterferometerSparse)
666+
667+
reference = np.asarray(inversion_mapping.data_vector)
668+
669+
np.testing.assert_allclose(
670+
np.asarray(inversion_sparse.data_vector),
671+
reference,
672+
rtol=1.0e-10,
673+
atol=1.0e-10 * np.abs(reference).max(),
674+
)
675+
676+
# The control: without the override the data vector is that of the unsubtracted visibilities.
677+
inversion_sparse_cached = aa.Inversion(
678+
dataset=aa.DatasetInterface(
679+
data=subtracted,
680+
noise_map=dataset.noise_map,
681+
grids=dataset.grids,
682+
transformer=dataset.transformer,
683+
sparse_operator=dataset_sparse.sparse_operator,
684+
),
685+
linear_obj_list=linear_obj_list,
686+
)
687+
688+
assert np.abs(
689+
np.asarray(inversion_sparse_cached.data_vector) - reference
690+
).max() > 1.0e-4 * np.abs(reference).max()
691+
692+
566693
def test__interferometer_sparse_operator__x2_mappers__identical_to_mapping():
567694
"""
568695
Two `Mapper` objects fitted simultaneously require the mapper-mapper off-diagonal block
@@ -866,3 +993,54 @@ def test__interferometer_sparse_operator__numpy_inversion_matches_jax_inversion(
866993
).max()
867994

868995
assert difference_sparse <= max(10.0 * difference_dense, 1.0e-14)
996+
997+
# A linear function list with no mapper (e.g. an MGE with no pixelization) takes the same
998+
# sparse path, with only the func-func curvature block and `Bᵀ d~` data vector.
999+
linear_obj = aa.m.MockLinearObjFuncList(
1000+
parameters=3,
1001+
mapping_matrix=np.random.default_rng(seed=seed).normal(
1002+
size=(mask.pixels_in_mask, 3)
1003+
),
1004+
)
1005+
1006+
inversion_np = aa.Inversion(
1007+
dataset=dataset_sparse, linear_obj_list=[linear_obj], xp=np
1008+
)
1009+
inversion_jax = aa.Inversion(
1010+
dataset=dataset_sparse, linear_obj_list=[linear_obj], xp=jnp
1011+
)
1012+
1013+
assert type(inversion_np) is aa.InversionInterferometerSparse
1014+
assert type(inversion_jax) is aa.InversionInterferometerSparse
1015+
1016+
for name in ("curvature_matrix", "data_vector"):
1017+
reference = np.asarray(getattr(inversion_jax, name))
1018+
1019+
np.testing.assert_allclose(
1020+
np.asarray(getattr(inversion_np, name)),
1021+
reference,
1022+
rtol=1.0e-10,
1023+
atol=1.0e-10 * np.abs(reference).max(),
1024+
err_msg=name,
1025+
)
1026+
1027+
# The unregularized func-list system is poorly conditioned, so the reconstruction is held to
1028+
# the same control as above: no worse than the dense inversion's NumPy / JAX spread.
1029+
difference_dense = np.abs(
1030+
np.asarray(
1031+
aa.Inversion(
1032+
dataset=dataset, linear_obj_list=[linear_obj], xp=np
1033+
).reconstruction
1034+
)
1035+
- np.asarray(
1036+
aa.Inversion(
1037+
dataset=dataset, linear_obj_list=[linear_obj], xp=jnp
1038+
).reconstruction
1039+
)
1040+
).max()
1041+
difference_sparse = np.abs(
1042+
np.asarray(inversion_np.reconstruction)
1043+
- np.asarray(inversion_jax.reconstruction)
1044+
).max()
1045+
1046+
assert difference_sparse <= max(10.0 * difference_dense, 1.0e-14)

‎test_autoarray/inversion/inversion/test_factory.py‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -490,6 +490,43 @@ def test__inversion_imaging__linear_obj_func_with_sparse_operator(
490490
)
491491

492492

493+
def test__inversion_interferometer__via_linear_obj_func_list__sparse_operator(
494+
interferometer_7_no_fft,
495+
):
496+
mask = interferometer_7_no_fft.real_space_mask
497+
498+
grid = aa.Grid2D.from_mask(mask=mask)
499+
500+
linear_obj = aa.m.MockLinearObjFuncList(
501+
parameters=2,
502+
grid=grid,
503+
mapping_matrix=np.random.default_rng(seed=1).normal(
504+
size=(mask.pixels_in_mask, 2)
505+
),
506+
)
507+
508+
inversion = aa.Inversion(
509+
dataset=interferometer_7_no_fft,
510+
linear_obj_list=[linear_obj],
511+
)
512+
513+
assert isinstance(inversion, aa.InversionInterferometerMapping)
514+
515+
# Unlike imaging, a func-list-only interferometer inversion (e.g. an MGE with no pixelization) uses the
516+
# sparse operator when the dataset has one.
517+
518+
inversion_sparse = aa.Inversion(
519+
dataset=interferometer_7_no_fft.apply_sparse_operator(use_jax=False),
520+
linear_obj_list=[linear_obj],
521+
)
522+
523+
assert type(inversion_sparse) is aa.InversionInterferometerSparse
524+
assert inversion_sparse.data_vector == pytest.approx(inversion.data_vector, 1.0e-8)
525+
assert inversion_sparse.curvature_matrix == pytest.approx(
526+
inversion.curvature_matrix, 1.0e-8
527+
)
528+
529+
493530
def test__inversion_interferometer__via_mapper(
494531
interferometer_7_no_fft,
495532
rectangular_mapper_7x7_3x3,

0 commit comments

Comments
 (0)