Progress bars

pyproject.toml

@@ -20,6 +20,7 @@ dependencies = [
   "click>=8.2.0",
   "more_itertools",
   "numpy>=1.23.5,<2.4",
+  "tqdm",
   "vcztools @ git+https://github.com/sgkit-dev/vcztools.git",
   "zarr>=3.1",
 ]

vczstore/append.py

@@ -3,12 +3,12 @@
 import numpy as np
 import zarr
 
-from vczstore.utils import variant_chunk_slices
+from vczstore.utils import variant_chunk_slices, variants_progress
 
 logger = logging.getLogger(__name__)
 
 
-def append(vcz1, vcz2):
+def append(vcz1, vcz2, *, show_progress=False):
     """Append vcz2 to vcz1 in place"""
     root1 = zarr.open(vcz1, mode="r+")
     root2 = zarr.open(vcz2, mode="r")
@@ -53,9 +53,11 @@ def append(vcz1, vcz2):
                 raise ValueError("unsupported number of array_dims")
 
     # append genotype fields
-    for variant_selection in variant_chunk_slices(root1):
-        for var in root1.keys():
-            if var.startswith("call_"):
-                root1[var][variant_selection, old_num_samples:new_num_samples, ...] = (
-                    root2[var][variant_selection, ...]
-                )
+    with variants_progress(n_variants1, "Append", show_progress) as pbar:
+        for v_sel in variant_chunk_slices(root1):
+            for var in root1.keys():
+                if var.startswith("call_"):
+                    root1[var][v_sel, old_num_samples:new_num_samples, ...] = root2[
+                        var
+                    ][v_sel, ...]
+            pbar.update(v_sel.stop - v_sel.start)

vczstore/cli.py

@@ -21,6 +21,14 @@ def show_work_summary(work_summary):
     click.echo(output)
 
 
