One window grid and one membership rule across the windowed engines by andrewkern · Pull Request #182 · kr-colab/pg_gpu

andrewkern · 2026-08-14T19:46:12Z

Closes #179. Closes #166.

The two issues sit on different axes of the same module. A windowed engine makes two
decisions: where the window grid starts, and which window each variant belongs to.
#179 was a grid problem in the iterator engine; #166 was a membership problem in the
fused engine. One commit per issue.

#179: iterator grid anchoring

WindowIterator._iter_bp_windows started the grid at the first variant and ignored
chrom_start / chrom_end, while every fused path anchors at the matrix bounds when
they are set. The engines tiled different grids for the same matrix whenever
chrom_start sits before the first variant: from_ts (chrom_start 0), a from_vcf
region load, or a get_subset_from_range subset. compute_region was the loudest
symptom -- its windows started at the first variant in the region, so neighbouring
region calls did not tile.

Fix: _bp_window_bounds applies the fused rule (matrix bounds when set, variants as
fallback) in both _iter_bp_windows and count_windows. count_windows now counts
the exact candidate starts the loop walks; the old formula under-counted whenever
step and window size differed (progress-bar total only).

#166: fused per-site window membership

bin_idx = searchsorted(we_gpu, positions) assigned each variant to one window, with
a boundary variant (pos == a window end == the next window's start) going to the
lower window. Everything else in the same call -- n_variants, the theta and
divergence kernels, the per-window slice stats -- is right-open [start, end). And
one window per variant means overlapping windows (step < window) silently dropped
each variant from every other window covering it. Affected: daf_hist, mu_sfs,
mean_nsl. The audit the issue asked for confirms the rest of the scatter set
(snp_dist_*, mu_var, zns, omega, mu_ld, dist_*) iterates
win_start[wi]:win_stop[wi] slices and was already correct.

Fix: membership is the consecutive run [k_lo, k_hi] -- k_hi the last window whose
start is <= pos, k_lo the first whose end is > pos. Both window families the engine
receives (contiguous bp_bins, uniform overlapping start/stop arrays) have
nondecreasing starts and ends, so the run is exact. Contiguous grids keep length-one
runs, so the common case keeps its old shape and cost. The chunked engine delegates
these stats to the fused function, so both paths are covered. This restores the
approach of the reverted df9716b and extends it to overlapping windows.

Tests

Nine new tests; eight fail on main (the ninth pins the unchanged no-bounds fallback):

per-window daf_hist / mu_sfs / mean_nsl equal the statistic recomputed from
that window's own [start, end) slice, on a boundary-heavy grid (every variant on
a window edge), on two- and three-deep overlapping grids, and through the direct
bp_bins branch;
the iterator anchors at chrom_start, its grid matches the fused engine, and two
consecutive compute_region calls tile without gaps or overlap.

pytest tests/ -n 10: 895 passed, 64 skipped. ruff: clean.

Notes

The fused engine reports every window; the iterator skips empty ones. Pre-existing
difference, unchanged here; the cross-engine test compares the non-empty rows.
A variant exactly at the grid tail (pos == the last window end, when the span is an
exact multiple of the step) falls outside all windows. Both engines agree on this
and did before; unchanged.
This will conflict lightly with the multiallelic branch, which reworks the daf /
mu_sfs internals next to the membership block (and where df9716b was reverted to
keep that branch scoped). The membership logic itself is orthogonal to those
reworks.

WindowIterator._iter_bp_windows started the window grid at the first variant and ignored chrom_start / chrom_end, while every fused windowed path anchors at the matrix bounds when they are set. The two engines tiled different grids for the same matrix whenever chrom_start sits before the first variant: a from_ts matrix (chrom_start 0), a from_vcf region load, or a get_subset_from_range subset. compute_region was the loudest symptom: its windows started at the first variant in the region and could run past the requested end, so neighbouring region calls did not tile. Use the matrix bounds when set, with the first and last variant as the fallback, in both _iter_bp_windows and count_windows. count_windows now counts the exact candidate starts the loop walks; the old formula under-counted whenever the step and window sizes differed. Tests: the iterator anchors at chrom_start, its grid matches the fused engine on the non-empty windows, the fallback without bounds is unchanged, and two consecutive compute_region calls tile without gaps or overlap.

