Windowed analysis supports different window_types

docs/source/changelog.rst

@@ -494,6 +494,14 @@ Bug fixes
   ``compute_ld_statistics`` also gained ``chunk_size`` and
   ``available_memory_bytes`` arguments to override the estimate; they do
   not change results, only memory use.
+* ``windowed_analysis`` ignored ``window_type='regions'`` and ``regions``
+  kwarg for statistics that went to the scattered or fused paths; only the
+  fallback used them. Now ``window_type`` defaults to ``'bp'`` and if 
+  ``'regions'`` is specified, it is sent to the proper engines. An error is
+  raised if no engine can compute the requested set of statistics or if 
+  ``regions`` is not passed alongside ``window_type='regions'`` and vice-versa.
+  An error is raise when ``window_type='snp'`` or ``'regions'`` for a 
+  ``StreamingHaplotypeMatrix``.
 
 New warnings
 ~~~~~~~~~~~~

pg_gpu/windowed_analysis.py

@@ -32,9 +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,
                           is_accessible=None):
@@ -432,8 +430,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 +440,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(
@@ -1198,6 +1198,7 @@ def _stream_windowed_analysis(streaming_hm, *, window_size, step_size,
 def windowed_analysis(haplotype_matrix: HaplotypeMatrix,
                      window_size: int = 50000,
                      step_size: Optional[int] = None,
+                     window_type: str = 'bp',
                      statistics: List[str] = ['pi'],
                      populations: Optional[List[str]] = None,
                      missing_data: str = 'include',
@@ -1214,7 +1215,17 @@ def windowed_analysis(haplotype_matrix: HaplotypeMatrix,
     haplotype_matrix : HaplotypeMatrix
         Input haplotype data
     window_size : int
-        Window size in base pairs
+        Window size in base pairs if `window_type = 'bp'` or number of variants
+        if `window_type='snp'`. Not used if `window_type = 'regions'`.
+    window_type : str
+        Type of windows: 'bp' (base pairs), 'regions', or 'snp'. If 'regions',
+        then provide an additional kwarg `regions` as a dictionary with the
+        form `{'chrom':[A, B, ...], 'start': [A_s, B_s, ...], 'end': [A_e, B_e, ...]}`. For
+        example, if you want the regions 1000-5000 and 8000-9000 on chromosome
+        1, then `regions = {'chrom': [1, 1], 'start': [1000, 8000], 'end': [5000, 9000]}`.
+        Any number of regions can be specified for any combination of chromosomes. Region
+        intervals are treated as half open, so a window spans [start, end). If 'snp' is given,
+        then `window_size` counts the number of variants.
     step_size : int, optional
         Step size. If None, uses window_size
     statistics : list
@@ -1247,8 +1258,17 @@ def windowed_analysis(haplotype_matrix: HaplotypeMatrix,
         of haplotypes with no missing call in each window (suffixed by
         population where the Garud columns are).
     """
+    valid_wt = ('bp', 'snp', 'regions')
+    if window_type not in valid_wt:
+        raise ValueError(f"Unknown window_type: {window_type!r}. Valid types are {valid_wt}.")
+    if (kwargs.get('regions') is not None) & (window_type != 'regions'):
+        raise ValueError("`regions` requested with incompatible `window_type`. Please pass `window_type = 'regions'`.")
+    if (window_type == 'regions') & (kwargs.get('regions') is None):
+        raise ValueError("`regions` kwarg must be provided for `window_type='regions'`.")
+    if missing_data not in ('include', 'exlcude'):
+        raise ValueError("Unknown value for `missing_data`. Please use 'include' or 'exclude'.")
     result = _windowed_analysis(
-        haplotype_matrix, window_size=window_size, step_size=step_size,
+        haplotype_matrix, window_size=window_size, step_size=step_size, window_type=window_type,
         statistics=statistics, populations=populations,
         missing_data=missing_data, span_normalize=span_normalize,
         accessible_bed=accessible_bed, chrom=chrom, **kwargs)
@@ -1293,6 +1313,7 @@ def _check_complete_haplotypes(df, matrix, populations, min_complete):
 def _windowed_analysis(haplotype_matrix: HaplotypeMatrix,
                       window_size: int = 50000,
                       step_size: Optional[int] = None,
+                      window_type: str = 'bp',
                       statistics: List[str] = ['pi'],
                       populations: Optional[List[str]] = None,
                       missing_data: str = 'include',
@@ -1303,6 +1324,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 +1358,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:
@@ -1349,6 +1375,7 @@ def _windowed_analysis(haplotype_matrix: HaplotypeMatrix,
                 haplotype_matrix,
                 window_size=window_size,
                 step_size=step_size,
+                window_type=window_type,
                 statistics=other_stats,
                 populations=populations,
                 missing_data=missing_data,
@@ -1374,7 +1401,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 +1464,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 +1476,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 +1502,24 @@ 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 +1554,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 +2125,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 +2227,14 @@ 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 +2623,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 +2683,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)