Add plink contig info

Author: benjeffery

bio2zarr/plink.py

@@ -5,7 +5,7 @@
 import numpy as np
 import zarr
 
-from bio2zarr import constants, schema, writer
+from bio2zarr import constants, core, schema, writer
 
 logger = logging.getLogger(__name__)
 
@@ -18,6 +18,9 @@ def __init__(self, path):
         self.samples = [schema.Sample(id=sample) for sample in self.bed.iid]
         self.num_samples = len(self.samples)
         self.root_attrs = {}
+        self.contigs = [
+            schema.Contig(id=str(chrom)) for chrom in np.unique(self.bed.chromosome)
+        ]
 
     def iter_alleles(self, start, stop, num_alleles):
         ref_field = self.bed.allele_1
@@ -49,6 +52,11 @@ def iter_genotypes(self, shape, start, stop):
             gt[values == 2] = [0, 0]  # Homozygous REF (0 in PLINK)
             yield gt, phased
 
+    def iter_contig(self, start, stop):
+        chrom_to_contig_index = {contig.id: i for i, contig in enumerate(self.contigs)}
+        for chrom in self.bed.chromosome[start:stop]:
+            yield chrom_to_contig_index[str(chrom)]
+
 
 def generate_schema(
     bed,
@@ -88,6 +96,15 @@ def generate_schema(
             chunks=[variants_chunk_size, 2],
             description=None,
         ),
+        schema.ZarrArraySpec.new(
+            vcf_field=None,
+            name="variant_contig",
+            dtype=core.min_int_dtype(0, len(np.unique(bed.chromosome))),
+            shape=[m],
+            dimensions=["variants"],
+            chunks=[variants_chunk_size],
+            description="Contig/chromosome index for each variant",
+        ),
         schema.ZarrArraySpec.new(
             vcf_field=None,
             name="call_genotype_phased",
@@ -155,8 +172,7 @@ def convert(
     )
     vzw.finalise(show_progress)
 
-    # TODO - index code needs variant_contig
-    # vzw.create_index()
+    vzw.create_index()
 
 
 # FIXME do this more efficiently - currently reading the whole thing

tests/test_plink.py

@@ -106,6 +106,16 @@ def test_variant_position(self, ds):
     def test_variant_allele(self, ds):
         nt.assert_array_equal(ds.variant_allele, [["A", "G"], ["T", "C"]])
 
+    def test_contig_id(self, ds):
+        """Test that contig identifiers are correctly extracted and stored."""
+        nt.assert_array_equal(ds.contig_id, ["1"])
+
+    def test_variant_contig(self, ds):
+        """Test that variant to contig mapping is correctly stored."""
+        nt.assert_array_equal(
+            ds.variant_contig, [0, 0]
+        )  # Both variants on chromosome 1
+
     def test_genotypes(self, ds):
         call_genotype = ds.call_genotype.values
         assert call_genotype.shape == (2, 10, 2)
@@ -196,6 +206,24 @@ def test_genotypes(self, ds):
         # FIXME not working
         nt.assert_array_equal(actual, expected)
 
+    def test_contig_arrays(self, ds):
+        """Test that contig arrays are correctly created and filled."""
+        # Verify contig_id exists and is a string array
+        assert hasattr(ds, "contig_id")
+        assert ds.contig_id.dtype == np.dtype("O")
+
+        # Verify variant_contig exists and contains integer indices
+        assert hasattr(ds, "variant_contig")
+        assert ds.variant_contig.dtype == np.dtype("int8")
+        assert ds.variant_contig.shape[0] == 100  # 100 variants
+
+        # Verify mapping between variant_contig and contig_id
+        # For each unique contig index, check at least one corresponding variant exists
+        for i, contig in enumerate(ds.contig_id.values):
+            assert np.any(
+                ds.variant_contig.values == i
+            ), f"No variants found for contig {contig}"
+
     # @pytest.mark.xfail
     @pytest.mark.parametrize(
         ("variants_chunk_size", "samples_chunk_size"),