Basic averaging

[AlecThomson]

Hi all!

Apologies for the (second) PR without a issue / discussion first. I had the code ready, so hopefully we can discuss here. I was inspired by both #98 and plot-ms's averaging to add some simple averaging. For now I've just added time and channel-wise averaging, using codex-africanus - hopefully an acceptable dependency addition. I think further averaging types will require custom averaging logic, or support upstream.

I elected to use an interface like this:

$ shadems --xaxis CHAN --yaxis DATA:amp --average CHAN:4 <msname>          # 4 channels per bin
$ shadems --xaxis TIME --yaxis DATA:amp --average TIME:60 <msname>         # 60-second time bins
$ shadems --xaxis FREQ --yaxis DATA:amp --average TIME:all <msname>        # collapse all time
$ shadems --xaxis TIME --yaxis DATA:amp --average TIME:60 --average CHAN:4 <msname>

Hopefully that's also in line with the project style as well. Overall it seems to work very nicely on the data I've played with.

[o-smirnov]

An excellent feature, thanks!

The codex dependency is perfectly acceptable, of course. It is an in-house library after all, and designed for exactly this sort of thing. @sjperkins will be happy.

I have only one improvement to suggest -- can we not check the bin size for a unit, and treat it as a quantity string if given? So e.g. --average TIME:60s --average CHAN:256MHz would be legit usage. And then if the bin is just an integer, treat it as "timeslots", "channels", etc.

[sjperkins]

Very cool @AlecThomson. I'll find some time to review this. codex is absolutely fine as a dependency.

[sjperkins]

@AlecThomson I haven't had time to review yet. Would it be possible to fix the test case failures? AFAICT they're failing due to an inconsistency in FLAG_ROW and FLAG.

[AlecThomson]

Thanks both! I think I've got the request CLI interface in place. I also fixed the CI issues, and also fixed up the behaviour of flags in averaging. I've added tests around these changes too

[sjperkins]

Note on authorship: this review was performed by Claude (Claude Code), and is posted through Simon Perkins' account only because that is the account with access — the analysis, experiments and conclusions below are Claude's, not Simon's. Please treat it as an automated review whose findings are worth verifying, not as a maintainer sign-off.

The overall shape of the integration looks good, and codex-africanus seems a reasonable dependency to take on here. There is one issue that probably needs fixing before merge, because it silently produces wrong output rather than failing.

The main issue: time_and_channel maps rows per dask row-chunk

africanus.averaging.dask.time_and_channel builds its row mapping via tc_row_mapper, which is a da.blockwise over ("row",) with adjust_chunks={"row": nan}:

def tc_row_mapper(time, interval, antenna1, antenna2, flag_row=None, time_bin_secs=1.0):
    """Create a dask row mapping structure for each row chunk"""
    return da.blockwise(
        np_tc_row_mapper, ("row",), ...
        adjust_chunks={"row": lambda x: np.nan}, ...)

So the numba row_mapper runs once per row chunk, and a time bin can never span a chunk boundary. get_plot_data reads its groups with chunks=dict(row=row_chunk_size) (default 100000), so this bites in practice.

Measured against codex-africanus 0.4.4, on 3 baselines × 20 timeslots of 10s:

row_chunk=60  TIME:all ->  3 out rows (correct),  interval=200s
row_chunk=30  TIME:all ->  6 out rows,            interval=100s
row_chunk=15  TIME:all -> 12 out rows,            interval=50s

row_chunk=60  TIME:5 -> 12 rows, slots-per-bin {5:12}     (correct)
row_chunk=20  TIME:5 -> 18 rows, slots-per-bin {5:9, 2:6, 1:3}
row_chunk=40  TIME:5 -> 15 rows, slots-per-bin {5:9, 3:2, 4:1, 2:2, 1:1}

Two consequences:

• --average TIME:all does not collapse the time axis. It collapses per chunk, giving nchunks points per baseline instead of one. The README documents otherwise.
• --average TIME:N gets a runt bin at every chunk boundary. For a MeerKAT MS (2016 baselines) that is 100000/2016 ≈ 49.6 timeslots per chunk, so an under-averaged partial bin roughly every 50 timeslots — which shows up as extra scatter in the plot.

average_group computes full_thresh and the overflow warning over the whole group, which reads as though the binning were global. It isn't.

Suggested fix: chunk rows on time-bin boundaries. Compute the group's TIME column (cheap, one column), find the row indices at which the bin index changes, and rechunk everything passed to time_and_channel to those boundaries. Rechunking to a single row chunk also works, but costs the memory of the unaveraged group.

The tests don't catch this — the fixture builds its arrays with chunks=a.shape, i.e. a single row chunk. Worth parametrising test_average_group_bin_sizes over row chunk size.

Other correctness issues

