test_zarr.py

"""Testing reading zarr file with zarr vs. dask."""

import random
import time
from contextlib import contextmanager

import click
import dask.array as da
import numpy as np

import zarr


@contextmanager
def time_block(label: str) -> None:
    start = time.perf_counter_ns()
    yield
    elapsed_ms = (time.perf_counter_ns() - start) / 1e6
    print(f"{elapsed_ms:6.3f}ms - {label}")


def test_reads(label, array, level):
    rows, cols = array.shape[-2:]
    for i in range(5):
        r = random.randint(0, rows - 1)
        c = random.randint(0, cols - 1)
        with time_block(f"level[{level}][{r}, {c}]"):
            np.asarray(array[0, 0, 0, r, c])
    print(f"  shape {array.shape}")


def test_zarr(path, level):
    with time_block("zarr.convenience.open"):
        infile = zarr.convenience.open(path)
    array = infile[level]
    test_reads("zarr", array, level)


def test_dask(path, level):
    with time_block("da.from_zarr"):
        array = da.from_zarr(path, level)
    test_reads("dask", array, level)


@click.command()
@click.argument('path', required=True)
@click.argument('method', required=True)
def main(path, method):

    for level in [7, 0]:
        print(f"\n{method} level {level} - {path}:")
        if method == "zarr":
            test_zarr(path, level)
        elif method == "dask":
            test_dask(path, level)
        else:
            ValueError(f"Unknown method {method}")


if __name__ == '__main__':
    main()