Merge pull request #19 from jeromekelleher/document-bgen-lims

Document bgen lims

Author: jeromekelleher

README.md

@@ -98,6 +98,27 @@ bgenix -g /tmp/bgen-mnt/sample.bgen -list
 fusermount3 -u /tmp/bgen-mnt
 ```
 
+#### Limitations: ploidy
+
+- **Mixed ploidy is not supported by `mount-bgen`.** The fixed-size BGEN
+  encoder used for random-access serving requires uniform ploidy across
+  every sample and variant in the view. Mounts whose region includes
+  mixed-ploidy chromosomes (typically X, Y, MT) open successfully and
+  serve `.sample` and `.bgen.bgi`, but the first `.bgen` read will fail
+  with `EIO`. Workaround: restrict the view to autosomes at mount time
+  (e.g. via the inherited `-r` / `-R` / `-t` / `-T` region filters), or
+  use the one-shot `vcztools view-bgen` CLI for full-file conversions
+  that include X / Y / MT — `view-bgen` uses the streaming
+  variable-size encoder which handles mixed ploidy correctly.
+- **Pure haploid VCZ is supported by `mount-bgen`** (the encoder emits a
+  uniform-haploid BGEN payload).
+- **`mount-plink` is diploid-only.** Pure haploid VCZ inputs (e.g.
+  mitochondrial-only stores) are rejected by the underlying encoder
+  with `EIO` on the first `.bed` read. Mixed-ploidy VCZ inputs serve
+  successfully, but haploid samples are encoded as homozygous for the
+  called allele — this matches the PLINK 1 BED format, which has no
+  haploid representation.
+
 ## Development
 
 ```bash

fs_tests/harness/fixtures.py

@@ -21,7 +21,7 @@
 @dataclasses.dataclass(frozen=True)
 class FixtureSpec:
     name: str
-    num_diploid_samples: int
+    num_samples: int
     sequence_length: int
     mutation_rate: float
     variants_chunk_size: int
