Skip to content

Commit e5efcf3

Browse files
Jammy2211claude
andcommitted
feat(interferometer): fits and save_attributes on array-free datasets (streaming phase 2, PyAutoGalaxy#638)
Mirror the autogalaxy guards in FitInterferometer (profile_visibilities None or a typed DatasetException without a transformer; profile_subtracted_visibilities None without data; inversion_with_data returns the inversion) and write dataset.fits through autogalaxy's interferometer_hdu_list_from, skipping transformer_class.json when there is no transformer. The aggregator loader is autogalaxy's and needs no change. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JZZksyZ8LTA4LLxoZjQMNF
1 parent 092897e commit e5efcf3

4 files changed

Lines changed: 244 additions & 19 deletions

File tree

‎autolens/interferometer/fit_interferometer.py‎

Lines changed: 36 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -127,15 +127,34 @@ def profile_image(self) -> aa.Array2D:
127127
return self.tracer.image_2d_from(grid=self.grids.lp, xp=self._xp)
128128

129129
@cached_property
130-
def profile_visibilities(self) -> aa.Visibilities:
130+
def profile_visibilities(self) -> Optional[aa.Visibilities]:
131131
"""
132132
Returns the visibilities of every light profile in the tracer, which are computed by performing a Fourier
133133
transform to the sum of light profile images.
134134
135135
If the tracer has no ordinary (non-linear) light profile (e.g. its light is entirely an MGE of linear
136136
Gaussians), the image is all zeros and the Fourier transform is skipped. This is decided structurally,
137137
so it is safe under `jax.jit`.
138+
139+
On an array-free dataset (built by `Interferometer.from_stream` / `from_sparse_terms`, which has no
140+
`uv_wavelengths` and so no transformer) there are no visibilities to compute: this returns `None` when
141+
the tracer has no ordinary light profile, and raises an `exc.DatasetException` when it does, because
142+
subtracting a light profile's visibilities needs the visibility arrays.
138143
"""
144+
if self.dataset.transformer is None:
145+
if _has_light_profile_non_linear(galaxies=self.tracer.galaxies):
146+
raise aa.exc.DatasetException(
147+
"This FitInterferometer's dataset is array-free (built by from_stream / "
148+
"from_sparse_terms) and has no visibilities or transformer, so the tracer's "
149+
"ordinary (non-linear) light profiles cannot be Fourier transformed and "
150+
"subtracted. An array-free dataset supports pixelization-only and linear-light "
151+
"fits; non-linear light profiles arrive in a later phase. Use the in-memory "
152+
"constructor (`Interferometer(data=..., noise_map=..., uv_wavelengths=..., ...)`) "
153+
"to fit them."
154+
)
155+
156+
return None
157+
139158
if _has_light_profile_non_linear(galaxies=self.tracer.galaxies):
140159
return self.dataset.transformer.visibilities_from(
141160
image=self.profile_image, xp=self._xp
@@ -146,12 +165,21 @@ def profile_visibilities(self) -> aa.Visibilities:
146165
)
147166

148167
@cached_property
149-
def profile_subtracted_visibilities(self) -> aa.Visibilities:
168+
def profile_subtracted_visibilities(self) -> Optional[aa.Visibilities]:
150169
"""
151170
Returns the interferometer dataset's visibilities with all transformed light profile images in the fit's
152171
tracer subtracted.
172+
173+
On an array-free dataset there are no visibilities, so this is `None`. `profile_visibilities` is
174+
evaluated first, so a fit with ordinary light profiles on an array-free dataset raises rather than
175+
silently fitting the unsubtracted sparse terms.
153176
"""
154-
return self.data - self.profile_visibilities
177+
profile_visibilities = self.profile_visibilities
178+
179+
if self.data is None:
180+
return None
181+
182+
return self.data - profile_visibilities
155183

156184
@property
157185
def _uses_precomputed_data_term(self) -> bool:
@@ -250,10 +278,14 @@ def inversion_with_data(self) -> Optional[aa.AbstractInversion]:
250278
from the visibilities in that case, so the data it fitted are `fit.data`: this returns a shallow copy of
251279
the inversion (solved first, so it shares the reconstruction and every other cached quantity) whose dataset interface carries
252280
`fit.data`. In every other case it returns `inversion` itself.
281+
282+
On an array-free dataset (built by `Interferometer.from_stream` / `from_sparse_terms`) `fit.data` is
283+
`None`, so there are no visibilities to carry and `inversion` itself is returned; output quantities
284+
that read the visibilities are unavailable on such a fit.
253285
"""
254286
inversion = self.inversion
255287

256-
if inversion is None or inversion.dataset.data is not None:
288+
if inversion is None or inversion.dataset.data is not None or self.data is None:
257289
return inversion
258290

259291
# Solve first, so the copy shares the reconstruction (and everything it cached) rather than repeating it.

‎autolens/interferometer/model/analysis.py‎

Lines changed: 11 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,11 @@
1717
from typing import Optional
1818

1919
from autonerves.dictable import to_dict
20-
from autonerves.fitsable import hdu_list_for_output_from
2120

2221
import autofit as af
2322
import autoarray as aa
2423
import autogalaxy as ag
24+
from autogalaxy.interferometer.model.analysis import interferometer_hdu_list_from
2525

2626
from autolens.analysis.analysis.dataset import AnalysisDataset
2727
from autolens.analysis.exceptions import raise_fit_exception
@@ -351,16 +351,9 @@ def save_attributes(self, paths: af.DirectoryPaths):
351351
"""
352352
super().save_attributes(paths=paths)
353353

354-
hdu_list = hdu_list_for_output_from(
355-
values_list=[
356-
self.dataset.real_space_mask.astype("float"),
357-
self.dataset.data.in_array,
358-
self.dataset.noise_map.in_array,
359-
self.dataset.uv_wavelengths,
360-
],
361-
ext_name_list=["mask", "data", "noise_map", "uv_wavelengths"],
362-
header_dict=self.dataset.real_space_mask.header_dict,
363-
)
354+
# An array-free dataset writes its `SparseTerms` instead of the visibility arrays (see
355+
# `autogalaxy.interferometer.model.analysis.interferometer_hdu_list_from`).
356+
hdu_list = interferometer_hdu_list_from(dataset=self.dataset)
364357

365358
# `dataset.fits` is written once per search, to the `image` folder, and is written
366359
# unconditionally (it is not gated on any visualization setting). The write is skipped
@@ -374,7 +367,10 @@ def save_attributes(self, paths: af.DirectoryPaths):
374367
if not dataset_path.exists():
375368
hdu_list.writeto(dataset_path, overwrite=True)
376369

377-
paths.save_json(
378-
"transformer_class",
379-
to_dict(self.dataset.transformer.__class__),
380-
)
370+
# An array-free dataset has no transformer; its transformer class name is recorded in
371+
# the `dataset.fits` header instead.
372+
if self.dataset.transformer is not None:
373+
paths.save_json(
374+
"transformer_class",
375+
to_dict(self.dataset.transformer.__class__),
376+
)

‎test_autolens/interferometer/model/test_analysis_interferometer.py‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -187,3 +187,104 @@ def test__shared_state_from__populates_mesh_geometry_fields(interferometer_7):
187187
shared = analysis.shared_state_from(instance=instance)
188188
assert shared.source_plane_mesh_grid is not None
189189
assert shared.image_plane_mesh_grid is not None
190+
191+
192+
class _FitStub:
193+
"""
194+
The part of a `PyAutoFit` aggregator `Fit` the interferometer loader reads: `value(name)` returning the
195+
`dataset.fits` HDU list written by `save_attributes` and the saved `transformer_class` json (if any).
196+
"""
197+
198+
def __init__(self, paths):
199+
self.paths = paths
200+
self.children = []
201+
202+
def value(self, name):
203+
from astropy.io import fits
204+
from autonerves.dictable import from_dict
205+
206+
if name == "dataset":
207+
return fits.open(self.paths.image_path / "dataset.fits")
208+
209+
if name == "transformer_class":
210+
path = self.paths._files_path / "transformer_class.json"
211+
212+
if not path.exists():
213+
return None
214+
215+
return from_dict(self.paths.load_json("transformer_class"))
216+
217+
return None
218+
219+
220+
def test__save_attributes__array_free_dataset__aggregator_round_trip(
221+
interferometer_7, tmp_path
222+
):
223+
"""
224+
`save_attributes` writes an array-free dataset's `SparseTerms` to `dataset.fits` and the aggregator loader
225+
(shared with autogalaxy) rebuilds it via `Interferometer.from_sparse_terms`, so a fit of the reloaded
226+
dataset with a tracer reproduces the original `log_evidence`.
227+
"""
228+
from astropy.io import fits
229+
230+
from autolens.aggregator import _interferometer_from
231+
232+
dataset = aa.Interferometer.from_stream(
233+
[
234+
(
235+
interferometer_7.uv_wavelengths,
236+
interferometer_7.data,
237+
interferometer_7.noise_map,
238+
)
239+
],
240+
real_space_mask=interferometer_7.real_space_mask,
241+
transformer_class=type(interferometer_7.transformer),
242+
)
243+
244+
paths = af.DirectoryPaths(name="array_free_round_trip", path_prefix=str(tmp_path))
245+
246+
analysis = al.AnalysisInterferometer(dataset=dataset, use_jax=False)
247+
analysis.save_attributes(paths=paths)
248+
249+
with fits.open(paths.image_path / "dataset.fits") as hdu_list:
250+
assert [hdu.name for hdu in hdu_list] == [
251+
"MASK",
252+
"NUFFT_PRECISION_OPERATOR",
253+
"DIRTY_IMAGE",
254+
"DIRTY_BEAM",
255+
"SPARSE_TERMS_SCALARS",
256+
]
257+
258+
assert not (paths._files_path / "transformer_class.json").exists()
259+
260+
dataset_reloaded = _interferometer_from(fit=_FitStub(paths=paths))[0]
261+
262+
assert dataset_reloaded.is_array_free
263+
assert (
264+
dataset_reloaded.sparse_operator.data_term
265+
== dataset.sparse_operator.data_term
266+
)
267+
assert (
268+
dataset_reloaded.sparse_operator.noise_normalization
269+
== dataset.sparse_operator.noise_normalization
270+
)
271+
272+
lens = al.Galaxy(
273+
redshift=0.5,
274+
mass=al.mp.Isothermal(centre=(0.0, 0.0), einstein_radius=1.0),
275+
)
276+
source = al.Galaxy(
277+
redshift=1.0,
278+
pixelization=al.Pixelization(
279+
mesh=al.mesh.RectangularUniform(shape=(3, 3)),
280+
regularization=al.reg.Constant(coefficient=1.0),
281+
),
282+
)
283+
tracer = al.Tracer(galaxies=[lens, source])
284+
285+
log_evidence = al.FitInterferometer(dataset=dataset, tracer=tracer).log_evidence
286+
log_evidence_reloaded = al.FitInterferometer(
287+
dataset=dataset_reloaded, tracer=tracer
288+
).log_evidence
289+
290+
assert log_evidence_reloaded == pytest.approx(log_evidence, rel=1.0e-8)

‎test_autolens/interferometer/test_fit_interferometer.py‎

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -733,3 +733,99 @@ def figure_of_merit_from(coefficient, xp):
733733
assert float(figure_of_merit_jax) == pytest.approx(
734734
figure_of_merit_numpy, rel=1.0e-8
735735
)
736+
737+
738+
def _array_free_dataset_from(dataset):
739+
"""
740+
The array-free counterpart of `dataset` (no visibilities, uv-wavelengths or transformer), built by
741+
streaming its visibilities through `Interferometer.from_stream` with the same transformer class.
742+
"""
743+
return aa.Interferometer.from_stream(
744+
[(dataset.uv_wavelengths, dataset.data, dataset.noise_map)],
745+
real_space_mask=dataset.real_space_mask,
746+
transformer_class=type(dataset.transformer),
747+
)
748+
749+
750+
def test__fit_figure_of_merit__array_free_dataset__pixelization_only__matches_in_memory_sparse(
751+
interferometer_7,
752+
):
753+
dataset_sparse = interferometer_7.apply_sparse_operator(use_jax=False)
754+
dataset_array_free = _array_free_dataset_from(interferometer_7)
755+
756+
assert dataset_array_free.is_array_free
757+
758+
tracer = _pixelized_source_tracer()
759+
760+
fit_sparse = al.FitInterferometer(dataset=dataset_sparse, tracer=tracer)
761+
fit_array_free = al.FitInterferometer(dataset=dataset_array_free, tracer=tracer)
762+
763+
assert fit_array_free._uses_precomputed_data_term
764+
assert fit_array_free.inversion.dataset.data is None
765+
assert isinstance(fit_array_free.inversion, aa.InversionInterferometerSparse)
766+
767+
assert fit_array_free.figure_of_merit == pytest.approx(
768+
fit_sparse.figure_of_merit, rel=1.0e-8
769+
)
770+
assert fit_array_free.log_evidence == pytest.approx(
771+
fit_sparse.log_evidence, rel=1.0e-8
772+
)
773+
774+
assert fit_array_free.profile_visibilities is None
775+
assert fit_array_free.profile_subtracted_visibilities is None
776+
assert fit_array_free.inversion_with_data is fit_array_free.inversion
777+
778+
779+
def test__fit_figure_of_merit__array_free_dataset__pixelization_only__jax_jit_matches_numpy(
780+
interferometer_7,
781+
):
782+
jax = pytest.importorskip("jax")
783+
import jax.numpy as jnp
784+
785+
dataset_sparse = interferometer_7.apply_sparse_operator(use_jax=False)
786+
dataset_array_free = _array_free_dataset_from(interferometer_7)
787+
788+
def figure_of_merit_from(coefficient, dataset, xp):
789+
fit = al.FitInterferometer(
790+
dataset=dataset,
791+
tracer=_pixelized_source_tracer(coefficient=coefficient),
792+
xp=xp,
793+
)
794+
795+
assert fit.inversion.dataset.data is None
796+
797+
return fit.figure_of_merit
798+
799+
figure_of_merit_sparse = figure_of_merit_from(
800+
coefficient=1.0, dataset=dataset_sparse, xp=np
801+
)
802+
figure_of_merit_numpy = figure_of_merit_from(
803+
coefficient=1.0, dataset=dataset_array_free, xp=np
804+
)
805+
806+
figure_of_merit_jax = jax.jit(
807+
lambda c: figure_of_merit_from(c, dataset=dataset_array_free, xp=jnp)
808+
)(1.0)
809+
810+
assert figure_of_merit_numpy == pytest.approx(figure_of_merit_sparse, rel=1.0e-8)
811+
assert float(figure_of_merit_jax) == pytest.approx(
812+
figure_of_merit_numpy, rel=1.0e-8
813+
)
814+
815+
816+
def test__fit_figure_of_merit__array_free_dataset__lens_light_profile__raises(
817+
interferometer_7,
818+
):
819+
dataset_array_free = _array_free_dataset_from(interferometer_7)
820+
821+
tracer = _pixelized_source_tracer(lens_light=True)
822+
823+
fit = al.FitInterferometer(dataset=dataset_array_free, tracer=tracer)
824+
825+
assert not fit._uses_precomputed_data_term
826+
827+
with pytest.raises(aa.exc.DatasetException):
828+
fit.profile_visibilities
829+
830+
with pytest.raises(aa.exc.DatasetException):
831+
al.FitInterferometer(dataset=dataset_array_free, tracer=tracer).figure_of_merit

0 commit comments

Comments
 (0)