visibility-tokenizer

Tokenization / embedding experiments for radio interferometric visibilities, working toward a foundation-model-style pretext task: what should a "token" be, and what does an unsupervised latent actually learn from raw visibilities?

This is a restructured, cleaned-up port of exploratory scripts (see "Current state" below for the full history). The physics findings are unchanged; the code is now a package with a real data feeder instead of a pile of top-level scripts that each reload and reshape the whole dataset.

Install

Environments are managed with pixi; pixi.lock is committed, so these are reproducible.

pixi run -e dev test          # core + pytest, runs the suite
pixi run -e dev vistok --help
pixi shell -e ms              # adds casatools, for `vistok extract`

Or with plain pip:

# core (torch, sklearn, etc.)
pip install -e .

# + casatools, to extract tiles from a Measurement Set
pip install -e ".[ms]"

# + pytest, for the test suite
pip install -e ".[dev]"

Python 3.12 (pixi pins it; the pip path works on 3.10+). Trains fine on Apple MPS.

Quick start

# 1. extract tiles from a Measurement Set (requires the `ms` extra)
vistok extract /path/to/data.ms tiles.npz

# 2. compute per-pixel statistical RFI flags (used as auxiliary model input)
vistok flag tiles.npz pixel_flags.npz

# 3. (optional) k-means pseudo-labels for the auxiliary class loss
vistok cluster tiles.npz labels.npz

# 4. train the masked-tile transformer
vistok train tiles.npz --pixel-flags pixel_flags.npz --labels labels.npz --epochs 12

Or drive it from Python for more control — see src/visibility_tokenizer/:

• config.py — TileGeometry, ModelConfig, TrainConfig dataclasses
• extract.py — MS -> tile npz
• features.py — window-grid physics features (torch, differentiable through the shift augmentation)
• flagging.py — statistical RFI flagging (tile-frame | observation-frame)
• dataset.py — VisTileDataset + GroupBatchSampler + collate_groups: the data feeder. Tiles are grouped by (spw, scan, chunk) because the model attends within a group, not across the whole dataset; batches are assembled and (optionally) shift-augmented on the fly rather than precomputed, so augmentation stays cheap without reloading the npz.
• model.py — MaskedTileTransformer
• train.py / evaluate.py — training loop and post-hoc analysis (embeddings, masked-prediction baselines, sigma calibration, persistent-stripe mining)

Run tests with pixi run -e dev test (small synthetic tiles, no MS required).

Current state (as of 2026-07-16)

What a "tile" is

One baseline x one spectral window x 8 integrations x 64 channels x 2 polarizations (RR, LL), normalized by one scalar per spw (median unflagged amplitude over the whole observation) so relative brightness across baselines/time survives.

What we've tried, in order

1. Uncalibrated calibrator spectra are trivially compressible — 95% variance in 2 of 958 PCA components; the latent is just visibility phase/delay. AE ties PCA (the structure is linear).
2. Per-tile reconstruction error works as an unsupervised coherence/SNR meter — separates calibrator from noise-dominated target tiles by 100-1000x with no labels.
3. Strong RFI detection works unsupervised (AE vs real FLAG column, AUC 0.98), but weak RFI (0.5-2x noise, persistent over one tile) is undetectable per-tile by anything, AE or MAD baseline — real spectral structure dominates the null distribution. Per-token detection has a hard floor around ~2x noise at 8x64-int tile size.
4. AE ≈ PCA everywhere tried — at this token size the structure is linear; nonlinear capacity buys nothing per-token. This pushed the design toward spending capacity on attention across tokens instead (cross-baseline stacking gain ~sqrt(N_baseline) for coherent weak RFI/sky signal) rather than a fancier per-token encoder.
5. Geometry conditioning + asinh amplitude compression removed most baseline-length leakage from the latent (R2 0.06). Heavy-tailed amplitudes need asinh/log compression before anything else works — raw MSE training lets a handful of monster tiles dominate the gradient.
6. Physics-informed derived features carry type information the learned latent misses — six then twelve then (window_features.py, the current version) 189 per-tile moments/delay-rate stats on a sliding window grid, fed as auxiliary content to a masked-tile attention model (MaskedTileTransformer). The window grid made physical axes far more linearly readable (R2 log-amp 0.94 vs 0.60) without moving the masked-reconstruction pretext loss — auxiliary features are cheap representation fuel even when they don't move the pretext loss itself.
7. Statistical RFI flagging (flagging.py) combines a short-memory per-tile frame and a long-memory whole-observation frame as a union — each frame normalizes away exactly what the other catches (always-on stripes vs minutes-long broadband bursts). Feeding these flags as token content + recon-loss weights (MASK_IN) roughly doubled the attention margin over zeros/group-mean baselines.
8. Heteroscedastic reconstruction (predict per-pixel mu + sigma, NLL): the learned sigma map reproduces the statistical flagger without ever seeing it (AUC 0.77, actually better than with the flags fed in as input, AUC 0.73) — gradient self-concentrates on predictable pixels. This suggests the explicit flagging step may be droppable in favor of just training with NLL loss.
9. Persistent-stripe mining from the reconstructed mean (evaluate.mine_persistent_stripes): sigma is a surprise detector, so it structurally misses always-on RFI (perfectly predictable = low sigma). Mining the model's own confidently-reconstructed mean for channels that are constant for hours across all baselines catches exactly this blind spot — narrowband structure invariant across baselines cannot be sky.
10. Class probing (linear + MLP heads on k-means pseudo-labels): MLP ≈ linear performance everywhere tried — class information the embedding does carry is linearly readable; more depth doesn't help. Polarized RFI classes are the consistent failure mode (pol ratio needs to be an explicit per-window feature, not just a global one).

What we're about to try

• Correlated weak-RFI injection (same channel, all baselines in a scan) + cross-baseline stacking of per-tile channel residuals, to demonstrate the sqrt(N_baseline) gain that's the whole motivation for attention across baseline-tokens rather than better per-token models.
• Rerun on calibrated / fringe-stopped data to see what the latent learns once the trivial phase/delay structure (finding 1) is gone — everything so far has been on uncalibrated or single-target data.
• A larger MeerKAT dataset is in preparation; extraction should port as-is via extract.py, but the per-spw scale/channel constants in evaluate.mine_persistent_stripes and the CLI defaults are currently tuned to the VLA G55 tutorial dataset and will need updating.
• rfi_sweep.py-style amplitude-vs-detectability sweep on clean data was planned but never run in the original exploratory repo.

Known deviations from the original exploratory scripts

• The k-means pseudo-labels used for the auxiliary classification loss originally depended on a separately-trained AE latent (derived_features2.py mixed window features with latent_dualpol.npz embeddings before clustering). cluster.py / vistok cluster ports this using window features alone (no external latent), which is reproducible from just the tile npz but won't match the original cluster assignments exactly.
• Plotting (UMAP latent panels, sigma/flag example grids, stripe-mining figures) was deliberately left out of src/ — evaluate.py returns plain arrays/dataclasses; make plots in a notebook or scratch script against those return values, so the library has no hard matplotlib/umap dependency in its core path.