Rename ConversionSpec to VcfZarrSchema

Author: jeromekelleher

bio2zarr/plink.py

@@ -4,6 +4,7 @@
 import numpy as np
 import zarr
 import bed_reader
+import numcodecs
 
 from . import core
 
@@ -82,11 +83,16 @@ def convert(
     chunks = [variants_chunk_size, samples_chunk_size]
     dimensions = ["variants", "samples"]
 
+    # TODO we should be reusing some logic from vcfzarr here on laying
+    # out the basic dataset, and using the schema generator. Currently
+    # we're not using the best Blosc settings for genotypes here.
+    default_compressor = numcodecs.Blosc(cname="zstd", clevel=7)
+
     a = root.array(
         "sample_id",
         bed.iid,
         dtype="str",
-        compressor=core.default_compressor,
+        compressor=default_compressor,
         chunks=(samples_chunk_size,),
     )
     a.attrs["_ARRAY_DIMENSIONS"] = ["samples"]
@@ -98,7 +104,7 @@ def convert(
         "variant_position",
         bed.bp_position,
         dtype=np.int32,
-        compressor=core.default_compressor,
+        compressor=default_compressor,
         chunks=(variants_chunk_size,),
     )
     a.attrs["_ARRAY_DIMENSIONS"] = ["variants"]
@@ -109,7 +115,7 @@ def convert(
         "variant_allele",
         alleles,
         dtype="str",
-        compressor=core.default_compressor,
+        compressor=default_compressor,
         chunks=(variants_chunk_size,),
     )
     a.attrs["_ARRAY_DIMENSIONS"] = ["variants", "alleles"]
@@ -121,7 +127,7 @@ def convert(
         dtype="bool",
         shape=list(shape),
         chunks=list(chunks),
-        compressor=core.default_compressor,
+        compressor=default_compressor,
     )
     a.attrs["_ARRAY_DIMENSIONS"] = list(dimensions)
 
@@ -132,7 +138,7 @@ def convert(
         dtype="i1",
         shape=list(shape),
         chunks=list(chunks),
-        compressor=core.default_compressor,
+        compressor=default_compressor,
     )
     a.attrs["_ARRAY_DIMENSIONS"] = list(dimensions)
 
@@ -141,7 +147,7 @@ def convert(
         dtype="bool",
         shape=list(shape),
         chunks=list(chunks),
-        compressor=core.default_compressor,
+        compressor=default_compressor,
     )
     a.attrs["_ARRAY_DIMENSIONS"] = list(dimensions)
 

bio2zarr/vcf.py

@@ -1098,11 +1098,9 @@ def _choose_compressor_settings(self):
 
 ZARR_SCHEMA_FORMAT_VERSION = "0.2"
 
-# RENAME to ZarrSchema
-
 
 @dataclasses.dataclass
-class ZarrConversionSpec:
+class VcfZarrSchema:
     format_version: str
     samples_chunk_size: int
     variants_chunk_size: int
@@ -1126,15 +1124,15 @@ def fromdict(d):
                 "Zarr schema format version mismatch: "
                 f"{d['format_version']} != {ZARR_SCHEMA_FORMAT_VERSION}"
             )
-        ret = ZarrConversionSpec(**d)
+        ret = VcfZarrSchema(**d)
         ret.columns = {
             key: ZarrColumnSpec(**value) for key, value in d["columns"].items()
         }
         return ret
 
     @staticmethod
     def fromjson(s):
-        return ZarrConversionSpec.fromdict(json.loads(s))
+        return VcfZarrSchema.fromdict(json.loads(s))
 
     @staticmethod
     def generate(pcvcf, variants_chunk_size=None, samples_chunk_size=None):
@@ -1280,7 +1278,7 @@ def fixed_field_spec(
                 )
             )
 
-        return ZarrConversionSpec(
+        return VcfZarrSchema(
             format_version=ZARR_SCHEMA_FORMAT_VERSION,
             samples_chunk_size=samples_chunk_size,
             variants_chunk_size=variants_chunk_size,
@@ -1673,7 +1671,7 @@ def service_completed_futures():
 
 def mkschema(if_path, out):
     pcvcf = PickleChunkedVcf.load(if_path)
-    spec = ZarrConversionSpec.generate(pcvcf)
+    spec = VcfZarrSchema.generate(pcvcf)
     out.write(spec.asjson())
 
 
@@ -1690,7 +1688,7 @@ def encode(
 ):
     pcvcf = PickleChunkedVcf.load(if_path)
     if schema_path is None:
-        schema = ZarrConversionSpec.generate(
+        schema = VcfZarrSchema.generate(
             pcvcf,
             variants_chunk_size=variants_chunk_size,
             samples_chunk_size=samples_chunk_size,
@@ -1700,7 +1698,7 @@ def encode(
         if variants_chunk_size is not None or samples_chunk_size is not None:
             raise ValueError("Cannot specify schema along with chunk sizes")
         with open(schema_path, "r") as f:
-            schema = ZarrConversionSpec.fromjson(f.read())
+            schema = VcfZarrSchema.fromjson(f.read())
     zarr_path = pathlib.Path(zarr_path)
     if zarr_path.exists():
         logger.warning(f"Deleting existing {zarr_path}")

tests/test_pcvcf.py

@@ -27,8 +27,8 @@ def test_mkschema(self, tmp_path, pcvcf):
         with open(schema_file, "w") as f:
             vcf.mkschema(pcvcf.path, f)
         with open(schema_file, "r") as f:
-            schema1 = vcf.ZarrConversionSpec.fromjson(f.read())
-        schema2 = vcf.ZarrConversionSpec.generate(pcvcf)
+            schema1 = vcf.VcfZarrSchema.fromjson(f.read())
+        schema2 = vcf.VcfZarrSchema.generate(pcvcf)
         assert schema1 == schema2
 
     def test_summary_table(self, pcvcf):
@@ -95,7 +95,7 @@ def pcvcf(self, tmp_path_factory):
 
     @pytest.fixture(scope="class")
     def schema(self, pcvcf):
-        return vcf.ZarrConversionSpec.generate(pcvcf)
+        return vcf.VcfZarrSchema.generate(pcvcf)
 
     @pytest.mark.parametrize(
         ("name", "dtype", "shape"),
@@ -165,8 +165,8 @@ def test_repr(self, pcvcf):
 
     def test_pos_repr(self, pcvcf):
         assert repr(pcvcf["POS"]).startswith(
-            "PickleChunkedVcfField(name=POS, partition_chunks=[8, 8, 8, 8, 8], path=")
-
+            "PickleChunkedVcfField(name=POS, partition_chunks=[8, 8, 8, 8, 8], path="
+        )
 
     def test_partition_record_index(self, pcvcf):
         nt.assert_array_equal(

tests/test_vcf.py

@@ -45,7 +45,7 @@ def test_zarr_schema_mismatch(self, schema, version):
         d = dict(schema)
         d["format_version"] = version
         with pytest.raises(ValueError, match="Zarr schema format version mismatch"):
-            vcf.ZarrConversionSpec.fromdict(d)
+            vcf.VcfZarrSchema.fromdict(d)
 
     @pytest.mark.parametrize("version", ["0.0", "1.0", "xxxxx", 0.1])
     def test_exploded_metadata_mismatch(self, tmpdir, exploded_path, version):