Skip to content

Make undefined Fst NaN, and fix multi-population single-stat column names #311

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
nspope merged 1 commit into from
Sep 28, 2026
Merged

Make undefined Fst NaN, and fix multi-population single-stat column names #311

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions docs/source/changelog.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -177,6 +177,10 @@ Read this section if you are comparing against older pg_gpu results.
Bug fixes
~~~~~~~~~

* ``fst_hudson``, ``fst_weir_cockerham``, ``fst_tskit``, and ``fst_nei``
returned ``0.0`` for an undefined ratio (no site with data in both
populations) while every windowed engine already used ``NaN`` for the
same case; they now return ``NaN`` too.
* Several ``HaplotypeMatrix``/``GenotypeMatrix`` methods that subset or
convert a matrix by variant (``get_subset``, ``get_subset_from_range``,
``restrict_to_biallelic``, ``restrict_to_segregating``,
Expand Down
27 changes: 14 additions & 13 deletions pg_gpu/divergence.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -199,7 +199,7 @@ def fst_hudson(haplotype_matrix: HaplotypeMatrix,
Returns
-------
float
Hudson's FST estimate
Hudson's FST estimate, or NaN if no site has data in both populations
"""
# Ensure data is on GPU if available
if haplotype_matrix.device == 'CPU':
Expand All @@ -209,7 +209,7 @@ def fst_hudson(haplotype_matrix: HaplotypeMatrix,
haplotype_matrix = haplotype_matrix.exclude_missing_sites(
populations=[pop1, pop2])
if haplotype_matrix.num_variants == 0:
return 0.0
return float('nan')

pop1_idx = _get_population_indices(haplotype_matrix, pop1)
pop2_idx = _get_population_indices(haplotype_matrix, pop2)
Expand All @@ -227,7 +227,7 @@ def fst_hudson(haplotype_matrix: HaplotypeMatrix,
valid_mask = den > 0
if cp.any(valid_mask):
return float((cp.sum(num[valid_mask]) / cp.sum(den[valid_mask])).get())
return 0.0
return float('nan')

@andrewkern andrewkern Sep 24, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

okay so now fst(method='hudson') disagrees with fst_nei and fst_tskit which still return 0.0 for this case - can we add those two as well to this PR?



def fst_tskit(haplotype_matrix: HaplotypeMatrix,
Expand Down Expand Up @@ -263,7 +263,7 @@ def fst_tskit(haplotype_matrix: HaplotypeMatrix,
Returns
-------
float
tskit's FST estimate
tskit's FST estimate, or NaN if no site has data in both populations
"""
if haplotype_matrix.device == 'CPU':
haplotype_matrix.transfer_to_gpu()
Expand All @@ -272,7 +272,7 @@ def fst_tskit(haplotype_matrix: HaplotypeMatrix,
haplotype_matrix = haplotype_matrix.exclude_missing_sites(
populations=[pop1, pop2])
if haplotype_matrix.num_variants == 0:
return 0.0
return float('nan')

pop1_idx = _get_population_indices(haplotype_matrix, pop1)
pop2_idx = _get_population_indices(haplotype_matrix, pop2)
Expand All @@ -289,7 +289,7 @@ def fst_tskit(haplotype_matrix: HaplotypeMatrix,
total = between_sum + within_sum
if total > 0:
return (between_sum - within_sum) / total
return 0.0
return float('nan')


def _pop_wc_stats(pop_haps, k):
Expand Down Expand Up @@ -432,7 +432,8 @@ def fst_weir_cockerham(haplotype_matrix,
Returns
-------
float
Weir & Cockerham's FST estimate
Weir & Cockerham's FST estimate, or NaN if no site has data in both
populations
"""

if hasattr(haplotype_matrix, 'device') and haplotype_matrix.device == 'CPU':
Expand All @@ -442,7 +443,7 @@ def fst_weir_cockerham(haplotype_matrix,
haplotype_matrix = haplotype_matrix.exclude_missing_sites(
populations=[pop1, pop2])
if haplotype_matrix.num_variants == 0:
return 0.0
return float('nan')

pop1_idx = _get_population_indices(haplotype_matrix, pop1)
pop2_idx = _get_population_indices(haplotype_matrix, pop2)
Expand All @@ -464,7 +465,7 @@ def fst_weir_cockerham(haplotype_matrix,
sum_abc = float(cp.sum(abc_site).get())
if sum_abc > 0:
return sum_a / sum_abc
return 0.0
return float('nan')


def fst_nei(haplotype_matrix: HaplotypeMatrix,
Expand Down Expand Up @@ -492,7 +493,7 @@ def fst_nei(haplotype_matrix: HaplotypeMatrix,
Returns
-------
float
Nei's GST estimate
Nei's GST estimate, or NaN if no site has data in both populations
"""
# Ensure data is on GPU if available
if haplotype_matrix.device == 'CPU':
Expand All @@ -502,7 +503,7 @@ def fst_nei(haplotype_matrix: HaplotypeMatrix,
haplotype_matrix = haplotype_matrix.exclude_missing_sites(
populations=[pop1, pop2])
if haplotype_matrix.num_variants == 0:
return 0.0
return float('nan')

pop1_idx = _get_population_indices(haplotype_matrix, pop1)
pop2_idx = _get_population_indices(haplotype_matrix, pop2)
Expand Down Expand Up @@ -536,13 +537,13 @@ def fst_nei(haplotype_matrix: HaplotypeMatrix,
valid_mask = (ht > 0) & (n1 > 0) & (n2 > 0)

if not cp.any(valid_mask):
return 0.0
return float('nan')

# Ratio-of-averages: sum(HT-HS) / sum(HT)
sum_ht = float(cp.sum(ht[valid_mask]).get())
sum_hs = float(cp.sum(hs[valid_mask]).get())
if sum_ht == 0:
return 0.0
return float('nan')
return (sum_ht - sum_hs) / sum_ht


Expand Down
36 changes: 33 additions & 3 deletions tests/test_divergence.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -422,11 +422,11 @@ def test_no_variation(self):
'pop2': list(range(20, 40))
}

# FST should be 0 (no variation to differentiate)
# FST is a ratio, so NaN when there's no data to normalize against.
fst_val = divergence.fst(matrix, 'pop1', 'pop2')
assert fst_val == 0.0
assert np.isnan(fst_val)

# Dxy should be 0
# Dxy is a sum, so 0 is well-defined.
dxy_val = divergence.dxy(matrix, 'pop1', 'pop2')
assert dxy_val == 0.0

Expand Down Expand Up @@ -458,3 +458,33 @@ def test_fst_tskit_masks_population_gap_sites(self):
reference = divergence.fst_tskit(reference_matrix, 'pop1', 'pop2')

assert np.isclose(gapped, reference, rtol=1e-9, atol=1e-12)

def test_fst_nan_when_undefined(self):
"""Every FST estimator returns NaN, not 0.0, when the ratio has no
denominator: no site with data in both populations."""
n_variants = 30
haplotypes = np.random.randint(0, 2, size=(20, n_variants))
positions = np.arange(n_variants) * 1000

matrix = HaplotypeMatrix(haplotypes, positions)
matrix.sample_sets = {
'pop1': list(range(10)),
'pop2': list(range(10, 20)),
}
# Every site missing in pop1 -- pop2 never has anything to pair with.
matrix.haplotypes[0:10, :] = -1

assert np.isnan(divergence.fst_hudson(matrix, 'pop1', 'pop2'))
assert np.isnan(divergence.fst_weir_cockerham(matrix, 'pop1', 'pop2'))
assert np.isnan(divergence.fst_tskit(matrix, 'pop1', 'pop2'))
assert np.isnan(divergence.fst_nei(matrix, 'pop1', 'pop2'))

# 'exclude' mode short-circuits to the same empty-matrix case.
assert np.isnan(divergence.fst_hudson(
matrix, 'pop1', 'pop2', missing_data='exclude'))
assert np.isnan(divergence.fst_weir_cockerham(
matrix, 'pop1', 'pop2', missing_data='exclude'))
assert np.isnan(divergence.fst_tskit(
matrix, 'pop1', 'pop2', missing_data='exclude'))
assert np.isnan(divergence.fst_nei(
matrix, 'pop1', 'pop2', missing_data='exclude'))
Loading