Skip to content

Commit 0cc47cd

Browse files
Jammy2211claude
authored andcommitted
feat(interferometer): from_fits accepts raise_error_dft_visibilities_limit
Thread the existing constructor kwarg through `Interferometer.from_fits` so callers loading >10,000-visibility datasets with TransformerDFT (e.g. when profiling the JAX-traceable DFT path at ALMA-scale before nufftax is wired in) can opt out of the safety guard without reconstructing the dataset by hand. Default behaviour unchanged. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent f75bc19 commit 0cc47cd

2 files changed

Lines changed: 53 additions & 0 deletions

File tree

autoarray/dataset/interferometer/dataset.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -139,6 +139,7 @@ def from_fits(
139139
noise_map_hdu=0,
140140
uv_wavelengths_hdu=0,
141141
transformer_class=TransformerNUFFT,
142+
raise_error_dft_visibilities_limit: bool = True,
142143
):
143144
"""
144145
Load an interferometer dataset from multiple .fits files.
@@ -176,6 +177,11 @@ def from_fits(
176177
transformer_class
177178
The class of the Fourier Transform which maps images from real space to Fourier space
178179
visibilities. Defaults to `TransformerNUFFT` for efficiency with large datasets.
180+
raise_error_dft_visibilities_limit
181+
If True (default), raise a `DatasetException` when ``transformer_class`` is
182+
`TransformerDFT` and the dataset has more than 10,000 visibilities. Set to False to
183+
opt into the slow DFT path at ALMA-scale (e.g. when profiling the JAX-traceable
184+
DFT path before a JIT-friendly NUFFT is available).
179185
180186
Returns
181187
-------
@@ -199,6 +205,7 @@ def from_fits(
199205
noise_map=noise_map,
200206
uv_wavelengths=uv_wavelengths,
201207
transformer_class=transformer_class,
208+
raise_error_dft_visibilities_limit=raise_error_dft_visibilities_limit,
202209
)
203210

204211
def apply_sparse_operator(

test_autoarray/dataset/interferometer/test_dataset.py

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,52 @@ def test__dirty_signal_to_noise_map__shape_native_matches_real_space_mask(
6464
).all()
6565

6666

67+
def test__from_fits__raise_error_dft_visibilities_limit__threads_kwarg(
68+
tmp_path, mask_2d_7x7
69+
):
70+
"""``from_fits`` must forward ``raise_error_dft_visibilities_limit`` to the
71+
``Interferometer`` constructor so callers loading large DFT-based datasets can opt out
72+
of the >10,000-visibility safety check (e.g. for profiling the JAX-traceable DFT path)."""
73+
from astropy.io import fits
74+
75+
n_visibilities = 10_001
76+
visibilities = np.ones((n_visibilities, 2), dtype=np.float64)
77+
noise_map = np.ones((n_visibilities, 2), dtype=np.float64)
78+
uv_wavelengths = np.zeros((n_visibilities, 2), dtype=np.float64)
79+
80+
data_path = tmp_path / "data.fits"
81+
noise_map_path = tmp_path / "noise_map.fits"
82+
uv_path = tmp_path / "uv_wavelengths.fits"
83+
84+
for arr, path in (
85+
(visibilities, data_path),
86+
(noise_map, noise_map_path),
87+
(uv_wavelengths, uv_path),
88+
):
89+
fits.PrimaryHDU(data=arr).writeto(path, overwrite=True)
90+
91+
with pytest.raises(aa.exc.DatasetException):
92+
aa.Interferometer.from_fits(
93+
data_path=data_path,
94+
noise_map_path=noise_map_path,
95+
uv_wavelengths_path=uv_path,
96+
real_space_mask=mask_2d_7x7,
97+
transformer_class=transformer.TransformerDFT,
98+
)
99+
100+
dataset = aa.Interferometer.from_fits(
101+
data_path=data_path,
102+
noise_map_path=noise_map_path,
103+
uv_wavelengths_path=uv_path,
104+
real_space_mask=mask_2d_7x7,
105+
transformer_class=transformer.TransformerDFT,
106+
raise_error_dft_visibilities_limit=False,
107+
)
108+
109+
assert dataset.uv_wavelengths.shape[0] == n_visibilities
110+
assert type(dataset.transformer) == transformer.TransformerDFT
111+
112+
67113
def test__from_fits__all_files_in_one_fits__load_using_different_hdus(mask_2d_7x7):
68114
dataset = aa.Interferometer.from_fits(
69115
real_space_mask=mask_2d_7x7,

0 commit comments

Comments
 (0)