Fixes to max_daf returns, under various "empty" scenarios

docs/dev/pipeline_map.toml

@@ -768,9 +768,9 @@ inputs = ["HaplotypeMatrix"]
 streaming = "generic_error"
 missing_data = ["include", "exclude"]
 span = "none"
-undefined = "0.0 with no derived allele; -1.0 under 'exclude' when no complete site carries a derived allele"
+undefined = "0.0 with no derived allele or no variant; under 'exclude', NaN where a site existed but none is complete"
 parallel = ["windowed_analysis._windowed_thetas_scatter", "windowed_analysis._fused_windowed_kernel_v2", "windowed_analysis.windowed_statistics_fused", "windowed_analysis.windowed_statistics_fused_chunked"]
-notes = "Scalar 'include' keeps sites with n_valid = 1; the windowed engines need n_valid >= 2."
+notes = "'include' needs n_valid >= 2 at a site, matching pi/theta_w and both windowed engines; 'exclude' tracks each window's pre-filter variant count to tell the two 0.0/NaN cases apart."
 
 [functions."diversity.haplotype_count"]
 # signature: (haplotype_matrix, population, missing_data)

docs/source/changelog.rst

@@ -199,6 +199,14 @@ Read this section if you are comparing against older pg_gpu results.
   - ``mu_var``: a window length that is not positive;
   - ``haplotype_diversity``: fewer than two haplotypes.
 
+  ``max_daf(missing_data='exclude')`` returned ``-1.0``, not even a
+  plausible-looking real value, when no site was complete; it is NaN now.
+  The windowed engines (``windowed_analysis``, ``windowed_statistics_fused``,
+  ``windowed_statistics_fused_chunked``) had the same gap one level up: a
+  window whose sites were all excluded gave ``0.0`` rather than NaN,
+  indistinguishable from a window with no variant to begin with. Both now
+  agree with the scalar function.
+
   With the default ``sigma_d2`` estimator, ``zns`` and ``omega`` leave
   out a pair that they cannot estimate (fewer than four haplotypes
   called at both sites, or a pi2 estimate of 0). Before, ``zns`` counted
@@ -212,6 +220,13 @@ Read this section if you are comparing against older pg_gpu results.
   ``omega``, ``mu_ld`` and ``dist_*`` columns skipped every window with
   fewer than four variants. Now a window with one to three variants
   gets the value that the scalar function gives for its variants.
+* The scalar ``max_daf(missing_data='include')`` now needs at least two
+  calls at a site to count it, matching ``pi``, ``theta_w``, and both
+  windowed engines, which already worked this way. A single stray call
+  of the derived allele used to be enough to set the whole-region value
+  to that call's own frequency, as high as 1.0; it now takes the best
+  value among sites with their own pair, the same answer the windowed
+  path already gave.
 * In ``windowed_statistics``, ``singletons`` now counts only alternate
   alleles seen once (it used to also count reference alleles seen
   once), and ``segregating_sites`` counts mutations.

pg_gpu/diversity.py

