Merge pull request #66 from jeromekelleher/some-polishing

Some polishing

Author: jeromekelleher

bio2zarr/cli.py

@@ -15,9 +15,21 @@
 )
 
 # TODO help text
-chunk_length = click.option("-l", "--chunk-length", type=int, default=None)
+chunk_length = click.option(
+    "-l",
+    "--chunk-length",
+    type=int,
+    default=None,
+    help="Chunk size in the variants dimension",
+)
 
-chunk_width = click.option("-w", "--chunk-width", type=int, default=None)
+chunk_width = click.option(
+    "-w",
+    "--chunk-width",
+    type=int,
+    default=None,
+    help="Chunk size in the samples dimension",
+)
 
 version = click.version_option(version=provenance.__version__)
 
@@ -83,13 +95,30 @@ def mkschema(if_path):
 @click.argument("if_path", type=click.Path())
 @click.argument("zarr_path", type=click.Path())
 @verbose
-@click.option("-s", "--schema", default=None)
-# TODO: these are mutually exclusive with schema, tell click this
+@click.option("-s", "--schema", default=None, type=click.Path(exists=True))
 @chunk_length
 @chunk_width
+@click.option(
+    "-V",
+    "--max-variant-chunks",
+    type=int,
+    default=None,
+    help=(
+        "Truncate the output in the variants dimension to have "
+        "this number of chunks. Mainly intended to help with "
+        "schema tuning."
+    ),
+)
 @worker_processes
 def encode(
-    if_path, zarr_path, verbose, schema, chunk_length, chunk_width, worker_processes
+    if_path,
+    zarr_path,
+    verbose,
+    schema,
+    chunk_length,
+    chunk_width,
+    max_variant_chunks,
+    worker_processes,
 ):
     """
     Encode intermediate format (see explode) to vcfzarr
@@ -101,6 +130,7 @@ def encode(
         schema,
         chunk_length=chunk_length,
         chunk_width=chunk_width,
+        max_v_chunks=max_variant_chunks,
         worker_processes=worker_processes,
         show_progress=True,
     )
@@ -132,6 +162,9 @@ def convert_vcf(vcfs, out_path, chunk_length, chunk_width, verbose, worker_proce
 @click.argument("vcfs", nargs=-1, required=True)
 @click.argument("out_path", type=click.Path())
 def validate(vcfs, out_path):
+    """
+    Development only, do not use. Will be removed before release.
+    """
     # FIXME! Will silently not look at remaining VCFs
     vcf.validate(vcfs[0], out_path, show_progress=True)
 
@@ -158,7 +191,9 @@ def vcf2zarr():
 @verbose
 @chunk_length
 @chunk_width
-def convert_plink(in_path, out_path, verbose, worker_processes, chunk_length, chunk_width):
+def convert_plink(
+    in_path, out_path, verbose, worker_processes, chunk_length, chunk_width
+):
     """
     In development; DO NOT USE!
     """

bio2zarr/core.py

@@ -23,13 +23,15 @@
 )
 
 
-def chunk_aligned_slices(z, n):
+def chunk_aligned_slices(z, n, max_chunks=None):
     """
     Returns at n slices in the specified zarr array, aligned
     with its chunks
     """
     chunk_size = z.chunks[0]
     num_chunks = int(np.ceil(z.shape[0] / chunk_size))
+    if max_chunks is not None:
+        num_chunks = min(num_chunks, max_chunks)
     slices = []
     splits = np.array_split(np.arange(num_chunks), min(n, num_chunks))
     for split in splits:
@@ -132,7 +134,7 @@ class ProgressConfig:
     units: str = ""
     title: str = ""
     show: bool = False
-    poll_interval: float = 0.001
+    poll_interval: float = 0.01
 
 
 # NOTE: this approach means that we cannot have more than one
@@ -175,7 +177,7 @@ def __init__(self, worker_processes=1, progress_config=None):
         self.progress_config = progress_config
         self.progress_bar = tqdm.tqdm(
             total=progress_config.total,
-            desc=f"{progress_config.title:>7}",
+            desc=f"{progress_config.title:>9}",
             unit_scale=True,
             unit=progress_config.units,
             smoothing=0.1,
@@ -186,6 +188,7 @@ def __init__(self, worker_processes=1, progress_config=None):
         self.progress_thread = threading.Thread(
             target=self._update_progress_worker,
             name="progress-update",
+            daemon=True,  # Avoids deadlock on exit in awkward error conditions
         )
         self.progress_thread.start()