[jaxengodfrey]

Fixes #335

windowed_analysis accepted window_type and regions through **kwargs, but only the per-window fallback read them. When the scatter or fused engines served the request they returned the default bp grid, with no error or warning — so window_type='regions', window_type='snp', and the regions argument were silently ignored for most statistics. This PR routes each window type to an engine that honors it and raises an error when no engine is able to.

Behavior changes

• window_type is now an explicit argument in windowed_analysis and defaults to bp.
• Error is raised for window_type='regions' or 'snp' on a StreamingHaplotypeMatrix.
• The scatter engine is guarded with window_type=='bp', since it only tiles a fixed grid.
• window_type='regions' is sent to the fused engines, possibly sent to the fallback if necessary.
  • The scatter-only statistics:'theta_l', 'zeng_e', 'zeng_dh', 'normalized_fay_wu_h' for now these raise NotImplementedError with regions because they are not currently implemented in the fallback.
  • The stats: 'mean_nsl', 'daf_hist', 'mu_sfs' require window families to have non-decreasing starts and ends, which would break with e.g. unsorted regions. The fallback doesn't currently implement them, so for now these raise NotImplementedError with window_type='regions'.
• window_type='snp' is sent to the fallback (or local_pca if present).
• _iter_region_windows in the fallback now records chrom_start/chrom_end.
• Fallback no longer defaults to window_type='bp' in WindowedAnalyzer.

Left for follow-up

• Deferred to Fused snp windows and remaining statistic coverage for custom window types #352:
  • Fused GPU 'snp' path: Supporting snp here means feeding variant-index ranges directly and reconstructing the bp outputs and the span from each window's variant positions consistently across every statistic family the engine emits. This felt a little out of scope and I wasn't confident in my ability to implement this.
  • Unsupported statistics with regions: See above and Lines 1465:1473 in windowed_analysis.

Tests

• tests/test_windowed_chunked_coverage.py: the fused engines (single and chunked) agree over arbitrary region edges: overlapping, nested, out-of-order, and one past the last variant for pi, tajimas_d, segregating_sites, fst_hudson and garud_h1.
• tests/test_windowed_analysis.py: a regions request returns on window per region and the fused engine agrees with the per-window fallback for single and two populations.
  • The existing test, TestWindowedGarudCorrectness was passing by coincidence. It specified window_type='snp' but windowed_analysis was silently ignoring this on the fused path, tiling a bp-grid instead. I hardened this test against this type of coincidence with a non-unit-stride positions and a window-count assertion that would only work for 'snp' windows. It no longer uses the fused path, but once 'snp' windows are supported on that path it will.
• tests/test_windowed_error_coverage.py: pins the deferred NotImplementedError statistics left for follow-up PRs. When one is later implemented, the test will fail, at which point the list can be pruned or the test removed entirely.

Related

• windowed_analysis raises "Unknown statistic" when a request mixes statistics that no single engine computes #347
• windowed_analysis ignores unknown keyword arguments on the fused and scatter engines #349 - window_type is removed from **kwargs in this PR, instead named explicitly in function signature. window_type and regions also removed from _LOCAL_PCA_ONLY_KWARGS.
• windowed_statistics is largely superceded by windowed_analysis, is documented inconsistently, and is another route for bugs #305

[jaxengodfrey]

@andrewkern or @nspope this is ready for review!

[andrewkern]

Thanks, this is a good fix. I ran the new tests and the windowed test files on a GPU, and they pass. I also compared the fused region output with the scalar functions on each region's subset, for unsorted, overlapping and nested regions: pi, theta_w, tajimas_d, zns, omega, mu_var and garud_h1 all agree. Some items:

1. regions= without window_type='regions' is still ignored. #335 reports this case too. The call below returns the default bp grid, with no error:

windowed_analysis(hm, window_size=20_000, regions=regions, statistics=['pi'])
# 5 rows, starts [1, 20001, 40001, 60001, 80001]

I suggest a ValueError when regions is given and window_type != 'regions'.

2. The new positional parameter breaks positional calls. window_type is now the second positional parameter, before step_size. So windowed_analysis(hm, 20_000, 10_000) used to set step_size=10_000, but now it raises ValueError: Unknown window_type: 10000. No caller in the repo uses this form, but users can. Please move window_type after the existing parameters, or make it keyword-only.

3. An empty region gives a row on one engine and no row on the other. The fused engine returns a row for a region with no variants. The fallback skips it (if len(variant_indices) > 0 in _iter_region_windows). So the number of rows for the same regions table depends on missing_data and on the statistics:

regions: [60000, 90000), [10000, 30000), [10000, 25000), [99990, 99999)   # last one empty
missing_data='include' -> 4 rows (fused)
missing_data='exclude' -> 3 rows (fallback)

Users will join the output back to their regions table, so a missing row is easy to miss. _region_case avoids empty regions, so the tests do not catch it. The bp fallback has the same behavior (see #347), so a follow-up is fine. But the fallback region iterator could also yield empty regions here.

4. Not caused by this PR: with n_total_sites set and no mask, the region pi differs between the fused engine and the fallback, also with no missing data (0.00152 against 0.000912 for one region). The fused engine divides by the region length, and the fallback divides by the whole matrix's n_total_sites, because get_subset keeps it. This is the general engine span problem, and a separate issue will track it. The test comment "both paths normalize by the region span" is true only when n_total_sites is not set.

5. Docs.

• The windowed_analysis docstring now lists window_type, but not regions. Please document:
  • the columns that regions needs;
  • that end is exclusive, so a window is [start, end);
  • that regions needs window_type='regions';
  • that window_size counts variants for 'snp'.
• Please add a changelog entry, because the output changes for users who passed these arguments before.

6. Small items.

• # Build equivalent bp_bins for _compute_window_ranges (near the pop1 = line) no longer belongs to any code; you can remove it.
• #TODO: treat region bounds as inclusive. If you want to change the convention, an issue is a better place for this. Region strings elsewhere use inclusive ends (parse_region), so the two forms can confuse users. Also, #TODO and #fused need a space after #.
• windowed_statistics_fused has an extra blank line after the new raise.
• The old guard if missing_data in ('include', 'exclude') also kept an unknown missing_data out of the scatter engines. Now missing_data='bogus' runs the scatter engine as 'include'. Other functions do not check this value either, but window_type is now checked at the top, so one more line there would also check missing_data.