from math import isfinite, sqrt
from jax import checkpoint, lax, vmap, random
import jax.numpy as jnp
from tabascal.components import Component, assert_attr_shape
from tabascal.dist import standard_normal
from tabascal.interferometry import fov_to_eff_diameter, max_ast_fringe_rate
from tabascal.fft_gp import FROM_DATA, latent_to_signal_init, latent_to_signal, latent_to_signal_dft_init, latent_to_signal_dft, signal_to_latent_init, signal_to_latent, knee_from_corr_scale, pow_spec_nd, rms_vis, validate_cutoff, validate_gp_cov
from tabascal.timing import measure_runtime
from tabascal.truth import read_true_vis_ast
#: The ``ast.gp_cov`` keys and what each may be, for
#: :func:`tabascal.fft_gp.validate_gp_cov`. The same keys the RFI prior takes,
#: except that the time scale is given as a field of view: a sky source's
#: coherence time is its fringe rate, which follows from how far off axis it is
#: and not from a number to pick, so the key asks for the angle.
_GP_COV_RULES = {
"std": "amplitude",
"corr_freq": "number",
"fov_deg": "number",
"gammas": "pair",
}
#: Both knees may be null, and mean different things by it. ``fov_deg`` unset is
#: the telescope's own primary beam, 2 * 1.22 * lambda / D from the dish
#: diameter in the measurement set, rather than a field of view chosen by hand.
#: ``corr_freq`` unset is no roll-off along the frequency axis at all, which
#: the comment above :class:`GPVisAst` explains -- where the RFI prior reads the
#: same null as half the observed band.
_GP_COV_OPTIONAL = ("fov_deg", "corr_freq")
#: ``ast.cutoff`` when it is left null. What every shipped configuration sets,
#: kept as the default now that the key may be omitted.
_DEFAULT_PK_CUTOFF = 1e-6
#: The power spectrum is used for its shape alone -- which modes survive
#: ``cutoff``, and their relative weight -- and the amplitude is applied by
#: normalising the mode standard deviations to ``std``. This is the ``p0`` the
#: shape is evaluated with; any positive value gives the same shape and the
#: same surviving modes, so it is 1.
_SHAPE_ONLY = 1.0
#: ``sqrt(E|z|^2)`` for the latent :meth:`GPVisAst.build_set_params` draws: a
#: real and an imaginary standard normal, so the complex latent carries twice
#: the variance of the unit circularly-symmetric one and the signal comes out
#: ``sqrt(2)`` wider than the mode variances alone. Dividing it out is what
#: makes ``std`` the width of the visibility rather than of its real part.
#:
#: ``math``, not ``jnp``: this module is imported before ``set_precision`` runs,
#: so a jax scalar here would be built in float32 and a double run would carry a
#: float32 sqrt(2) into its prior.
_LATENT_WIDTH = sqrt(2.0)
#: ``corr_freq: null`` is no roll-off along the frequency axis:
#: :func:`~tabascal.fft_gp.knee_from_corr_scale` returns an infinite knee and
#: the power spectrum is flat there.
#:
#: On a single channel it makes no difference what the knee is at all -- the
#: only delay mode is zero -- which is why every shipped configuration leaves it
#: unset.
[docs]
class GPVisAst(Component):
# Accumulates into vis_ast, which Model zeroes before the components run.
required_inputs = {
"vis_ast": ("n_bl", "n_freq", "n_time"),
}
output_shapes = {
"vis_ast": ("n_bl", "n_freq", "n_time"),
}
# Add parameter specifications
parameters = {
"ast_k_r_base": ("n_bl", "n_k_freq_ast", "n_k_time_ast"),
"ast_k_i_base": ("n_bl", "n_k_freq_ast", "n_k_time_ast"),
}
[docs]
def setup(self, config):
"""All validation and error-prone operations here"""
try:
# Store only what's needed for forward computation
self.n_time = config.n_time
self.n_bl = config.n_bl
self.n_freq = config.n_freq
self.int_time = config.int_time
self.chan_width = config.chan_width
self.dish_d = config.dish_d
self.uvw = config.uvw
self.dec = config.phase_centre["dec"]
self.freqs = config.freqs
self.times = config.times
# Checked rather than indexed: these went straight to the Fourier
# machinery, so a negative gamma, a string, or a cutoff of 1 -- which
# cuts every mode -- surfaced from inside fft_gp with a message about
# array shapes, if it surfaced at all. The same validator serves the
# RFI prior; the sections differ only in which keys are live.
gp_cov = validate_gp_cov(
config.args["ast"].get("gp_cov"),
"ast",
_GP_COV_RULES,
optional=_GP_COV_OPTIONAL,
)
config.args["ast"]["gp_cov"] = gp_cov
self.std = gp_cov["std"]
if self.std == FROM_DATA:
self.std = self._std_from_data(
config.vis_obs, config.estimator_flags
)
self.gammas = gp_cov["gammas"]
self.fov_deg = gp_cov["fov_deg"]
# Outside the covariance block: it sets how many modes are fitted,
# not what the prior believes, and since the spectrum is normalised
# after the cut it does not change the width either.
self.pk_cutoff = validate_cutoff(
config.args["ast"].get("cutoff"), "ast.cutoff", _DEFAULT_PK_CUTOFF
)
config.args["ast"]["cutoff"] = self.pk_cutoff
# corr_freq is the bandwidth over which the sky stays correlated;
# the knee is the delay conjugate to it. Shared with the RFI prior
# so the two sections cannot drift; null is no roll-off.
self.corr_freq = gp_cov["corr_freq"]
# float(), so the knee is one type whether or not corr_freq was set
# and is computed at host precision: the helper returns a jax scalar
# for a value and a Python float for None, and under x32 the former
# would carry the reciprocal in float32 for no reason.
self.k0_freq = float(knee_from_corr_scale(self.corr_freq))
self.freq_pad_factor = config.args["ast"]["freq_pad_factor"]
self.time_pad_factor = config.args["ast"]["time_pad_factor"]
# null is a setting, not a missing value: every baseline in one
# step, which is the padded transform at full width. int() alone
# would turn 1.9 into 1 without a word, and the finiteness test
# comes before it because yaml spells .inf and .nan, on which int()
# raises with a message about floats rather than about the key.
block_size = config.args["ast"].get("baseline_block_size", "auto")
if block_size not in ("auto", None) and (
isinstance(block_size, bool)
or not isinstance(block_size, (int, float))
or not isfinite(block_size)
or block_size != int(block_size)
or block_size < 1
):
raise ValueError(
"ast.baseline_block_size is the number of baselines "
"transformed per scan step: a whole number of at least 1, "
"'auto' to size it from the padded grid, or null for a "
f"single block over every baseline, got {block_size!r}."
)
self.baseline_block_size_setting = block_size
self.xs = [self.freqs, self.times]
self.pad_factors = [self.freq_pad_factor, self.time_pad_factor]
self.ss_factors = [1, 1]
# Do expensive setup operations once
self._compute_gp_params()
self._resolve_baseline_block_size()
self._compute_prior_params(config.args["ast"]["mean"], config.vis_obs)
if config.args["plots"]["truth"] or config.args["ast"]["init"] == "truth":
self._compute_true_params(
config.args["data"]["zarr_path"], config.args["data"]["data_col"]
)
self._compute_init_params(config.args["ast"]["init"], config.vis_obs)
self._set_outputs()
# Validate dimensions
self._validate_dimensions()
except Exception as e:
raise RuntimeError(f"GPVisAst setup failed: {e}")
#: How much of a padded Fourier grid one scan step may build, in bytes. The
#: block is sized from this under ``baseline_block_size: auto``. Three arrays
#: of that size are live at the peak -- the pad, the shift and the transform
#: -- so the budget is the transform's transient, not the model's.
#:
#: It exists because a block that does not bind costs scan steps for nothing:
#: on a single-channel observation, whose padded grid is (1, 180), a fixed
#: block of 128 split 4560 baselines into 36 steps and cost 13 % of the
#: optimiser's time on the cheaper of the two RFI kernels, with no memory to
#: show for it. At this budget that observation runs in one step, and a wide
#: band -- where the grid is the largest array in the model -- still blocks.
_BLOCK_BUDGET_BYTES = 64 * 1024**2
def _transform_bytes_per_baseline(self):
"""What one baseline costs the transform at its peak, in bytes.
The padded Fourier grid: the pad, the shift and the inverse transform
are each of that size. Split out because a transform that does not
build it has a different budget -- see :class:`GPVisAstDFT`.
"""
padded = 1
for dim, (lo, hi) in zip((self.n_k_freq_ast, self.n_k_time_ast), self.pads):
padded *= dim + lo + hi
return padded * jnp.zeros((), dtype=complex).dtype.itemsize
def _resolve_baseline_block_size(self):
"""Turn the ``baseline_block_size`` setting into a block, or ``None``."""
setting = self.baseline_block_size_setting
if setting != "auto":
self.baseline_block_size = None if setting is None else int(setting)
return
per_baseline = self._transform_bytes_per_baseline()
block = max(1, self._BLOCK_BUDGET_BYTES // max(1, per_baseline))
# A block at or above the axis is one step over all of it, which is what
# the whole grid fitting the budget should mean.
self.baseline_block_size = min(int(block), self.n_bl)
[docs]
def build_set_params(self):
n_bl = self.n_bl
n_k_freq_ast = self.n_k_freq_ast
n_k_time_ast = self.n_k_time_ast
def set_params(params):
params["ast_k_r_base"] = standard_normal(
"ast_k_r_base", (n_bl, n_k_freq_ast, n_k_time_ast)
)
params["ast_k_i_base"] = standard_normal(
"ast_k_i_base", (n_bl, n_k_freq_ast, n_k_time_ast)
)
return params
return set_params
[docs]
def build_constants(self):
return {
"sigma_ast_k": self.sigma_ast_k,
"mu_ast_k": self.mu_ast_k,
}
[docs]
def build_forward(self):
"""Return pure, JIT-compatible function
The baseline axis is walked in blocks of ``ast.baseline_block_size``
rather than vmapped whole. ``latent_to_signal`` pads the latent block up
to the padded k-grid, shifts, inverse-transforms and crops back, so a
vmap over every baseline holds ``(n_bl, n_freq_pad, n_time_pad)`` three
times over -- at the default padding each axis is about twice the data
one, so about four times the elements of the visibilities, and the crop
throws all of that away. The scan replaces ``n_bl`` in that shape with
the block, which is the whole of the term this component contributes to
peak memory.
There is deliberately no ``checkpoint`` on the body, unlike the RFI
components' scans. The chain from the latent parameters to ``vis_ast`` is
affine -- an elementwise ``sigma * base + mu``, then pad, shift, ifftn and
crop, every one of them linear -- so reverse mode is its transpose and
stores no primal intermediates to begin with. Issue #153 records the
measurement; a remat here would be a no-op today and a silent
recomputation if the chain ever stopped being linear.
The affine transform runs inside the body, on the block, so that neither
it nor the padded grid is ever formed for every baseline at once.
Splitting the baseline axis across devices was measured alongside this
and is deliberately not here: on the astronomical benchmark the scan
takes the value-and-gradient peak from 1008 MB to 163.5 MB on its own,
and adding the split takes it back up to 174.5 MB, because the gather
the visibilities need on the way out makes the backward pass scatter a
cotangent the forward has just collected. See issue #209 for the change
that would remove the padded grid altogether.
"""
prefix = self.prefix
forward_transform = self.forward_transform
block_size = self.baseline_block_size
block_signal = self.build_block_signal()
def blocked_vis(ast_k_base, sigma_ast_k, mu_ast_k):
# The size comes off the array rather than off the config, so the
# body says what it walks.
n_bl = ast_k_base.shape[0]
n_block = max(1, n_bl if block_size is None else min(block_size, n_bl))
n_pad = -n_bl % n_block
if n_block >= n_bl:
# One step: the scan would stack a single block and reshape it
# back, which is a copy for nothing. This is the path a padded
# grid that fits the budget takes, and it is the transform as it
# was before the scan existed.
return block_signal(
forward_transform(ast_k_base, sigma_ast_k, mu_ast_k)
)
def block_vis(carry, block):
k_base, sigma, mu = block
return carry, block_signal(
forward_transform(k_base, sigma, mu)
)
# The padding baselines of the last block carry a zero latent, a zero
# sigma and a zero mu, so they transform to zero; the slice below
# drops them, and being linear it drops their gradient with them.
def bl_blocks(x):
padded = jnp.pad(x, ((0, n_pad), (0, 0), (0, 0)))
return jnp.reshape(padded, (-1, n_block) + x.shape[1:])
_, vis_ast = lax.scan(
block_vis,
None,
tuple(bl_blocks(x) for x in (ast_k_base, sigma_ast_k, mu_ast_k)),
)
# lax.scan stacks along axis 0, which is the baseline axis of the
# result already, so the blocks are flattened rather than transposed.
return jnp.reshape(vis_ast, (-1,) + vis_ast.shape[2:])[:n_bl]
def forward(params, state, constants):
# Pure JAX operations only
sigma_ast_k = constants[f"{prefix}/sigma_ast_k"]
mu_ast_k = constants[f"{prefix}/mu_ast_k"]
ast_k_base = params["ast_k_r_base"] + 1.0j * params["ast_k_i_base"]
vis_ast = blocked_vis(ast_k_base, sigma_ast_k, mu_ast_k)
state = {**state, "vis_ast": state["vis_ast"] + vis_ast}
return state
return forward
[docs]
def build_block_signal(self):
"""The latent-to-signal transform of one block of baselines.
Split out so a subclass can replace the transform and nothing else:
the priors, the initialisation and the blocking above are the same
whichever way the modes are turned into a signal.
"""
pads = self.pads
ss_idxs = self.ss_idxs
transform = vmap(latent_to_signal, (0, None, None), 0)
return lambda Y_block: transform(Y_block, pads, ss_idxs)
[docs]
def validate_and_test(self):
"""Call this before using in JIT context"""
pass
def _compute_gp_params(self):
if self.fov_deg:
# fov_deg is the full field of view (diameter) out to the first null;
# the effective diameter makes the beam radius in
# max_ast_fringe_rate equal to fov_deg / 2.
eff_dish_d = float(fov_to_eff_diameter(self.fov_deg, jnp.min(self.freqs)))
else:
eff_dish_d = self.dish_d
# One maximum fringe rate per baseline; time and frequency are reduced
# inside max_ast_fringe_rate.
self.ast_fr = max_ast_fringe_rate(
self.uvw, self.dec, self.freqs, eff_dish_d
)
self.k0_time = self.ast_fr
self.k0s = [self.k0_freq, self.k0_time.max()]
ns = [self.n_freq, self.n_time]
dxs = [self.chan_width, self.int_time]
# 1.0, not the amplitude: both of these want the *shape* of the power
# spectrum, and use it only to decide which modes survive `cutoff`. That
# test is `pk > cutoff * pk.max(axis=i)` on each axis, and `pow_spec_nd`
# applies p0 as one final multiply, so both sides scale together and the
# surviving set is the same for any positive p0 -- exactly so in double,
# and in single for everything except a spectrum that underflows, which
# is what passing 1.0 rather than the amplitude rules out. The amplitude
# is applied afterwards, by normalising sigma.
self.pk, self.ks, self.pads, self.ss_idxs = latent_to_signal_init(
ns,
dxs,
self.pad_factors,
self.ss_factors,
_SHAPE_ONLY,
self.k0s,
self.gammas,
self.pk_cutoff,
)
# Pre-compute slicing indices for JIT-compatible latent extraction
self.latent_idxs, _ = signal_to_latent_init(
ns,
dxs,
self.pad_factors,
_SHAPE_ONLY,
self.k0s,
self.gammas,
self.pk_cutoff,
)
self.signal_to_latent = lambda vis_ast: vmap(signal_to_latent, (0, None, None), 0)(vis_ast, self.pad_factors, self.latent_idxs)
print("\nAST specs")
print(f"(d_freq, d_time): ({dxs[0]:.3e}, {dxs[1]:.3e})")
print(f"(n_freq, n_time): ({self.n_freq}, {self.n_time})")
print(f"(n_k_fq, n_k_tm): {self.pk.shape}")
self.n_k_freq_ast, self.n_k_time_ast = self.pk.shape
def sigma(k0, std):
"""Mode standard deviations that make ``std`` the width of vis_ast.
The modes are independent, so the variance of the signal they build
is the sum of theirs. Dividing the shape by its own sum makes that
sum one; the remaining ``sqrt(2)`` is the latent.
:meth:`build_set_params` draws ``ast_k_r_base`` and ``ast_k_i_base``
as two independent standard normals, so the complex latent has
``E|z|^2 = 2`` rather than 1, and the visibility comes out
``sqrt(2)`` wider than the mode variances alone would say.
With it, ``std`` is the width of the *complex* visibility about
its prior mean, ``sqrt(E|vis_ast - mu|^2)``. At the default
``ast.mean: 0`` the mean is zero and that is exactly
``rms|vis_ast|``, so "set it to the amplitude you see in a clean
channel" is literally true rather than true after dividing by
``sqrt(2)``. That is the property the key exists to have. Under
``ast.mean: data`` the prior is centred on the observed
visibilities and ``std`` is the scatter allowed around them, not
the total amplitude -- ``rms|vis_ast|`` is then
``sqrt(std^2 + |mu|^2)``.
``rfi.gp_cov.std`` names the same quantity in the same units, at
``rfi.mean: 0``. It gets there differently, since it normalises
``rfi_A``, which the visibility is quadratic in -- and a non-zero
``rfi.mean`` adds power to the RFI visibility rather than merely
shifting it, so the two stop matching there.
Dividing by the sum rather than by the mode count is what makes
the width independent of ``cutoff`` and ``gammas``: those decide
which modes are fitted and how they are weighted relative to each
other, neither of which is a statement about how bright the sky is.
Per baseline, because ``k0`` is that baseline's own maximum fringe
rate: normalising each one separately is what makes the configured
width the width on every baseline rather than on an average one.
"""
shape = pow_spec_nd(self.ks, _SHAPE_ONLY, [self.k0_freq, k0], self.gammas)
return (std / _LATENT_WIDTH) * jnp.sqrt(shape / jnp.sum(shape))
# std is a scalar when it was configured and one per baseline when it
# was measured, so it is broadcast to the baseline axis and mapped
# alongside the knee. The width is per baseline either way; the two
# cases differ only in whether the baselines were given the same one.
std_bl = jnp.broadcast_to(jnp.asarray(self.std), jnp.shape(self.k0_time))
self.sigma_ast_k = vmap(sigma, (0, 0), 0)(self.k0_time, std_bl)
# validate_gp_cov accepts any finite positive float, but sigma is
# built in the run's own precision: std = 1e-50 flushes it to zero in
# float32 and std = 1e40 overflows it, and either way the first thing
# that divides by sigma -- inv_transform, encoding the initial sky --
# produces non-finite parameters. Shape validation does not look at
# values, so without this the run starts and fails later somewhere that
# says nothing about std.
if not jnp.all(jnp.isfinite(self.sigma_ast_k)) or jnp.any(
self.sigma_ast_k <= 0
):
reported = (
self.std if jnp.ndim(self.std) == 0
else f"{float(jnp.min(self.std)):.4g}..{float(jnp.max(self.std)):.4g}"
)
raise ValueError(
f"ast.gp_cov.std ({reported}) is not representable in this "
f"run's precision: the mode standard deviations it gives come "
f"out {'non-finite' if not jnp.all(jnp.isfinite(self.sigma_ast_k)) else 'zero'}. "
f"std is a visibility amplitude in Jy, so it should be within "
f"a few orders of magnitude of the data; check the units it "
f"was set from, or run in double precision."
)
@measure_runtime
def _compute_true_params(self, zarr_path, data_col):
true_vis_ast = read_true_vis_ast(zarr_path, data_col)
self.true_ast_k = self.signal_to_latent(true_vis_ast)
self.true_ast_k_base = self.inv_transform(
self.true_ast_k, self.sigma_ast_k, self.mu_ast_k
)
def _compute_prior_params(self, prior_type: str, vis_obs):
if prior_type == "data":
print("Using data for AST prior mean")
self.mu_ast_k = self._compute_data_est(vis_obs)
elif prior_type in ["zeros", 0]:
print("Using zeros for AST prior mean")
self.mu_ast_k = jnp.zeros(
(self.n_bl, self.n_k_freq_ast, self.n_k_time_ast), dtype=complex
)
else:
raise ValueError(f"Provided prior type: {prior_type} is not valid. Choose from (data, zeros, 0).")
def _set_outputs(self):
self.state_outputs = {
"vis_ast": jnp.zeros((self.n_bl, self.n_freq, self.n_time), dtype=complex),
}
def _std_from_data(self, vis_obs, estimator_flags):
"""``rms|V|`` per baseline over the samples nothing has flagged.
``std`` is defined as that quantity, so this is not a proxy for the
prior width -- it is the prior width, read off the data instead of
guessed. Per baseline because that is what the data offers; the model
then has one width per baseline rather than one for all of them.
The measurement is :func:`tabascal.fft_gp.rms_vis`, shared with
``rfi.gp_cov.std: data``. What is here is the policy: which samples this
prior wants, and what to say when the mask cannot give them.
The mask is ``estimator_flags`` -- everything known to be bad, the
MS's own flags and the samples no gain table could calibrate -- and
not the likelihood's. Those answer different questions and this is
where the difference shows: a strong emitter flagged by some other
task is data tabascal is here to recover, so ``data.flags: false``
keeps it in the fit, and it is still the wrong place to measure a
clean sky amplitude from. An uncalibratable visibility is worse: it
carries a unity gain where its neighbours were divided by a real one,
so it is not even on the same flux scale.
**It measures whatever is in the unflagged data, including RFI.** That
is the point of taking only unflagged samples, and it is why an MS that
flags nothing gets a warning rather than a silent estimate: on the
shipped 8A simulation, whose RFI is unflagged because modelling it is
the job, this returns about 11 Jy where the true sky is under 3. Where
the RFI *is* flagged, the estimate is the sky. See GitHub #220 for
making it use the RFI model's own view of which samples are
contaminated, which is what would fix the unflagged case.
"""
keep = ~jnp.asarray(estimator_flags)
if not bool(jnp.any(keep)):
raise ValueError(
"ast.gp_cov.std: data has nothing to measure -- every "
"visibility is flagged. Set a width in Jy instead."
)
if not bool(jnp.any(~keep)):
print(
"Warning: ast.gp_cov.std: data is measuring every "
"visibility, because nothing flags any of them. Whatever "
"RFI is in them is in the prior width too, so it will be "
"wider than the sky by however much RFI there is. Flag the "
"contaminated samples in the MS -- this reads those flags "
"whatever data.flags says, so flagging them does not stop "
"tabascal fitting them -- or set ast.gp_cov.std to a width "
"in Jy."
)
# The measurement itself is rfi.gp_cov.std: data's as well; only the mask and
# the reduction differ. See fft_gp.rms_vis.
std = rms_vis(vis_obs, keep, "ast.gp_cov.std", per_baseline=True)
print(
f"Using ast.gp_cov.std from data: {float(jnp.min(std)):.4g} to "
f"{float(jnp.max(std)):.4g} Jy across baselines "
f"(median {float(jnp.median(std)):.4g})"
)
return std
def forward_transform(self, base_params, sigma, mu):
params = sigma * base_params + mu
return params
def inv_transform(self, params, sigma, mu):
base_params = (params - mu) / sigma
return base_params
def _compute_data_est(self, vis_obs):
est_ast_k = self.signal_to_latent(vis_obs)
return est_ast_k
def _compute_init_params(self, init_type: str, vis_obs):
if init_type == "data":
print("Using data for AST init")
self.init_ast_k = self._compute_data_est(vis_obs)
elif init_type == "prior":
print("Using prior mean for AST init")
self.init_ast_k = self.mu_ast_k
elif init_type == "truth":
print("Using truth for AST init")
self.init_ast_k = self.true_ast_k
elif init_type == "sample":
print("Using prior sample for AST init")
prior_sample = random.normal(
random.PRNGKey(1),
(self.n_bl, self.n_k_freq_ast, self.n_k_time_ast),
dtype=complex,
)
self.init_ast_k = self.forward_transform(
prior_sample, self.sigma_ast_k, self.mu_ast_k
)
elif init_type == "zeros":
print("Using zeros for AST init")
# The zero signal put through the same encoding `data` uses, rather
# than the zero latent written straight in, so `zeros` keeps meaning
# "no sky" independently of what the encoding does. One
# (n_freq, n_time) plane rather than the full (n_bl, n_freq, n_time)
# array: the transform is per-baseline, so every baseline encodes to
# the same latent, and encoding all of them would allocate a padded
# FFT of the whole visibility array to get it.
zeros_k = signal_to_latent(
jnp.zeros((self.n_freq, self.n_time), dtype=complex),
self.pad_factors,
self.latent_idxs,
)
self.init_ast_k = jnp.broadcast_to(zeros_k, (self.n_bl, *zeros_k.shape))
else:
raise ValueError(f"Provided init type: {init_type} is not valid. Choose from (data, prior, truth, sample, zeros).")
self.init_ast_k_base = self.inv_transform(
self.init_ast_k, self.sigma_ast_k, self.mu_ast_k
)
self.init_params = {
"ast_k_r": self.init_ast_k.real,
"ast_k_i": self.init_ast_k.imag,
}
self.init_params_base = {
"ast_k_r_base": self.init_ast_k_base.real,
"ast_k_i_base": self.init_ast_k_base.imag,
}
def _validate_dimensions(self):
"""Ensure all setup operations completed successfully"""
ast_shape = (self.n_bl, self.n_k_freq_ast, self.n_k_time_ast)
assert_attr_shape(self, "mu_ast_k", ast_shape)
assert_attr_shape(self, "sigma_ast_k", ast_shape)
assert_attr_shape(self, "init_ast_k", ast_shape)
assert_attr_shape(self, "init_ast_k_base", ast_shape)
[docs]
def radec_to_lmn(ra, dec, ra0, dec0):
"""Direction cosines of sources at ``(ra, dec)`` about a phase centre, in radians.
Returns ``(l, m, n, n - 1)``. ``n - 1`` is returned alongside ``n`` because it is the
quantity the w term actually needs and the two cannot both be computed accurately
from one expression.
``n`` is the exact spherical form ``sin(d) sin(d0) + cos(d) cos(d0) cos(da)`` rather
than ``sqrt(1 - l^2 - m^2)``. The square root is the cosine of the angular distance
only on the near hemisphere: it is unsigned, so it folds a source more than 90
degrees from the phase centre back onto the near side instead of giving it the
negative ``n`` it has.
``n - 1`` uses the haversine identity ``n - 1 = -2 h`` with
``h = sin^2((d - d0)/2) + cos(d) cos(d0) sin^2(da/2)``, the haversine of the angular
distance. ``h`` runs over ``[0, 1]`` across the whole sphere — 0 at the phase centre,
1 at the antipode — so ``n - 1`` runs over ``[-2, 0]`` and ``n < 0`` exactly when
``h > 1/2``. Subtracting a nearby ``n``
from 1 cancels catastrophically: at a 40 arcsec offset ``1 - n ~ 2e-8``, which in
single precision is below the spacing of the floats either expression lands on, so
the difference comes out as exactly zero and the w term disappears. The haversine
form never forms the difference, so it keeps full relative accuracy at any offset.
"""
dra = ra - ra0
l = jnp.cos(dec) * jnp.sin(dra)
m = jnp.sin(dec) * jnp.cos(dec0) - jnp.cos(dec) * jnp.sin(dec0) * jnp.cos(dra)
n = jnp.sin(dec) * jnp.sin(dec0) + jnp.cos(dec) * jnp.cos(dec0) * jnp.cos(dra)
hav = (
jnp.sin((dec - dec0) / 2) ** 2
+ jnp.cos(dec) * jnp.cos(dec0) * jnp.sin(dra / 2) ** 2
)
return l, m, n, -2.0 * hav
[docs]
class GPVisAstDFT(GPVisAst):
""":class:`GPVisAst` with the padded grid never formed.
The prior, the initialisation and the blocking over baselines are the
parent's; only the step from the surviving Fourier modes to the
visibilities differs. :func:`~tabascal.fft_gp.latent_to_signal` pads those
modes back up to the padded grid, inverse-transforms the whole of it and
crops -- at the default padding, about four times the elements of the
visibilities formed and thrown away, per block, in the forward pass and
again in reverse. The chain is linear and separable, so the same values
come out of one small matrix per axis
(:func:`~tabascal.fft_gp.latent_to_signal_dft`), and the largest thing
formed is the size of the visibilities.
The modes of this component are set by the array's geometry rather than by
a choice, so the matrices are fixed at setup and shared by every baseline.
Select it where the parent would go::
model:
components:
- ast_vis:GPVisAstDFT
"""
def _compute_gp_params(self):
super()._compute_gp_params()
# The same modes and the same spectrum as the parent -- checked below
# against its own pads and crop -- with the transform as matrices.
pk, ks, self.dft_mats, self.dft_first_axis = latent_to_signal_dft_init(
[self.n_freq, self.n_time],
[self.chan_width, self.int_time],
self.pad_factors,
self.ss_factors,
_SHAPE_ONLY,
self.k0s,
self.gammas,
self.pk_cutoff,
)
if pk.shape != self.pk.shape:
raise RuntimeError(
"The DFT transform kept a different set of Fourier modes than "
f"the FFT one: {pk.shape} against {self.pk.shape}. They are "
"built from the same power spectrum and cutoff, so this is a "
"bug rather than a configuration error."
)
n_out = [m.shape[1] for m in self.dft_mats]
if n_out != [self.n_freq, self.n_time]:
raise RuntimeError(
f"The DFT transform outputs {n_out} samples where the "
f"visibilities are ({self.n_freq}, {self.n_time})."
)
def _resolve_baseline_block_size(self):
"""``auto`` is a single step over every baseline.
The block exists to bound the padded Fourier grid, and this component
builds none: its only transient is the intermediate of the two matrix
products, which is the modes in one axis and the visibilities in the
other. What the scan does still cost is the stack of its steps'
outputs, so on this transform blocking is worse on both counts.
Measured on one GH200, 2016 baselines at 150 integrations, value and
gradient together:
========= ================= ================== =================
channels FFT, best block FFT, single step DFT, single step
========= ================= ================== =================
8 273.9 MB, 1.0 ms 441.2 MB, 1.0 ms 218.6 MB, 0.2 ms
32 828.6 MB, 4.0 ms 1711.2 MB, 3.3 ms 683.8 MB, 0.4 ms
128 2386.4 MB, 14.7 ms 5926.7 MB, 12.7 ms 1799.1 MB, 1.1 ms
========= ================= ================== =================
An explicit ``baseline_block_size`` is still honoured: it is the way
to trade this back if some other part of a model wants the room.
"""
setting = self.baseline_block_size_setting
if setting != "auto":
self.baseline_block_size = None if setting is None else int(setting)
return
self.baseline_block_size = self.n_bl
[docs]
def build_block_signal(self):
mats = self.dft_mats
first_axis = self.dft_first_axis
transform = vmap(latent_to_signal_dft, (0, None, None), 0)
return lambda Y_block: transform(Y_block, mats, first_axis)
[docs]
class DiscreteSkyVis(Component):
"""Visibilities of a discrete sky by direct DFT, with the full w term.
Reads ``ast_radec`` (n_src, 2) radians, ``ast_I`` (n_src, n_freq) Jy and
``ast_shape`` (n_src, 3) radians from the state — see
:class:`~tabascal.components.ast_signal.FixedDiscreteSky`, which must be listed
before this component — and ACCUMULATES into ``vis_ast`` using the visibility
equation
V(u,v,w) = sum_k I_k G_k(u,v) exp(-2i pi (u l_k + v m_k + w (n_k - 1)) / lambda)
For a discrete sky the direct sum is exact, gridless, differentiable, and unaffected
by field of view or baseline length; "discrete" is the set of sources, not their
size, so ``ImageSkyVis`` is reserved for a sky carried as an image.
``I_k`` enters undivided: there is no ``1 / n``. The RIME integrand carries ``B / n``
because ``dOmega = dl dm / n``, but a source of integrated flux ``S`` is
``B = S delta_Omega`` and ``delta_Omega = n delta(l) delta(m)``, so the Jacobian
cancels and a source contributes its catalogue flux exactly, in every direction.
``(u, v, w)`` is the ANTENNA2 - ANTENNA1 baseline the equation above is written for,
which is ``ast.uvw_sign`` times the UVW column. Which baseline that column holds is a
property of whatever wrote the data, so the sign is a config key; the default,
``(-1, -1, -1)``, is the convention tab-sim writes — see ``setup``.
``G_k`` is the uv-plane envelope of an elliptical Gaussian source,
G(u,v) = exp(-pi^2 / (4 ln 2) * (a^2 u'^2 + b^2 v'^2))
for FWHM ``a`` (major) and ``b`` (minor) in radians, with ``(u', v')`` the baseline
in wavelengths rotated into the source frame::
u' = u sin(phi) + v cos(phi) along the major axis
v' = u cos(phi) - v sin(phi) along the minor axis
``phi`` is the position angle in the radio convention, measured from north (the m
axis) through east (the l axis), so ``phi = 0`` puts the major axis north-south and a
north-south baseline is the one that resolves the source out. A zero FWHM gives
``G = 1`` exactly, so points and Gaussians are the same code path.
It accumulates rather than assigns, so it composes with the astronomical GP: with
both listed, ``vis_ast`` is the GP plus the fixed sources. ``Model`` zeroes
``vis_ast`` before the forward chain runs, so the two may be listed in either order.
The source axis is walked in blocks of ``ast.source_block_size`` with
:func:`jax.lax.scan`, and the block body is rematerialised: the delay array is
(n_bl, n_time, n_src), which for a real catalogue is the largest array in the model,
and blocking replaces ``n_src`` in that shape with the block size at the cost of
recomputing each block in the backward pass.
Sources more than 90 degrees from the phase centre are modelled, not rejected: only
``n - 1`` enters the phase and :func:`radec_to_lmn` computes it exactly over the
whole sphere, so nothing here breaks down at ``n <= 0``. Such a source is almost
always a catalogue mistake rather than a real one, so
:class:`~tabascal.components.ast_signal.FixedDiscreteSky` warns about it at setup
and leaves the decision to the caller.
A fixed sky exists to make a per-antenna gain identifiable (see issue #124); the flux
scale it fixes the gain against is only physical if the data are calibrated to Jy.
"""
required_inputs = {
"ast_radec": ("n_src", 2),
"ast_I": ("n_src", "n_freq"),
"ast_shape": ("n_src", 3),
"vis_ast": ("n_bl", "n_freq", "n_time"),
}
output_shapes = {
"vis_ast": ("n_bl", "n_freq", "n_time"),
}
parameter_shapes = {}
[docs]
def setup(self, config):
try:
self.n_bl = config.n_bl
self.n_freq = config.n_freq
self.n_time = config.n_time
# config.uvw is (n_time, n_bl, 3) as read_ms gives it; the DFT below is
# written baseline-first, to match the (n_bl, n_freq, n_time) visibilities.
self.uvw = jnp.swapaxes(jnp.asarray(config.uvw), 0, 1) # (n_bl, n_time, 3)
self.freqs = jnp.asarray(config.freqs)
self.phase_centre_ra = jnp.deg2rad(config.phase_centre["ra"])
self.phase_centre_dec = jnp.deg2rad(config.phase_centre["dec"])
# Per-axis uvw sign toggles (u, v, w), applied before the exponent below,
# from ast.uvw_sign.
#
# The measurement equation's exp(-2i pi b.(s - s0) / lambda) is written for
# b = ANTENNA2 - ANTENNA1, and which of the two baselines a UVW column
# holds is a property of whatever wrote it, not something we can read off
# the data. The default negates because tab-sim writes
# bl_uvw = ants_uvw[a1] - ants_uvw[a2] and tabascal forms its own baselines
# the same way throughout (interferometry.py:
# bl_u = ants_u[:, a1] - ants_u[:, a2]), so negating is what makes this
# agree with tab-sim's astro_vis on the very visibilities the model is fit
# to -- but an MS from another toolchain may well carry the opposite
# convention, and a wrong sign is a sky mirrored through the phase centre,
# exactly the corruption a gain solved against a fixed sky would absorb.
sign = config.args["ast"].get("uvw_sign", (-1, -1, -1))
entries = sign if isinstance(sign, (list, tuple)) else ()
if len(entries) != 3 or not all(
not isinstance(s, bool) and isinstance(s, (int, float)) and abs(s) == 1
for s in entries
):
raise ValueError(
"ast.uvw_sign is the sign applied to each of the u, v and w axes "
"of the UVW column: a sequence of three values, each exactly "
f"+1 or -1, e.g. [-1, -1, -1], got {sign!r}."
)
self.uvw_sign = jnp.asarray([float(s) for s in entries])
# int() alone would turn 1.9 into 1 without a word, which is a hundredfold
# slowdown dressed up as a valid setting.
block_size = config.args["ast"].get("source_block_size", 128)
if (
isinstance(block_size, bool)
or not isinstance(block_size, (int, float))
or block_size != int(block_size)
or block_size < 1
):
raise ValueError(
"ast.source_block_size is the number of sources handled per scan "
f"step: a whole number of at least 1, got {block_size!r}."
)
self.source_block_size = int(block_size)
self._set_outputs()
except Exception as e:
raise RuntimeError(f"{self.__class__.__name__} setup failed: {e}")
[docs]
def build_constants(self):
return {
"uvw": self.uvw,
"freqs": self.freqs,
"ra0": self.phase_centre_ra,
"dec0": self.phase_centre_dec,
"uvw_sign": self.uvw_sign,
}
[docs]
def build_forward(self):
prefix = self.prefix
block_size = self.source_block_size
vis_shape = (self.n_freq, self.n_bl, self.n_time)
C = 299792458.0
# exp(-4 ln 2 x^2 / a^2) on the sky transforms to exp(-pi^2 a^2 u^2 / (4 ln 2)).
gauss_uv = jnp.pi**2 / (4 * jnp.log(2.0))
def forward(params, state, constants):
uvw = constants[f"{prefix}/uvw"] * constants[f"{prefix}/uvw_sign"]
freqs = constants[f"{prefix}/freqs"] # (n_freq,)
ra0 = constants[f"{prefix}/ra0"]
dec0 = constants[f"{prefix}/dec0"]
ra = state["ast_radec"][:, 0] # (n_src,)
dec = state["ast_radec"][:, 1]
shape = state["ast_shape"] # (n_src, 3)
l, m, _, n_minus_1 = radec_to_lmn(ra, dec, ra0, dec0) # (n_src,)
lmn = jnp.stack([l, m, n_minus_1], axis=-1) # (n_src, 3)
# No 1 / n. The RIME integrand carries B / n because dOmega = dl dm / n, but
# a source of integrated flux S is B = S delta_Omega with
# delta_Omega = n delta(l) delta(m), so the Jacobian cancels and the source
# contributes S exactly, in any direction. Catalogue fluxes -- OSKAR's
# included -- are integrated fluxes, so ast_I goes in as it stands.
weights = state["ast_I"] # (n_src, n_freq)
u = uvw[..., 0, None] # (n_bl, n_time, 1)
v = uvw[..., 1, None]
def block_vis(lmn_b, weights_b, shape_b):
# Geometric path-length delay per (baseline, time, source), in metres.
tau = jnp.einsum("btx,sx->bts", uvw, lmn_b)
fwhm_maj, fwhm_min, pa = shape_b[:, 0], shape_b[:, 1], shape_b[:, 2]
u_rot = u * jnp.sin(pa) + v * jnp.cos(pa) # along the major axis
v_rot = u * jnp.cos(pa) - v * jnp.sin(pa) # along the minor axis
# -log G, in metres^2; scaled to wavelengths^2 by (freq / c)^2 per
# channel below. Zero for a point source, so exp() leaves it alone.
log_envelope = gauss_uv * (
(fwhm_maj * u_rot) ** 2 + (fwhm_min * v_rot) ** 2
)
# vmap over frequency to avoid materialising a 4D (bl, time, src, freq)
# array.
def vis_at_freq(freq, weights_f):
k = freq / C
exponent = -log_envelope * k**2 - 2.0j * jnp.pi * tau * k
return jnp.sum(jnp.exp(exponent) * weights_f, axis=-1) # (n_bl, n_time)
return vmap(vis_at_freq)(freqs, weights_b.T) # (n_freq, n_bl, n_time)
n_src = lmn.shape[0]
n_block = min(block_size, n_src)
n_pad = -n_src % n_block
# Padding sources are (l, m, n - 1) = 0 -- the phase centre -- at zero flux,
# so they contribute exactly zero rather than merely something small.
pad = lambda x: jnp.pad(x, ((0, n_pad), (0, 0)))
blocks = tuple(
pad(x).reshape(-1, n_block, x.shape[1]) for x in (lmn, weights, shape)
)
def accumulate(vis, block):
return vis + block_vis(*block), None
vis_dtype = jnp.result_type(uvw, freqs, weights, jnp.complex64)
vis, _ = lax.scan(
checkpoint(accumulate), jnp.zeros(vis_shape, vis_dtype), blocks
)
return {**state, "vis_ast": state["vis_ast"] + vis.transpose(1, 0, 2)}
return forward
def _set_outputs(self):
self.state_outputs = {
"vis_ast": jnp.zeros((self.n_bl, self.n_freq, self.n_time), dtype=complex),
}