@@ -37,7 +37,7 @@ class FixtureSpec:
 
 SMALL = FixtureSpec(
     name="small",
-    num_diploid_samples=10,
+    num_samples=10,
     sequence_length=10_000,
     mutation_rate=1e-3,
     variants_chunk_size=7,
@@ -47,7 +47,7 @@ class FixtureSpec:
 
 MEDIUM = FixtureSpec(
     name="medium",
-    num_diploid_samples=200,
+    num_samples=200,
     sequence_length=10_000_000,
     mutation_rate=1e-3,
     variants_chunk_size=1000,
@@ -57,7 +57,7 @@ class FixtureSpec:
 
 LARGE = FixtureSpec(
     name="large",
-    num_diploid_samples=500,
+    num_samples=500,
     sequence_length=20_000_000,
     mutation_rate=2e-3,
     variants_chunk_size=2000,
@@ -81,7 +81,7 @@ def get_or_build(
         "building fixture %s under %s (samples=%d seq_len=%d mut_rate=%g)",
         spec.name,
         spec_dir,
-        spec.num_diploid_samples,
+        spec.num_samples,
         spec.sequence_length,
         spec.mutation_rate,
     )
@@ -90,7 +90,7 @@ def get_or_build(
     build_started = time.monotonic()
     fixture = helpers.simulate_vcz(
         out_dir=spec_dir,
-        num_diploid_samples=spec.num_diploid_samples,
+        num_samples=spec.num_samples,
         sequence_length=spec.sequence_length,
         mutation_rate=spec.mutation_rate,
         variants_chunk_size=spec.variants_chunk_size,

tests/conftest.py

@@ -26,7 +26,7 @@ def fx_small_vcz(fx_session_dir) -> helpers.VczFixture:
     """Tiny VCZ: ~10 diploid samples, multiple variant chunks for boundary testing."""
     return helpers.simulate_vcz(
         out_dir=fx_session_dir / "small",
-        num_diploid_samples=10,
+        num_samples=10,
         sequence_length=10_000,
         mutation_rate=1e-3,
         variants_chunk_size=7,
@@ -41,7 +41,7 @@ def fx_medium_vcz(fx_session_dir) -> helpers.VczFixture:
     """Mid-sized VCZ: ~50 diploid samples × hundreds of variants for app tests."""
     return helpers.simulate_vcz(
         out_dir=fx_session_dir / "medium",
-        num_diploid_samples=50,
+        num_samples=50,
         sequence_length=200_000,
         mutation_rate=1e-4,
         variants_chunk_size=50,
@@ -56,7 +56,7 @@ def fx_singleton_vcz(fx_session_dir) -> helpers.VczFixture:
     """Single-variant VCZ. Tiny edge case: chunk count 1, file size minimal."""
     return helpers.simulate_vcz(
         out_dir=fx_session_dir / "singleton",
-        num_diploid_samples=4,
+        num_samples=4,
         sequence_length=1000,
         mutation_rate=2e-3,
         variants_chunk_size=10,
@@ -76,7 +76,7 @@ def fx_multiallelic_vcz(fx_session_dir) -> helpers.VczFixture:
     """
     return helpers.simulate_vcz(
         out_dir=fx_session_dir / "multiallelic",
-        num_diploid_samples=10,
+        num_samples=10,
         sequence_length=10_000,
         mutation_rate=1e-2,
         variants_chunk_size=10,
@@ -85,3 +85,37 @@ def fx_multiallelic_vcz(fx_session_dir) -> helpers.VczFixture:
         seed=53,
         biallelic=False,
     )
+
+
+@pytest.fixture(scope="session")
+def fx_haploid_vcz(fx_session_dir) -> helpers.VczFixture:
+    """Tiny haploid VCZ: ``call_genotype`` shape ``(V, S, 1)``."""
+    return helpers.simulate_vcz(
+        out_dir=fx_session_dir / "haploid",
+        num_samples=10,
+        sequence_length=10_000,
+        mutation_rate=1e-3,
+        variants_chunk_size=7,
+        samples_chunk_size=10,
+        name="haploid",
+        seed=67,
+        ploidy=1,
+    )
+
+
+@pytest.fixture(scope="session")
+def fx_mixed_ploidy_vcz(fx_small_vcz, fx_session_dir) -> helpers.VczFixture:
+    """Diploid-shaped VCZ with half the samples flipped to effectively haploid.
+
+    Derived from :func:`fx_small_vcz` by writing ``-2`` into
+    ``call_genotype[:, :S//2, 1]``. The resulting ``call_genotype``
+    still has shape ``(V, S, 2)`` but the BgenEncoder's chunk-level
+    uniform-ploidy check raises ``NotImplementedError`` on the first
+    read, mirroring real X / Y / MT data alongside diploid autosomes.
+    """
+    return helpers.mutate_to_mixed_ploidy(
+        fx_small_vcz.path,
+        fx_session_dir / "mixed_ploidy",
+        name="mixed_ploidy",
+        haploid_sample_fraction=0.5,
+    )

tests/helpers.py

@@ -1,6 +1,7 @@
 """Test helpers shared across the biofuse test suite."""
 
 import pathlib
+import shutil
 from dataclasses import dataclass
 
 import bio2zarr.tskit as bio2zarr_tskit
@@ -61,25 +62,29 @@ def _keep_first_mutation_per_site(ts):
 def simulate_vcz(
     out_dir: pathlib.Path,
     *,
-    num_diploid_samples: int,
+    num_samples: int,
     sequence_length: float,
     mutation_rate: float,
     variants_chunk_size: int,
     samples_chunk_size: int,
     name: str = "sim",
     seed: int = 1,
     biallelic: bool = True,
+    ploidy: int = 2,
 ) -> VczFixture:
     """Simulate a tree sequence and convert it to VCZ on disk.
 
     Returns a VczFixture pointing at a directory-format VCZ. The caller
     owns cleanup of out_dir. By default the fixture is biallelic: any
     site with multiple mutations keeps only its first mutation. Pass
     ``biallelic=False`` to keep recurrent mutations and exercise the
-    multi-allelic rejection path.
+    multi-allelic rejection path. ``ploidy=1`` yields a VCZ with
+    ``call_genotype`` of shape ``(V, num_samples, 1)``; the default
+    ``ploidy=2`` yields ``(V, num_samples, 2)``.
     """
     ts = msprime.sim_ancestry(
-        samples=num_diploid_samples,
+        samples=num_samples,
+        ploidy=ploidy,
         sequence_length=sequence_length,
         recombination_rate=1e-8,
         random_seed=seed,
@@ -102,7 +107,6 @@ def simulate_vcz(
     )
 
     num_variants = ts.num_sites
-    num_samples = ts.num_samples // 2
     return VczFixture(
         path=vcz_path,
         num_samples=num_samples,
@@ -111,3 +115,41 @@ def simulate_vcz(
         variants_chunk_size=variants_chunk_size,
         samples_chunk_size=samples_chunk_size,
     )
+
+
+def mutate_to_mixed_ploidy(
+    src_vcz: pathlib.Path,
+    dst_dir: pathlib.Path,
+    *,
+    name: str = "mixed",
+    haploid_sample_fraction: float = 0.5,
+) -> VczFixture:
+    """Copy a diploid VCZ and convert half its samples to effectively haploid.
+
+    Writes the haploid sentinel ``-2`` into slot 1 of ``call_genotype`` for
+    the first ``int(S * haploid_sample_fraction)`` samples across all
+    variants. This is the canonical VCZ representation for samples with
+    lower ploidy within a fixed-width genotype array — it is what an X /
+    Y / MT chromosome looks like alongside diploid autosomes in the same
+    store.
+    """
+    dst_dir.mkdir(parents=True, exist_ok=True)
+    dst_vcz = dst_dir / f"{name}.vcz"
+    shutil.copytree(src_vcz, dst_vcz)
+
+    store = zarr.open(dst_vcz, mode="r+")
+    call_genotype = store["call_genotype"]
+    num_variants, num_samples, _ = call_genotype.shape
+    num_haploid = int(num_samples * haploid_sample_fraction)
+    arr = call_genotype[:]
+    arr[:, :num_haploid, 1] = -2
+    call_genotype[:] = arr
+
+    return VczFixture(
+        path=dst_vcz,
+        num_samples=num_samples,
+        num_variants=num_variants,
+        num_biallelic_sites=_count_biallelic_sites(dst_vcz),
+        variants_chunk_size=call_genotype.chunks[0],
+        samples_chunk_size=call_genotype.chunks[1],
+    )

tests/test_bgen_apps.py

@@ -7,6 +7,8 @@
 encoder's bytes.
 """
 
+import contextlib
+import errno
 import os
 import pathlib
 import shutil
@@ -81,6 +83,22 @@ async def _arun(cmd) -> None:
     await trio.run_process(cmd, capture_stdout=True, capture_stderr=True, check=True)
 
 
+@contextlib.asynccontextmanager
+async def _mount_bgen(tmp_path, vcz):
+    """Mount ``vcz`` as a BGEN fileset; yield ``(mnt, basename)``."""
+    mnt = tmp_path / "mnt"
+    mnt.mkdir()
+    basename = vcz.path.stem
+    sock_path = tmp_path / "bgen.sock"
+    async with await encoder_client.EncoderClient.start(
+        str(vcz.path), sock_path, formats.BGEN_SPEC
+    ) as client:
+        ops = encoder_ops.EncoderOps(client, basename, formats.BGEN_SPEC)
+        async with fuse_adapter.mount(ops, str(mnt)):
+            await _wait_for_mount(mnt)
+            yield mnt, basename
+
+
 class TestBgenBytes:
     async def test_full_bgen_matches_encoder(self, fx_mounted_bgen):
         mnt, basename, expected, _ = fx_mounted_bgen
@@ -207,6 +225,56 @@ async def test_statfs_via_os_statvfs(self, fx_mounted_bgen):
         assert info.f_namemax >= 255
 
 
+class TestHaploidAndMixedPloidy:
+    """Pins the mount-level behaviour for non-diploid VCZ inputs.
+
+    See ``tests/test_formats.py::TestBgenSpecHaploid`` and
+    ``TestBgenSpecMixedPloidy`` for the encoder-layer contracts these
+    end-to-end tests exercise.
+    """
+
+    async def test_haploid_bgen_matches_encoder(self, tmp_path, fx_haploid_vcz):
+        """Pure haploid is fully supported by BgenEncoder; the mounted
+        ``.bgen`` bytes must match the in-process encoder's output."""
+        expected = _encoder_bytes(fx_haploid_vcz.path)
+        async with _mount_bgen(tmp_path, fx_haploid_vcz) as (mnt, basename):
+            data = await trio.to_thread.run_sync((mnt / f"{basename}.bgen").read_bytes)
+        assert data == expected
+
+    @needs_plink2
+    async def test_haploid_bgen_plink2_freq(self, tmp_path, fx_haploid_vcz):
+        """``plink2 --bgen`` reads the mounted haploid view end-to-end."""
+        async with _mount_bgen(tmp_path, fx_haploid_vcz) as (mnt, basename):
+            out = tmp_path / "p2_freq_hap"
+            await _arun(
+                [
+                    PLINK2,
+                    "--bgen",
+                    str(mnt / f"{basename}.bgen"),
+                    "ref-first",
+                    "--sample",
+                    str(mnt / f"{basename}.sample"),
+                    "--freq",
+                    "--out",
+                    str(out),
+                ]
+            )
+            assert out.with_suffix(".afreq").exists()
+
+    async def test_mixed_ploidy_bgen_read_raises_eio(
+        self, tmp_path, fx_mixed_ploidy_vcz
+    ):
+        """The fixed-size BGEN encoder cannot represent mixed ploidy.
+        The mount and the static sidecars (``.sample``, ``.bgen.bgi``)
+        still serve, but the first ``.bgen`` read fails with EIO."""
+        async with _mount_bgen(tmp_path, fx_mixed_ploidy_vcz) as (mnt, basename):
+            await trio.to_thread.run_sync((mnt / f"{basename}.sample").read_bytes)
+            await trio.to_thread.run_sync((mnt / f"{basename}.bgen.bgi").read_bytes)
+            with pytest.raises(OSError, match="Input/output error") as excinfo:
+                await trio.to_thread.run_sync((mnt / f"{basename}.bgen").read_bytes)
+            assert excinfo.value.errno == errno.EIO
+
+
 class TestAccessTrace:
     """Captures access patterns for inspection in phase-2 design."""