Skip to content

Commit ffd13ba

Browse files
author
James Nightingale
authored
Merge pull request #601 from PyAutoLabs/feature/streaming-p5-cubes-phase-centre
feat: per-channel SparseTerms sum and phase-centre shifts in sparse_terms_from_chunks (streaming P5)
2 parents 176f61e + 8a2332f commit ffd13ba

4 files changed

Lines changed: 678 additions & 8 deletions

File tree

‎autoarray/dataset/interferometer/dataset.py‎

Lines changed: 25 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import logging
22
import numpy as np
3-
from typing import Optional
3+
from typing import Optional, Tuple
44

55
from autonerves.fitsable import ndarray_via_fits_from
66
from autonerves import cached_property
@@ -338,6 +338,7 @@ def from_stream(
338338
use_jax: bool = False,
339339
show_progress: bool = False,
340340
batch_size: int = 128,
341+
phase_centre: Optional[Tuple[float, float]] = None,
341342
) -> "Interferometer":
342343
"""
343344
Build an array-free `Interferometer` by accumulating a stream of visibility chunks,
@@ -360,6 +361,12 @@ def from_stream(
360361
Passed to `sparse_terms_from_chunks`.
361362
batch_size
362363
The number of source-pixel columns processed per batch by the sparse operator.
364+
phase_centre
365+
The `(y, x)` phase-centre shift in arcseconds applied to every chunk's
366+
visibilities, `d' = d * exp(+2 pi i (u * x0 + v * y0))` (radians), so a source at
367+
`(y0, x0)` lands at the image origin; recorded as `sparse_terms.phase_centre`.
368+
`None` applies no shift (recorded as `(0.0, 0.0)`). See
369+
`sparse_terms_from_chunks`.
363370
364371
Raises
365372
------
@@ -379,6 +386,7 @@ def from_stream(
379386
chunk_k=chunk_k,
380387
use_jax=use_jax,
381388
show_progress=show_progress,
389+
phase_centre=phase_centre,
382390
)
383391

384392
return cls.from_sparse_terms(
@@ -670,7 +678,10 @@ def apply_sparse_operator_from_chunks(
670678
The number of source-pixel columns processed per batch by the sparse operator.
671679
accumulator_kwargs
672680
Passed to `sparse_terms_from_chunks` (`transformer_class`, `method`, `eps`,
673-
`chunk_size`, `chunk_k`, `use_jax`, `show_progress`). When not given,
681+
`chunk_size`, `chunk_k`, `use_jax`, `show_progress`). `phase_centre` is rejected:
682+
it would shift the operator's terms but not this dataset's retained `data`, so the
683+
two would describe different phase centres; use `Interferometer.from_stream`
684+
(array-free, no retained data) to stream with a phase-centre shift. When not given,
674685
`transformer_class`, `eps` and `chunk_size` follow this dataset's transformer (the
675686
same defaults `psf_precision_operator_from` takes), so the accumulated terms match
676687
`apply_sparse_operator()`.
@@ -684,8 +695,19 @@ def apply_sparse_operator_from_chunks(
684695
Raises
685696
------
686697
exc.DatasetException
687-
If any chunk has unequal real and imaginary noise sigma.
698+
If any chunk has unequal real and imaginary noise sigma, or a `phase_centre` is
699+
passed.
688700
"""
701+
if accumulator_kwargs.get("phase_centre") is not None:
702+
raise exc.DatasetException(
703+
"Interferometer.apply_sparse_operator_from_chunks does not take a "
704+
"`phase_centre`: it would shift the sparse operator's dirty image and "
705+
"`data_term` but not this dataset's retained `data`, so the two would describe "
706+
"different phase centres and give contradictory likelihoods. Use "
707+
"`Interferometer.from_stream(..., phase_centre=...)`, which builds an "
708+
"array-free dataset with no retained visibilities."
709+
)
710+
689711
if disable_jax() and accumulator_kwargs.get("use_jax", False):
690712
accumulator_kwargs["use_jax"] = False
691713

‎autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py‎

