diff --git a/pyproject.toml b/pyproject.toml index 590d7d2..71785fa 100644 --- a/pyproject.toml +++ b/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", ] diff --git a/uv.lock b/uv.lock index 8627741..94a5e68 100644 --- a/uv.lock +++ b/uv.lock @@ -1552,6 +1552,7 @@ dependencies = [ { name = "click" }, { name = "more-itertools" }, { name = "numpy" }, + { name = "tqdm" }, { name = "vcztools" }, { name = "zarr" }, ] @@ -1597,6 +1598,7 @@ requires-dist = [ { name = "icechunk", marker = "extra == 'icechunk'" }, { name = "more-itertools" }, { name = "numpy", specifier = ">=1.23.5,<2.4" }, + { name = "tqdm" }, { name = "vcztools", git = "https://github.com/sgkit-dev/vcztools.git" }, { name = "zarr", specifier = ">=3.1" }, ] diff --git a/vczstore/append.py b/vczstore/append.py index c090e17..a494fbf 100644 --- a/vczstore/append.py +++ b/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) diff --git a/vczstore/cli.py b/vczstore/cli.py index 9df7deb..4d07b72 100644 --- a/vczstore/cli.py +++ b/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() diff --git a/vczstore/normalise.py b/vczstore/normalise.py index be83ff6..5df633b 100644 --- a/vczstore/normalise.py +++ b/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) diff --git a/vczstore/remove.py b/vczstore/remove.py index 63e502f..2a1ea99 100644 --- a/vczstore/remove.py +++ b/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) diff --git a/vczstore/utils.py b/vczstore/utils.py index a730937..04dd6c6 100644 --- a/vczstore/utils.py +++ b/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.