+progress = click.option(
+    "-P /-Q",
+    "--progress/--no-progress",
+    default=True,
+    help="Show progress bars (default: show)",
+)
+
+
 zarr_backend_storage = click.option(
     "--zarr-backend-storage",
     type=str,
@@ -32,8 +40,9 @@ def show_work_summary(work_summary):
 @click.command()
 @click.argument("vcz1", type=click.Path())
 @click.argument("vcz2", type=click.Path())
+@progress
 @zarr_backend_storage
-def append(vcz1, vcz2, zarr_backend_storage):
+def append(vcz1, vcz2, progress, zarr_backend_storage):
     """Append vcz2 to vcz1 in place"""
     if zarr_backend_storage == "icechunk":
         from vczstore.icechunk_utils import icechunk_transaction
@@ -42,23 +51,25 @@ def append(vcz1, vcz2, zarr_backend_storage):
     else:
         cm = nullcontext(vcz1)
     with cm as vcz1:
-        append_function(vcz1, vcz2)
+        append_function(vcz1, vcz2, show_progress=progress)
 
 
 @click.command()
 @click.argument("vcz1", type=click.Path())
 @click.argument("vcz2", type=click.Path())
 @click.argument("vcz2_norm", type=click.Path())
-def normalise(vcz1, vcz2, vcz2_norm):
+@progress
+def normalise(vcz1, vcz2, vcz2_norm, progress):
     """Normalise variants in vcz2 with respect to vcz1 and write to vcz2_norm"""
-    normalise_function(vcz1, vcz2, vcz2_norm)
+    normalise_function(vcz1, vcz2, vcz2_norm, show_progress=progress)
 
 
 @click.command()
 @click.argument("vcz", type=click.Path())
 @click.argument("sample_id", type=str)
+@progress
 @zarr_backend_storage
-def remove(vcz, sample_id, zarr_backend_storage):
+def remove(vcz, sample_id, progress, zarr_backend_storage):
     """Remove a sample from vcz and overwrite with missing data"""
     if zarr_backend_storage == "icechunk":
         from vczstore.icechunk_utils import icechunk_transaction
@@ -67,7 +78,7 @@ def remove(vcz, sample_id, zarr_backend_storage):
     else:
         cm = nullcontext(vcz)
     with cm as vcz:
-        remove_function(vcz, sample_id)
+        remove_function(vcz, sample_id, show_progress=progress)
 
 
 @click.command()

vczstore/normalise.py

@@ -6,17 +6,17 @@
 from vcztools.retrieval import VczReader
 from vcztools.utils import array_dims, search
 
-from vczstore.utils import missing_val, variant_chunk_slices
+from vczstore.utils import missing_val, variant_chunk_slices, variants_progress
 
 
-def normalise(vcz1, vcz2, vcz2_norm):
+def normalise(vcz1, vcz2, vcz2_norm, show_progress=False):
     """Normalise variants in vcz2 with respect to vcz1 and write to vcz2_norm.
 
     vcz1, vcz2, vcz2_norm are paths or Zarr stores. Variants in vcz1 not present
     in vcz2 are filled with missing values.
     """
     index, remap_alleles, allele_mappings, updated_allele_mappings = index_variants(
-        vcz1, vcz2
+        vcz1, vcz2, show_progress=show_progress
     )
 
     if len(updated_allele_mappings) > 0:
@@ -96,31 +96,33 @@ def normalise(vcz1, vcz2, vcz2_norm):
     allele_mappings_list = list(allele_mappings.values())
 
     # copy variant-chunked arrays
-    for i, v_sel in enumerate(variant_chunk_slices(root1)):
-        for var, arr in root2.arrays():
-            if array_dims(arr)[0] != "variants":
-                continue
-            # copy genotype fields from vcz2, everything else from vcz1
-            if var.startswith("call_"):
-                arr = root2[var]
-                chunk_n = v_sel.stop - v_sel.start
-                shape = (chunk_n,) + arr.shape[1:]
-                data = np.full(shape, fill_value=missing_val(arr), dtype=arr.dtype)
-
-                match_sl = slice(match_starts[i], match_starts[i + 1])
-                local_idx = match_idx[match_sl] - v_sel.start
-                data[local_idx, ...] = arr[match_sl, ...]
-
-                if var == "call_genotype":
-                    remap_sl = slice(remap_starts[i], remap_starts[i + 1])
-                    local_remap_idx = remap_idx[remap_sl] - v_sel.start
-                    chunk_maps = allele_mappings_list[remap_sl]
-                    remap_genotypes(data, local_remap_idx, chunk_maps)
-
-                norm_root[var][v_sel] = data
-            else:
-                arr = root1[var]
-                norm_root[var][v_sel] = arr[v_sel, ...]
+    with variants_progress(n_variants, "Write", show_progress) as pbar:
+        for i, v_sel in enumerate(variant_chunk_slices(root1)):
+            for var, arr in root2.arrays():
+                if array_dims(arr)[0] != "variants":
+                    continue
+                # copy genotype fields from vcz2, everything else from vcz1
+                if var.startswith("call_"):
+                    arr = root2[var]
+                    chunk_n = v_sel.stop - v_sel.start
+                    shape = (chunk_n,) + arr.shape[1:]
+                    data = np.full(shape, fill_value=missing_val(arr), dtype=arr.dtype)
+
+                    match_sl = slice(match_starts[i], match_starts[i + 1])
+                    local_idx = match_idx[match_sl] - v_sel.start
+                    data[local_idx, ...] = arr[match_sl, ...]
+
+                    if var == "call_genotype":
+                        remap_sl = slice(remap_starts[i], remap_starts[i + 1])
+                        local_remap_idx = remap_idx[remap_sl] - v_sel.start
+                        chunk_maps = allele_mappings_list[remap_sl]
+                        remap_genotypes(data, local_remap_idx, chunk_maps)
+
+                    norm_root[var][v_sel] = data
+                else:
+                    arr = root1[var]
+                    norm_root[var][v_sel] = arr[v_sel, ...]
+            pbar.update(v_sel.stop - v_sel.start)
 
 
 def remap_genotypes(gt, indices, mappings):
@@ -139,7 +141,7 @@ def remap_genotypes(gt, indices, mappings):
                     gt[i, j, k] = mapping[val]
 
 
-def index_variants(vcz1, vcz2):
+def index_variants(vcz1, vcz2, *, show_progress=False):
     """Construct an index for variants of vcz2 that are in vcz1.
 
     Returns:
@@ -170,27 +172,30 @@ def index_variants(vcz1, vcz2):
     fields = ["variant_contig", "variant_position", "variant_allele"]
     it1 = VczReader(vcz1).variants(fields=fields)
     it2 = peekable(VczReader(vcz2).variants(fields=fields))
-    for i, variant in enumerate(it1):
-        v = it2.peek(None)
-        if v is None:
-            # it2 is exhausted - leave rest of index as False
-            break
-        matched, mapping, updated = variants_match(variant, v)
-        if matched:
-            # advance it2 and continue to next variant in it1
-            next(it2)
-            index[i] = True
-            if mapping is not None:
-                remap_alleles[i] = True
-                allele_mappings[i] = mapping
-            if updated is not None:
-                updated_allele_mappings[i] = updated
-        else:
-            if cmp_variant_site(variant, v) > 0:
-                raise ValueError(
-                    "Variant in vcz2 not found in vcz1 (or vcz2 is out of order): "
-                    f"{variant_repr(v)}"
-                )
+
+    with variants_progress(n_variants, "Index", show_progress) as pbar:
+        for i, variant in enumerate(it1):
+            v = it2.peek(None)
+            if v is None:
+                # it2 is exhausted - leave rest of index as False
+                break
+            matched, mapping, updated = variants_match(variant, v)
+            if matched:
+                # advance it2 and continue to next variant in it1
+                next(it2)
+                index[i] = True
+                if mapping is not None:
+                    remap_alleles[i] = True
+                    allele_mappings[i] = mapping
+                if updated is not None:
+                    updated_allele_mappings[i] = updated
+            else:
+                if cmp_variant_site(variant, v) > 0:
+                    raise ValueError(
+                        "Variant in vcz2 not found in vcz1 (or vcz2 is out of order): "
+                        f"{variant_repr(v)}"
+                    )
+            pbar.update()
 
     # when it1 is exhausted error if there are any variants left in it2
     v = it2.peek(None)

vczstore/remove.py

@@ -4,14 +4,15 @@
 import zarr
 from vcztools.utils import array_dims, search
 
-from vczstore.utils import missing_val, variant_chunk_slices
+from vczstore.utils import missing_val, variant_chunk_slices, variants_progress
 
 logger = logging.getLogger(__name__)
 
 
-def remove(vcz, sample_id):
+def remove(vcz, sample_id, *, show_progress=False):
     """Remove a sample from vcz and overwrite with missing data"""
     root = zarr.open(vcz, mode="r+")
+    n_variants = root["variant_contig"].shape[0]
     all_samples = root["sample_id"][:]
 
     # find index of sample to remove
@@ -22,12 +23,14 @@ def remove(vcz, sample_id):
 
     # overwrite sample data
     root["sample_id"][sample_selection] = ""
-    for variant_selection in variant_chunk_slices(root):
-        for var in root.keys():
-            arr = root[var]
-            if (
-                var.startswith("call_")
-                and array_dims(arr)[0] == "variants"
-                and array_dims(arr)[1] == "samples"
-            ):
-                arr[variant_selection, sample_selection, ...] = missing_val(arr)
+    with variants_progress(n_variants, "Remove", show_progress) as pbar:
+        for v_sel in variant_chunk_slices(root):
+            for var in root.keys():
+                arr = root[var]
+                if (
+                    var.startswith("call_")
+                    and array_dims(arr)[0] == "variants"
+                    and array_dims(arr)[1] == "samples"
+                ):
+                    arr[v_sel, sample_selection, ...] = missing_val(arr)
+        pbar.update(v_sel.stop - v_sel.start)

vczstore/utils.py

@@ -1,6 +1,7 @@
 from collections.abc import Callable
 from typing import Any
 
+import tqdm
 from vcztools.constants import FLOAT32_MISSING, INT_MISSING, STR_MISSING
 
 
@@ -29,6 +30,17 @@ def variant_chunk_slices(root):
         yield slice(start, end)
 
 
+def variants_progress(n_variants, title, show_progress=False):
+    return tqdm.tqdm(
+        total=n_variants,
+        desc=f"{title:>8}",
+        unit_scale=True,
+        unit="vars",
+        smoothing=0.1,
+        disable=not show_progress,
+    )
+
+
 def merge_lists(l1: list, l2: list, *, key: Callable[[Any], Any] = lambda x: x) -> list:
     """Merge two lists preserving the relative order of elements from both inputs.