diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst index f5c3ac6a..397c5a9d 100644 --- a/docs/source/changelog.rst +++ b/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 scatter 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 raised when ``window_type='snp'`` or ``'regions'`` for a + ``StreamingHaplotypeMatrix``. New warnings ~~~~~~~~~~~~ diff --git a/pg_gpu/windowed_analysis.py b/pg_gpu/windowed_analysis.py index db0d2e4f..8f9a57e6 100644 --- a/pg_gpu/windowed_analysis.py +++ b/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( @@ -1205,6 +1206,7 @@ def windowed_analysis(haplotype_matrix: HaplotypeMatrix, accessible_bed: str = None, chrom: str = None, min_complete_haplotypes: Optional[int] = None, + window_type: str = 'bp', **kwargs) -> pd.DataFrame: """ Convenience function for windowed analysis. @@ -1214,7 +1216,8 @@ 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'`. step_size : int, optional Step size. If None, uses window_size statistics : list @@ -1236,6 +1239,15 @@ def windowed_analysis(haplotype_matrix: HaplotypeMatrix, missing call (the ``n_complete_haplotypes`` column). When it is not given, a ``FewCompleteHaplotypesWarning`` reports when many windows have few complete haplotypes, where H1 and H12 come out too low. + window_type : str + Type of windows: 'bp' (base pairs), 'regions', or 'snp'. If 'regions', + then provide an additional kwarg `regions` as a pandas DataFrame with the + form `pd.DataFrame({'start': [s1, s2, ...], 'end': [e1, e2, ...]})`. For + example, if you want the regions 1000-5000 and 8000-9000 on the chromosome, + then `regions = pd.DataFrame({'start': [1000, 8000], 'end': [5000, 9000]})`. + Any number of regions can be specified. 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. **kwargs Additional arguments passed to WindowedAnalyzer @@ -1247,11 +1259,20 @@ 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) and (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', 'exclude'): + 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, statistics=statistics, populations=populations, missing_data=missing_data, span_normalize=span_normalize, - accessible_bed=accessible_bed, chrom=chrom, **kwargs) + accessible_bed=accessible_bed, chrom=chrom, window_type=window_type, **kwargs) # A local-PCA request returns a LocalPCAResult; its scalar columns sit # in .windows. frame = getattr(result, 'windows', result) @@ -1299,10 +1320,16 @@ def _windowed_analysis(haplotype_matrix: HaplotypeMatrix, span_normalize=True, accessible_bed: str = None, chrom: str = None, + window_type: str = 'bp', **kwargs) -> pd.DataFrame: """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 +1359,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: @@ -1355,6 +1382,7 @@ def _windowed_analysis(haplotype_matrix: HaplotypeMatrix, span_normalize=span_normalize, accessible_bed=None, # already applied above chrom=chrom, + window_type=window_type, **{k: v for k, v in kwargs.items() if k not in _LOCAL_PCA_ONLY_KWARGS}, ) @@ -1374,7 +1402,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 +1465,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 +1477,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 +1503,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 +1555,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 +2126,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 +2228,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 +2624,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 +2684,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 diff --git a/tests/test_windowed_analysis.py b/tests/test_windowed_analysis.py index d019e77b..58e001bc 100644 --- a/tests/test_windowed_analysis.py +++ b/tests/test_windowed_analysis.py @@ -86,7 +86,6 @@ def test_region_windows(self): haplotypes = np.random.randint(0, 2, size=(10, n_variants)) regions = pd.DataFrame({ - 'chrom': [1, 1, 1], 'start': [10000, 30000, 70000], 'end': [20000, 50000, 90000] }) @@ -1656,8 +1655,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 +1666,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 +1680,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 +1788,25 @@ 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({ + '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 +1873,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) diff --git a/tests/test_windowed_chunked_coverage.py b/tests/test_windowed_chunked_coverage.py index e9eb4496..5164abd3 100644 --- a/tests/test_windowed_chunked_coverage.py +++ b/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) diff --git a/tests/test_windowed_error_coverage.py b/tests/test_windowed_error_coverage.py index 351d2393..627c95ec 100644 --- a/tests/test_windowed_error_coverage.py +++ b/tests/test_windowed_error_coverage.py @@ -3,6 +3,7 @@ missing-regions raises, and the fused engine's CPU->GPU transfer. """ import numpy as np +import pandas as pd import pytest from pg_gpu import HaplotypeMatrix @@ -68,3 +69,31 @@ def test_fused_transfers_cpu_matrix(): statistics=("pi",), population="p1") assert m.device == "GPU" assert np.isfinite(result["pi"][0]) + + +# ── forward-looking: window_type/statistic combinations deferred to later ── +# Each case below raises NotImplementedError today and marks work the issue +# leaves for the future. When one is implemented its dispatch guard goes away +# and this test fails -- the signal to drop that case here and from the guard. + +# Scatter-only theta estimators / neutrality tests and per-site stats. The +# fused regions path can't serve them, and they aren't implemented in the +# per-window fallback either, so the dispatch raises. +_REGIONS_UNSUPPORTED = ["theta_l", "zeng_e", "zeng_dh", "normalized_fay_wu_h", + "mean_nsl", "daf_hist", "mu_sfs"] +# snp windows route to the per-window fallback, so stats that aren't +# implemented there still raise. +_SNP_UNSUPPORTED = ["theta_l", "zeng_e", "theta_h", "max_daf", "mean_nsl", + "zns"] + + +@pytest.mark.parametrize("window_type,stat", ( + [("regions", s) for s in _REGIONS_UNSUPPORTED] + + [("snp", s) for s in _SNP_UNSUPPORTED])) +def test_deferred_window_type_stats_raise(window_type, stat): + m = _matrix(n_var=40) + regions = (pd.DataFrame({"start": [100], "end": [3000]}) + if window_type == "regions" else None) + with pytest.raises(NotImplementedError): + windowed_analysis(m, window_size=10, window_type=window_type, + statistics=[stat], regions=regions)