The fused engine assigned each variant to a window with searchsorted on the window ENDS. Two failures follow. A variant sitting exactly on a boundary (pos == a window end == the next window's start) lands in the lower window, while n_variants, the theta / divergence kernels and the per-window slice stats all treat a window as right-open [start, end) -- so daf_hist / mu_sfs / mean_nsl counted a boundary variant in one window and everything else counted it in the next. And a single searchsorted assigns exactly one window per variant, so with overlapping windows (step < window) these stats dropped the variant from every other window that covers it. Replace the assignment with the member run [k_lo, k_hi]: k_hi is the last window whose start is <= pos, k_lo the first whose end is > pos. Both window families the engine receives (contiguous bp_bins and uniform overlapping start/stop arrays) have nondecreasing starts and ends, so the member windows of a variant are exactly that consecutive run. For contiguous windows the run has length one and the arrays keep their old shape and cost. The chunked engine delegates these stats to the fused function, so both paths are covered. Tests: every window's daf_hist / mu_sfs / mean_nsl equals the statistic recomputed from that window's own [start, end) variant slice, on boundary-heavy and overlapping grids, through both the windowed_analysis path and the direct bp_bins branch.

The grid-anchoring rule and the arange grid construction each existed in four copies across the engines; the two fixes in this branch had added the fourth. Extract _bp_grid_bounds and _bp_window_grid and call them from the iterator, both scatter paths, _build_scatter_indices, and the fused dispatcher, so the engines agree by construction rather than by test. count_windows returns the grid length instead of re-deriving it. The iterator now finds each window's variants with one searchsorted pair instead of a full boolean scan per candidate window. The scan cost was O(n_windows * n_var) and grew badly for a subset matrix whose chromosome bounds are much wider than its variant span. In the fused per-site block, drop the always-true upper-bound clause (site_win <= k_hi < n_windows by construction), and replace the tiling closure with zero-copy broadcast views: the consumers mask and scatter straight through the (n_per_var, n_var) views, which removes up to three tiled copies per call on overlapping grids. The per-site stat list and the daf bin count become module constants shared with the chunked delegation set and the tests. Tests: one shared leading-gap matrix builder, the four membership scenarios as one parametrized test, and expected values built once in the fixture from the engine's own binning constant.

The candidate-row count came from a reduction over the per-variant run lengths, which forced a device-to-host sync and needed an empty-input guard because CuPy's max rejects empty arrays. Coverage depth is a property of the window grid alone: it is maximized at a window start, and the window arrays are already on the host. Two numpy searchsorted calls give the depth with no device work and no guard. The depth can exceed what the variants need when they avoid the deep-coverage spots; site_ok masks the unused rows, so an upper bound is what the block wants anyway, and the allocation shape no longer depends on where the variants fall.

nspope

Looks great, left a couple comments on edge cases, and spotting one inclusive/exclusive endpoint bug (I think this is a bug, depends how window interval is defined)

nspope · 2026-08-17T18:19:13Z

+        win_starts, win_stops = _bp_window_grid(
+            chrom_start, chrom_end, window_size, step_size)
        # Build equivalent bp_bins for _compute_window_ranges
        bp_bins = np.concatenate([win_starts, [win_stops[-1]]])


Very minor edge case, but this'll error out ambiguously on a size-0 grid (which is reachable if chrom_start == chrom_end). Might be worth adding a if n_windows == 0: guard.

Here's a MWE:

import numpy as np from pg_gpu import HaplotypeMatrix, windowed_analysis pos = np.array([0, 50, 100]) hap = np.tile(np.array([0, 1, 0, 1], dtype=np.int8)[:, None], (1, 3)) hm = HaplotypeMatrix(hap, pos) # a one-base span should not crash (but currently throws IndexError) windowed_analysis(hm.get_subset_from_range(50, 51), window_size=100, step_size=100, statistics=['pi'])

nspope · 2026-08-17T18:44:53Z

+    chrom_end; each window spans window_size. Shared by every windowed
+    engine so the grids agree by construction.
+    """
+    win_starts = np.arange(int(chrom_start), int(chrom_end), step_size,


So I think chrom_end is intended to be inclusive, but the arange here treats it as exclusive. This allows SNPs to vanish on a window boundary. For example:

import numpy as np from pg_gpu import HaplotypeMatrix, windowed_analysis pos = np.array([0, 50, 100]) hap = np.tile(np.array([0, 1, 0, 1], dtype=np.int8)[:, None], (1, 3)) hm = HaplotypeMatrix(hap, pos) # every variant should land in some window df = windowed_analysis(hm, window_size=100, step_size=100, statistics=['pi']) assert df['n_variants'].sum() == hm.num_variants, f"expected {hm.num_variants}, got {df['n_variants'].sum()}"

nspope · 2026-08-17T18:53:22Z

+        assert int(it['start'].iloc[0]) == 0
+        assert all(int(s) % 10_000 == 0 for s in it['start'])
+
+    def test_iterator_grid_matches_fused(self):


I think it'd be worth: (a) checking multiple stats here, (b) doing overlapping windows

nspope · 2026-08-17T18:57:01Z

                                        err_msg=f"H2H1 window {w}")
+
+
+class TestFusedPerSiteWindowMembership:


This might by a bit orthogonal to the point of this PR, but I think we need a more through two-pop scatter-vs-fused comparison. The existing test_fused_twopop_missing_matches_scatter rebuilds windows by hand under a specific missing-data setup. So I think it'd be worth adding some comparisons on an overlapping grid here (i.e. for da/dxy/fst/fst_hudson).

nspope · 2026-08-17T18:59:22Z

+                                       atol=1e-12)
+
+
+class TestEngineGridAgreement:


just brainstorming some edge cases to check with tests (feel free to skip):

single variant windows

zero-window grid (that's the index error I commented on above)

that different loaders->windowed_analysis result in the same output

andrewkern added 3 commits August 14, 2026 12:44

andrewkern mentioned this pull request Aug 14, 2026

Unify window membership: _build_scatter_indices and the fused per-site block are two implementations of one rule #183

Open

andrewkern requested a review from nspope August 14, 2026 20:02

nspope reviewed Aug 17, 2026

View reviewed changes

Navigation Menu

Uh oh!

One window grid and one membership rule across the windowed engines - #182

One window grid and one membership rule across the windowed engines#182
andrewkern wants to merge 4 commits into
mainfrom
fix/window-grid-and-membership

andrewkern commented Aug 14, 2026

Uh oh!

nspope left a comment

Uh oh!

nspope Aug 17, 2026

Uh oh!

nspope Aug 17, 2026

Uh oh!

nspope Aug 17, 2026

Uh oh!

nspope Aug 17, 2026

Uh oh!

nspope Aug 17, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		err_msg=f"H2H1 window {w}")


		class TestFusedPerSiteWindowMembership:

Uh oh!

Conversation

andrewkern commented Aug 14, 2026

#179: iterator grid anchoring

#166: fused per-site window membership

Tests

Notes

Uh oh!

nspope left a comment

Choose a reason for hiding this comment

Uh oh!

nspope Aug 17, 2026

Choose a reason for hiding this comment

Uh oh!

nspope Aug 17, 2026

Choose a reason for hiding this comment

Uh oh!

nspope Aug 17, 2026

Choose a reason for hiding this comment

Uh oh!

nspope Aug 17, 2026

Choose a reason for hiding this comment

Uh oh!

nspope Aug 17, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants