Add plink contig info

Author: benjeffery

bio2zarr/plink.py

@@ -5,7 +5,7 @@
 import numpy as np
 import zarr
 
-from bio2zarr import constants, vcz
+from bio2zarr import constants, core, vcz
 
 logger = logging.getLogger(__name__)
 
@@ -18,6 +18,9 @@ def __init__(self, path):
         self.samples = [vcz.Sample(id=sample) for sample in self.bed.iid]
         self.num_samples = len(self.samples)
         self.root_attrs = {}
+        self.contigs = [
+            vcz.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
@@ -32,6 +35,11 @@ def iter_alleles(self, start, stop, num_alleles):
             alleles[1 : 1 + len(alt)] = alt
             yield alleles
 
+    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 iter_field(self, field_name, shape, start, stop):
         assert field_name == "position"  # Only position field is supported from plink
         yield from self.bed.bp_position[start:stop]
@@ -88,6 +96,15 @@ def generate_schema(
                 chunks=[schema_instance.variants_chunk_size, 2],
                 description=None,
             ),
+            vcz.ZarrArraySpec.new(
+                vcf_field=None,
+                name="variant_contig",
+                dtype=core.min_int_dtype(0, len(np.unique(self.bed.chromosome))),
+                shape=[m],
+                dimensions=["variants"],
+                chunks=[schema_instance.variants_chunk_size],
+                description="Contig/chromosome index for each variant",
+            ),
             vcz.ZarrArraySpec.new(
                 vcf_field=None,
                 name="call_genotype_phased",
@@ -159,9 +176,7 @@ def convert(
         show_progress=show_progress,
     )
     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"),