[andrewkern]

Fixes #330.

Problem

Under missing_data='exclude', windowed_analysis sent only Garud-only requests to the fused engine. Every other request went to the per-window fallback, which does not know these statistics, so the call raised Unknown statistic:

• zns, omega, mu_ld, mu_var, mu_sfs, daf_hist, the snp_dist_* and dist_* statistics, and mean_nsl, alone or mixed;
• theta_h, singletons, fay_wu_h and max_daf when requested with a Garud statistic, or with two populations and no two-population statistic. windowed_analysis rejects the diploS/HIC statistics and mean_nsl under missing_data='exclude' #330 did not list these four.

Called directly, windowed_statistics_fused and windowed_statistics_fused_chunked silently ignored 'exclude' for every statistic except the Garud family and omega. The values and n_variants were the same as under 'include'.

Changes

• Pre-filter in both fused engines. Under 'exclude', both engines restrict the matrix to its complete sites before they index windows (_apply_exclude).
  • A site's completeness does not depend on the window, so this equals dropping each window's incomplete sites. The window grid comes from the bin edges and does not move.
  • Completeness is judged over the rows the request reads: both populations when a two-population statistic is requested, else the single population, else every row.
  • After the filter the engine runs in 'include' mode. So the chunked engine's hand-off to the unchunked engine does not filter again over a different set of rows.
  • A Garud-only request skips the filter, because garud_h_windows drops incomplete columns itself without a copy. That path, including its n_variants, is unchanged.
• Routing. Under 'exclude', windowed_analysis uses the fused engine for any request that contains a statistic the fallback cannot compute. The set comes from the fallback's own statistic tables (_fallback_stats), not from a separate list. Every request that this routes differently raised on main; requests that worked on main keep their engine.
• Cheaper mask. exclude_missing_sites finds incomplete columns with column_has_missing: one pass and one byte per site, with no matrix-sized boolean.
• One constant. The five copies of the two-population statistic names in windowed_analysis.py now share _TWO_POP_STATS.

Known limits

• Memory. The filter copies the complete sites, so under 'exclude' an engine holds one extra copy of the matrix next to the caller's. At large sizes the chunked engine can run out of memory under 'exclude' where 'include' works. Before this PR those requests raised at every size. Measured on an 80 GB A100 with pi + mu_sfs: 'exclude' works up to a 46.6 GiB matrix and runs out of memory at 51.3 GiB, while 'include' works at 55.9 GiB. Fused engines copy the matrix under missing_data='exclude', which defeats the chunked engine's memory bound #333 tracks the fix: carry a keep index instead of copying.
• Different site sets in mixed requests. With two populations under 'exclude', adding zns to ['pi', 'fst'] changes pi. The scatter engine judges completeness per statistic; the fused engine uses one site set for the whole request, as its allele-cap filter already does. This is part of Engines disagree on the undefined-FST sentinel and on single-population columns in mixed requests #200.

Tests

tests/test_windowed_exclude.py, 25 tests:

• Each affected statistic under 'exclude' against the scalar function on the window's complete sites, with n_variants and the window starts.
• A mixed request against the single-statistic requests, and pi against the complete sites.
• daf_hist against diversity.daf_histogram.
• theta_h, singletons, fay_wu_h and max_daf with a Garud statistic and with two populations.
• Both fused engines called directly, against 'include' on the complete sites.
• The chunked and unchunked engines agree when some rows belong to neither population. This test fails if the hand-off filters a second time.
• Completeness judged within one population and within two.
• The grid does not move when an end site is incomplete.
• The streaming path against the scalar functions.

24 of the 25 tests fail on main. The chunked-against-unchunked test passes on main only because both engines ignore 'exclude' there.

The full suite passes: 1957 passed, 81 skipped, 15 xfailed. ruff check pg_gpu/ tests/ examples/ passes.

[nspope]

looks good! couple minor points

[nspope]

Hmm ... so if statistics contains something that's not a Garud stat, then this doesn't fire, and the call goes through to the exclude_missing_sites below? That doesn't seem quite right?

[nspope]

for example,

import numpy as np
from pg_gpu import HaplotypeMatrix
from pg_gpu.windowed_analysis import windowed_analysis

N_HAP, N_VAR, SPACING, WIN = 40, 40, 10, 200
rng = np.random.default_rng(0)
hap = (rng.random((N_HAP, N_VAR)) < 0.3).astype(np.int8)
hap[:, 5] = np.where(np.arange(N_HAP) < 20, hap[:, 5], -1)  # site 5 missing in rows 20-39 only
pos = np.arange(N_VAR, dtype=np.int64) * SPACING
hm = HaplotypeMatrix(hap.copy(), pos)
df_garud_only = windowed_analysis(
    hm, window_size=WIN, step_size=WIN, 
    statistics=["garud_h1"],
    missing_data="exclude",
)
df_garud_plus = windowed_analysis(
    hm, window_size=WIN, step_size=WIN, 
    statistics=["garud_h1", "zns"],
    missing_data="exclude",
)
print("n_variants garud-only:", df_garud_only["n_variants"].to_numpy())
print("n_variants garud+zns: ", df_garud_plus["n_variants"].to_numpy())
n_variants garud-only: [20 20]
n_variants garud+zns:  [19 20]

[andrewkern]

ah-- you are right. fixed this so the fused engine now counts the complete sites when it builds the window columns and the fallback sets the complete count on the window, so the n_variants column and the n_variants statistic agree.

with the latest commit your example gives

n_variants garud-only: [19 20]                                                                                                   │
n_variants garud+zns:  [19 20]

[nspope]

this is happening before the pop1/pop2 parameter check, so we don't raise the right error if pops aren't supplied to a raw call of the engine:

import numpy as np
from pg_gpu import HaplotypeMatrix
from pg_gpu.windowed_analysis import windowed_statistics_fused

hap = np.zeros((40, 40), dtype=np.int8)
hm = HaplotypeMatrix(hap, np.arange(40, dtype=np.int64) * 10)
edges = np.arange(0, 401, 200)
windowed_statistics_fused(hm, bp_bins=edges, statistics=("fst",), missing_data="exclude")
# TypeError: 'NoneType' object is not iterable

(should raise the value error indicating pops are required for Fst)

[andrewkern]

Good catch. Fixed this so pop1/pop2 check now runs at the start of _apply_exclude before the filter

[nspope]

Looks good, merging