Windowed analysis supports different window_types

pg_gpu/windowed_analysis.py

@@ -32,8 +32,7 @@
 # Kwargs that the 'local_pca' / 'local_pca_jackknife' dispatch consumes but
 # scalar-stat paths don't accept. Filtered out before the recursive call.
 _LOCAL_PCA_ONLY_KWARGS = frozenset(
-    {'k', 'population', 'batch_size', 'window_type', 'regions',
-     'n_blocks', 'aggregate'})
+    {'k', 'population', 'batch_size', 'n_blocks', 'aggregate'})
 
 
 def _compute_window_bases(haplotype_matrix, win_starts, win_stops,
@@ -432,8 +431,8 @@ def _iter_region_windows(self) -> Iterator[WindowData]:
             raise ValueError("Regions must be provided for region window type")
 
         for window_id, region in self.params.regions.iterrows():
-            start = region['start']
-            end = region['end']
+            start = int(region['start'])
+            end = int(region['end'])
             center = (start + end) // 2
 
             # Find variants in region
@@ -442,6 +441,8 @@ def _iter_region_windows(self) -> Iterator[WindowData]:
 
             if len(variant_indices) > 0:
                 window_matrix = self.matrix.get_subset(variant_indices)
+                window_matrix.chrom_start = start
+                window_matrix.chrom_end = end - 1 #TODO: treat region bounds as inclusive.
                 self._attach_window_mask(window_matrix, start, end)
 
                 yield WindowData(
@@ -1197,6 +1198,7 @@ def _stream_windowed_analysis(streaming_hm, *, window_size, step_size,
 
 def windowed_analysis(haplotype_matrix: HaplotypeMatrix,
                      window_size: int = 50000,
+                     window_type: str = 'bp',
                      step_size: Optional[int] = None,
                      statistics: List[str] = ['pi'],
                      populations: Optional[List[str]] = None,
@@ -1215,6 +1217,8 @@ def windowed_analysis(haplotype_matrix: HaplotypeMatrix,
         Input haplotype data
     window_size : int
         Window size in base pairs
+    window_type : str
+        Type of windows: 'bp' (base pairs), 'regions', or 'snp'.
     step_size : int, optional
         Step size. If None, uses window_size
     statistics : list
@@ -1247,8 +1251,10 @@ def windowed_analysis(haplotype_matrix: HaplotypeMatrix,
         of haplotypes with no missing call in each window (suffixed by
         population where the Garud columns are).
     """
+    if window_type not in ('bp', 'snp', 'regions'):
+        raise ValueError(f"Unknown window_type: {window_type!r}")
     result = _windowed_analysis(
-        haplotype_matrix, window_size=window_size, step_size=step_size,
+        haplotype_matrix, window_size=window_size, window_type=window_type, step_size=step_size,
         statistics=statistics, populations=populations,
         missing_data=missing_data, span_normalize=span_normalize,
         accessible_bed=accessible_bed, chrom=chrom, **kwargs)
@@ -1292,6 +1298,7 @@ def _check_complete_haplotypes(df, matrix, populations, min_complete):
 
 def _windowed_analysis(haplotype_matrix: HaplotypeMatrix,
                       window_size: int = 50000,
+                      window_type: str = 'bp',
                       step_size: Optional[int] = None,
                       statistics: List[str] = ['pi'],
                       populations: Optional[List[str]] = None,
@@ -1303,6 +1310,11 @@ def _windowed_analysis(haplotype_matrix: HaplotypeMatrix,
     """The body of ``windowed_analysis``, without its final checks."""
     from .streaming_matrix import StreamingHaplotypeMatrix
     if isinstance(haplotype_matrix, StreamingHaplotypeMatrix):
+        if window_type in ('regions', 'snp'):
+            raise NotImplementedError(
+                "window_type='regions' or 'snp' is not available "
+                "on a StreamingHaplotypeMatrix."
+            )
         return _stream_windowed_analysis(
             haplotype_matrix,
             window_size=window_size, step_size=step_size,
@@ -1332,7 +1344,7 @@ def _windowed_analysis(haplotype_matrix: HaplotypeMatrix,
             'batch_size': kwargs.get('batch_size', None),
             'window_size': window_size,
             'step_size': step_size,
-            'window_type': kwargs.get('window_type', 'bp'),
+            'window_type': window_type,
             'regions': kwargs.get('regions', None),
         }
         if want_jackknife:
@@ -1348,6 +1360,7 @@ def _windowed_analysis(haplotype_matrix: HaplotypeMatrix,
             scalar_df = _windowed_analysis(
                 haplotype_matrix,
                 window_size=window_size,
+                window_type=window_type,
                 step_size=step_size,
                 statistics=other_stats,
                 populations=populations,
@@ -1374,7 +1387,7 @@ def _windowed_analysis(haplotype_matrix: HaplotypeMatrix,
     scatter_twopop = _TWO_POP_STATS
     requested = set(statistics)
 
-    if missing_data in ('include', 'exclude'):
+    if window_type == 'bp':
         # Pure single-pop request
         if requested <= scatter_single and len(populations or []) <= 1:
             result = _windowed_thetas_scatter(
@@ -1437,6 +1450,10 @@ def _windowed_analysis(haplotype_matrix: HaplotypeMatrix,
     can_fuse = ((missing_data == 'include' or requested <= fused_garud
                  or bool(requested & fused_only))
                 and requested <= fused_all
+                # The fused engine tiles bp grids and custom regions; snp
+                # windows are variant-index ranges that only the per-window
+                # fallback builds, so they must skip the fused path.
+                and window_type in ('bp', 'regions')
                 # A two-population statistic with anything other than exactly
                 # two populations must not reach the fused engine, which
                 # computes the first pair only and names the column after the
@@ -1445,6 +1462,24 @@ def _windowed_analysis(haplotype_matrix: HaplotypeMatrix,
                 and (not (requested & fused_two)
                      or len(populations or []) == 2))
 
+    if window_type == 'regions':
+        # theta_l, zeng_e, zeng_dh, normalized_fay_wu_h, mean_nsl,
+        # daf_hist, mu_sfs currently not supported with window_type = 'regions'.
+        unsupported_regions = {'theta_l', 'zeng_e', 'zeng_dh', 'normalized_fay_wu_h',
+                                'mean_nsl', 'daf_hist', 'mu_sfs'}
+        if requested & unsupported_regions:
+            raise NotImplementedError(f"{sorted(requested & unsupported_regions)} are not "
+            "supported for window_type='regions'."
+            )
+
+    #fused engine can't handle 'snp' windows yet, only local_pca and fallback.
+    if window_type == 'snp':
+        unsupported_snp = {s for s in requested if isinstance(s, str)} - _fallback_stats()
+        if unsupported_snp:
+            raise NotImplementedError(
+                f"{sorted(unsupported_snp)} not yet supported for window_type='snp'."
+            )
+
     if can_fuse:
         if haplotype_matrix.device == 'CPU':
             haplotype_matrix.transfer_to_gpu()
@@ -1453,16 +1488,25 @@ def _windowed_analysis(haplotype_matrix: HaplotypeMatrix,
             positions = positions.get()
         positions = np.asarray(positions)
 
-        chrom_start, chrom_end = _bp_grid_bounds(haplotype_matrix, positions)
+        if window_type == 'regions':
+            regions = kwargs.get('regions')
+            if regions is None:
+                raise ValueError("Regions must be provided for region window type")
+            win_starts = np.asarray(regions['start'], dtype=np.float64)
+            win_stops = np.asarray(regions['end'], dtype=np.float64)
+            bp_bins = None
+        else:
+            chrom_start, chrom_end = _bp_grid_bounds(haplotype_matrix, positions)
+
+            # Build window start/stop arrays (supports overlapping windows)
+            win_starts, win_stops = _bp_window_grid(
+                chrom_start, chrom_end, window_size, step_size)
+            # A zero-span matrix (one variant, no bounds) tiles no windows.
+            bp_bins = np.concatenate([win_starts, [win_stops[-1]]])
 
-        # Build window start/stop arrays (supports overlapping windows)
-        win_starts, win_stops = _bp_window_grid(
-            chrom_start, chrom_end, window_size, step_size)
-        # A zero-span matrix (one variant, no bounds) tiles no windows.
         if len(win_starts) == 0:
             return pd.DataFrame()
         # Build equivalent bp_bins for _compute_window_ranges
-        bp_bins = np.concatenate([win_starts, [win_stops[-1]]])
 
         pop1 = populations[0] if populations and len(populations) >= 1 else None
         pop2 = populations[1] if populations and len(populations) >= 2 else None
@@ -1497,7 +1541,7 @@ def _windowed_analysis(haplotype_matrix: HaplotypeMatrix,
 
     # Fallback: per-window Python loop
     analyzer = WindowedAnalyzer(
-        window_type='bp',
+        window_type=window_type,
         window_size=window_size,
         step_size=step_size,
         statistics=statistics,
@@ -2068,7 +2112,7 @@ def _apply_exclude(haplotype_matrix, missing_data, statistics, population,
 
 
 def windowed_statistics_fused(haplotype_matrix: HaplotypeMatrix,
-                              bp_bins,
+                              bp_bins=None,
                               statistics=('pi', 'theta_w', 'tajimas_d',
                                           'segregating_sites', 'singletons'),
                               population=None,
@@ -2170,12 +2214,15 @@ def windowed_statistics_fused(haplotype_matrix: HaplotypeMatrix,
         n_windows = len(ws_gpu)
         win_start = cp.searchsorted(positions, ws_gpu, side='left').astype(cp.int64)
         win_stop = cp.searchsorted(positions, we_gpu, side='left').astype(cp.int64)
-    else:
+    elif bp_bins is not None:
         bp_bins_gpu = cp.asarray(bp_bins, dtype=cp.float64)
         n_windows = len(bp_bins_gpu) - 1
         win_start, win_stop = _compute_window_ranges(positions, bp_bins_gpu)
         ws_gpu = bp_bins_gpu[:-1]
         we_gpu = bp_bins_gpu[1:]
+    else:
+        raise ValueError('bp_bins and _win_starts/_win_stops must not both be None.')
+
 
     ws_cpu = ws_gpu.get() if hasattr(ws_gpu, 'get') else ws_gpu
     we_cpu = we_gpu.get() if hasattr(we_gpu, 'get') else we_gpu
@@ -2564,7 +2611,7 @@ def _expand(v):
 
 
 def windowed_statistics_fused_chunked(haplotype_matrix: HaplotypeMatrix,
-                                      bp_bins,
+                                      bp_bins=None,
                                       statistics=('pi', 'theta_w', 'tajimas_d'),
                                       population=None,
                                       pop1=None,
@@ -2624,12 +2671,14 @@ def windowed_statistics_fused_chunked(haplotype_matrix: HaplotypeMatrix,
         n_windows = len(ws_gpu)
         win_start = cp.searchsorted(positions, ws_gpu, side='left').astype(cp.int64)
         win_stop = cp.searchsorted(positions, we_gpu, side='left').astype(cp.int64)
-    else:
+    elif bp_bins is not None:
         bp_bins_gpu = cp.asarray(bp_bins, dtype=cp.float64)
         n_windows = len(bp_bins_gpu) - 1
         win_start, win_stop = _compute_window_ranges(positions, bp_bins_gpu)
         ws_gpu = bp_bins_gpu[:-1]
         we_gpu = bp_bins_gpu[1:]
+    else:
+        raise ValueError('bp_bins and _win_starts/_win_stops must not both be None.')
 
     ws_cpu = ws_gpu.get() if hasattr(ws_gpu, 'get') else ws_gpu
     we_cpu = we_gpu.get() if hasattr(we_gpu, 'get') else we_gpu

tests/test_windowed_analysis.py

@@ -1656,8 +1656,8 @@ def test_window_id_is_monotonic(self, hm):
 
 class TestWindowedGarudCorrectness:
     """Anchor windowed Garud H against scikit-allel's per-window
-    garud_h. ``pg_gpu`` runs the windowed pass on the GPU via a fused
-    kernel; this confirms each window's H1 / H12 / H123 / H2H1 lines
+    garud_h. ``pg_gpu`` runs the windowed pass on the fallback engine;
+    this confirms each window's H1 / H12 / H123 / H2H1 lines
     up with the canonical implementation."""
 
     @pytest.mark.parametrize("seed", [1, 2, 3])
@@ -1667,9 +1667,10 @@ def test_matches_allel_per_window(self, seed):
         n_hap = 32
         n_var = 600
         window_size_var = 50  # variants per window
-        # Use unit-stride positions so window_size_var-many variants
-        # land in each tile window.
-        positions = np.arange(n_var, dtype=np.int64)
+        # Spread positions on a non-unit stride so a 50-bp grid and a 50-variant snp
+        # grid diverge. Previous check passed even though the fused engine was silently using
+        # a bp window_type, not snp.
+        positions = (np.arange(n_var) * 100).astype(np.int64)
         chrom_len = int(positions[-1]) + 1
         haps = rng.integers(0, 2, (n_hap, n_var), dtype=np.int8)
 
@@ -1680,6 +1681,7 @@ def test_matches_allel_per_window(self, seed):
                                window_type="snp",
                                statistics=["garud_h1", "garud_h12",
                                             "garud_h123", "garud_h2h1"])
+        assert len(df) == n_var // window_size_var   # 12 genuine snp windows
 
         # Each tile window is exactly window_size_var variants on the
         # contiguous range; allel.garud_h expects (n_var, n_hap).
@@ -1787,6 +1789,26 @@ def test_direct_bp_bins_branch(self, data):
                                daf_bin, f"bp bin {wi}")
 
 
+def _region_case(seed=0):
+    """A matrix plus custom regions (unsorted, overlapping/nested, gapped, all
+    populated) for the regions-dispatch agreement tests. Every region holds at
+    least one variant so the fallback -- which skips empty regions -- yields a
+    row per region too, aligned with the fused engine's dense output.
+    """
+    rng = np.random.RandomState(seed)
+    n_var = 400
+    pos = np.sort(rng.choice(np.arange(1, 1_000_000), n_var, replace=False))
+    hap = rng.randint(0, 2, (40, n_var)).astype(np.int8)
+    hm = HaplotypeMatrix(hap, pos, 1, 999_999)
+    hm.sample_sets = {'p1': list(range(20)), 'p2': list(range(20, 40))}
+    regions = pd.DataFrame({
+        'chrom': [1, 1, 1, 1],
+        'start': [600_000, 100_000, 100_000, 400_000],
+        'end':   [900_000, 300_000, 250_000, 500_000],
+    })
+    return hm, regions
+
+
 class TestEngineGridAgreement:
     """Both windowed engines tile the same grid for the same matrix,
     anchored at the matrix chromosome bounds when they are set."""
@@ -1853,6 +1875,49 @@ def test_compute_region_tiles_from_region_start(self):
         assert int(r1['n_variants'].sum()) == in1
         assert int(r2['n_variants'].sum()) == in2
 
+    def test_region_windows_single_pop_fused_vs_fallback(self):
+        # window_type='regions' must honour the regions (one window each, not
+        # the default bp grid -- the reported bug) and route single-pop stats
+        # through the fused engine, whose raw values match the per-window
+        # fallback. pi's agreement also confirms both paths normalize by the
+        # region span (the second fault).
+        hm, regions = _region_case()
+        stats = ['pi', 'tajimas_d', 'segregating_sites', 'garud_h1']
+        fused = windowed_analysis(hm, window_type='regions', regions=regions,
+                                  statistics=stats)
+        fb = WindowedAnalyzer(window_type='regions', regions=regions,
+                              statistics=stats, progress_bar=False).compute(hm)
+        # Regions honoured, not the default bp grid.
+        assert len(fused) == len(regions)
+        assert list(fused['start'].astype(int)) == list(regions['start'])
+        assert list(fused['end'].astype(int)) == list(regions['end'])
+        # Same window order in both engines; raw values agree.
+        for k in stats:
+            np.testing.assert_allclose(fused[k].to_numpy(dtype=float),
+                                       fb[k].to_numpy(dtype=float),
+                                       rtol=1e-9, atol=1e-11, equal_nan=True,
+                                       err_msg=k)
+
+    def test_region_windows_two_pop_fused_vs_fallback(self):
+        # Two-pop stats over regions: the fused engine (bare column names) and
+        # the fallback (suffixed names) agree, and dxy's agreement confirms the
+        # region span normalization on the two-pop path too.
+        hm, regions = _region_case(seed=1)
+        fused = windowed_analysis(hm, window_type='regions', regions=regions,
+                                  statistics=['fst_hudson', 'dxy'],
+                                  populations=['p1', 'p2'])
+        fb = WindowedAnalyzer(window_type='regions', regions=regions,
+                              statistics=['fst_hudson', 'dxy'],
+                              populations=['p1', 'p2'],
+                              progress_bar=False).compute(hm)
+        assert len(fused) == len(regions)
+        for bare, suffixed in (('fst_hudson', 'fst_hudson_p1_p2'),
+                               ('dxy', 'dxy_p1_p2')):
+            np.testing.assert_allclose(fused[bare].to_numpy(dtype=float),
+                                       fb[suffixed].to_numpy(dtype=float),
+                                       rtol=1e-9, atol=1e-11, equal_nan=True,
+                                       err_msg=bare)
+
 
 def _host(a):
     return a.get() if hasattr(a, "get") else np.asarray(a)

tests/test_windowed_chunked_coverage.py

@@ -62,6 +62,30 @@ def _compare(hm, statistics, monkeypatch, window_size=10_000,
                                    atol=1e-14, equal_nan=True, err_msg=k)
 
 
+def _region_edges():
+    """Custom window edges the dispatcher builds for ``window_type='regions'``:
+    overlapping, nested, out-of-order, gapped, and one past the last variant
+    (empty). No grid structure, so each engine must bin every window on its own.
+    """
+    starts = np.array([40_000, 10_000, 10_000, 55_000, 200_000], dtype=np.float64)
+    stops = np.array([70_000, 30_000, 25_000, 60_000, 210_000], dtype=np.float64)
+    return starts, stops
+
+
+def _compare_edges(hm, statistics, monkeypatch, starts, stops, **kwargs):
+    """Like ``_compare`` but drives the engines from explicit window edges
+    (``bp_bins=None``) instead of a bp grid."""
+    common = dict(statistics=tuple(statistics), bp_bins=None,
+                  _win_starts=starts, _win_stops=stops, **kwargs)
+    single = windowed_statistics_fused(hm, **common)
+    monkeypatch.setattr(_memutil, "estimate_fused_chunk_size",
+                        lambda n, memory_fraction=0.35: _CHUNK)
+    chunked = windowed_statistics_fused_chunked(hm, **common)
+    for k in statistics:
+        np.testing.assert_allclose(single[k], chunked[k], rtol=1e-10,
+                                   atol=1e-14, equal_nan=True, err_msg=k)
+
+
 # ── two-population: fst_hudson column + per_base=False dxy/da ───────────
 def test_twopop_fst_hudson_and_perbase_false(twopop_hm, monkeypatch):
     _compare(twopop_hm, ["fst", "fst_hudson", "dxy", "da"], monkeypatch,
@@ -109,6 +133,21 @@ def test_sparse_windows_empty_overlap(twopop_hm, monkeypatch):
              pop1="p1", pop2="p2")
 
 
+# ── custom region windows: overlapping / nested / unsorted / empty ─────
+def test_region_windows_fused_vs_chunked(twopop_hm, monkeypatch):
+    # The dispatcher feeds window_type='regions' as explicit _win_starts /
+    # _win_stops with no grid structure, so the chunked engine must bin and
+    # accumulate arbitrary windows the same as the single-shot engine: pi
+    # (per-base accumulation + span), tajimas_d / segregating_sites (single-pop
+    # kernels), fst_hudson (two-pop kernel), and garud_h1 (delegated to the
+    # single-shot engine -- this checks the delegation forwards the edges).
+    starts, stops = _region_edges()
+    _compare_edges(twopop_hm,
+                   ["pi", "tajimas_d", "segregating_sites", "fst_hudson",
+                    "garud_h1"],
+                   monkeypatch, starts, stops, pop1="p1", pop2="p2")
+
+
 # ── population validation (before the chunk loop) ──────────────────────
 def test_chunked_requires_pop1_pop2(twopop_hm):
     grid = _grid(twopop_hm, 10_000, 10_000)