Lines changed: 93 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1926,6 +1926,12 @@ class SparseTerms:
19261926
The NUFFT precision the precision operator was built with.
19271927
transformer_class_name
19281928
The class name of the transformer that formed the dirty image and beam.
1929+
phase_centre
1930+
The `(y, x)` phase-centre shift in arcseconds the visibilities were re-centred on
1931+
before forming the dirty image (`sparse_terms_from_chunks(phase_centre=...)`, which
1932+
records `(0.0, 0.0)` when no shift is applied). `None` means not recorded. Terms with
1933+
different phase centres have dirty images referred to different sky origins and
1934+
cannot be summed.
19291935
"""
19301936

19311937
nufft_precision_operator: np.ndarray
@@ -1940,13 +1946,15 @@ class SparseTerms:
19401946
origin: Optional[tuple] = None
19411947
eps: Optional[float] = None
19421948
transformer_class_name: Optional[str] = None
1949+
phase_centre: Optional[Tuple[float, float]] = None
19431950

19441951
def __add__(self, other: "SparseTerms") -> "SparseTerms":
19451952
"""
19461953
The field-wise sum of two `SparseTerms`.
19471954
19481955
Raises `exc.InversionException` if a provenance field (`shape_native`,
1949-
`pixel_scales`, `origin`, `eps`) is recorded on both sides with different values.
1956+
`pixel_scales`, `origin`, `eps`, `phase_centre`, `transformer_class_name`) is recorded on
1957+
both sides with different values.
19501958
The result carries, for each provenance field, the recorded value from either side
19511959
(the left one when both are recorded), so an unrecorded operand never erases the
19521960
provenance of a recorded one.
@@ -1965,14 +1973,23 @@ def __add__(self, other: "SparseTerms") -> "SparseTerms":
19651973
f"{other.nufft_precision_operator.shape} differ."
19661974
)
19671975

1968-
for name in ("shape_native", "pixel_scales", "origin", "eps"):
1976+
for name in (
1977+
"shape_native",
1978+
"pixel_scales",
1979+
"origin",
1980+
"eps",
1981+
"phase_centre",
1982+
"transformer_class_name",
1983+
):
19691984
value_self = getattr(self, name)
19701985
value_other = getattr(other, name)
19711986

19721987
if value_self is None or value_other is None:
19731988
continue
19741989

1975-
if isinstance(value_self, float) or isinstance(value_other, float):
1990+
if isinstance(value_self, str) or isinstance(value_other, str):
1991+
differ = str(value_self) != str(value_other)
1992+
elif isinstance(value_self, float) or isinstance(value_other, float):
19761993
differ = float(value_self) != float(value_other)
19771994
else:
19781995
differ = tuple(value_self) != tuple(value_other)
@@ -1981,8 +1998,8 @@ def __add__(self, other: "SparseTerms") -> "SparseTerms":
19811998
raise exc.InversionException(
19821999
"SparseTerms can only be added when accumulated with the same provenance: "
19832000
f"`{name}` is {value_self!r} on one and {value_other!r} on the other. "
1984-
"Terms accumulated on different real-space masks or NUFFT accuracies do "
1985-
"not describe the same operator."
2001+
"Terms accumulated on different real-space masks, NUFFT accuracies, "
2002+
"phase centres or transformers do not describe the same operator."
19862003
)
19872004

19882005
return SparseTerms(
@@ -2003,8 +2020,23 @@ def __add__(self, other: "SparseTerms") -> "SparseTerms":
20032020
transformer_class_name=_recorded(
20042021
self.transformer_class_name, other.transformer_class_name
20052022
),
2023+
phase_centre=_recorded(self.phase_centre, other.phase_centre),
20062024
)
20072025

2026+
def __radd__(self, other) -> "SparseTerms":
2027+
"""
2028+
Support `sum(list_of_terms)`, which starts from the integer `0`: `0 + terms` is
2029+
`terms`. Any other left operand is not supported.
2030+
"""
2031+
if (
2032+
isinstance(other, (int, float))
2033+
and not isinstance(other, bool)
2034+
and other == 0
2035+
):
2036+
return self
2037+
2038+
return NotImplemented
2039+
20082040

20092041
def _recorded(value_left, value_right):
20102042
"""
@@ -2042,6 +2074,7 @@ def sparse_terms_from_chunks(
20422074
chunk_k: int = 2048,
20432075
use_jax: bool = False,
20442076
show_progress: bool = False,
2077+
phase_centre: Optional[Tuple[float, float]] = None,
20452078
) -> SparseTerms:
20462079
"""
20472080
Accumulate the `SparseTerms` of an interferometer dataset one chunk of visibilities at a
@@ -2068,6 +2101,35 @@ def sparse_terms_from_chunks(
20682101
`K` may differ between chunks; empty chunks are skipped. Chunks are consumed once, in
20692102
order, and only one is referenced at a time.
20702103
2104+
Multi-channel (MFS) terms
2105+
-------------------------
2106+
Because every field is a sum over visibilities, the terms of several channels (or any
2107+
partition of the visibilities) accumulated separately on the same mask sum to the terms
2108+
of all of them accumulated together: `sum(per_channel_terms)` (via `SparseTerms.__add__`
2109+
/ `__radd__`) is the multi-frequency-synthesis (MFS) terms, equal to one accumulation
2110+
over every channel's chunks to summation order.
2111+
2112+
Phase-centre shift
2113+
------------------
2114+
With `phase_centre=(y0, x0)` (arcseconds, autoarray `(y, x)` order like a mask `origin`),
2115+
each chunk's visibilities are multiplied by the unit phase
2116+
2117+
d' = d * exp(+2 pi i (u * l0 + v * m0)), l0 = x0, m0 = y0 in radians,
2118+
2119+
before any data-dependent term is formed. The forward transform is
2120+
`V(u, v) = sum I exp(-2 pi i (u x + v y))`, so this re-centres the phase centre onto
2121+
`(y0, x0)`: a source at `(y0, x0)` lands at the image origin. The shifted visibilities
2122+
form both the dirty image and `data_term`, so every data-dependent term describes the same
2123+
visibilities. The precision operator, dirty beam, `sum_weights`, `noise_normalization` and
2124+
`n_vis` depend on the baselines and sigmas alone and are bit-identical to the unshifted
2125+
accumulation. `data_term` is invariant under a unit phase only when the real and imaginary
2126+
sigmas are exactly equal; `check_noise_map_real_imag_equal` accepts them equal to a relative
2127+
tolerance, so with slightly unequal sigmas it differs from the unshifted value at that
2128+
tolerance (it is still the correct `data_term` of the shifted data). The shift is recorded
2129+
as `SparseTerms.phase_centre` provenance (`(0.0, 0.0)` when no shift is applied), so terms
2130+
with different phase centres -- including shifted and unshifted ones -- refuse to be
2131+
summed.
2132+
20712133
Parameters
20722134
----------
20732135
chunks
@@ -2082,6 +2144,10 @@ def sparse_terms_from_chunks(
20822144
method, eps, chunk_size, chunk_k, use_jax, show_progress
20832145
Passed to `nufft_precision_operator_from` for each chunk (`chunk_size` is the NUFFT
20842146
builder's own inner chunk, a memory ceiling within a chunk). `eps=None` is `1e-12`.
2147+
phase_centre
2148+
The `(y, x)` phase-centre shift in arcseconds applied to every chunk's visibilities
2149+
before forming the dirty image (see "Phase-centre shift" above). `None` applies no
2150+
shift and records `(0.0, 0.0)`.
20852151
20862152
Returns
20872153
-------
@@ -2105,6 +2171,18 @@ def sparse_terms_from_chunks(
21052171
if eps is None:
21062172
eps = 1.0e-12
21072173

2174+
shift = phase_centre is not None
2175+
2176+
if not shift:
2177+
phase_centre = (0.0, 0.0)
2178+
else:
2179+
from astropy import units
2180+
2181+
phase_centre = (float(phase_centre[0]), float(phase_centre[1]))
2182+
arcsec_to_rad = units.arcsec.to(units.rad)
2183+
m0 = phase_centre[0] * arcsec_to_rad
2184+
l0 = phase_centre[1] * arcsec_to_rad
2185+
21082186
terms = None
21092187

21102188
for uv_wavelengths, data, noise_map in chunks:
@@ -2135,6 +2213,15 @@ def sparse_terms_from_chunks(
21352213
uv_wavelengths=uv_wavelengths, real_space_mask=real_space_mask
21362214
)
21372215

2216+
# The phase-centre shift is applied to the data before every data-dependent term (the
2217+
# dirty image and `data_term`), so all of them describe the same shifted visibilities.
2218+
# The precision operator, dirty beam, `sum_weights` and `noise_normalization` depend on
2219+
# the baselines and sigmas alone.
2220+
if shift:
2221+
data = data * np.exp(
2222+
2j * np.pi * (uv_wavelengths[:, 0] * l0 + uv_wavelengths[:, 1] * m0)
2223+
)
2224+
21382225
# The same arguments `Interferometer.psf_precision_operator_from` builds from the
21392226
# dataset's own transformer, so the per-chunk operators sum to the dataset's.
21402227
nufft_precision_operator = np.asarray(
@@ -2192,6 +2279,7 @@ def sparse_terms_from_chunks(
21922279
origin=tuple(real_space_mask.origin),
21932280
eps=float(eps),
21942281
transformer_class_name=type(transformer).__name__,
2282+
phase_centre=phase_centre,
21952283
)
21962284

21972285
terms = chunk_terms if terms is None else terms + chunk_terms

0 commit comments

Comments
 (0)