@@ -1430,13 +1430,16 @@ def max_daf(haplotype_matrix: HaplotypeMatrix,
     haplotype_matrix : HaplotypeMatrix
     population : str or list, optional
     missing_data : str
-        'include' - per-site n_valid for frequency
+        'include' - a site needs at least two calls to contribute, same as
+        ``pi``/``theta_w``/``heterozygosity_expected``; fewer contribute
+        nothing, since a frequency from one call is not a useful estimate.
         'exclude' - only sites with no missing data
 
     Returns
     -------
     float
-        Maximum DAF in [0, 1].
+        Maximum DAF in [0, 1], or NaN when ``'exclude'`` leaves no complete
+        site.
     """
 
     if population is not None:
@@ -1457,10 +1460,12 @@ def max_daf(haplotype_matrix: HaplotypeMatrix,
     freq = ac[:, 1:].astype(cp.float64) / cp.maximum(n_valid[:, None], 1.0)
     present = ac[:, 1:] > 0
     if missing_data == 'exclude':
-        keep = present & (n_valid_i[:, None] == matrix.haplotypes.shape[0])
-        freq = cp.where(keep, freq, -1.0)
+        complete = n_valid_i == matrix.haplotypes.shape[0]
+        if not bool(cp.any(complete)):
+            return float('nan')
+        freq = cp.where(present & complete[:, None], freq, 0.0)
     else:
-        freq = cp.where(present & (n_valid_i[:, None] > 0), freq, 0.0)
+        freq = cp.where(present & (n_valid_i[:, None] >= 2), freq, 0.0)
     return float(cp.max(freq).get())
 
 

pg_gpu/windowed_analysis.py

@@ -847,6 +847,7 @@ def _windowed_thetas_scatter(haplotype_matrix, window_size, step_size,
     if chrom_end <= chrom_start:
         return pd.DataFrame()
 
+    pos_cpu_pre_exclude = pos_cpu
     if missing_data == 'exclude':
         # Keep the already-anchored window grid even when exclusion removes
         # every site: the scatter machinery renders empty windows as 0/NaN
@@ -883,6 +884,13 @@ def scatter_sum(values):
     # Compute requested per-variant contributions and scatter
     n_variants_per_window = np.bincount(
         k_safe[contains], minlength=n_windows)
+    # A window that lost every site to 'exclude' is undefined for max_daf
+    # (NaN); one with no variant to begin with stays 0.0.
+    excluded_to_empty = None
+    if missing_data == 'exclude' and 'max_daf' in statistics:
+        lo = np.searchsorted(pos_cpu_pre_exclude, win_starts, side='left')
+        hi = np.searchsorted(pos_cpu_pre_exclude, win_stops, side='left')
+        excluded_to_empty = (hi > lo) & (n_variants_per_window == 0)
     results = _init_window_results(
         chrom, win_starts, win_stops,
         n_variants=n_variants_per_window)
@@ -955,6 +963,8 @@ def scatter_sum(values):
     if 'singletons' in stats_set:
         results['singletons'] = sing_count.get()
     if 'max_daf' in stats_set:
+        if excluded_to_empty is not None:
+            max_daf_arr = np.where(excluded_to_empty, np.nan, max_daf_arr)
         results['max_daf'] = max_daf_arr
 
     # Composite stats from raw theta sums
@@ -2265,6 +2275,14 @@ def windowed_statistics_fused(haplotype_matrix: HaplotypeMatrix,
     if haplotype_matrix.device == 'CPU':
         haplotype_matrix.transfer_to_gpu()
 
+    # Pre-filter positions: a window exclusion emptied (NaN for max_daf)
+    # looks identical to one with no variant to begin with (0.0) once
+    # filtering below removes the distinction.
+    orig_missing_data = missing_data
+    orig_positions_gpu = haplotype_matrix.positions
+    if not isinstance(orig_positions_gpu, cp.ndarray):
+        orig_positions_gpu = cp.asarray(orig_positions_gpu)
+
     haplotype_matrix, missing_data = _apply_exclude(
         haplotype_matrix, missing_data, statistics, population, pop1, pop2)
 
@@ -2317,6 +2335,12 @@ def windowed_statistics_fused(haplotype_matrix: HaplotypeMatrix,
     else:
         raise ValueError('bp_bins and _win_starts/_win_stops must not both be None.')
 
+    excluded_to_empty = None
+    if orig_missing_data == 'exclude' and 'max_daf' in statistics:
+        orig_lo = cp.searchsorted(orig_positions_gpu, ws_gpu, side='left')
+        orig_hi = cp.searchsorted(orig_positions_gpu, we_gpu, side='left')
+        excluded_to_empty = ((orig_hi > orig_lo) & (win_stop == win_start)).get()
+
     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
     # Raw variant-index span is the fallback n_variants; single-pop kernels
@@ -2399,6 +2423,8 @@ def windowed_statistics_fused(haplotype_matrix: HaplotypeMatrix,
                                         per_base)
 
         if 'max_daf' in statistics:
+            if excluded_to_empty is not None:
+                max_daf = np.where(excluded_to_empty, np.nan, max_daf)
             results['max_daf'] = max_daf
 
     # two-pop stats via fused kernel
@@ -2740,6 +2766,14 @@ def windowed_statistics_fused_chunked(haplotype_matrix: HaplotypeMatrix,
     if haplotype_matrix.device == 'CPU':
         haplotype_matrix.transfer_to_gpu()
 
+    # Pre-filter positions: a window exclusion emptied (NaN for max_daf)
+    # looks identical to one with no variant to begin with (0.0) once
+    # filtering below removes the distinction.
+    orig_missing_data = missing_data
+    orig_positions_gpu = haplotype_matrix.positions
+    if not isinstance(orig_positions_gpu, cp.ndarray):
+        orig_positions_gpu = cp.asarray(orig_positions_gpu)
+
     haplotype_matrix, missing_data = _apply_exclude(
         haplotype_matrix, missing_data, statistics, population, pop1, pop2)
 
@@ -2785,6 +2819,12 @@ def windowed_statistics_fused_chunked(haplotype_matrix: HaplotypeMatrix,
     else:
         raise ValueError('bp_bins and _win_starts/_win_stops must not both be None.')
 
+    excluded_to_empty = None
+    if orig_missing_data == 'exclude' and 'max_daf' in statistics:
+        orig_lo = cp.searchsorted(orig_positions_gpu, ws_gpu, side='left')
+        orig_hi = cp.searchsorted(orig_positions_gpu, we_gpu, side='left')
+        excluded_to_empty = ((orig_hi > orig_lo) & (win_stop == win_start)).get()
+
     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
     # Raw variant-index span is the fallback n_variants; single-pop kernels
@@ -2910,6 +2950,8 @@ def windowed_statistics_fused_chunked(haplotype_matrix: HaplotypeMatrix,
                                         per_base)
 
         if 'max_daf' in statistics:
+            if excluded_to_empty is not None:
+                max_daf_arr = np.where(excluded_to_empty, np.nan, max_daf_arr)
             results['max_daf'] = max_daf_arr
 
     # ── Two-pop stats via chunked fused kernel ───────────────────────────

tests/test_diversity.py

@@ -668,6 +668,68 @@ def test_haplotype_diversity_with_population(self):
         assert abs(h_div_all - h_div_allel_all) < 1e-10
 
 
+class TestMaxDafExclude:
+    """max_daf(missing_data='exclude') must give NaN, not a sentinel value
+    outside [0, 1], when no site is complete."""
+
+    def test_no_complete_site_is_nan(self):
+        # Every site has exactly one missing call, so none is complete.
+        h = np.array([[0, 1, -1], [1, -1, 0], [-1, 0, 1], [1, 1, 0]],
+                     dtype=np.int8)
+        hm = HaplotypeMatrix(h, np.array([1, 11, 21]), 1, 30)
+        assert np.isnan(diversity.max_daf(hm, missing_data='exclude'))
+
+    def test_complete_but_monomorphic_site_is_zero_not_nan(self):
+        # One complete site, but it carries no derived allele -- there is
+        # data to measure, unlike the no-complete-site case above, so the
+        # correct max DAF is 0.0.
+        h = np.array([[0, 0, -1], [0, 1, 0], [0, -1, 1], [0, 1, 0]],
+                     dtype=np.int8)
+        hm = HaplotypeMatrix(h, np.array([1, 11, 21]), 1, 30)
+        assert diversity.max_daf(hm, missing_data='exclude') == 0.0
+
+
+class TestMaxDafIncludeNeedsOwnPair:
+    """max_daf(missing_data='include') must ignore a site with only one
+    call, matching pi/theta_w/heterozygosity_expected and both windowed
+    engines, instead of letting a single stray call set the whole-region
+    max to that call's own frequency."""
+
+    def _matrix(self):
+        # Site 0: every haplotype called, DAF 3/6 = 0.5. Site 1: one call
+        # (derived), n_valid=1 -- must be ignored, not read as DAF 1.0.
+        # Site 2: 5 calls, DAF 1/5 = 0.2 -- the best of the eligible sites.
+        h = np.array([[1, 1, 1],
+                      [-1, 1, 0],
+                      [-1, 0, 0],
+                      [-1, 0, 0],
+                      [-1, 0, 0],
+                      [-1, 0, 0]], dtype=np.int8)
+        return HaplotypeMatrix(h, np.array([1, 11, 21]), 1, 30)
+
+    def test_single_call_site_ignored(self):
+        hm = self._matrix()
+        np.testing.assert_allclose(diversity.max_daf(hm), 1.0 / 3.0,
+                                   rtol=1e-12)
+
+    def test_matches_scatter_and_fused_engines(self):
+        from pg_gpu.windowed_analysis import (
+            windowed_analysis, windowed_statistics_fused,
+        )
+        hm = self._matrix()
+        ref = diversity.max_daf(hm)
+
+        scatter = windowed_analysis(hm, window_size=30,
+                                    statistics=['max_daf'])
+        np.testing.assert_allclose(scatter['max_daf'].iloc[0], ref, rtol=1e-12)
+
+        hm.transfer_to_gpu()
+        bp_bins = np.array([0.0, 30.0])
+        fused = windowed_statistics_fused(hm, bp_bins=bp_bins,
+                                          statistics=('max_daf',))
+        np.testing.assert_allclose(fused['max_daf'][0], ref, rtol=1e-12)
+
+
 def test_diversity_stats_key_order_matches_request():
     """The result dict lists keys in the caller's order, whatever the
     process hash seed, so DataFrame columns built from it are stable."""

tests/test_windowed_exclude.py

@@ -235,6 +235,46 @@ def test_exclude_keeps_the_grid_when_an_end_site_is_incomplete():
     np.testing.assert_array_equal(inc["end"], exc["end"])
 
 
+def test_max_daf_nan_when_exclude_empties_a_window_with_data():
+    # Window 0 has two sites, each with a missing call -- 'exclude' drops
+    # both, so nothing is left to measure: NaN, not 0.0. Window 1 has no
+    # variant at all, a different case (nothing was ever there), and stays
+    # 0.0. The scatter and fused engines both filter the whole matrix before
+    # windowing, which loses the distinction unless each window's pre-filter
+    # variant count is checked too.
+    hap = np.array([[0, 1],
+                    [1, -1],
+                    [-1, 0],
+                    [1, 1]], dtype=np.int8)
+    pos = np.array([10, 20])
+    hm = HaplotypeMatrix(hap, pos, 0, 200)
+
+    scatter = windowed_analysis(hm, window_size=100, step_size=100,
+                                statistics=['max_daf'], missing_data='exclude',
+                                span_normalize=False)
+    assert scatter['n_variants'].tolist() == [0, 0]
+    assert np.isnan(scatter['max_daf'].iloc[0])
+    assert scatter['max_daf'].iloc[1] == 0.0
+
+    bp = np.array([0.0, 100.0, 200.0])
+    fused = windowed_statistics_fused(hm, bp_bins=bp, statistics=('max_daf',),
+                                      missing_data='exclude', per_base=False)
+    np.testing.assert_array_equal(fused['n_variants'], [0, 0])
+    assert np.isnan(fused['max_daf'][0])
+    assert fused['max_daf'][1] == 0.0
+
+    chunked = windowed_statistics_fused_chunked(
+        hm, bp_bins=bp, statistics=('max_daf',), missing_data='exclude',
+        per_base=False)
+    np.testing.assert_array_equal(chunked['n_variants'], [0, 0])
+    assert np.isnan(chunked['max_daf'][0])
+    assert chunked['max_daf'][1] == 0.0
+
+    # Pinned against the scalar function directly, not just a mutual check.
+    window0 = HaplotypeMatrix(hap, pos, 0, 100)
+    assert np.isnan(diversity.max_daf(window0, missing_data='exclude'))
+
+
 def test_streaming_exclude_matches_scalar(tmp_path):
     hm = simulate_hm(n_samples=20, seq_length=50_000, seed=42, mutation_model='binary')
     hap = cp.asnumpy(hm.haplotypes).astype(np.int8)