Adding phased property to HaplotypeMatrix

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:

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
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

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
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

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,
     )
 
 

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,
     )
 
 

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,

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_<tag>`` 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):