diff --git a/pyproject.toml b/pyproject.toml index b21262a..02ed879 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,6 +30,8 @@ xarray = [ "xarray>=2026.04", "zarr~=3.1.6", "zarrs>=0.2", + "bottleneck", + "fsspec", "pint", "unum", ] @@ -39,6 +41,7 @@ dev = [ "polars", "pyarrow", "fsspec", + "bottleneck", "tqdm", "xarray>=2026.04", "zarr~=3.1.6", diff --git a/tests/test_zarr_mapreduce_engine.py b/tests/test_zarr_mapreduce_engine.py new file mode 100644 index 0000000..315e767 --- /dev/null +++ b/tests/test_zarr_mapreduce_engine.py @@ -0,0 +1,153 @@ +"""End-to-end smoke test for the ported zarr map-reduce ENGINE foundation. + +Proves the reconciled storage glue (WorkflowConfig / WorkflowPaths / Substore / +XarrayStoragePartition workflow methods) and the read-side zarr_utils helpers +work against a workflow store actually produced by viva's XArrayEmitter — i.e. +the multi-lineage, multi-generation ``experiment_id=/variant=/lineage_seed=`` +layout the engine walks. + +The full ``ZarrMapReduce`` driver (~12 abstract hooks) is exercised by the +concrete moving-average pipeline downstream; here we cover the foundation: +substore discovery + async array read. +""" + +import asyncio +from pathlib import Path + +import pytest + +pytest.importorskip("xarray") +pytest.importorskip("zarr") + +from bigraph_schema import allocate_core # noqa: E402 + +from viva_emitters.xarray_emitter.storage import ( # noqa: E402 + WorkflowConfig, + WorkflowPaths, + Substore, + XarrayStoragePartition, +) +from viva_emitters.xarray_emitter.zarr_utils import ( # noqa: E402 + get_async_group, +) + +EXPERIMENT_ID = "e" +N_LINEAGES = 2 + + +def _emitter_config(store_root, agent_id, variant, lineage_seed): + return { + "emit": {}, + "out_uri": store_root, + "strategy": "colony", + "emit_root": ["agents", agent_id], + "transducer": { + "predicate": [[{"subsample": {"interval": 1}}]], + "buffer": {"size": 3}, + }, + "view": [ + { + "root": ("listeners",), + "metadata": False, + "variables": { + "mass": [{"path": "listeners/mass", "dtype": " num_variants == 1 (the baseline, variant=0) + } + + +def test_workflow_paths_locate_discovers_real_viva_substores(workflow_store): + wc = WorkflowConfig.build(_config(workflow_store["out_dir"])) + assert wc.is_uri is False + + wp = WorkflowPaths.locate(wc) + assert len(wp) == N_LINEAGES + subs = sorted((s.variant, s.lineage) for s in wp) + assert subs == [("variant=0", "lineage_seed=0"), ("variant=0", "lineage_seed=1")] + + +def test_from_substore_round_trips(workflow_store): + wc = WorkflowConfig.build(_config(workflow_store["out_dir"])) + sub = Substore("variant=0", "lineage_seed=1") + part = XarrayStoragePartition.from_substore(wc, sub, generation=2) + assert part.variant == 0 + assert part.lineage_seed == 1 + assert part.generation == 2 + assert part.agent_id == "00" + assert str(part.independent_path) == "experiment_id=e/variant=0/lineage_seed=1" + + +def test_zarr_utils_async_read_over_substore(workflow_store): + from zarr.api.asynchronous import open_group + + wc = WorkflowConfig.build(_config(workflow_store["out_dir"])) + wp = WorkflowPaths.locate(wc) + sub = next(iter(wp)) + sub_path = f"{Path(wp.root).name}/{sub}" + + async def read(): + root = await open_group(store=workflow_store["store_root"], mode="r") + g = await get_async_group(root, sub_path) + async for name, node in g.members(max_depth=None): + if hasattr(node, "shape") and name.endswith("generation=2"): + data = await node.getitem(Ellipsis) + return name, tuple(data.shape), data[:3].tolist() + return None + + result = asyncio.run(read()) + assert result is not None, "no generation=2 array read from substore" + _name, shape, head = result + assert shape == (6,) + assert head == [10.0, 11.0, 12.0] diff --git a/viva_emitters/xarray_emitter/mapreduce.py b/viva_emitters/xarray_emitter/mapreduce.py new file mode 100644 index 0000000..d81b0d2 --- /dev/null +++ b/viva_emitters/xarray_emitter/mapreduce.py @@ -0,0 +1,320 @@ + +""" +Basic building blocks for `map-reduce`_ pipelines. + +This module defines `closures`_ for parallel/asynchronous *maps* (with hashable +keys) and abstract/concrete *reduce* operations. Backend-specific pipeline +implementations, which use these building blocks, should be placed into separate +modules, e.g., :py:mod:`~ecoli.analysis.xarray_emitter.zarr_mapreduce`. + +.. _map-reduce: https://en.wikipedia.org/wiki/MapReduce +.. _closures: https://en.wikipedia.org/wiki/Closure_(computer_programming) +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from asyncio import Future, TaskGroup +from collections.abc import ( + AsyncGenerator, + Callable, + Coroutine, + Generator, + Iterable, + Mapping, +) +from dataclasses import astuple, dataclass +from inspect import isasyncgen, isgenerator +from itertools import chain +from multiprocessing.pool import Pool +from typing import Any, Self, cast + +from bottleneck import nanmax, nanmin +from numpy import ndarray, vstack +from numpy.typing import NDArray + +# ============================================================================== +# concurrency utils +# ============================================================================== + + +@dataclass(slots=True, frozen=True) +class PoolDictItem[KeyT, InputT, ResultT]: + """ + A constructor of `closures`_ intended for *parallel execution* by + :py:meth:`~multiprocessing.pool.Pool.imap_unordered`. + + This object holds the definition of the computation and of the (unordered) + output key, while :py:func:`.pool_dict` provides the execution context. + """ + + #: Output key as a function of the input value. + key: Callable[[InputT], KeyT] + #: Output value as a function of the input value. + closure: Callable[[InputT], ResultT] + + def __call__(self, arg: InputT) -> tuple[KeyT, ResultT]: + """ + Execute the closure for a given input value. + """ + return (self.key(arg), self.closure(arg)) + + +def pool_dict[KeyT, InputT, ResultT]( + worker: PoolDictItem[KeyT, InputT, ResultT], + args: Iterable[InputT], + *, + n_procs: int, n_items: int +) -> dict[KeyT, ResultT]: + """ + Create :py:class:`.PoolDictItem` closures for a concrete set of input + values, and execute them *in parallel* using + :py:meth:`~multiprocessing.pool.Pool.imap_unordered`. + """ + assert isinstance(worker, PoolDictItem) + assert isinstance(n_procs, int) and n_procs > 0 + if n_procs == 1: + return dict(map(worker, args)) + else: + with Pool(processes=n_procs) as pool: + return dict(pool.imap_unordered(worker, args, + chunksize=n_items // n_procs)) + + +# ------------------------------------------------------------------------------ + + +type AnyGenerator[ResultT] = Generator[ResultT] | AsyncGenerator[ResultT] + + +@dataclass(slots=True, frozen=True) +class AsyncDictItem[KeyT, ResultT]: + """ + A constructor of `closures`_ intended for *asynchronous execution* by a + :py:class:`~asyncio.TaskGroup`. + + This object holds the definition of the computation and of the (unordered) + output key, while :py:func:`.async_dict` provides the execution context. + """ + + #: Output key. + key: KeyT + #: :py:class:`asyncio.Task` name. + name: str + #: Coroutine producing the output value. + closure: Coroutine[None, None, ResultT] + + async def item(self) -> tuple[KeyT, ResultT]: + """ + Associate the closure with its output key. + """ + return (self.key, await self.closure) + + def create_task(self, g: TaskGroup) -> Future: + """ + Assign the closure to a :py:class:`asyncio.TaskGroup`. + """ + return g.create_task(self.item(), name=self.name) + + +async def async_dict[KeyT, ResultT]( + items: AnyGenerator[AsyncDictItem[KeyT, ResultT]] +) -> dict[KeyT, ResultT]: + r""" + Execute a set of :py:class:`.AsyncDictItem`\ s in a + :py:class:`~asyncio.TaskGroup`. + """ + async with TaskGroup() as g: + if isgenerator(items): + futures = [i.create_task(g) for i in items] + elif isasyncgen(items): + futures = [i.create_task(g) async for i in items] + else: + raise TypeError(str(type(items))) + return dict(f.result() for f in futures) + + +# ============================================================================== +# reduce operators +# ============================================================================== + + +@dataclass(slots=True, frozen=True) +class Reducer(ABC): + """ + A wrapper type for reduce operations inside a `map-reduce`_ pipeline. + """ + + @classmethod + @abstractmethod + def reduce(cls, args: list[Self], /) -> Self: + """ + Apply the reduce operation to a set of instances of this type. + """ + ... + + @abstractmethod + def extract(self) -> Any: + """ + Extract the reduced value from the wrapper type. + """ + ... + + +# ------------------------------------------------------------------------------ + + +@dataclass(slots=True, frozen=True) +class ReducerList(Reducer): + """ + A composite reducer, with component reducers indexed by an integer. + """ + + #: Component reducers. + reducers: list[Reducer] + + def __post_init__(self) -> None: + assert isinstance(self.reducers, list) + assert all(isinstance(r, Reducer) for r in self.reducers) + + @classmethod + def reduce(cls, args: list[Self], /) -> Self: + reducers = [ + [a.reducers[i] for a in args] + for i in range(len(args[0].reducers))] + return cls([type(r[0]).reduce(r) for r in reducers]) + + def extract(self) -> list[Any]: + return [r.extract() for r in self.reducers] + + +@dataclass(slots=True, frozen=True) +class ReducerMapping[KeyT](Reducer): + """ + A composite reducer, with component reducers indexed by a hashable key. + """ + + #: Component reducers. + reducers: Mapping[KeyT, Reducer] + + def __post_init__(self) -> None: + assert isinstance(self.reducers, dict) + assert all(isinstance(r, Reducer) for r in self.reducers.values()) + + @classmethod + def reduce(cls, args: list[Self], /) -> Self: + reducers: list[list[Reducer]] = [[a.reducers[k] for a in args] + for k in args[0].reducers] + return cls(dict(zip(args[0].reducers.keys(), + (type(r[0]).reduce(r) for r in reducers)))) + + def extract(self) -> dict[KeyT, Any]: + return {k: r.extract() for (k, r) in self.reducers.items()} + + +@dataclass(slots=True, frozen=True) +class ReducerDistribute[KeyT1, KeyT2](ReducerMapping): + """ + A composite reducer, with two nested levels of :py:class:`ReducerMapping` + that can be `distributed`_/exchanged. + + This is a useful operation for explicitly translating between different data + hierarchies, e.g., when the input hierarchy is optimised for parallel + computation, but is transposed relative to the desired output hierarchy. + + .. _distributed: https://en.wikipedia.org/wiki/Distributive_property + """ + + #: Component reducers. + reducers: Mapping[KeyT1, ReducerMapping[KeyT2]] + + def __post_init__(self) -> None: + assert isinstance(self.reducers, dict) + assert all(isinstance(r, ReducerMapping) for r in self.reducers.values()) + + @classmethod + def reduce(cls, args: list[Self], /) -> Self: + return cls(cast( + Mapping[KeyT1, ReducerMapping[KeyT2]], + ReducerMapping.reduce( + [ReducerMapping(a.reducers) for a in args] + ).reducers)) + + def distribute(self) -> ReducerDistribute[KeyT2, KeyT1]: + """ + Exchange the two nested levels of :py:class:`ReducerMapping`, and apply + :py:meth:`Reducer.reduce` to the output's inner level. + """ + transpose: dict[KeyT2, dict[KeyT1, list[Reducer]]] = {} + for (k1, r1) in self.reducers.items(): + for (k2, r2) in r1.reducers.items(): + transpose.setdefault(k2, {}).setdefault(k1, []).append(r2) + result: dict[KeyT2, dict[KeyT1, Reducer]] = {} + for (k2, x2) in transpose.items(): + for (k1, x1) in x2.items(): + result.setdefault(k2, {})[k1] = type(x1[0]).reduce(x1) + return ReducerDistribute[KeyT2, KeyT1]({k2: ReducerMapping[KeyT1](m2) + for (k2, m2) in result.items()}) + + +# ------------------------------------------------------------------------------ + + +@dataclass(kw_only=True, slots=True, frozen=True) +class MinMax(Reducer): + """ + A primitive reducer that computes the minimum and maximum values across its + inputs, which are assumed to be 1-D arrays with identical shapes. + """ + + mins: NDArray + maxs: NDArray + + def __post_init__(self) -> None: + assert isinstance(self.mins, ndarray) + assert isinstance(self.maxs, ndarray) + assert self.mins.ndim == 1 + assert self.maxs.ndim == 1 + assert self.mins.shape == self.maxs.shape + + @classmethod + def reduce(cls, args: list[Self], /) -> Self: + return cls(mins=nanmin(vstack([a.mins for a in args]), axis=0), + maxs=nanmax(vstack([a.maxs for a in args]), axis=0)) + + def extract(self) -> tuple[NDArray, NDArray]: + return astuple(self) + + +@dataclass(slots=True, frozen=True) +class Chain[ElemT](Reducer): + """ + A primitive reducer that simply accumulates its inputs into a list. + """ + + string: list[ElemT] + + def __post_init__(self) -> None: + assert isinstance(self.string, list) + + @classmethod + def reduce(cls, args: list[Self], /) -> Self: + return cls(list(chain.from_iterable(a.string for a in args))) + + def extract(self) -> list[ElemT]: + return self.string + + +@dataclass(slots=True, frozen=True) +class Unique[ElemT](Chain): + """ + A primitive reducer that checks whether all of its inputs are identical. + """ + + def extract(self) -> list[ElemT]: + """ + Return the unique input element as a singleton list. + """ + fst = self.string[0] + assert all(el == fst for el in self.string[1:]) + return [fst] diff --git a/viva_emitters/xarray_emitter/storage.py b/viva_emitters/xarray_emitter/storage.py index 1cf0697..565f400 100644 --- a/viva_emitters/xarray_emitter/storage.py +++ b/viva_emitters/xarray_emitter/storage.py @@ -7,8 +7,12 @@ from __future__ import annotations -from dataclasses import dataclass, field, fields +import json +from collections.abc import Callable, Generator, Mapping +from dataclasses import astuple, dataclass, field, fields from functools import cached_property +from glob import glob +from os.path import join from pathlib import Path from typing import Any, Self, TYPE_CHECKING @@ -43,6 +47,13 @@ TIME_VAR_DTYPE = np.dtype(np.float32) """ Data type for :py:attr:`.XarrayStoragePartition.time_var_name`. """ +EXPERIMENT_PREFIX = "experiment_id=" +""" Prefix for a part of :py:attr:`.XarrayStoragePartition.independent_path`. """ +VARIANT_PREFIX = "variant=" +""" Prefix for a part of :py:attr:`.XarrayStoragePartition.independent_path`. """ +LINEAGE_PREFIX = "lineage_seed=" +""" Prefix for a part of :py:attr:`.XarrayStoragePartition.independent_path`. """ + # ============================================================================== # Xarray storage layout @@ -65,6 +76,28 @@ def cast(cls, partition: StoragePartition) -> Self: return cls(**{f.name: getattr(partition, f.name) for f in fields(partition) if f.init}) + @classmethod + def from_substore( + cls, config: WorkflowConfig, substore: Substore, generation: int + ) -> Self: + """ + Reconstruct the partition for a single generation of a located + :py:class:`.Substore`. Used by analysis pipelines walking a workflow + store. ``agent_id`` is set to ``"0" * generation`` so the base + :py:class:`.StoragePartition` derives the matching ``generation``. + """ + assert isinstance(config, WorkflowConfig) + assert isinstance(substore, Substore) + assert isinstance(generation, int) + return cls( + experiment_id=config.sim["experiment_id"], + variant=int(substore.variant.removeprefix(VARIANT_PREFIX)), + lineage_seed=int(substore.lineage.removeprefix(LINEAGE_PREFIX)), + agent_id="0" * generation) + + def __hash__(self): + return hash(tuple(getattr(self, f.name) for f in fields(self))) + # ~~~~~~~~~~~~~~~~~ # @cached_property @@ -88,6 +121,23 @@ def independent_path(self) -> Path: return Path(*(f"{k}={getattr(self, k)}" for k in ["experiment_id", "variant", "lineage_seed"])) + @staticmethod + def independent_path_glob(experiment_id: str) -> str: + """ + Glob pattern matching every independent substore path of a workflow. + """ + return join(f"{EXPERIMENT_PREFIX}{experiment_id}", + f"{VARIANT_PREFIX}[0-9]*", + f"{LINEAGE_PREFIX}[0-9]*") + + @staticmethod + def get_variant(path: str) -> str: + return path.split(sep="/")[1] + + @staticmethod + def get_lineage(path: str) -> str: + return path.split(sep="/")[2] + # ~~~~~~~~~~~~~~~~~ # @cached_property @@ -122,6 +172,10 @@ def time_coo_name(self) -> str: """ return f"{TIME_COO_PREFIX}{self.sim_id}" + @staticmethod + def is_time_coo_name(key: str) -> bool: + return key.startswith(TIME_COO_PREFIX) + @cached_property def time_var_name(self) -> str: """ @@ -151,9 +205,24 @@ def success_attr_name(self) -> str: # ============================================================================== +type VariablePath = str type VariableEncoding = dict[str, Any] +def var_name(path: VariablePath) -> VariablePath: + """ + Extract the variable name from a full variable path. + """ + return path.rsplit("/", maxsplit=1)[-1] + + +def coo_path(path: VariablePath) -> VariablePath: + """ + Compute the coordinate path associated with a variable path. + """ + return join(path, VariableSpec.var_coo_name(var_name(path))) + + # ============================================================================== @@ -367,3 +436,227 @@ def alloc_var(self, buf_size: int, /) -> Dataset: """ return Dataset(data_vars={ self.datavar_name: (self.dim_names, self.zeros(buf_size))}) + + +# ============================================================================== +# workflow configuration +# ============================================================================== + + +@dataclass(kw_only=True, frozen=True) +class WorkflowConfig: + """ + Simulation and analysis workflow configuration, parsed directly from JSON. + + This object is intended for use by analysis pipelines (e.g. the + :py:mod:`.zarr_mapreduce` engine). + """ + + #: Full workflow configuration. + sim: dict[str, Any] + #: Indicates whether the workflow store is remote or local. + is_uri: bool + #: Variant parameters, parsed from :py:attr:`.sim`. + variants: list[dict[str, Any]] + + @classmethod + def load( + cls, path: str | Path, *, + variant_parser: Callable[[Any], list[dict[str, Any]]] | None = None + ) -> Self: + """ + Calls: :py:meth:`.build`. + """ + with open(path, "r") as f: + return cls.build(json.load(f), variant_parser=variant_parser) + + @classmethod + def build( + cls, config: dict[str, Any], *, + variant_parser: Callable[[Any], list[dict[str, Any]]] | None = None + ) -> Self: + """ + Parse and validate high-level information about the simulation workflow. + + Called by: :py:meth:`.load`. + + The variant specification is application-specific, so the parser is + *injected* rather than imported: viva-emitters is a generic library and + must not depend on any downstream simulator's variant machinery. Pass + ``variant_parser`` (e.g. v2ecoli's ``create_variants.parse_variants``) + to populate :py:attr:`.variants`; when omitted, :py:attr:`.variants` is + left empty (sufficient for locating substores and running map-reduce + pipelines that do not resolve variant parameters). + """ + assert config["experiment_id"] + assert config["emitter"] == "xarray" + assert any(map(config["emitter_arg"].__contains__, ["out_dir", "out_uri"])) + + variants = config["variants"] + assert isinstance(variants, dict) + assert len(variants) <= 1 + variants = (variant_parser(next(iter(variants.values()))) + if variants and variant_parser is not None else + []) + return cls(is_uri="out_uri" in config["emitter_arg"], + sim=config, variants=variants) + + # ~~~~~~~~~~~~~~~~~ # + + def variant_params(self, variant: str) -> dict[str, Any]: + """ + Translate a variant string into variant parameters. + + Requires that a ``variant_parser`` was supplied to :py:meth:`.build`. + + Calls: :py:meth:`.variant_index`. + """ + ix = self.variant_index(variant) + ix -= int(not self.sim.get("skip_baseline", False)) + return self.variants[ix] if ix >= 0 else {} + + @staticmethod + def variant_index(variant: str) -> int: + """ + Interpret a variant string as a pointer into :py:attr:`.variants`. + + Called by: :py:meth:`.variant_params`. + """ + ix = int(variant.lstrip(VARIANT_PREFIX)) + assert variant == f"{VARIANT_PREFIX}{ix}" + return ix + + +# ============================================================================== +# workflow storage layout +# ============================================================================== + + +@dataclass(slots=True, kw_only=True, frozen=True) +class WorkflowPaths: + """ + Efficiently locate all independent substore paths within a simulation + workflow that was emitted using the :py:class:`.XarrayStoragePartition` path + scheme. + + This object is intended for use by analysis pipelines. + """ + + #: Root path to the persistent store of a simulation workflow. + root: str + #: Tree of located :py:attr:`.XarrayStoragePartition.independent_path`\ s. + substores: dict[str, list[str]] + + def __post_init__(self) -> None: + assert self.root.rsplit(sep="/", maxsplit=1)[-1].startswith(EXPERIMENT_PREFIX) + for (variant, lineages) in self.substores.items(): + assert variant.startswith(VARIANT_PREFIX) + assert all(lin.startswith(LINEAGE_PREFIX) for lin in lineages) + + # ~~~~~~~~~~~~~~~~~ # + + def __len__(self) -> int: + """ + Number of independent substores. + """ + return sum(map(len, self.substores.values())) + + def __iter__(self) -> Generator[Substore]: + """ + Iterator over :py:attr:`.substores`. + """ + for (variant, lineages) in self.substores.items(): + for lineage in lineages: + yield Substore(variant, lineage) + + # ~~~~~~~~~~~~~~~~~ # + + @classmethod + def locate(cls, config: WorkflowConfig) -> Self: + """ + Find all substore paths using a single ``glob()`` call to the file + system. + """ + # load config + assert isinstance(config, WorkflowConfig) + experiment_id = config.sim["experiment_id"] + emitter = config.sim["emitter_arg"] + num_variants = len(config.variants) + num_variants += int(not config.sim.get("skip_baseline", False)) + num_lineages = config.sim["n_init_sims"] + + # find workflow store + if config.is_uri: + # fsspec is only needed for remote (URI) stores; import lazily so + # local analysis needs no extra dependency. + from fsspec import get_fs_token_paths + store_path = join(emitter["out_uri"], experiment_id, "store") + fs, _, store_path = get_fs_token_paths(store_path) + else: + store_path = join(emitter["out_dir"], experiment_id, "store") + assert Path(store_path).exists() + + # find independent substore paths + substore_glob = XarrayStoragePartition.independent_path_glob(experiment_id) + substores = cls.group( + fs.glob(join(store_path, substore_glob)) + if config.is_uri else + glob(substore_glob, root_dir=store_path)) + + # check consistency with workflow config + assert set(substores.keys()) == {f"{VARIANT_PREFIX}{v}" + for v in range(num_variants)} + for lineages in substores.values(): + assert len(lineages) == num_lineages + + return cls(root=join(store_path, f"{EXPERIMENT_PREFIX}{experiment_id}"), + substores=substores) + + @staticmethod + def group(paths: list[str]) -> dict[str, list[str]]: + """ + Reassemble a partition hierarchy from a flat list of substore paths. + """ + substores: dict[str, list[str]] = {} + for path in paths: + substores.setdefault( + XarrayStoragePartition.get_variant(path), [] + ).append(XarrayStoragePartition.get_lineage(path)) + return substores + + +# ------------------------------------------------------------------------------ + + +@dataclass(slots=True, frozen=True) +class Substore: + """ + Hashable identifier of an independent substore. + """ + + variant: str + lineage: str + + def __post_init__(self) -> None: + assert self.variant.startswith(VARIANT_PREFIX) + assert self.lineage.startswith(LINEAGE_PREFIX) + + def __str__(self) -> str: + return join(*astuple(self)) + + @classmethod + def identity(cls, path: Self) -> Self: + return path + + @staticmethod + def groupby_variant[ResultT]( + results: Mapping[Substore, ResultT] + ) -> Mapping[str, Mapping[str, ResultT]]: + """ + Group a hash map over substore identifiers into ``variant``/``lineage`` + levels. + """ + grouped: dict[str, dict[str, ResultT]] = {} + for (s, res) in results.items(): + grouped.setdefault(s.variant, {})[s.lineage] = res + return grouped diff --git a/viva_emitters/xarray_emitter/zarr_mapreduce.py b/viva_emitters/xarray_emitter/zarr_mapreduce.py new file mode 100644 index 0000000..958796f --- /dev/null +++ b/viva_emitters/xarray_emitter/zarr_mapreduce.py @@ -0,0 +1,940 @@ + +""" +Abstract base classes defining a `map-reduce`_ framework, which is intended +for executing complex numerical analysis pipelines on multi-variant, multi-seed, +multi-generation simulation workflows that were stored using the +:py:class:`.AsyncZarrBufferWriter` backend of the :py:class:`.XarrayEmitter`. + +This framework is designed to take into account the full (in-)dependence +assumptions underlying the :ref:`storage ` and :ref:`variable +` layouts of the :py:class:`.XarrayEmitter`, and to leverage +both :py:mod:`multiprocessing` parallelism and the +:py:mod:`~zarr.api.asynchronous` Zarr API. + +Concrete analysis pipelines should be defined in separate modules, by +subclassing the following abstract classes: :py:class:`.ZarrMapReduceConfig`, +:py:class:`.ZarrMapReduceResult`, :py:class:`.ZarrMapReduce`, and, when +appropriate, :py:class:`.ZarrMapReducePlotConfig` and +:py:class:`.ZarrMapReducePlot`. + +:py:meth:`ZarrMapReduceConfig.make_cli_parser` defines analysis-agnostic CLI +:py:class:`~argparse.ArgumentParser` arguments, including profiling +(`cProfile`_/`memray`_) and debugging options, as well as controls over whether +the numerical pipeline results are stored into or loaded from a ZIP file and +whether they are directly post-processed, e.g., by rendering plots. + +.. _map-reduce: https://en.wikipedia.org/wiki/MapReduce +.. _cProfile: https://docs.python.org/3/library/profile.html#module-cProfile +.. _memray: https://bloomberg.github.io/memray/getting_started.html +""" + +from __future__ import annotations + +import pathlib +from abc import ABC, abstractmethod +from argparse import ArgumentParser, Namespace, RawDescriptionHelpFormatter +from dataclasses import astuple, dataclass, field +from functools import partial +from os.path import join +from pathlib import Path +from typing import TYPE_CHECKING, Any, ClassVar, Self, cast + +import zarr +from zarr.api.asynchronous import open_group +from zarr.core.sync import sync + +from .storage import ( + Substore, + WorkflowConfig, + WorkflowPaths, + XarrayStoragePartition, + coo_path, + var_name, +) +from .view import ForestView +from .zarr_utils import get_async_array, get_async_group + +from .mapreduce import AsyncDictItem, PoolDictItem, async_dict, pool_dict + +if TYPE_CHECKING: + + from zarr.core.buffer import NDArrayLikeOrScalar + from zarr.core.group import AsyncGroup + from zarr.core.indexing import Selector + + from .storage import VariablePath + + from .mapreduce import Reducer + +# ============================================================================== +# analysis configuration +# ============================================================================== + + +@dataclass(kw_only=True, slots=True) +class ZarrMapReduceConfig(ABC): + """ + Configuration of a :py:class:`ZarrMapReduce` pipeline, excluding + post-processing steps such as plot rendering. + """ + + # workflow config + # ~~~~~~~~~~~~~~~~~ # + #: Name of the analysis pipeline. + #: Used to access :py:attr:`.WorkflowConfig.sim`. + name: ClassVar[str] + #: Variable names and coordinate descriptors used by the analysis pipeline. + #: Validated against :py:attr:`.WorkflowConfig.sim`. + variables: dict[VariablePath, Any] = field(init=False) + + # file I/O + # ~~~~~~~~~~~~~~~~~ # + #: File path for pipeline result data (Zarr ZIP file). + result_file: pathlib.Path = field(init=False) + #: File path for plots of pipeline results (e.g., SVG file). + figure_file: pathlib.Path = field(init=False) + #: Store the pipeline result data, instead of post-processing it. + save_data: bool = False + #: Load result data, instead of executing the pipeline. + load_data: bool = False + + # compute resources + # ~~~~~~~~~~~~~~~~~ # + #: Number of :py:mod:`multiprocessing` processes. + n_procs: int + #: Number of Zarr threads per :py:mod:`multiprocessing` process. + n_threads: int + + # ~~~~~~~~~~~~~~~~~ # + #: Print debug information and reduce the plot size. + debug: bool = False + #: Execute pipeline inside `cProfile`_ (``n_procs == 1``). + prof_time: bool = False + #: Execute pipeline inside `memray`_. + prof_mem: bool = False + #: File path for profiling output (relative to ``./``). + prof_file: pathlib.Path | None = None + + # ~~~~~~~~~~~~~~~~~ # + + def __post_init__(self) -> None: + assert isinstance(self.save_data, bool) + assert isinstance(self.load_data, bool) + assert not (self.save_data and self.load_data) + + assert isinstance(self.n_procs, int) and self.n_procs > 0 + assert isinstance(self.n_threads, int) and self.n_threads > 0 + if self.load_data: + if self.n_procs > 1: + raise ValueError("Set `--procs 1` for `--load`.") + if self.n_threads > 1: + raise ValueError("Set `--threads 1` for `--load`.") + if self.prof_file is not None: + raise ValueError("No profiling support for `--load`.") + + assert isinstance(self.prof_time, bool) + assert isinstance(self.prof_mem, bool) + assert isinstance(self.prof_file, Path | None) + match (self.prof_time, self.prof_mem, self.prof_file): + case (True, True, _): + raise TypeError("Choose one profiler at a time.") + case (True, False, _) if self.n_procs > 1: + raise ValueError("Set `--procs 1` for `--prof-time`.") + case (True, False, f) | (False, True, f) if f == Path(): + raise ValueError("Set `--prof-file` when profiling.") + case (False, False, f) if f is not None: + raise ValueError("Only set `--prof-file` when profiling.") + + assert isinstance(self.debug, bool) + + # ~~~~~~~~~~~~~~~~~ # + + @staticmethod + def make_cli_parser(description: str, /) -> ArgumentParser: + """ + This :py:class:`~argparse.ArgumentParser` may be further customised by a + concrete pipeline. + """ + parser = ArgumentParser( + prog=Path(__file__).name, description=description, + formatter_class=RawDescriptionHelpFormatter) + parser.add_argument( + "config", type=str, + help="simulation configuration JSON (relative to `configs/`)") + parser.add_argument( + "-p", "--procs", type=int, metavar="P", required=True, + help="number of `multiprocessing` processes") + parser.add_argument( + "-t", "--threads", type=int, metavar="T", required=True, + help="number of `zarr` threads per `multiprocessing` process") + parser.add_argument( + "--save", action="store_true", + help="save analysis result to ZIP without post-processing") + parser.add_argument( + "--load", action="store_true", + help="load analysis result from ZIP and post-process (--procs 1)") + parser.add_argument( + "--debug", action="store_true", + help="print debug information and reduce plot size") + parser.add_argument( + "--prof-time", action="store_true", + help="run analysis inside `cProfile` (--procs 1)") + parser.add_argument( + "--prof-mem", action="store_true", + help="run analysis inside `memray`") + parser.add_argument( + "--prof-file", type=str, metavar="F", + help="profiling output (relative to `./`)") + return parser + + @classmethod + def from_cli_parser(cls, args: Namespace, /, **kwargs) -> Self: + """ + Interpret the CLI arguments returned by the parser from + :py:meth:`.make_cli_parser`. + """ + assert isinstance(args, Namespace) + return cls( + save_data=args.save, load_data=args.load, debug=args.debug, + n_procs=args.procs, n_threads=args.threads, + prof_time=args.prof_time, prof_mem=args.prof_mem, + prof_file=args.prof_file, + **kwargs) + + # ~~~~~~~~~~~~~~~~~ # + + def cfg_error(self, msg: str, path: str, /) -> str: + return (f"{msg}:\n " + f"\"analysis_options.zarr_mapreduce.{self.name}{path}\"") + + def validate(self, workflow: WorkflowConfig, /) -> None: + """ + Parse and validate the analysis configuration. + + Called by: :py:meth:`ZarrMapReduce.__init__`. + + Calls: :py:meth:`.validate_generic`, :py:meth:`.validate_specific`. + """ + # load analysis config from workflow JSON + assert isinstance(workflow, WorkflowConfig) + analysis = (workflow.sim + .get("analysis_options", {}).get("zarr_mapreduce", {}) + .get(self.name, {})) + if not (isinstance(analysis, dict) and analysis): + raise KeyError(self.cfg_error("Missing analysis configuration", "")) + for k in ["paths", "variables", "parameters"]: + if not analysis.get(k, {}): + raise KeyError(self.cfg_error("Missing analysis configuration", + f".{k}")) + for k in ["result", "figure"]: + if not (isinstance(p := analysis["paths"].get(k), str) and p): + raise KeyError(self.cfg_error("Missing analysis configuration", + f".paths.{k}")) + if k == "result" and not p.endswith(".zip"): + raise ValueError(self.cfg_error("Invalid analysis configuration", + f".paths.{k}")) + setattr(self, f"{k}_file", Path(p)) + variables = analysis["variables"] + if not all(p == str(Path(p)) for p in variables): + raise ValueError(self.cfg_error("Invalid analysis configuration", + ".variables")) + self.variables = variables + + # validate against simulation config + self.validate_generic(workflow) + self.validate_specific(workflow) + + def validate_generic(self, workflow: WorkflowConfig, /) -> None: + """ + Analysis-agnostic validation checks of the *analysis configuration* + against the *simulation configuration*. + """ + view = ForestView.from_dict(workflow.sim["emitter_arg"]["view"]) + leaves = {str(lf.path) for lf in view.leaves.values()} + if not set(self.variables.keys()).issubset(leaves): + raise KeyError("Analysis requires variables missing from " + "the emitter configuration.") + + @abstractmethod + def validate_specific(self, workflow: WorkflowConfig, /) -> None: + """ + Analysis-specific validation checks of the *analysis configuration* + against the *simulation configuration*. + """ + ... + + def zarr_config(self) -> None: + """ + Apply configuration options related to the Zarr library. + """ + zarr.config.update({ + "async": {"concurrency": self.n_threads}, + "threading": {"max_workers": self.n_threads}}) + assert zarr.config.get("async.concurrency") == self.n_threads + assert zarr.config.get("threading.max_workers") == self.n_threads + + +# ============================================================================== +# analysis result +# ============================================================================== + + +@dataclass(kw_only=True, frozen=True) +class ZarrMapReduceResult[AnalysisConfigT: ZarrMapReduceConfig](ABC): + """ + Output type of a :py:class:`ZarrMapReduce` pipeline, with dedicated + serialisation/deserialisation methods. + + .. note:: + + This object should contain the full information required by + post-processing tasks, including metadata. + """ + + @staticmethod + @abstractmethod + def zarr_config(cfg: AnalysisConfigT, /) -> None: + """ + Apply configuration options related to the Zarr library. + """ + ... + + @abstractmethod + def to_zarr(self, cfg: AnalysisConfigT, /) -> None: + """ + Serialise the pipeline results into a Zarr ZIP file. + + Calls: :py:meth:`.zarr_config`. + """ + ... + + @classmethod + @abstractmethod + def from_zarr(self, cfg: AnalysisConfigT, /) -> Self: + """ + Deserialise the pipeline results from a Zarr ZIP file. + + Calls: :py:meth:`.zarr_config`. + """ + ... + + @abstractmethod + def to_pandas( + self, workflow_cfg: WorkflowConfig, + /, max_rows_per_panel: int + ) -> Any: + r""" + Export the pipeline results into a *grid* of `long-form`_ + :py:class:`~pandas.DataFrame`\ s suitable for plotting with + `Vega-Altair`_. + + Called by: :py:meth:`.ZarrMapReducePlot.import_data`. + + .. note:: + The issue of `dataset size`_ is controlled by the following measures: + + - Each panel in the grid is exported as a separate + :py:class:`~pandas.DataFrame`, rather than using `repeated`_ or + `faceted`_ charts. + - The size of each output :py:class:`~pandas.DataFrame` is limited by + :py:attr:`.ZarrMapReducePlotConfig.max_rows_per_panel`, and + `Vega-Altair`_ is configured accordingly. + + .. _long-form: https://altair-viz.github.io/user_guide/data.html#long-form-vs-wide-form-data + .. _Vega-Altair: https://altair-viz.github.io/index.html + .. _dataset size: https://altair-viz.github.io/user_guide/large_datasets.html#large-datasets + .. _repeated: https://altair-viz.github.io/user_guide/compound_charts.html#repeated-charts + .. _faceted: https://altair-viz.github.io/user_guide/compound_charts.html#faceted-charts + """ + ... + + +# ============================================================================== +# analysis pipeline +# ============================================================================== + + +@dataclass(kw_only=True, slots=True, frozen=True) +class CoordDescriptor: + """ + Analysis-related metadata and subarray `selector`_ for a single simulation + variable. + + Constructed by: :py:meth:`.ZarrMapReduce.select_variable_coordinate` and + :py:meth:`.ZarrMapReduce.select_time_coordinate`. + + .. _selector: https://zarr.readthedocs.io/en/stable/user-guide/arrays/#advanced-indexing + """ + + #: Labels of coordinate dimensions. Used for plotting. + dim_names: list[str] + #: A `selector`_ applicable to the Zarr array representing a variable. + selector: Selector + #: Unit annotation. + unit: str | None + + def __post_init__(self) -> None: + assert isinstance(self.dim_names, list) + assert all(isinstance(d, str) for d in self.dim_names) + assert isinstance(self.unit, str | None) + + def __eq__(self, other, /) -> bool: + return (self.unit, self.dim_names) == (other.unit, other.dim_names) + + +# ------------------------------------------------------------------------------ + + +class ZarrMapReduce[ + AnalysisConfigT: ZarrMapReduceConfig, + ResultT: ZarrMapReduceResult, + PlotConfigT: ZarrMapReducePlotConfig, +](ABC): + """ + Driver class for map-reduce pipelines on Zarr stores produced by the + :py:class:`.XarrayEmitter`. + + This abstract class defines a specialised data flow, which takes into + account the full (in-)dependence assumptions underlying the :ref:`storage + ` and :ref:`variable ` layouts of the + :py:class:`.XarrayEmitter`, and which leverages both + :py:mod:`multiprocessing` parallelism and the + :py:mod:`~zarr.api.asynchronous` Zarr API. The data flow is organised into a + method call hierarchy as follows:: + + .compute() + ├─ .mapreduce_workflow() + │ ├─ [parallel] .sync_mapreduce_substore() + │ │ └─ [async] .mapreduce_substore() + │ │ ├─ [async] .load_substore() + │ │ │ └─ [async] zarr.open_group() + │ │ ├─ .select_partitions() ; (abstract) + │ │ ├─ [async] .select_substore_coordinates() + │ │ │ ├─ [async] zarr.getitem() + │ │ │ └─ [async] .select_variable_coordinate() ; (abstract) + │ │ ├─ [async] .map_substore() + │ │ │ └─ [async] .mapreduce_partition() + │ │ │ ├─ [async] zarr.getitem() + │ │ │ ├─ [async] .select_time_coordinate() ; (abstract) + │ │ │ ├─ [async] .map_partition() + │ │ │ │ └─ [async] .mapreduce_variable() + │ │ │ │ ├─ [async] zarr.getitem() + │ │ │ │ ├─ [async] .get_array_selection() ; (abstract) + │ │ │ │ └─ [async] .reduce_variable() ; (abstract) + │ │ │ └─ [async] .reduce_partition() ; (abstract) + │ │ └─ .reduce_substore() ; (abstract) + │ └─ .reduce_workflow() ; (abstract) + └─ .post_process() ; (abstract) + + Concrete pipelines are defined by subclasses overloading the abstract + methods, which are provided with full access to the locally relevant + configuration options and Zarr objects at each level of the pipeline. For a + complete example, see :py:class:`.MovingAvgPipeline`. + + .. note:: + + This design is motivated by the following considerations: + + - Within a *workflow*, ``mapreduce`` operations on independent substores + can be processed in parallel, up to the final top level ``reduce`` + operation. + - Within each *independent substore* (lineage), the relevant partitions + (generations) and variable coordinate subsets can be selected, in an + analysis-specific way, at the top level without accessing all child + arrays. Such selections help avoid unnecessary data transfers and + decompressions of child arrays. + - Similarly, within each *partition*, the relevant time coordinate subset + can be selected, in an analysis-specific way, at the top level. + - Once these selections are made, the retrieval of relevant child array + subsets and the evaluation of their local ``reduce`` operations can be + performed asynchronously, thus mitigating the latency cost of child + arrays. + - By performing intermediate ``reduce`` operations at the partition level + and at the substore level, data marshalling can be further reduced. + """ + + def __init__( + self, workflow_cfg: WorkflowConfig, analysis_cfg: AnalysisConfigT, + plot_cfg: PlotConfigT + ) -> None: + assert isinstance(workflow_cfg, WorkflowConfig) + assert issubclass(self.config_type, ZarrMapReduceConfig) + assert isinstance(analysis_cfg, self.config_type) + assert isinstance(plot_cfg, ZarrMapReducePlotConfig) + self.workflow_cfg = workflow_cfg + self.analysis_cfg = analysis_cfg + self.analysis_cfg.validate(self.workflow_cfg) + self.analysis_cfg.zarr_config() + self.plot_cfg = plot_cfg + self.store = WorkflowPaths.locate(self.workflow_cfg) + + @property + @abstractmethod + def config_type(self) -> type[AnalysisConfigT]: + """ + Expected type of configuration data class. + """ + ... + + @property + @abstractmethod + def result_type(self) -> type[ResultT]: + """ + Expected type of pipeline result. + """ + ... + + # ~~~~~~~~~~~~~~~~~ # + # workflow level + # ~~~~~~~~~~~~~~~~~ # + + def compute(self) -> ResultT: + """ + Driver method. + + If a profiler is set in the :py:attr:`ZarrMapReduceConfig` constructor + argument, then the map-reduce pipeline is wrapped inside a profiler, and + post-processing is skipped. Otherwise, the map-reduce pipeline is + executed and post-processed. + + In addition, :py:attr:`.ZarrMapReduceConfig.save_data` and + :py:attr:`.ZarrMapReduceConfig.load_data` can be used to skip repeated + executions of the map-reduce pipeline, e.g., while the post-processing + code is under active development. + + Calls: :py:meth:`.mapreduce_workflow` and :py:meth:`.post_process`. + """ + if (config := self.analysis_cfg).prof_time: + # execute analysis pipeline inside time profiler + from cProfile import Profile + with Profile() as prof: + result = self.mapreduce_workflow() + prof.dump_stats(cast(Path, config.prof_file)) + + elif config.prof_mem: + # execute analysis pipeline inside memory profiler + try: + from memray import Tracker + except ImportError: + raise ImportError("Install the [prof] extra for profiling.") + with Tracker(file_name=cast(Path, config.prof_file), + native_traces=True, follow_fork=True): + result = self.mapreduce_workflow() + + else: + # execute analysis pipeline or load its result + result = (self.result_type.from_zarr(self.analysis_cfg) + if self.analysis_cfg.load_data else + self.mapreduce_workflow()) + # store or post-process pipeline result + if self.analysis_cfg.save_data: + result.to_zarr(self.analysis_cfg) + else: + self.post_process(self.workflow_cfg, self.analysis_cfg, result, + self.plot_cfg) + + assert isinstance(result, self.result_type) + return result + + @classmethod + @abstractmethod + def post_process( + cls, workflow_cfg: WorkflowConfig, analysis_cfg: AnalysisConfigT, + result: ResultT, plot_cfg: PlotConfigT + ) -> None: + """ + Downstream uses of map-reduce pipeline results, e.g., plot rendering. + + Called by: :py:meth:`.compute`. + """ + ... + + def mapreduce_workflow(self) -> ResultT: + """ + Entry point for the global ``mapreduce`` data flow. + + This method spawns a pool of parallel processes for executing + ``mapreduce`` operations at the level of independent substores, and then + performs the final global ``reduce`` operation. + + Called by: :py:meth:`.compute`. + + Calls: :py:meth:`.sync_mapreduce_substore`, :py:meth:`.reduce_workflow`. + """ + worker = PoolDictItem( + Substore.identity, + partial(self.sync_mapreduce_substore, + self.workflow_cfg, self.analysis_cfg, self.store.root)) + return self.reduce_workflow(pool_dict( + worker, iter(self.store), + n_items=len(self.store), n_procs=self.analysis_cfg.n_procs)) + + @abstractmethod + def reduce_workflow( + self, workflow_map: dict[Substore, Reducer], / + ) -> ResultT: + """ + Execute the final, global ``reduce`` operation. + + Called by: :py:meth:`.mapreduce_workflow`. + """ + ... + + # ~~~~~~~~~~~~~~~~~ # + # substore level + # ~~~~~~~~~~~~~~~~~ # + + @classmethod + def sync_mapreduce_substore( + cls, workflow_cfg: WorkflowConfig, analysis_cfg: AnalysisConfigT, + root: str, substore: Substore, / + ) -> Reducer: + """ + Entry point for the ``mapreduce`` data flow at the level of an + independent substore. + + This method is merely a synchronisation barrier for + :py:meth:`.mapreduce_substore`. + + Called by: :py:meth:`.mapreduce_workflow`. + + Calls: :py:meth:`.mapreduce_substore`. + """ + return sync( + cls.mapreduce_substore(workflow_cfg, analysis_cfg, root, substore)) + + @classmethod + async def mapreduce_substore( + cls, workflow_cfg: WorkflowConfig, analysis_cfg: AnalysisConfigT, + root: str, substore: Substore, / + ) -> Reducer: + """ + Access the root of an independent substore, select relevant partition + subtrees, coordinate variables and dimensions, and execute ``mapreduce`` + operations up to the level of the independent substore. + + Called by: :py:meth:`.sync_mapreduce_substore`. + + Calls: :py:meth:`.load_substore`, :py:meth:`.select_partitions`, + :py:meth:`.select_substore_coordinates`, :py:meth:`.map_substore`, + :py:meth:`.reduce_substore`. + """ + group, partitions = await cls.load_substore(workflow_cfg, root, substore) + var_dscrs = await cls.select_substore_coordinates(analysis_cfg, group) + return cls.reduce_substore( + analysis_cfg, var_dscrs, + await cls.map_substore( + analysis_cfg, group, + cls.select_partitions(analysis_cfg, substore, partitions), + {v: d.selector for (v, d) in var_dscrs.items()})) + + @classmethod + async def load_substore( + cls, config: WorkflowConfig, root: str, substore: Substore, / + ) -> tuple[zarr.AsyncGroup, list[XarrayStoragePartition]]: + """ + Access the root of an independent substore, and detect its partitions. + + Called by: :py:meth:`.mapreduce_substore`. + """ + group: AsyncGroup = await open_group( + root, path=join(*astuple(substore)), + mode="r", use_consolidated=True) + t_coords: set[str] = { + key async for key in group.array_keys() + if XarrayStoragePartition.is_time_coo_name(key)} + if (n_gens := len(t_coords)) > config.sim["generations"]: + raise ValueError( + f"Unexpected number of generations in \"{substore}\": {n_gens}") + partitions = [ + XarrayStoragePartition.from_substore(config, substore, g) + for g in range(1, 1 + n_gens)] + if t_coords != {p.time_coo_name for p in partitions}: + raise ValueError( + f"Unexpected generations in \"{substore}\": {t_coords}") + return (group, partitions) + + @staticmethod + @abstractmethod + def select_partitions( + cfg: AnalysisConfigT, substore: Substore, + partitions: list[XarrayStoragePartition], / + ) -> list[XarrayStoragePartition]: + """ + Select relevant partitions within an independent substore. + + Called by: :py:meth:`.mapreduce_substore`. + """ + ... + + @classmethod + async def select_substore_coordinates( + cls, cfg: AnalysisConfigT, group: zarr.AsyncGroup, / + ) -> dict[VariablePath, CoordDescriptor]: + """ + Select relevant coordinate variables and dimensions within an + independent substore. + + Called by: :py:meth:`.mapreduce_substore`. + + Calls: :py:meth:`.select_variable_coordinate`. + """ + return await async_dict( + AsyncDictItem(path, path, + cls.select_variable_coordinate( + cfg, path, + ((await get_async_group(group, path)) + .attrs.get(var_name(path))), + await get_async_array(group, coo_path(path)), + match_coo)) + for (path, match_coo) in cfg.variables.items()) + + @staticmethod + @abstractmethod + async def select_variable_coordinate( + cfg: AnalysisConfigT, path: VariablePath, attr: Any, + coo: zarr.AsyncArray, match_coo: str, / + ) -> CoordDescriptor: + """ + Select relevant coordinate dimensions within a selected variable. + + Called by: :py:meth:`.select_substore_coordinates`. + """ + ... + + @classmethod + async def map_substore( + cls, cfg: AnalysisConfigT, + group: zarr.AsyncGroup, partitions: list[XarrayStoragePartition], + var_ixs: dict[VariablePath, Selector], / + ) -> dict[XarrayStoragePartition, Reducer]: + """ + Execute ``mapreduce`` operations for all relevant partitions within an + independent substore. + + Called by: :py:meth:`.mapreduce_substore`. + + Calls: :py:meth:`.mapreduce_partition`. + """ + return await async_dict( + AsyncDictItem(p, str(p.generation), + cls.mapreduce_partition(cfg, group, p, var_ixs)) + for p in partitions) + + @staticmethod + @abstractmethod + def reduce_substore( + cfg: AnalysisConfigT, + var_dscrs: dict[VariablePath, CoordDescriptor], + substore_map: dict[XarrayStoragePartition, Reducer], / + ) -> Reducer: + """ + Execute the ``reduce`` operation at the level of an independent + substore. + + Called by: :py:meth:`.mapreduce_substore`. + """ + ... + + # ~~~~~~~~~~~~~~~~~ # + # partition level + # ~~~~~~~~~~~~~~~~~ # + + @classmethod + async def mapreduce_partition( + cls, cfg: AnalysisConfigT, group: zarr.AsyncGroup, + p: XarrayStoragePartition, var_ixs: dict[VariablePath, Selector], / + ) -> Reducer: + """ + Select the relevant time coordinate dimensions within a relevant + partition, and execute ``mapreduce`` operations up to the level of the + relevant partition. + + Called by: :py:meth:`.map_substore`. + + Calls: :py:meth:`.select_time_coordinate`, :py:meth:`.map_partition`, + :py:meth:`.reduce_partition`. + """ + time_coo = await get_async_array(group, p.time_coo_name) + time_var = await get_async_array(group, p.time_var_name) + time_dscr = await cls.select_time_coordinate( + cfg, p, cast(str, group.attrs[p.time_var_name]), + time_coo, time_var) + return await cls.reduce_partition( + cfg, time_dscr, time_coo, time_var, + await cls.map_partition(cfg, group, p, time_dscr.selector, var_ixs)) + + @staticmethod + @abstractmethod + async def select_time_coordinate( + cfg: AnalysisConfigT, p: XarrayStoragePartition, attr: str, + time_coo: zarr.AsyncArray, time_var: zarr.AsyncArray, / + ) -> CoordDescriptor: + """ + Select relevant time coordinate dimensions within a relevant partition. + + Called by: :py:meth:`.mapreduce_partition`. + """ + ... + + @classmethod + async def map_partition( + cls, cfg: AnalysisConfigT, group: zarr.AsyncGroup, + p: XarrayStoragePartition, + time_ix: Selector, var_ixs: dict[VariablePath, Selector], / + ) -> dict[VariablePath, Reducer]: + """ + Execute ``mapreduce`` operations for all relevant variables within a + relevant partition. + + Called by: :py:meth:`.mapreduce_partition`. + + Calls: :py:meth:`.mapreduce_variable`. + """ + return await async_dict( + AsyncDictItem(path, path, + cls.mapreduce_variable( + cfg, group, p, path, time_ix, var_ix)) + for (path, var_ix) in var_ixs.items()) + + @staticmethod + @abstractmethod + async def reduce_partition( + cfg: AnalysisConfigT, time_dscr: CoordDescriptor, + time_coo: zarr.AsyncArray, time_var: zarr.AsyncArray, + partition_map: dict[VariablePath, Reducer], / + ) -> Reducer: + """ + Execute the ``reduce`` operation at the level of a relevant partition, + using both the time coordinate and the inner reduced variables. + + Called by: :py:meth:`.mapreduce_partition`. + """ + ... + + # ~~~~~~~~~~~~~~~~~ # + # variable level + # ~~~~~~~~~~~~~~~~~ # + + @classmethod + async def mapreduce_variable( + cls, cfg: AnalysisConfigT, group: zarr.AsyncGroup, + p: XarrayStoragePartition, path: VariablePath, + time_ix: Selector, var_ix: Selector, / + ) -> Reducer: + """ + Retrieve the relevant subset of a relevant simulation variable, and + execute the local ``reduce`` operation for it. + + Called by: :py:meth:`.map_partition`. + + Calls: :py:meth:`.get_array_selection`, :py:meth:`.reduce_variable`. + """ + return await cls.reduce_variable( + cfg, path, + await cls.get_array_selection( + cfg, path, time_ix, var_ix, + await get_async_array(group, join(path, p.dynamic_suffix)))) + + @staticmethod + @abstractmethod + async def get_array_selection( + cfg: AnalysisConfigT, path: VariablePath, + time_ix: Selector, var_ix: Selector, var: zarr.AsyncArray, / + ) -> NDArrayLikeOrScalar: + """ + Load and decompress the relevant subset of a relevant variable. + + Called by: :py:meth:`.mapreduce_variable`. + + .. note:: + + What fraction of the underlying chunked array will actually be loaded + into memory and decompressed depends on the ordering of dimensions, on + chunk sizes, on the filter and compression codecs, as well as on Zarr + library internals that may change over time. + + When preparing large simulation and analysis workflows, in particular + for analyses that require only small fractions of stored arrays, users + may benefit from empirically testing the configuration-specific + resource usage by profiling sample analyses, see :py:meth:`.compute`. + """ + ... + + @staticmethod + @abstractmethod + async def reduce_variable( + cfg: AnalysisConfigT, path: VariablePath, data: NDArrayLikeOrScalar, / + ) -> Reducer: + """ + Execute the ``reduce`` operation at the level of a single relevant + variable, without using the time coordinate yet. + + Called by: :py:meth:`.mapreduce_variable`. + """ + ... + + +# ============================================================================== +# analysis plot +# ============================================================================== + + +@dataclass(kw_only=True, slots=True) +class ZarrMapReducePlotConfig(ABC): + """ + Configuration of a :py:class:`ZarrMapReducePlot`. This object is expected to + mostly contain analysis-specific rendering options. + """ + + #: Used by :py:meth:`.ZarrMapReduceResult.to_pandas`. + max_rows_per_panel: int = 20_000 + + +# ------------------------------------------------------------------------------ + + +@dataclass(slots=True) +class ZarrMapReducePlot[ + AnalysisConfigT: ZarrMapReduceConfig, + ResultT: ZarrMapReduceResult, + PlotConfigT: ZarrMapReducePlotConfig +](ABC): + """ + Container for the full information required to produce analysis-specific + plots. + """ + + workflow_cfg: WorkflowConfig + analysis_cfg: AnalysisConfigT + result: ResultT + plot_cfg: PlotConfigT + plot_data: Any = None + + # ~~~~~~~~~~~~~~~~~ # + + def __post_init__(self) -> None: + assert isinstance(self.workflow_cfg, WorkflowConfig) + assert isinstance(self.analysis_cfg, ZarrMapReduceConfig) + assert isinstance(self.result, ZarrMapReduceResult) + assert isinstance(self.plot_cfg, ZarrMapReducePlotConfig) + assert isinstance(self.plot_cfg.max_rows_per_panel, int) + self.import_data() + + def import_data(self) -> None: + """ + Load plot data in a format suitable for the rendering engine. + + Called by: ``.__post_init__()``. + + Calls: :py:meth:`.ZarrMapReduceResult.to_pandas`. + """ + self.plot_data = self.result.to_pandas(self.workflow_cfg, + self.plot_cfg.max_rows_per_panel) + + @abstractmethod + def render(self) -> Any: + """ + Driver method. Renders and stores plots. + """ + ... diff --git a/viva_emitters/xarray_emitter/zarr_utils.py b/viva_emitters/xarray_emitter/zarr_utils.py new file mode 100644 index 0000000..8c760f9 --- /dev/null +++ b/viva_emitters/xarray_emitter/zarr_utils.py @@ -0,0 +1,228 @@ + +""" +Utilities for controlling low-level Zarr internals (read side). + +Ported from vEcoli's ``ecoli.library.xarray_emitter.zarr_utils`` (PR #414). Only +the store-access and codec-parsing helpers used by the map-reduce engine +(:py:mod:`.zarr_mapreduce`) and by downstream analysis pipelines are kept here; +the consolidated-metadata management functions live in +:py:mod:`viva_emitters.xarray_emitter.zarr_writer` (viva-emitters keeps its own +diverged implementation of those). +""" + +from lzma import FILTER_LZMA2, FORMAT_RAW +from typing import Any, Literal, cast + +from numpy.typing import NDArray +from zarr.abc.codec import Codec +from zarr.abc.numcodec import Numcodec +from zarr.core.array import ( + Array, + AsyncArray, + _parse_chunk_encoding_v2, + default_compressors_v3, +) +from zarr.core.dtype import parse_dtype +from zarr.core.group import AsyncGroup, Group +from zarr.core.indexing import BlockIndex +from zarr.core.metadata import v2, v3 +from zarr.errors import UnstableSpecificationWarning, ZarrUserWarning + +from .storage import VariableEncoding +from .utils import WarningFilter, filter_warnings + +# ============================================================================== +# Zarr warnings +# ============================================================================== + + +zarr_warnings: dict[str, WarningFilter] = { + "consolidated_metadata": WarningFilter( + module="zarr.api.asynchronous", + category=ZarrUserWarning, + message="Consolidated metadata.*Zarr format 3", + action="ignore"), + "string": WarningFilter( + module="zarr.core.dtype.npy.string", + category=UnstableSpecificationWarning, + message=".*data type.*Zarr V3", + action="ignore"), + "numcodecs": WarningFilter( + module="zarr.codecs.numcodecs", + category=ZarrUserWarning, + message=".*Numcodecs codecs.*Zarr version 3 specification", + action="ignore"), + "zarrs": WarningFilter( + module="zarrs.pipeline", + category=UserWarning, + message="Array is unsupported by ZarrsCodecPipeline", + action="ignore") +} + + +# ============================================================================== +# Zarr codecs +# ============================================================================== + + +ZARR_FILTERS: dict[str, dict[int, list[dict[str, Any]]]] = { + "delta": { + 2: [{"id": "delta", "dtype": None}], + 3: [{"name": "numcodecs.delta", "configuration": {"dtype": None}}] + }, + "num": { + 2: [], + 3: [] + }, +} +""" +Default filter codecs for :py:meth:`.AsyncZarrBufferWriter.var_codecs`, as a +function of the array category and the Zarr format. +""" +ZARR_COMPRESSORS: dict[str, dict[int, list[dict[str, Any]]]] = { + "delta": { + 2: [{"id": "blosc", "cname": "zstd", "clevel": 6, + "shuffle": -1, "blocksize": 0}], + 3: [{"name": "blosc", "configuration": { + "cname": "zstd", "clevel": 6, + "typesize": None, "shuffle": None, "blocksize": 0}}] + }, + "num": { + 2: [{"id": "lzma", "check": -1, "preset": None, + "format": FORMAT_RAW, + "filters": [{"id": FILTER_LZMA2, "preset": 6}]}], + 3: [{"name": "numcodecs.lzma", + "configuration": {"format": FORMAT_RAW, + "filters": [{"id": FILTER_LZMA2, "preset": 6}]}}] + }, +} +""" +Default compression codecs for :py:meth:`.AsyncZarrBufferWriter.var_codecs`, as +a function of the array category and the Zarr format. +""" + + +# ------------------------------------------------------------------------------ + + +def parse_codecs( + zarr_format: Literal[2, 3], /, *, + codecs: dict[str, Any] | None = None, + category: str | None = None, + dtype: str | None = None +) -> VariableEncoding: + """ + Translate a Vivarium JSON configuration into a codec specification that is + interpretable by the Zarr API, leveraging Zarr internal functions for + parsing. + + Used by: :py:meth:`.AsyncZarrBufferWriter.coo_codecs`, + :py:meth:`.AsyncZarrBufferWriter.var_codecs`. + """ + filters: tuple[Codec | Numcodec, ...] | None + compressors: tuple[Codec | Numcodec | None, ...] + + # dispatch on spec type + z = zarr_format + if codecs or (category is not None): + # fetch non-default config + if codecs: + # fetch custom JSON config + assert category is None + assert dtype is None + _filters = codecs.get(f"filters_v{z}", []) + _compressors = codecs.get(f"compressors_v{z}", []) + if not (_filters or _compressors): + raise KeyError( + f"Missing arguments:\n " + f"{{\"filters_v{z}\": ..., \"compressors_v{z}\": ...}}") + elif category is not None: + # fetch library preset, supply data type information + _filters = ZARR_FILTERS[category][z] + _compressors = ZARR_COMPRESSORS[category][z] + assert dtype is not None + for f in _filters: + if z == 2: + f["dtype"] = dtype + else: + f["configuration"]["dtype"] = dtype + # parse non-default config + with filter_warnings(list(zarr_warnings.values())): + if z == 2: + filters = v2.parse_filters(_filters) + compressors = tuple(map(v2.parse_compressor, _compressors)) + else: + filters = v3.parse_codecs(_filters) + compressors = v3.parse_codecs(_compressors) + else: + # fetch default config, supply data type information + assert dtype is not None + _dtype = parse_dtype(dtype, zarr_format=z) + if z == 2: + filters, compressor = _parse_chunk_encoding_v2( + filters="auto", compressor="auto", dtype=_dtype + ) + compressors = (compressor,) + else: + filters = () + compressors = default_compressors_v3(_dtype) + + return {"filters": filters, "compressors": compressors} + + +# ============================================================================== +# Zarr store access +# ============================================================================== + + +def get_group(group: Group, path: str) -> Group: + """ + Access a Zarr store path known to be a :py:class:`~zarr.Group`. + """ + return cast(Group, group[path]) + + +def get_array(group: Group, path: str) -> Array: + """ + Access a Zarr store path known to be an :py:class:`~zarr.Array`. + """ + return cast(Array, group[path]) + + +def get_ndarray(group: Group, path: str) -> NDArray: + """ + Access a Zarr store path known to be an :py:class:`~zarr.Array`, and + retrieve the uncompressed array data. + """ + return cast(NDArray, get_array(group, path)[:]) + + +def get_rectilinear_ndarray(group: Group, path: str) -> list: + """ + Given a Zarr array that was stored using a `rectilinear chunk grid`_, return + a list over its block views. + + .. _rectilinear chunk grid: https://zarr.readthedocs.io/en/stable/user-guide/examples/rectilinear_chunks/ + """ + arr = get_array(group, path) + blk_view: BlockIndex = arr.blocks + blk_ix = [range(len(dim)) for dim in arr.write_chunk_sizes] + def get_blocks(outer: tuple[int,...], inner: list[range]): + return ([get_blocks(outer + (i,), inner[1:]) for i in inner[0]] + if inner else + blk_view[outer]) + return get_blocks((), blk_ix) + + +async def get_async_group(group: AsyncGroup, path: str) -> AsyncGroup: + """ + Access a Zarr store path known to be an :py:class:`~zarr.AsyncGroup`. + """ + return cast(AsyncGroup, await group.getitem(path)) + + +async def get_async_array(group: AsyncGroup, path: str) -> AsyncArray: + """ + Access a Zarr store path known to be an :py:class:`~zarr.AsyncArray`. + """ + return cast(AsyncArray, await group.getitem(path))