From f2e75d313f82157567099cd6ca232ebcff5d03ee Mon Sep 17 00:00:00 2001 From: Murillo R Date: Thu, 29 Oct 2026 16:47:44 -0700 Subject: [PATCH] adding phased property to HaplotypeMatrix --- docs/source/api.rst | 1 + docs/source/changelog.rst | 13 ++ docs/source/quickstart.rst | 45 ++++-- pg_gpu/_gpu_genotype_prep.py | 3 +- pg_gpu/_utils.py | 2 + pg_gpu/genotype_matrix.py | 3 +- pg_gpu/haplotype_matrix.py | 65 +++++++-- pg_gpu/streaming_matrix.py | 22 ++- pg_gpu/windowed_analysis.py | 3 +- tests/test_gpu_genotype_prep.py | 15 +- tests/test_phasing.py | 233 ++++++++++++++++++++++++++++++++ tests/test_streaming.py | 2 +- 12 files changed, 368 insertions(+), 39 deletions(-) create mode 100644 tests/test_phasing.py diff --git a/docs/source/api.rst b/docs/source/api.rst index 644ad4ee..04d43f61 100644 --- a/docs/source/api.rst +++ b/docs/source/api.rst @@ -44,6 +44,7 @@ diploid genotypes; haploid and polyploid stores are rejected. .. autoclass:: pg_gpu.streaming_matrix.StreamingHaplotypeMatrix :members: iter_gpu_chunks, materialize, num_haplotypes, num_variants, chrom, chrom_start, chrom_end, sample_sets, align_bp, + phased, set_phased, compute_ld_statistics_gpu_single_pop, compute_ld_statistics_gpu_two_pops :show-inheritance: diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst index f5c3ac6a..470347c6 100644 --- a/docs/source/changelog.rst +++ b/docs/source/changelog.rst @@ -4,6 +4,19 @@ Changelog Unreleased ---------- +HaplotypeMatrix now tracks phasing +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Added ``phased`` attribute to ``HaplotypeMatrix`` (which defaults to +``False``). This tracks whether the haplotypes in the matrix +represent true phased haplotypes or not. Each loader exposes a +``phased`` keyword argument: ``from_vcf`` and ``from_zarr`` default +to ``False``; ``from_ts`` defaults to ``True``. It is necessary to +call ``set_phased()`` to change it after construction. + +A future release will use this flag to warn when an unphased matrix +is passed to a phase-dependent statistic. + Sites with more than two alleles are now handled correctly ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/docs/source/quickstart.rst b/docs/source/quickstart.rst index c302991a..930fbc89 100644 --- a/docs/source/quickstart.rst +++ b/docs/source/quickstart.rst @@ -7,22 +7,32 @@ Basic Usage Loading Data ~~~~~~~~~~~~ -The primary container, ``HaplotypeMatrix``, holds *phased* haplotypes -(one row per haplotype, two rows per diploid sample). When loaded from a -diploid VCF, the two haploid copies of each sample become two rows; the -loader does not check whether your data are actually phased, so if the -VCF contains unphased calls the resulting matrix is best thought of as a -random-phasing of the genotypes. For unphased / diploid-aware analyses, -use ``GenotypeMatrix`` (see "Phased to Unphased" below). +The primary data structure, ``HaplotypeMatrix``, holds haplotypes (one row +per haplotype, with variants as columns). When loaded from a diploid +VCF, the two haploid copies of each sample become two adjacent rows. + +The ``phased`` flag records whether the rows represent true phased +haplotypes or a pseudo-phased (arbitrary assignment of alleles within +each individual). The loaders default to ``phased=False`` for VCF and +Zarr, but ``from_ts`` defaults to ``phased=True`` for tree sequences. +You can override it at load time or flip it later with ``set_phased()``. +Note some analyses use genotype dosages (0, 1, 2 or -1 for missing), in which case +you should use ``GenotypeMatrix`` (see "HaplotypeMatrix vs GenotypeMatrix" below). .. code-block:: python from pg_gpu import HaplotypeMatrix # From VCF file (sample names are stored automatically). Diploid VCFs - # are loaded as 2*n_samples haplotypes -- treated as phased. + # are loaded as 2*n_samples haplotypes -- treated as unphased by default.. h = HaplotypeMatrix.from_vcf("data.vcf.gz") + # If the VCF is phased, you can set it at load time: + h = HaplotypeMatrix.from_vcf("data.vcf.gz", phased=True) + + # Or change it after loading: + h = HaplotypeMatrix.from_vcf("data.vcf.gz").set_phased() + # Load a region; both ends are inclusive, as in samtools # (requires tabix index) h = HaplotypeMatrix.from_vcf("data.vcf.gz", region="chr1:1000000-2000000") @@ -527,18 +537,21 @@ directly: statistics=('pi', 'theta_w', 'tajimas_d'), ) -Phased to Unphased -~~~~~~~~~~~~~~~~~~ +HaplotypeMatrix vs GenotypeMatrix +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -Use ``GenotypeMatrix`` for unphased diploid genotypes (entries 0 / 1 / 2 -counting alt alleles). Many functions auto-dispatch on input type, so -the same call works on either container: +In ``HaplotypeMatrix``, rows are haplotypes (0/1 per allele, two rows +per diploid individual). In ``GenotypeMatrix``, rows are individuals +(0/1/2 dosage). Use ``GenotypeMatrix`` only for statistics which need +dosage encoding. Many functions auto-dispatch on input type, so the same +call works on either a ``HaplotypeMatrix`` or a ``GenotypeMatrix``. .. code-block:: python from pg_gpu import GenotypeMatrix, distance_stats - # Convert from haploid (pairs consecutive haplotypes) + # Convert to dosages: each pair of consecutive haplotype rows is + # summed into one individual's 0/1/2 alternate-allele count. gm = GenotypeMatrix.from_haplotype_matrix(h) # Same functions work on both types @@ -549,6 +562,10 @@ the same call works on either container: # Distance distribution moments var, skew, kurt = distance_stats.dist_moments(gm) +Note that ``gm.to_haplotype_matrix()`` always produces an unphased +(``phased=False``) matrix, because the conversion assigns heterozygous +sites as ``(0, 1)`` arbitrarily. + Custom Theta via FrequencySpectrum ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/pg_gpu/_gpu_genotype_prep.py b/pg_gpu/_gpu_genotype_prep.py index 4e906c9b..89364fcf 100644 --- a/pg_gpu/_gpu_genotype_prep.py +++ b/pg_gpu/_gpu_genotype_prep.py @@ -30,6 +30,7 @@ def build_haplotype_matrix(gt, pos, *, + phased, chrom_start=None, chrom_end=None, sample_sets=None, n_total_sites=None, samples=None, accessible_mask=None, @@ -83,7 +84,7 @@ def build_haplotype_matrix(gt, pos, *, chrom_start=chrom_start, chrom_end=chrom_end, sample_sets=sample_sets, n_total_sites=n_total_sites, samples=samples, accessible_mask=accessible_mask, - allow_empty=allow_empty + allow_empty=allow_empty, phased=phased, ) diff --git a/pg_gpu/_utils.py b/pg_gpu/_utils.py index ffa8ac62..9e329f1f 100644 --- a/pg_gpu/_utils.py +++ b/pg_gpu/_utils.py @@ -78,6 +78,7 @@ def get_population_matrix(matrix, population: Union[str, list], UserWarning, stacklevel=2) cls = GenotypeMatrix if is_genotype else HaplotypeMatrix data = matrix.genotypes if is_genotype else matrix.haplotypes + hm_kw = {} if is_genotype else {'phased': matrix.phased} return cls( data[pop_indices, :], matrix.positions, @@ -87,6 +88,7 @@ def get_population_matrix(matrix, population: Union[str, list], n_total_sites=matrix.n_total_sites, allow_empty=allow_empty, **extra, + **hm_kw, ) diff --git a/pg_gpu/genotype_matrix.py b/pg_gpu/genotype_matrix.py index 0691368e..5241df06 100644 --- a/pg_gpu/genotype_matrix.py +++ b/pg_gpu/genotype_matrix.py @@ -445,7 +445,8 @@ def to_haplotype_matrix(self): sample_sets=new_sample_sets, samples=self.samples, fields=slice_fields(self.fields, xp.arange(n_var), - self._accessible_idx)) + self._accessible_idx), + phased=False) @classmethod def from_vcf(cls, path, include_invariant=False, accessible_bed=None, diff --git a/pg_gpu/haplotype_matrix.py b/pg_gpu/haplotype_matrix.py index 0b9c0201..e9a46261 100644 --- a/pg_gpu/haplotype_matrix.py +++ b/pg_gpu/haplotype_matrix.py @@ -212,10 +212,18 @@ def _decide_streaming_mode(zarr_path, region, streaming, pop_assignment, class HaplotypeMatrix: """Haplotype matrix for population genetics analysis. - Stores phased haplotype data with variant positions and population - labels. Supports GPU-accelerated computation of diversity, divergence, + Stores haplotype data with variant positions and population labels. + Supports GPU-accelerated computation of diversity, divergence, selection, and LD statistics. + Phasing + ------- + The ``phased`` flag records whether the haplotype rows represent true + phased haplotypes or not. It defaults to ``False`` on the raw + constructor and on ``from_vcf`` / ``from_zarr``, and to ``True`` on + ``from_ts`` (tree-sequence haplotypes are phased by definition). Use + ``set_phased()`` to change it to phased after construction. + Row order --------- Rows are gametes, and for diploid data the two gametes of one sample are @@ -263,6 +271,7 @@ def __init__(self, accessible_mask=None, fields: dict = None, allow_empty: bool = False, + phased: bool = False, ): # ``allow_empty`` is set only by the streaming chunk builder, which # yields zero-variant chunks over empty regions (e.g. an acrocentric @@ -305,6 +314,7 @@ def __init__(self, # requested at load time. Shape disambiguates per-variant vs # per-genotype. self.fields = dict(fields) if fields else {} + self._phased = bool(phased) if accessible_mask is not None and not isinstance(accessible_mask, AccessibleMask): accessible_mask = resolve_accessible_mask( @@ -420,6 +430,16 @@ def has_accessible_mask(self): """Whether an accessible site mask is attached.""" return self.accessible_mask is not None + @property + def phased(self): + """Whether the haplotype rows represent true phased haplotypes.""" + return self._phased + + def set_phased(self, phased: bool = True): + """Mark this matrix as phased or unphased. Returns ``self``.""" + self._phased = bool(phased) + return self + def set_accessible_mask(self, mask_or_path, chrom=None): """Attach an accessible site mask (non-destructive). @@ -541,7 +561,8 @@ def transfer_to_cpu(self): @classmethod def from_vcf(cls, path: str, region: str = None, samples: list = None, include_invariant: bool = False, - accessible_bed: str = None, fields: list = None): + accessible_bed: str = None, fields: list = None, + phased: bool = False): """Construct a HaplotypeMatrix from a VCF file. Parameters @@ -566,10 +587,14 @@ def from_vcf(cls, path: str, region: str = None, ``(n_var, n_samples)`` array. Tags missing from the VCF are warned and dropped silently. + phased : bool + Whether the data are phased haplotypes. Default ``False``: + the loader does not inspect the VCF phase bit. + Returns ------- HaplotypeMatrix - Phased haplotype data with sample names stored. + Haplotype data with sample names stored. """ from ._warnings import _maybe_memory_warn _maybe_memory_warn(path, region=region) @@ -627,7 +652,7 @@ def from_vcf(cls, path: str, region: str = None, n_total_sites = num_variants if include_invariant else None hm = cls(haplotypes, positions, chrom_start, chrom_end, n_total_sites=n_total_sites, samples=sample_names, - fields=qc_fields) + fields=qc_fields, phased=phased) if accessible_bed is not None: hm.set_accessible_mask(accessible_bed, chrom=chrom) return hm @@ -640,7 +665,8 @@ def from_zarr(cls, path: str, region: str = None, chunk_bp: int = 1_500_000, prefetch: int = 1, backend: str = "auto", - fields: list = None): + fields: list = None, + phased: bool = False): """Construct a HaplotypeMatrix from a Zarr store. Supports both VCZ (bio2zarr) and scikit-allel zarr layouts. @@ -713,6 +739,8 @@ def from_zarr(cls, path: str, region: str = None, ``(n_var,)`` array, ``call_`` as ``(n_var, n_samples)``. Tags missing from the store are warned and dropped silently. Not supported on the streaming path. + phased : bool + Whether the data are phased haplotypes. Default ``False``. Returns ------- @@ -739,6 +767,7 @@ def from_zarr(cls, path: str, region: str = None, path, region=region, pop_assignment=pop_assignment, chunk_bp=chunk_bp, prefetch=prefetch, backend=backend, accessible_bed=accessible_bed, + phased=phased, ) # 'auto' and 'never' both want eager when the matrix fits; 'auto' @@ -762,15 +791,16 @@ def from_zarr(cls, path: str, region: str = None, path, region=region, pop_assignment=pop_assignment, chunk_bp=chunk_bp, prefetch=prefetch, source=source, backend=backend, accessible_bed=accessible_bed, + phased=phased, ) return cls._build_eager(path, region=region, accessible_bed=accessible_bed, pop_assignment=pop_assignment, - fields=fields) + fields=fields, phased=phased) @classmethod def _build_eager(cls, path, *, region, accessible_bed, pop_assignment, - fields=None): + fields=None, phased=False): from .zarr_io import read_genotypes, read_qc_fields from ._gpu_genotype_prep import build_haplotype_matrix @@ -795,6 +825,7 @@ def _build_eager(cls, path, *, region, accessible_bed, pop_assignment, chrom_start=start, chrom_end=stop, samples=list(sample_names) if sample_names is not None else None, + phased=phased, ) if fields: @@ -831,7 +862,7 @@ def _build_eager(cls, path, *, region, accessible_bed, pop_assignment, @classmethod def _build_streaming(cls, path, *, region, pop_assignment, chunk_bp, prefetch, source=None, backend="auto", - accessible_bed=None): + accessible_bed=None, phased=False): from .streaming_matrix import ( HostChunkFetcher, KvikioChunkFetcher, StreamingHaplotypeMatrix, _pick_chunk_fetcher, @@ -862,6 +893,7 @@ def _build_streaming(cls, path, *, region, pop_assignment, chunk_bp, source, fetcher, chunk_bp=chunk_bp, prefetch=prefetch, accessible_mask=accessible_mask, + phased=phased, ) def to_zarr(self, zarr_path: str, format: str = 'vcz', @@ -1015,7 +1047,8 @@ def load_pop_file(self, pop_assignment, pops: list = None): def from_ts(cls, ts: tskit.TreeSequence, device: str = 'CPU', include_invariant: bool = False, accessible_bed: str = None, - chrom: str = None) -> 'HaplotypeMatrix': + chrom: str = None, + phased: bool = True) -> 'HaplotypeMatrix': """ Create a HaplotypeMatrix from a tskit.TreeSequence. @@ -1027,6 +1060,8 @@ def from_ts(cls, ts: tskit.TreeSequence, device: str = 'CPU', invariant sites analytically (no extra rows stored). accessible_bed: Path to a BED file defining accessible regions. chrom: Chromosome name for BED file filtering. + phased: Whether the data are phased. Default ``True`` + (tree-sequence haplotypes are phased by definition). Returns: HaplotypeMatrix: A new HaplotypeMatrix instance @@ -1072,7 +1107,7 @@ def from_ts(cls, ts: tskit.TreeSequence, device: str = 'CPU', n_total_sites = int(ts.sequence_length) if include_invariant else None hm = cls(haplotypes, positions, chrom_start, chrom_end, - n_total_sites=n_total_sites) + n_total_sites=n_total_sites, phased=phased) if accessible_bed is not None: hm.set_accessible_mask(accessible_bed, chrom=chrom) @@ -1156,7 +1191,8 @@ def __repr__(self): first_pos = self.positions[0] if self.positions.size > 0 else None last_pos = self.positions[-1] if self.positions.size > 0 else None return (f"HaplotypeMatrix(shape={self.shape}, " - f"first_position={first_pos}, last_position={last_pos})") + f"first_position={first_pos}, last_position={last_pos}, " + f"phased={self._phased})") def _empty_subset(self) -> "HaplotypeMatrix": """Build a copy of this matrix that holds no variants. @@ -1175,6 +1211,7 @@ def _empty_subset(self) -> "HaplotypeMatrix": result._hap_filtered = None result._pos_filtered = None result._device = self._device + result._phased = self._phased result._sample_sets = self._sample_sets result.chrom_start = self.chrom_start result.chrom_end = self.chrom_end @@ -1228,6 +1265,7 @@ def get_subset(self, positions) -> "HaplotypeMatrix": n_total_sites=self.n_total_sites, samples=self.samples, fields=slice_fields(self.fields, positions, self._accessible_idx), + phased=self._phased, ) def get_subset_from_range(self, low: int, high: int) -> "HaplotypeMatrix": @@ -1277,6 +1315,7 @@ def get_subset_from_range(self, low: int, high: int) -> "HaplotypeMatrix": samples=self.samples, accessible_mask=sliced_mask, fields=slice_fields(self.fields, indices, self._accessible_idx), + phased=self._phased, ) def _keep_sites(self, keep) -> "HaplotypeMatrix": @@ -1301,6 +1340,7 @@ def _keep_sites(self, keep) -> "HaplotypeMatrix": accessible_mask=self.accessible_mask, samples=self.samples, fields=slice_fields(self.fields, keep, self._accessible_idx), + phased=self._phased, ) def restrict_to_biallelic(self, *, warn_context=None) -> "HaplotypeMatrix": @@ -1701,6 +1741,7 @@ def filter(self, variants=None, genotypes=None, sample_sets=self._sample_sets, samples=self.samples, fields=new_fields, + phased=self._phased, ) def summarize_missing_data(self): diff --git a/pg_gpu/streaming_matrix.py b/pg_gpu/streaming_matrix.py index 9d526fe3..8c494f40 100644 --- a/pg_gpu/streaming_matrix.py +++ b/pg_gpu/streaming_matrix.py @@ -723,6 +723,23 @@ class StreamingHaplotypeMatrix(_StreamingMatrixBase): ``materialize`` a sub-region for repeated calls on the same data. """ + def __init__(self, source, fetcher, chunk_bp, prefetch, *, + phased=False, accessible_mask=None, n_total_sites=None): + super().__init__(source, fetcher, chunk_bp, prefetch, + accessible_mask=accessible_mask, + n_total_sites=n_total_sites) + self._phased = bool(phased) + + @property + def phased(self): + """Whether the haplotype rows represent true phased haplotypes.""" + return self._phased + + def set_phased(self, phased: bool = True): + """Mark this stream as phased or unphased. Returns ``self``.""" + self._phased = bool(phased) + return self + @property def num_haplotypes(self): return self._source.num_haplotypes @@ -749,13 +766,14 @@ def _build_chunk(self, gt, pos, **kwargs): # re-validation -- np.unique per population per chunk adds up to # minutes over a biobank-scale walk. sets = kwargs.pop('sample_sets', None) - m = build_haplotype_matrix(gt, pos, allow_empty = True, **kwargs) + m = build_haplotype_matrix(gt, pos, allow_empty=True, + phased=self._phased, **kwargs) if sets is not None: m._sample_sets = sets return m def _repr_sample_axis(self): - return f"num_haplotypes={self.num_haplotypes}" + return f"num_haplotypes={self.num_haplotypes}, phased={self._phased}" def compute_ld_statistics_gpu_single_pop(self, bp_bins, raw=False, chunk_size='auto'): diff --git a/pg_gpu/windowed_analysis.py b/pg_gpu/windowed_analysis.py index db0d2e4f..2574e2ff 100644 --- a/pg_gpu/windowed_analysis.py +++ b/pg_gpu/windowed_analysis.py @@ -2539,7 +2539,8 @@ def _expand(v): if need_winmat: win_mat = HaplotypeMatrix(matrix.haplotypes[:, s:e], - matrix.positions[s:e]) + matrix.positions[s:e], + phased=matrix.phased) if 'omega' in stat_arrays: # Feed already-biallelic input so omega's own restriction # drops nothing and stays silent; warned once for the scan diff --git a/tests/test_gpu_genotype_prep.py b/tests/test_gpu_genotype_prep.py index f6f2b694..c9f69578 100644 --- a/tests/test_gpu_genotype_prep.py +++ b/tests/test_gpu_genotype_prep.py @@ -36,7 +36,7 @@ class TestBuildHaplotypeMatrix: def test_layout_matches_load_pop_file_convention(self, synthetic_gt): gt, pos = synthetic_gt - hm = build_haplotype_matrix(gt, pos) + hm = build_haplotype_matrix(gt, pos, phased=False) haps_host = cp.asnumpy(hm.haplotypes) np.testing.assert_array_equal(haps_host, _reference_reshape(gt)) @@ -44,7 +44,7 @@ def test_missing_rows_preserved(self, synthetic_gt): # downstream kernels use 'include' / 'exclude' missing-data modes, # so the helper must keep -1 cells rather than dropping them. gt, pos = synthetic_gt - hm = build_haplotype_matrix(gt, pos) + hm = build_haplotype_matrix(gt, pos, phased=False) haps_host = cp.asnumpy(hm.haplotypes) # rows 3, 17, 42 should be all -1 across haplotypes. assert (haps_host[:, 3] == -1).all() @@ -53,14 +53,14 @@ def test_missing_rows_preserved(self, synthetic_gt): def test_positions_passed_through(self, synthetic_gt): gt, pos = synthetic_gt - hm = build_haplotype_matrix(gt, pos) + hm = build_haplotype_matrix(gt, pos, phased=False) pos_host = cp.asnumpy(hm.positions) np.testing.assert_array_equal(pos_host, pos) def test_kwargs_forwarded(self, synthetic_gt): gt, pos = synthetic_gt hm = build_haplotype_matrix( - gt, pos, + gt, pos, phased=True, chrom_start=1000, chrom_end=20000, samples=[f"s{i}" for i in range(50)], sample_sets={"all": list(range(100))}, @@ -72,20 +72,21 @@ def test_kwargs_forwarded(self, synthetic_gt): def test_accepts_device_inputs(self, synthetic_gt): gt, pos = synthetic_gt - hm = build_haplotype_matrix(cp.asarray(gt), cp.asarray(pos)) + hm = build_haplotype_matrix(cp.asarray(gt), cp.asarray(pos), + phased=False) np.testing.assert_array_equal(cp.asnumpy(hm.haplotypes), _reference_reshape(gt)) def test_mismatched_shapes_raise(self, synthetic_gt): gt, pos = synthetic_gt with pytest.raises(ValueError, match="disagree on n_var"): - build_haplotype_matrix(gt, pos[:50]) + build_haplotype_matrix(gt, pos[:50], phased=False) def test_wrong_ploidy_raises(self): gt = np.zeros((10, 5, 3), dtype=np.int8) pos = np.arange(10, dtype=np.int64) with pytest.raises(ValueError, match="must have shape"): - build_haplotype_matrix(gt, pos) + build_haplotype_matrix(gt, pos, phased=False) def _reference_dosage(gt): diff --git a/tests/test_phasing.py b/tests/test_phasing.py new file mode 100644 index 00000000..48cb1654 --- /dev/null +++ b/tests/test_phasing.py @@ -0,0 +1,233 @@ +"""Tests for the ``phased`` flag on HaplotypeMatrix and StreamingHaplotypeMatrix.""" + +import numpy as np +import pytest +import cupy as cp + +from pg_gpu.haplotype_matrix import HaplotypeMatrix +from pg_gpu.genotype_matrix import GenotypeMatrix +from pg_gpu._gpu_genotype_prep import build_haplotype_matrix + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _small_hm(phased=False): + """Minimal HaplotypeMatrix for flag-propagation tests.""" + hap = np.array([[0, 1, 0], [1, 0, 1], [0, 0, 1], [1, 1, 0]], dtype=np.int8) + pos = np.array([100, 200, 300], dtype=np.int64) + return HaplotypeMatrix(hap, pos, chrom_start=0, chrom_end=400, phased=phased) + + +# --------------------------------------------------------------------------- +# Constructor defaults and overrides +# --------------------------------------------------------------------------- + +class TestConstructorDefaults: + + def test_raw_constructor_default_false(self): + hm = _small_hm() + assert hm.phased is False + + def test_raw_constructor_explicit_true(self): + hm = _small_hm(phased=True) + assert hm.phased is True + + def test_from_vcf_default_false(self, simple_vcf_file): + hm = HaplotypeMatrix.from_vcf(simple_vcf_file) + assert hm.phased is False + + def test_from_vcf_override_true(self, simple_vcf_file): + hm = HaplotypeMatrix.from_vcf(simple_vcf_file, phased=True) + assert hm.phased is True + + def test_from_ts_default_true(self): + msprime = pytest.importorskip("msprime") + ts = msprime.sim_ancestry(4, sequence_length=1000, random_seed=1) + ts = msprime.sim_mutations(ts, rate=1e-3, random_seed=1) + hm = HaplotypeMatrix.from_ts(ts) + assert hm.phased is True + + def test_from_ts_override_false(self): + msprime = pytest.importorskip("msprime") + ts = msprime.sim_ancestry(4, sequence_length=1000, random_seed=1) + ts = msprime.sim_mutations(ts, rate=1e-3, random_seed=1) + hm = HaplotypeMatrix.from_ts(ts, phased=False) + assert hm.phased is False + + +# --------------------------------------------------------------------------- +# set_phased +# --------------------------------------------------------------------------- + +class TestSetPhased: + + def test_set_phased_flips_to_true(self): + hm = _small_hm(phased=False) + ret = hm.set_phased() + assert hm.phased is True + assert ret is hm + + def test_set_phased_explicit_false(self): + hm = _small_hm(phased=True) + hm.set_phased(False) + assert hm.phased is False + + def test_set_phased_returns_self(self): + hm = _small_hm() + assert hm.set_phased(True) is hm + + +# --------------------------------------------------------------------------- +# __repr__ +# --------------------------------------------------------------------------- + +class TestRepr: + + def test_phased_in_repr(self): + hm = _small_hm(phased=True) + assert "phased=True" in repr(hm) + + def test_unphased_in_repr(self): + hm = _small_hm(phased=False) + assert "phased=False" in repr(hm) + + +# --------------------------------------------------------------------------- +# Propagation through derived matrices +# --------------------------------------------------------------------------- + +class TestPropagation: + + @pytest.mark.parametrize("phased", [True, False]) + def test_get_subset(self, phased): + hm = _small_hm(phased=phased) + sub = hm.get_subset(cp.array([0, 1])) + assert sub.phased is phased + + @pytest.mark.parametrize("phased", [True, False]) + def test_get_subset_empty(self, phased): + hm = _small_hm(phased=phased) + sub = hm.get_subset(cp.array([], dtype=np.int64)) + assert sub.phased is phased + + @pytest.mark.parametrize("phased", [True, False]) + def test_get_subset_from_range(self, phased): + hm = _small_hm(phased=phased) + sub = hm.get_subset_from_range(50, 250) + assert sub.phased is phased + + @pytest.mark.parametrize("phased", [True, False]) + def test_get_subset_from_range_empty(self, phased): + hm = _small_hm(phased=phased) + sub = hm.get_subset_from_range(500, 600) + assert sub.phased is phased + + @pytest.mark.parametrize("phased", [True, False]) + def test_keep_sites(self, phased): + hm = _small_hm(phased=phased) + hm.transfer_to_gpu() + sub = hm._keep_sites(cp.array([0])) + assert sub.phased is phased + + @pytest.mark.parametrize("phased", [True, False]) + def test_restrict_to_biallelic(self, phased): + hm = _small_hm(phased=phased) + sub = hm.restrict_to_biallelic() + assert sub.phased is phased + + @pytest.mark.parametrize("phased", [True, False]) + def test_restrict_to_segregating(self, phased): + hm = _small_hm(phased=phased) + sub = hm.restrict_to_segregating() + assert sub.phased is phased + + @pytest.mark.parametrize("phased", [True, False]) + def test_filter(self, phased): + hm = _small_hm(phased=phased) + sub = hm.filter() + assert sub.phased is phased + + @pytest.mark.parametrize("phased", [True, False]) + def test_get_population_matrix(self, phased): + from pg_gpu._utils import get_population_matrix + hm = _small_hm(phased=phased) + hm.sample_sets = {"pop": [0, 1, 2, 3]} + pop_hm = get_population_matrix(hm, "pop") + assert pop_hm.phased is phased + + +# --------------------------------------------------------------------------- +# GM round-trip gives phased=False +# --------------------------------------------------------------------------- + +class TestGenotypeRoundTrip: + + def test_to_haplotype_matrix_always_unphased(self): + hm = _small_hm(phased=True) + gm = GenotypeMatrix.from_haplotype_matrix(hm) + hm2 = gm.to_haplotype_matrix() + assert hm2.phased is False + + +# --------------------------------------------------------------------------- +# build_haplotype_matrix requires phased +# --------------------------------------------------------------------------- + +class TestBuildHaplotypeMatrixRequired: + + def test_missing_phased_raises_typeerror(self): + gt = np.zeros((10, 5, 2), dtype=np.int8) + pos = np.arange(10, dtype=np.int64) + with pytest.raises(TypeError): + build_haplotype_matrix(gt, pos) + + def test_explicit_phased_works(self): + gt = np.zeros((10, 5, 2), dtype=np.int8) + pos = np.arange(10, dtype=np.int64) + hm = build_haplotype_matrix(gt, pos, phased=True) + assert hm.phased is True + + +# --------------------------------------------------------------------------- +# Zarr round-trip (eager and streaming) +# --------------------------------------------------------------------------- + +class TestZarrPhased: + + @pytest.fixture + def zarr_store(self, tmp_path, simple_vcf_file): + hm = HaplotypeMatrix.from_vcf(simple_vcf_file, phased=True) + path = str(tmp_path / "test.vcz.zarr") + hm.to_zarr(path, contig_name="1") + return path + + @pytest.mark.parametrize("phased", [True, False]) + def test_eager_from_zarr(self, zarr_store, phased): + hm = HaplotypeMatrix.from_zarr(zarr_store, streaming="never", + phased=phased) + assert hm.phased is phased + + @pytest.mark.parametrize("phased", [True, False]) + def test_streaming_iter_chunks(self, zarr_store, phased): + stream = HaplotypeMatrix.from_zarr(zarr_store, streaming="always", + phased=phased) + for _left, _right, chunk in stream.iter_gpu_chunks(): + assert chunk.phased is phased + + @pytest.mark.parametrize("phased", [True, False]) + def test_streaming_materialize(self, zarr_store, phased): + stream = HaplotypeMatrix.from_zarr(zarr_store, streaming="always", + phased=phased) + mat = stream.materialize() + assert mat.phased is phased + + def test_streaming_set_phased_propagates_to_chunks(self, zarr_store): + stream = HaplotypeMatrix.from_zarr(zarr_store, streaming="always", + phased=False) + assert stream.phased is False + stream.set_phased(True) + assert stream.phased is True + for _left, _right, chunk in stream.iter_gpu_chunks(): + assert chunk.phased is True diff --git a/tests/test_streaming.py b/tests/test_streaming.py index 5d6d98f7..0676e8f5 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -374,7 +374,7 @@ def pop_cols(self): return self.inner.pop_cols flaky = FlakySource(source) fetcher = HostChunkFetcher(flaky) smatrix = StreamingHaplotypeMatrix(flaky, fetcher, chunk_bp=5_000, - prefetch=1) + prefetch=1, phased=False) with pytest.raises(RuntimeError, match="synthetic producer error"): for _ in smatrix.iter_gpu_chunks(): pass