Non-DATA axis columns disappear from the rebuilt dataset. vis_columns only matches *DATA/*SPECTRUM, and avg_group is rebuilt with only TIME, INTERVAL, UVW, FLAG, FLAG_ROW, ANTENNA1/2 plus vis_columns. But shadems accepts any column in ms.valid_columns as an axis, so e.g.

shadems -y TIME_CENTROID --average TIME:60 <ms>     # or EXPOSURE, or a custom column

dies in DataAxis.get_column_data with NameError: column TIME_CENTROID not found in group. Both TIME_CENTROID and EXPOSURE are native time_and_channel parameters (time_centroid=, exposure=), so passing them through is nearly free. Failing that, a parser.error up front for unsupported axis columns would at least be honest.

WEIGHT_SPECTRUM is consumed twice. If it is an axis, c.endswith("SPECTRUM") puts it in vis_columns and the hasattr(group, "WEIGHT_SPECTRUM") branch passes it as weight_spectrum — so you get the weight-weighted mean of the weights, while the properly averaged res.weight_spectrum is discarded. SIGMA_SPECTRUM is likewise averaged as if it were a visibility rather than in quadrature. Suggest excluding both from vis_columns and taking them from res.

A count of timeslots is not honoured when INTERVAL is shorter than the time step. time_bin_secs = time_value * dt with dt = interval.mean(), but africanus bins on the time span:

elif time[r] + half_int - bin_low > time_bin_secs:

Whenever the dump spacing exceeds INTERVAL — correlator overheads, dropped dumps, quite common — you get fewer timeslots per bin than requested. On 10s-spaced 8s dumps:

asked for TIME:4 timeslots (-> 32.0s) -> slots per bin [3, 3, 3, 3]
asked for TIME:6 timeslots (-> 48.0s) -> slots per bin [5, 5, 2]

Using the median spacing of the unique times instead of mean INTERVAL is exact, since span(N) = (N-1)·Δt + interval ≤ N·Δt always holds (as interval ≤ Δt), while span(N+1) > N·Δt.

full_thresh for all is a floating-point knife-edge. (tmax - tmin) + dt is computed differently from africanus's internal span (tmax + interval/2) - (tmin - interval/2), so the rounding need not agree — and dt being a mean makes it outright wrong for varying INTERVAL. One ULP flips the answer:

TIME:all threshold 200.0          -> 3 out rows (expected 3)
TIME:all threshold 199.999999999  -> 6 out rows

time_bin_secs = np.inf is robust (verified: 3 rows, interval 200s) and also removes two of the three .compute() calls.

Whole-group materialisation. computed = dask.compute(out)[0] pulls the entire averaged group into memory. Groups are (FIELD_ID, DATA_DESC_ID, SCAN_NUMBER), and with --average CHAN:2 there is no row reduction at all, so a large single-scan MS means roughly half of it resident — which undoes shadems' streaming design. The comment calls the group "small", which it need not be.

Worth noting that dataframe_factory in dask_utils.py already has a have_nan_chunks path, so staying lazy may be closer than it looks; the blockers are the explicit row/chan coords and value.size in the point count.

Minor

• Three separate .compute() calls per group (time.min(), time.max(), interval.mean()), each re-reading columns from disk. Collapse into one dask.compute(...) — moot for all if np.inf is used.
• --average TIME:60S (capital S) fails, because casacore parses S as siemens. Worth normalising the unit or mentioning it in the error message.
• avg_warned is keyed on axis name only, so once one group has warned about bin-size overflow, later groups with fewer channels or a shorter time range silently fall back to all with no warning.
• Unrelated behaviour changes are bundled in: plotting a flag column with --noflags now hard-errors where it previously only warned, and the # -*- coding: future_fstrings -*- lines are dropped. Both seem reasonable, but they deserve a mention in the PR description.
• -x ROW under averaging now plots a bin index rather than an MS row number. (-x CHAN is already relative to the channel selection, so that one stays consistent.)

What's right

• _bin_mean_freqs exactly reproduces africanus's channel_mapper binning, partial last bin included. Verified numerically — 8 channels at chan_bin_size=3 gives [1, 4, 6.5] from both africanus and _bin_mean_freqs.
• Folding FLAG_ROW into FLAG and letting africanus re-derive a consistent flag_row is the right way around RowMapperError, and the comment explains why.
• Grouping by FIELD_ID / DATA_DESC_ID / SCAN_NUMBER is exactly the partitioning africanus's duplicate-(TIME, ANTENNA1, ANTENNA2) check asks for.
• The zero-width range padding (_pad_range) and the max(..., 2) discrete-canvas fix are genuinely needed, and nicely commented.
• The use_flags / noflags split is subtle but correct — use_flags is computed before the axis loop mutates options.noflags, which is what makes the flag-column-plus-averaging case work.

[sjperkins]

@AlecThomson Again, thanks for taking this on. I ended up getting Claude to review it (see above) as my knowledge of ShadeMS's internals is rather stale.

I'm trying to keep the human element in this comment as walls of LLM-generated text can sometimes be a bit overbearing in my experience. I'm interested in whether you find it useful (feel free to push back!)