diff --git a/.gitignore b/.gitignore index cf55a3998..514d3deff 100644 --- a/.gitignore +++ b/.gitignore @@ -69,4 +69,4 @@ trace-* test_sherlock/ # SMS API # -.hpc_env \ No newline at end of file +.hpc_env diff --git a/configs/test_configs/moving_avg_analysis.json b/configs/test_configs/moving_avg_analysis.json new file mode 100644 index 000000000..030d77023 --- /dev/null +++ b/configs/test_configs/moving_avg_analysis.json @@ -0,0 +1,63 @@ +{ + "experiment_id": "moving_avg", + "sim_data_path": "out/kb/simData.cPickle", + "suffix_time": false, + "fail_at_max_duration": true, + "generations": 6, + "n_init_sims": 10, + "skip_baseline": false, + "variants": { + "condition": {"condition": {"value": ["with_aa", "acetate"]}} + }, + "emitter": "xarray", + "emitter_arg": { + "out_dir": "out/experiments", + "transducer": { + "predicate": [ + [{"subsample": {"interval": 3}}] + ], + "buffer": {"size": 200} + }, + "writer": { + "threaded": true, + "buffers_per_chunk": 1, + "backend": "zarr", + "backend_config": { + "format": 3, + "async.concurrency": 3, + "threading.max_workers": 3 + } + }, + "view": [ + { + "root": ["listeners"], + "variables": { + "fba_results": { + "reaction_fluxes": [{ + "path": "metabolic_fluxes", + "unit": "[mmol/L.s]", + "dtype": "` or the + :py:mod:`ecoli.library.xarray_emitter`. The Parquet emitter is unable to serialize many of the objects contained in process updates (e.g., nested lists of inconsistent depth like ``[[1], 2]``). + ------------- Initial State ------------- diff --git a/doc/conf.py b/doc/conf.py index 9abcc5c31..89683f985 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -75,17 +75,43 @@ ("py:class", "numpy.int64"), ("py:class", "numpy.int32"), ("py:class", "numpy.bool_"), + ("py:class", "numpy._typing._array_like._ScalarType_co"), ("py:class", "duckdb.duckdb.DuckDBPyConnection"), + ("py:class", "pandas.core.series.Series"), ("py:class", "pandas.core.frame.DataFrame"), ("py:class", "polars.series.series.Series"), ("py:class", "polars.dataframe.frame.DataFrame"), ("py:class", "polars.datatypes.classes.DataTypeClass"), + ("py:class", "altair.vegalite.v5.api.LayerChart"), + ("py:class", "altair.vegalite.v5.api.VConcatChart"), # No docs for fsspec ("py:class", "fsspec.spec.AbstractFileSystem"), ("py:class", "fsspec.core.OpenFile"), # Silence warning in ecoli.processes.environment.field_timeline.FieldTimeline ("py:class", "vivarium.processes.timeline.TimelineProcess"), ("py:class", "concurrent.futures._base.Future"), + # Type annotations using library internals + ("py:class", "_asyncio.Future"), + ("py:class", "asyncio.taskgroups.TaskGroup"), + ("py:class", "concurrent.futures._base.Executor"), + ("py:class", "unittest.mock._patch"), + ("py:class", "xarray.backends.common.ArrayWriter"), + ("py:class", "xarray.core.treenode.NodePath"), + ("py:class", "zarr.core.group.ConsolidatedMetadata"), + ("py:class", "zarr.core._tree.TreeRepr"), + # Sphinx does not recognize type parameters in generic classes + ("py:class", "ArrT"), ("py:class", "NodeT"), ("py:class", "StoreT"), + ("py:class", "ElemT"), + ("py:class", "KeyT"), ("py:class", "KeyT1"), ("py:class", "KeyT2"), + ("py:class", "InputT"), ("py:class", "ResultT"), + ("py:class", "AnalysisConfigT"), ("py:class", "PlotConfigT"), + # Sphinx does not recognize type aliases + ("py:type", "VariableEncoding"), ("py:class", "VariableEncoding"), + ("py:type", "VariablePath"), ("py:class", "VariablePath"), + ("py:type", "AnyAsyncArray"), ("py:class", "AnyAsyncArray"), + ("py:type", "AnyGenerator"), ("py:class", "AnyGenerator"), + ("py:type", "Selector"), ("py:class", "Selector"), + ("py:type", "NDArrayLikeOrScalar"), ("py:class", "NDArrayLikeOrScalar"), ] @@ -115,15 +141,17 @@ # -- sphinx.ext.intersphinx options -- intersphinx_mapping = { "python": ("https://docs.python.org/3", None), - "vivarium": ( - "https://vivarium-core.readthedocs.io/en/latest/", - None, - ), + "pytest": ("https://docs.pytest.org/en/latest", None), + "vivarium": ("https://vivarium-core.readthedocs.io/en/latest/", None), "numpy": ("https://numpy.org/doc/stable", None), + "xarray": ("https://docs.xarray.dev/en/latest", None), + "zarr": ("https://zarr.readthedocs.io/en/latest", None), "matplotlib": ("https://matplotlib.org/stable/", None), "pandas": ("http://pandas.pydata.org/pandas-docs/dev", None), "polars": ("https://docs.pola.rs/api/python/stable", None), "sympy": ("https://docs.sympy.org/latest", None), + "pint": ("https://pint.readthedocs.io/en/stable", None), + "unum": ("https://unum.readthedocs.io/en/stable", None), } @@ -143,8 +171,24 @@ ] # Move typehints from signature into description autodoc_typehints = "description" -# Concatenate class and __init__ docstrings -autoclass_content = "both" +# Only use the class’s docstring. __init__ docstrings are now listed separately. +autoclass_content = "class" +# Default options for all autodoc directives. +autodoc_default_options = { + "member-order": "bysource", + "private-members": True, + "special-members": ( + # object + "__init__, __del__, __call__" + ), + "exclude-members": ( + # abc.ABC + "_abc_impl, " + # enum.Flag + "_flag_mask_, _singles_mask_, _all_bits_, _boundary_, _inverted_, " + "_generate_next_value_" + ) +} # Remove domain objects (e.g. functions, classes, attributes) from # table of contents toc_object_entries = False diff --git a/doc/experiments.rst b/doc/experiments.rst index 9e6a423e2..2fe680faa 100644 --- a/doc/experiments.rst +++ b/doc/experiments.rst @@ -396,6 +396,9 @@ Here are some general rules to remember when writing your own JSON config files: without dividing, this results in a more informative error message instead of a Nextflow error about missing daughter cell states. + +.. _experiment_output: + ------ Output ------ diff --git a/doc/stores.rst b/doc/stores.rst index d1a379355..4a28d433d 100644 --- a/doc/stores.rst +++ b/doc/stores.rst @@ -368,6 +368,8 @@ any number of attributes for all active (``_entryState`` is 1) unique molecules of a given type (e.g. RNA, active RNAP, etc.). +.. _listeners: + --------- Listeners --------- diff --git a/doc/workflows.rst b/doc/workflows.rst index b73ab6624..67f55bf52 100644 --- a/doc/workflows.rst +++ b/doc/workflows.rst @@ -330,6 +330,7 @@ it has access to. folder, you can just create stub files in the appropriate folders that simply import the ``plot`` function from a primary analysis script. + .. _analysis_config: Configuration @@ -502,6 +503,9 @@ Refer to :ref:`/output.rst` for more information about how to use DuckDB to read and analyze simulation output inside analysis scripts. + +.. _workflows: + --------- Workflows --------- @@ -706,7 +710,7 @@ is a list workflow behaviors enabled in our model to handle unexpected errors. The only exceptions are changes to resource allocation options (e.g., ``SIM_MEM``), allowing users to retry failed jobs with higher resource limits without triggering re-execution of already completed jobs. If you want to change other options, you - must launch a new workflow. + must launch a new workflow. .. _output: diff --git a/ecoli/__init__.py b/ecoli/__init__.py index 673b8391f..24e2c5c2f 100644 --- a/ecoli/__init__.py +++ b/ecoli/__init__.py @@ -6,6 +6,7 @@ ) from ecoli.library.parquet_emitter import ParquetEmitter +from ecoli.library.xarray_emitter.emitter import XarrayEmitter from ecoli.library.schema import ( divide_binomial, divide_bulk, @@ -39,6 +40,7 @@ faulthandler.enable() emitter_registry.register("parquet", ParquetEmitter) +emitter_registry.register("xarray", XarrayEmitter) # register :term:`updaters` inverse_updater_registry.register("accumulate", inverse_update_accumulate) diff --git a/ecoli/analysis/xarray_emitter/__init__.py b/ecoli/analysis/xarray_emitter/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/ecoli/analysis/xarray_emitter/mapreduce.py b/ecoli/analysis/xarray_emitter/mapreduce.py new file mode 100644 index 000000000..d81b0d293 --- /dev/null +++ b/ecoli/analysis/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/ecoli/analysis/xarray_emitter/zarr_mapreduce.py b/ecoli/analysis/xarray_emitter/zarr_mapreduce.py new file mode 100644 index 000000000..b73bf3a2d --- /dev/null +++ b/ecoli/analysis/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 ecoli.library.xarray_emitter.storage import ( + Substore, + WorkflowConfig, + WorkflowPaths, + XarrayStoragePartition, + coo_path, + var_name, +) +from ecoli.library.xarray_emitter.view import ForestView +from ecoli.library.xarray_emitter.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 ecoli.library.xarray_emitter.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/ecoli/composites/ecoli_master_tests.py b/ecoli/composites/ecoli_master_tests.py index ad60a4e7b..c30cfa051 100644 --- a/ecoli/composites/ecoli_master_tests.py +++ b/ecoli/composites/ecoli_master_tests.py @@ -436,7 +436,6 @@ def test_emit_unique(parquet_out_dir): sim.config["emitter_arg"] = {"out_dir": parquet_out_dir} sim.build_ecoli() sim.run() - sim.ecoli_experiment.emitter.finalize() unique_molecules = sim.ecoli_experiment.state["agents"]["0"]["unique"].inner.keys() history_sql, _, _ = dataset_sql(parquet_out_dir, [sim.experiment_id]) @@ -476,7 +475,6 @@ def test_emit_paths(parquet_out_dir): sim.config["emitter_arg"] = {"out_dir": parquet_out_dir} sim.build_ecoli() sim.run() - sim.ecoli_experiment.emitter.finalize() history_sql, _, _ = dataset_sql(parquet_out_dir, [sim.experiment_id]) conn = create_duckdb_conn() diff --git a/ecoli/experiments/ecoli_master_sim.py b/ecoli/experiments/ecoli_master_sim.py index fce81295f..2495a561f 100644 --- a/ecoli/experiments/ecoli_master_sim.py +++ b/ecoli/experiments/ecoli_master_sim.py @@ -24,15 +24,16 @@ import numpy as np from fsspec import open as fsspec_open from vivarium.core.engine import Engine -from vivarium.core.composer import deep_merge +from vivarium.core.composer import Composite, deep_merge from vivarium.core.process import Process from vivarium.core.serialize import deserialize_value, serialize_value from vivarium.library.dict_utils import deep_merge_check from vivarium.library.topology import inverse_topology from vivarium.library.topology import assoc_path, get_in from ecoli.library.logging_tools import write_json -from wholecell.utils.filepath import cloud_path_join +from ecoli.library.parquet_emitter import BufferedEmitter import ecoli.composites.ecoli_master +from wholecell.utils.filepath import cloud_path_join # Environment composer for spatial environment sim import ecoli.composites.environment.lattice @@ -42,7 +43,6 @@ from ecoli.processes.registries import topology_registry from configs import CONFIG_DIR_PATH -from ecoli.library.parquet_emitter import ParquetEmitter from ecoli.library.schema import not_a_process from wholecell.utils.filepath import ROOT_PATH @@ -256,6 +256,7 @@ def __init__( self.parser.add_argument( "--experiment_id", action="store", + type=str, help=( "ID for this experiment. A UUID will be generated if " 'this argument is not used and "experiment_id" is null ' @@ -346,8 +347,7 @@ def __init__( "--variant", action="store", help="Name of variant." ) self.parser.add_argument( - "--lineage_seed", - action="store", + "--lineage_seed", action="store", type=int, help="Seed used for first cell in lineage.", ) self.parser.add_argument( @@ -468,15 +468,19 @@ def __init__(self, config: dict[str, Any]): # Keep track of base experiment id # in case multiple simulations are run with suffix_time = True. - self.experiment_id_base = config["experiment_id"] + self.experiment_id: str + self.experiment_id_base: str = config["experiment_id"] self.config = config - self.ecoli = None - """vivarium.core.composer.Composite: Contains the fully instantiated - processes, steps, topologies, and flow necessary to run simulation. - Generated by + self.emitter_config: dict[str, Any] = {} + + self.ecoli: Composite + """ + Contains the fully instantiated processes, steps, topologies, and flow + necessary to run simulation. Generated by :py:meth:`~ecoli.experiments.ecoli_master_sim.EcoliSim.build_ecoli` and cleared when :py:meth:`~ecoli.experiments.ecoli_master_sim.EcoliSim.run` - is called to potentially free up memory after division.""" + is called to potentially free up memory after division. + """ self.generated_initial_state = None """dict: Fully populated initial state for simulation. Generated by :py:meth:`~ecoli.experiments.ecoli_master_sim.EcoliSim.build_ecoli` and @@ -746,60 +750,81 @@ def build_ecoli(self): initial_environment, self.generated_initial_state ) - def update_experiment(self, time_to_update: float = 0.0): + def update_experiment( + self, time_to_update: float = 0.0, finalize: bool = True + ) -> None: """ - Runs the E. coli simulation for a specified amount of time. If the - simulation reaches a division event and ``config['generations']`` is set, - it will save the daughter cell states to JSON files in the directory - specified by ``config['daughter_outdir']``. Also creates a file - ``division_time.sh`` that, when executed, sets the environment variable - ``division_time`` to the time at which division occurred (used in - Nextflow workflow runs). + Run the E. coli simulation for a specified amount of time. If the + simulation reaches a division event during this time and + ``config['generations']`` is set, then :py:meth:`~.persist_generation` + will be called and the Python interpreter will be terminated. + + Called by: :py:meth:`.run` or :py:meth:`.save_states`. """ try: - self.ecoli_experiment.update(time_to_update) + success = False + if time_to_update > 0: + self.ecoli_experiment.update(time_to_update) except DivisionDetected: - state = self.ecoli_experiment.state.get_value(condition=not_a_process) - assert len(state["agents"]) == 2 - # Daughter state should include all of the additional - # non-agent state (e.g. environment state) - non_agent_state = {k: v for k, v in state.items() if k != "agents"} - for i, (agent_id, agent_state) in enumerate(state["agents"].items()): - prepare_save_state(agent_state) - daughter_filename = f"daughter_state_{i}.json" - daughter_path = cloud_path_join(self.daughter_outdir, daughter_filename) - write_json( - daughter_path, - {**non_agent_state, "agents": {agent_id: agent_state}}, - ) - # Write daughter state URI to local file for Nextflow to read - with open(f"daughter_state_{i}_uri.txt", "w") as f: - f.write(daughter_path) - print( - f"Divided at t = {self.ecoli_experiment.global_time} after " - f"{self.ecoli_experiment.global_time - self.initial_global_time} sec." + success = True + self.persist_generation() + finally: + # Don't start new I/O operations during a manual shutdown. + if not isinstance(sys.exception(), KeyboardInterrupt): + if isinstance(emitter := self.ecoli_experiment.emitter, BufferedEmitter): + # Finish writing buffered emits to persistent storage, + # unless called inside the `.save_states()` loop. + if finalize: + emitter.finalize(success=success) + if success: + # Exit so that `.run()` does not raise `TimeLimitError`. + sys.exit() + + def persist_generation(self, *, mock: bool = False) -> None: + """ + Upon reaching cell division, save the daughter cell states to JSON files + in the directory specified by ``config['daughter_outdir']``. Also, + create a file ``division_time.sh`` that, when executed, sets the + environment variable ``division_time`` to the time at which division + occurred, as expected by + ``runscripts/nextflow/sim.nf::{simGen0,sim}.output``. + + Called by: :py:meth:`~.update_experiment`. + + Args: + mock: This argument exists solely for testing purposes, and has the + effect of reducing costly file system operations. + """ + state = self.ecoli_experiment.state.get_value(condition=not_a_process) + assert len(state["agents"]) == 2 + # Daughter state should include all of the additional + # non-agent state (e.g. environment state) + non_agent_state = {k: v for k, v in state.items() if k != "agents"} + for i, (agent_id, agent_state) in enumerate(state["agents"].items()): + if mock and i: + # skip export of second daughter state + continue + prepare_save_state(agent_state) + daughter_filename = f"daughter_state_{i}.json" + daughter_path = cloud_path_join(self.daughter_outdir, daughter_filename) + write_json( + daughter_path, + {**non_agent_state, "agents": {agent_id: agent_state}}, ) + # Write daughter state URI to local file for Nextflow to read + with open(f"daughter_state_{i}_uri.txt", "w") as f: + f.write(daughter_path) + print( + f"Divided at t = {self.ecoli_experiment.global_time} after " + f"{self.ecoli_experiment.global_time - self.initial_global_time} sec." + ) + if not mock: # Nextflow workflows will source division time to determine # initial global time to use for daughter cells with open("division_time.sh", "w") as f: f.write(f"export division_time={self.ecoli_experiment.global_time}") - # Tell Parquet emitter that simulation was successful - if isinstance(self.ecoli_experiment.emitter, ParquetEmitter): - self.ecoli_experiment.emitter.success = True - self.ecoli_experiment.emitter.finalize() - # Exit so that EcoliSim.run() does not raise TimeLimitError - sys.exit() - except: # noqa: E722 - # Finish writing any buffered emits to Parquet files if the simulation - # encounters any error (including KeyboardInterrupt) - # We use a bare except instead of finally because we don't want to - # run finalize() every time update_experiment is called to advance to - # save times in save_states() - if isinstance(self.ecoli_experiment.emitter, ParquetEmitter): - self.ecoli_experiment.emitter.finalize() - raise - - def save_states(self): + + def save_states(self) -> None: """ Runs the simulation while saving the states of specific timesteps to files named ``data/vivecoli_t{time}.json``. Invoked by @@ -807,6 +832,10 @@ def save_states(self): if ``config['save'] == True``. State is saved as a JSON that can be reloaded into a simulation as described in :py:meth:`~ecoli.composites.ecoli_master.Ecoli.initial_state`. + + Called by: :py:meth:`.run`. + + Calls: :py:meth:`~.update_experiment`. """ for time in self.save_times: if time > self.max_duration: @@ -820,7 +849,7 @@ def save_states(self): time_to_next_save = self.save_times[i] else: time_to_next_save = self.save_times[i] - self.save_times[i - 1] - self.update_experiment(time_to_next_save) + self.update_experiment(time_to_next_save, finalize=False) time_elapsed = self.save_times[i] state = self.ecoli_experiment.state.get_value(condition=not_a_process) if self.divide: @@ -831,18 +860,19 @@ def save_states(self): write_json("data/vivecoli_t" + str(time_elapsed) + ".json", state) print("Finished saving the state at t = " + str(time_elapsed)) time_remaining = self.max_duration - self.save_times[-1] - if time_remaining: - self.update_experiment(time_remaining) + self.update_experiment(time_remaining) - def run(self): - """Create and run an EcoliSim experiment. If the simulation reaches + def run(self) -> None: + """ + Create and run an EcoliSim experiment. If the simulation reaches the maximum duration specified by ``config['max_duration']``, it will raise a :py:class:`~ecoli.experiments.ecoli_master_sim.TimeLimitError` if ``config['fail_at_max_duration']`` is ``True``. + Calls: :py:meth:`~.update_experiment` or :py:meth:`~.save_states`. + .. WARNING:: - Run :py:meth:`~ecoli.experiments.ecoli_master_sim.EcoliSim.build_ecoli` - before calling :py:meth:`~ecoli.experiments.ecoli_master_sim.EcoliSim.run`! + Run :py:meth:`~.build_ecoli` before calling :py:meth:`~.run`! """ if self.ecoli is None: raise RuntimeError( @@ -858,16 +888,13 @@ def run(self): if self.emitter_arg is not None: for key, value in self.emitter_arg.items(): self.emitter_config[key] = value - if self.emitter == "parquet": - if ("out_dir" not in self.emitter_config) and ( - "out_uri" not in self.emitter_config - ): - raise RuntimeError( - "Must provide out_dir or out_uri" - " as emitter argument for parquet emitter." - ) + if self.emitter in ["parquet", "xarray"]: + if not any(map(self.emitter_config.__contains__, + ["out_dir", "out_uri"])): + raise KeyError( + "Must provide `out_dir` or `out_uri` in `emitter_arg`.") else: - raise RuntimeError( + raise TypeError( "Emitter option must be a string" " representing the emitter type with any additional config" " options under the emitter_arg key." @@ -905,6 +932,10 @@ def run(self): f" != {parse.quote_plus(self.experiment_id)}" ) experiment_config["experiment_id"] = self.experiment_id + # Ensure that `suffix_time` is in effect for all duplicates + # of `experiment_id` + assert metadata["experiment_id"] == self.experiment_id_base + metadata["experiment_id"] = self.experiment_id experiment_config["profile"] = self.profile # Since unique numpy updater is an class method, internal @@ -919,12 +950,11 @@ def run(self): self.ecoli_experiment = Engine(**experiment_config) # Only emit designated stores if specified - if self.config["emit_paths"]: - self.ecoli_experiment.state.set_emit_values([tuple()], False) - self.ecoli_experiment.state.set_emit_values( - self.emit_stores, - True, - ) + if isinstance(emitter := self.ecoli_experiment.emitter, BufferedEmitter): + emitter.reset_emit_flags( + engine=self.ecoli_experiment, + agent=("agents", self.agent_id), + emit_stores=self.emit_stores) # Clean up unnecessary references self.generated_initial_state = None @@ -1094,11 +1124,6 @@ def main(): ecoli_sim = EcoliSim.from_cli() ecoli_sim.build_ecoli() ecoli_sim.run() - # When max_duration is specified, a simulation can finish without dividing - # or raising an Exception. In this case, we still want to finalize the - # Parquet emitter to ensure all buffered data is written to file. - if isinstance(ecoli_sim.ecoli_experiment.emitter, ParquetEmitter): - ecoli_sim.ecoli_experiment.emitter.finalize() if __name__ == "__main__": diff --git a/ecoli/library/emitter.py b/ecoli/library/emitter.py new file mode 100644 index 000000000..455fd4af6 --- /dev/null +++ b/ecoli/library/emitter.py @@ -0,0 +1,167 @@ + +""" +Extensions to the :py:class:`~vivarium.core.emitter.Emitter` interface, as used +by :py:class:`.ParquetEmitter` and :py:class:`.XarrayEmitter`. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Callable +from concurrent.futures import Executor, Future +from dataclasses import dataclass, field, replace +from typing import Any, Self +from urllib import parse +from warnings import warn + +from vivarium.core.emitter import Emitter +from vivarium.core.engine import Engine +from vivarium.core.types import HierarchyPath + +# ============================================================================== + + +class BlockingExecutor(Executor): + + def __init__(self, *args) -> None: + assert not len(args) + super().__init__() + + def submit(self, fn: Callable, /, *args, **kwargs) -> Future: + """ + Run a function in the current thread, and return a + :py:class:`~concurrent.futures.Future` that is already done. + """ + future: Future = Future() + try: + result = fn(*args, **kwargs) + future.set_result(result) + except Exception as e: # noqa: BLE001 + future.set_exception(e) + return future + + def shutdown(self, wait=True, *, cancel_futures=False) -> None: + pass + + +# ============================================================================== + + +@dataclass(eq=True, kw_only=True, slots=True) +class StoragePartition: + """ + Metadata determining the relative storage location for the simulation + outputs of a single-generation :py:class:`.EcoliSim`, inside a hive + partition or hierarchical store (see :ref:`parquet_emitter`). + """ + + experiment_id: str + variant: int + lineage_seed: int + generation: int = field(init=False) + agent_id: str + + def __post_init__(self) -> None: + assert isinstance(self.experiment_id, str) + assert isinstance(self.variant, int) + assert isinstance(self.lineage_seed, int) + assert isinstance(self.agent_id, str) + self.generation = len(self.agent_id) + assert self.generation > 0 + + @property + def parent(self) -> Self: + """ + Metadata of the mother cell in the same cell lineage. + """ + return replace(self, agent_id=self.agent_id[:-1]) + + +# ============================================================================== + + +class BufferedEmitter(Emitter, ABC): + """ + An extension to the :py:class:`~vivarium.core.emitter.Emitter` interface + that buffers emitted simulation data before writing it to persistent + storage. In particular, this interface is used by + :py:meth:`.EcoliSim.update_experiment` and + :py:meth:`.EngineProcess.next_update`. + + .. warning:: + :py:meth:`~.finalize` must be explicitly called in a + ``try ... finally ...`` block around the call to + :py:meth:`vivarium.core.engine.Engine.update`, in order to ensure that + all buffered emits are written out when the simulation terminates for + any reason. + """ + + def __init__(self) -> None: + """ + .. warning:: + This method should be called **at the end** of a subclass + ``__init__()``. + """ + self.finalized: bool = False + """ + Flag set by :py:meth:`.finalize` after writing the last buffer. + """ + + @abstractmethod + def reset_emit_flags( + self, *, + engine: Engine, agent: HierarchyPath, emit_stores: list[HierarchyPath] + ) -> None: + """ + Reconfigure the simulation engine to avoid futile data marshalling, by + suppressing all default emissions and enabling only stores that were + explicitly requested by this emitter's configuration. + + Called by: :py:meth:`.EcoliSim.run` or + :py:meth:`.EngineProcess.create_emitter`. + """ + ... + + def extract_partition(self, metadata: dict[str, Any], /) -> StoragePartition: + """ + Define the current :py:class:`StoragePartition` from the simulation + metadata received via :py:meth:`!Engine._emit_configuration`. + """ + return StoragePartition( + experiment_id=parse.quote_plus( + metadata.get("experiment_id", "default")), + variant=int(metadata.get("variant", 0)), + lineage_seed=int(metadata.get("lineage_seed", 0)), + agent_id=metadata.get("agent_id", "1")) + + def finalize(self, *, success: bool = False) -> None: + """ + Emit the partially filled buffer at the end of a single-generation + simulation. + + Args: + success: Indicates whether the simulation reached a + :py:exc:`.DivisionDetected` event. + """ + if self.finalized: + raise RuntimeError( + f"`{type(self).__name__}.finalize()` was already called.") + assert isinstance(success, bool) + self._finalize(success=success) + self.finalized = True + + @abstractmethod + def _finalize(self, *, success: bool) -> None: + """ + Called by: :py:meth:`.finalize`. + """ + ... + + def __del__(self) -> None: + """ + When a successfully initialised :py:class:`.BufferedEmitter` instance is + destroyed, check that its last batch has been flushed by the simulation + loop. + """ + if not getattr(self, "finalized", True): + warn(f"\n `{type(self).__name__}.finalize()` was never called.") diff --git a/ecoli/library/parquet_emitter.py b/ecoli/library/parquet_emitter.py index 0d85ac967..2430fae5a 100644 --- a/ecoli/library/parquet_emitter.py +++ b/ecoli/library/parquet_emitter.py @@ -1,7 +1,9 @@ + import os import fnmatch from concurrent.futures import Future, ThreadPoolExecutor -from typing import Any, Callable, cast, Mapping, Optional +from dataclasses import asdict +from typing import Any, Callable, Mapping, Optional, cast, final from urllib import parse import duckdb @@ -12,7 +14,15 @@ from fsspec.core import filesystem, url_to_fs, OpenFile from fsspec.spec import AbstractFileSystem from tqdm import tqdm -from vivarium.core.emitter import Emitter + +from vivarium.core.types import HierarchyPath +from vivarium.core.engine import Engine + +from .emitter import BlockingExecutor, BufferedEmitter + + +# ============================================================================== + METADATA_PREFIX = "output_metadata__" """ @@ -62,6 +72,9 @@ """uint32 is 2x smaller than int64 for values between 0 - 4,294,967,295.""" +# ============================================================================== + + def json_to_parquet( emit_dict: dict[str, np.ndarray | list[pl.Series]], outfile: str, @@ -839,22 +852,11 @@ def pl_dtype_from_ndarray(arr: np.ndarray) -> pl.DataType: return pl_dtype -class BlockingExecutor: - def submit(self, fn: Callable, *args, **kwargs) -> Future: - """ - Run function in the current thread and return a Future that - is already done. - """ - future: Future = Future() - try: - result = fn(*args, **kwargs) - future.set_result(result) - except Exception as e: - future.set_exception(e) - return future +# ============================================================================== -class ParquetEmitter(Emitter): +@final +class ParquetEmitter(BufferedEmitter): """ Emit data to a Parquet dataset. Note that :py:meth:`~.finalize` must be explicitly called in a ``try...finally`` block around the call to @@ -919,20 +921,27 @@ def __init__(self, config: dict[str, Any]) -> None: } else: self.emit_prefixes = None - - def finalize(self): - """Convert remaining batched emits to Parquet at sim shutdown - and mark sim as successful if ``success`` flag was set. In vEcoli, - this is done by :py:class:`~ecoli.experiments.ecoli_master_sim.EcoliSim` - upon reaching division. + super().__init__() + + def reset_emit_flags( + self, *, + engine: Engine, agent: HierarchyPath, emit_stores: list[HierarchyPath] + ) -> None: + assert engine.emitter is self + if emit_stores: + engine.state.set_emit_value(emit=False, path=()) + engine.state.set_emit_values(emit=True, paths=emit_stores) + + def _finalize(self, *, success: bool): + """ + Convert remaining batched emits to Parquet at sim shutdown and mark sim + as successful if ``success`` flag was set. """ # Wait for last batch to finish writing self.last_batch_future.result() # Flush any remaining buffered emits to Parquet outfile = os.path.join( - self.out_uri, - self.experiment_id, - "history", + self.out_uri, self.experiment_id, "history", self.partitioning_path, f"{self.num_emits}.pq", ) @@ -944,11 +953,9 @@ def finalize(self): self.buffered_emits, outfile, self.pl_types, self.filesystem ) # Hive-partitioned directory that only contains successful sims - if self.success: + if success: success_file = os.path.join( - self.out_uri, - self.experiment_id, - "success", + self.out_uri, self.experiment_id, "success", self.partitioning_path, "s.pq", ) @@ -996,21 +1003,10 @@ def emit(self, data: dict[str, Any]): data = {**data["data"].pop("metadata", {}), **data["data"]} data["time"] = data.get("initial_global_time", 0.0) # Manually create filepaths with hive partitioning - agent_id = data.get("agent_id", "1") - quoted_experiment_id = parse.quote_plus( - data.get("experiment_id", "default") - ) - partitioning_keys = { - "experiment_id": quoted_experiment_id, - "variant": data.get("variant", 0), - "lineage_seed": data.get("lineage_seed", 0), - "generation": len(agent_id), - "agent_id": agent_id, - } - self.experiment_id = quoted_experiment_id - self.partitioning_path = os.path.join( - *(f"{k}={v}" for k, v in partitioning_keys.items()) - ) + partition = self.extract_partition(data) + self.partitioning_path = os.path.join(*( + f"{k}={v}" for (k, v) in asdict(partition).items())) + self.experiment_id = partition.experiment_id data = flatten_dict(data) config_emit: dict[str, Any] = {} config_schema: dict[str, pl.DataType] = {} @@ -1024,9 +1020,7 @@ def emit(self, data: dict[str, Any]): config_emit[k] = v config_schema[k] = v.dtype outfile = os.path.join( - self.out_uri, - self.experiment_id, - "configuration", + self.out_uri, self.experiment_id, "configuration", self.partitioning_path, "config.pq", ) @@ -1046,7 +1040,8 @@ def emit(self, data: dict[str, Any]): ) # Delete any sim output files in final filesystem history_outdir = os.path.join( - self.out_uri, self.experiment_id, "history", self.partitioning_path + self.out_uri, self.experiment_id, "history", + self.partitioning_path ) try: self.filesystem.delete(history_outdir, recursive=True) @@ -1159,9 +1154,7 @@ def emit(self, data: dict[str, Any]): # If last batch of emits failed, exception should be raised here self.last_batch_future.result() outfile = os.path.join( - self.out_uri, - self.experiment_id, - "history", + self.out_uri, self.experiment_id, "history", self.partitioning_path, f"{self.num_emits}.pq", ) diff --git a/ecoli/library/test_parquet_emitter.py b/ecoli/library/test_parquet_emitter.py index 68c7451cb..ef8ac875b 100644 --- a/ecoli/library/test_parquet_emitter.py +++ b/ecoli/library/test_parquet_emitter.py @@ -1,7 +1,5 @@ import os import re -import tempfile -import shutil import duckdb import numpy as np import polars as pl @@ -254,7 +252,7 @@ def test_union_pl_dtypes(self): pl.UInt32, ) == pl.List(pl.List(pl.List(pl.UInt32))) - def test_quote_columns(self): + def test_quote_columns(self, tmp_path): """Test quote_columns handles special characters correctly.""" # Test single string with special characters assert quote_columns("simple") == '"simple"' @@ -291,116 +289,114 @@ def test_quote_columns(self): assert quote_columns([]) == [] # Test that quoted columns actually work in DuckDB queries with weird column names - with tempfile.TemporaryDirectory() as tmp_path: - test_file = os.path.join(tmp_path, "weird_cols.parquet") - # Create test data with columns containing special characters - test_data = pl.DataFrame( - { - "simple": [1, 2, 3], - "with spaces": [4, 5, 6], - "with-hyphens": [7, 8, 9], - "with[brackets]": [10, 11, 12], - "with/slashes": [13, 14, 15], - 'has"quote': [16, 17, 18], - "dot.name": [19, 20, 21], - "colon:name": [22, 23, 24], - } - ) - test_data.write_parquet(test_file, statistics=False) - - conn = create_duckdb_conn() - - # Test selecting individual columns with special characters - for col in test_data.columns: - quoted_col = quote_columns(col) - result = conn.sql(f"SELECT {quoted_col} FROM '{test_file}'").pl() - assert result.shape == (3, 1) - assert result.columns[0] == col - expected_values = test_data[col].to_list() - assert result[col].to_list() == expected_values - - # Test selecting multiple columns at once - weird_cols = ["with spaces", "with-hyphens", "with[brackets]", 'has"quote'] - quoted_cols = ", ".join(quote_columns(weird_cols)) - result = conn.sql(f"SELECT {quoted_cols} FROM '{test_file}'").pl() - assert result.shape == (3, 4) - for col in weird_cols: - assert col in result.columns - assert result[col].to_list() == test_data[col].to_list() - - # Test that using WHERE clause works with quoted columns - quoted_space_col = quote_columns("with spaces") - result = conn.sql( - f"SELECT * FROM '{test_file}' WHERE {quoted_space_col} > 4" - ).pl() - assert result.shape == (2, 8) - assert result["with spaces"].to_list() == [5, 6] - - # Test aggregation with quoted columns - quoted_bracket_col = quote_columns("with[brackets]") - result = conn.sql( - f"SELECT AVG({quoted_bracket_col}) as avg_val FROM '{test_file}'" - ).pl() - assert result["avg_val"][0] == 11.0 - - # Test ORDER BY with quoted columns - quoted_slash_col = quote_columns("with/slashes") - result = conn.sql( - f"SELECT {quoted_slash_col} FROM '{test_file}' ORDER BY {quoted_slash_col} DESC" - ).pl() - assert result["with/slashes"].to_list() == [15, 14, 13] - - def test_list_columns(self): + test_file = os.path.join(tmp_path, "weird_cols.parquet") + # Create test data with columns containing special characters + test_data = pl.DataFrame( + { + "simple": [1, 2, 3], + "with spaces": [4, 5, 6], + "with-hyphens": [7, 8, 9], + "with[brackets]": [10, 11, 12], + "with/slashes": [13, 14, 15], + 'has"quote': [16, 17, 18], + "dot.name": [19, 20, 21], + "colon:name": [22, 23, 24], + } + ) + test_data.write_parquet(test_file, statistics=False) + + conn = create_duckdb_conn() + + # Test selecting individual columns with special characters + for col in test_data.columns: + quoted_col = quote_columns(col) + result = conn.sql(f"SELECT {quoted_col} FROM '{test_file}'").pl() + assert result.shape == (3, 1) + assert result.columns[0] == col + expected_values = test_data[col].to_list() + assert result[col].to_list() == expected_values + + # Test selecting multiple columns at once + weird_cols = ["with spaces", "with-hyphens", "with[brackets]", 'has"quote'] + quoted_cols = ", ".join(quote_columns(weird_cols)) + result = conn.sql(f"SELECT {quoted_cols} FROM '{test_file}'").pl() + assert result.shape == (3, 4) + for col in weird_cols: + assert col in result.columns + assert result[col].to_list() == test_data[col].to_list() + + # Test that using WHERE clause works with quoted columns + quoted_space_col = quote_columns("with spaces") + result = conn.sql( + f"SELECT * FROM '{test_file}' WHERE {quoted_space_col} > 4" + ).pl() + assert result.shape == (2, 8) + assert result["with spaces"].to_list() == [5, 6] + + # Test aggregation with quoted columns + quoted_bracket_col = quote_columns("with[brackets]") + result = conn.sql( + f"SELECT AVG({quoted_bracket_col}) as avg_val FROM '{test_file}'" + ).pl() + assert result["avg_val"][0] == 11.0 + + # Test ORDER BY with quoted columns + quoted_slash_col = quote_columns("with/slashes") + result = conn.sql( + f"SELECT {quoted_slash_col} FROM '{test_file}' ORDER BY {quoted_slash_col} DESC" + ).pl() + assert result["with/slashes"].to_list() == [15, 14, 13] + + def test_list_columns(self, tmp_path): """Test list_columns retrieves column names correctly.""" - with tempfile.TemporaryDirectory() as tmp_path: - # Create test Parquet file with known columns - test_file = os.path.join(tmp_path, "test.parquet") - test_data = pl.DataFrame( - { - "col_a": [1, 2, 3], - "col_b": [4.0, 5.0, 6.0], - "listeners__mass__cell_mass": [7.0, 8.0, 9.0], - "listeners__mass__dry_mass": [10.0, 11.0, 12.0], - "listeners__growth__instantaneous_growth_rate": [0.1, 0.2, 0.3], - "bulk": [[1, 2], [3, 4], [5, 6]], - } - ) - test_data.write_parquet(test_file, statistics=False) + # Create test Parquet file with known columns + test_file = os.path.join(tmp_path, "test.parquet") + test_data = pl.DataFrame( + { + "col_a": [1, 2, 3], + "col_b": [4.0, 5.0, 6.0], + "listeners__mass__cell_mass": [7.0, 8.0, 9.0], + "listeners__mass__dry_mass": [10.0, 11.0, 12.0], + "listeners__growth__instantaneous_growth_rate": [0.1, 0.2, 0.3], + "bulk": [[1, 2], [3, 4], [5, 6]], + } + ) + test_data.write_parquet(test_file, statistics=False) - conn = create_duckdb_conn() - subquery = f"SELECT * FROM '{test_file}'" + conn = create_duckdb_conn() + subquery = f"SELECT * FROM '{test_file}'" - # Test getting all columns - all_cols = list_columns(conn, subquery) - assert len(all_cols) == 6 - assert "col_a" in all_cols - assert "col_b" in all_cols - assert "listeners__mass__cell_mass" in all_cols + # Test getting all columns + all_cols = list_columns(conn, subquery) + assert len(all_cols) == 6 + assert "col_a" in all_cols + assert "col_b" in all_cols + assert "listeners__mass__cell_mass" in all_cols - # Test pattern matching with glob patterns - listener_cols = list_columns(conn, subquery, "listeners__*") - assert len(listener_cols) == 3 - assert all(col.startswith("listeners__") for col in listener_cols) + # Test pattern matching with glob patterns + listener_cols = list_columns(conn, subquery, "listeners__*") + assert len(listener_cols) == 3 + assert all(col.startswith("listeners__") for col in listener_cols) - # Test pattern matching for specific listener - mass_cols = list_columns(conn, subquery, "listeners__mass__*") - assert len(mass_cols) == 2 - assert "listeners__mass__cell_mass" in mass_cols - assert "listeners__mass__dry_mass" in mass_cols + # Test pattern matching for specific listener + mass_cols = list_columns(conn, subquery, "listeners__mass__*") + assert len(mass_cols) == 2 + assert "listeners__mass__cell_mass" in mass_cols + assert "listeners__mass__dry_mass" in mass_cols - # Test pattern that matches nothing - no_match = list_columns(conn, subquery, "nonexistent__*") - assert len(no_match) == 0 + # Test pattern that matches nothing + no_match = list_columns(conn, subquery, "nonexistent__*") + assert len(no_match) == 0 - # Test pattern with single character wildcard - col_pattern = list_columns(conn, subquery, "col_?") - assert len(col_pattern) == 2 - assert "col_a" in col_pattern - assert "col_b" in col_pattern + # Test pattern with single character wildcard + col_pattern = list_columns(conn, subquery, "col_?") + assert len(col_pattern) == 2 + assert "col_a" in col_pattern + assert "col_b" in col_pattern - # Test exact match pattern - exact = list_columns(conn, subquery, "bulk") - assert exact == ["bulk"] + # Test exact match pattern + exact = list_columns(conn, subquery, "bulk") + assert exact == ["bulk"] def compare_nested(a: list, b: list) -> bool: @@ -420,21 +416,16 @@ def compare_nested(a: list, b: list) -> bool: class TestParquetEmitter: - @pytest.fixture - def temp_dir(self): - """Create a temporary directory for testing.""" - tmp = tempfile.mkdtemp() - yield tmp - shutil.rmtree(tmp) - def test_initialization(self, temp_dir): + def test_initialization(self, tmp_path): """Test ParquetEmitter initialization with different configs.""" # Test with out_dir - emitter = ParquetEmitter({"out_dir": temp_dir}) + emitter = ParquetEmitter({"out_dir": tmp_path}) emitter.experiment_id = "test_exp" emitter.partitioning_path = "path/to/output" - assert emitter.out_uri == os.path.abspath(temp_dir) + assert emitter.out_uri == os.path.abspath(tmp_path) assert emitter.batch_size == 400 + emitter.finalized = True # Test with out_uri and custom batch size emitter = ParquetEmitter({"out_uri": "gs://bucket/path", "batch_size": 100}) @@ -442,10 +433,11 @@ def test_initialization(self, temp_dir): emitter.partitioning_path = "path/to/output" assert emitter.out_uri == "gs://bucket/path" assert emitter.batch_size == 100 + emitter.finalized = True - def test_emit_configuration(self, temp_dir): + def test_emit_configuration(self, tmp_path): """Test emitting configuration data.""" - emitter = ParquetEmitter({"out_dir": temp_dir}) + emitter = ParquetEmitter({"out_dir": tmp_path}) # Setup ThreadPoolExecutor mock future = Future() @@ -466,6 +458,7 @@ def test_emit_configuration(self, temp_dir): } emitter.emit(config_data) + emitter.finalized = True # Verify partitioning path assert emitter.experiment_id == "test_exp" @@ -477,9 +470,9 @@ def test_emit_configuration(self, temp_dir): args, _ = emitter.executor.submit.call_args assert args[0] == json_to_parquet - def test_emit_simulation_data(self, temp_dir): + def test_emit_simulation_data(self, tmp_path): """Test emitting simulation data with various types.""" - emitter = ParquetEmitter({"out_dir": temp_dir, "batch_size": 2}) + emitter = ParquetEmitter({"out_dir": tmp_path, "batch_size": 2}) # Configuration emit to initialize variables config_data = { @@ -542,6 +535,7 @@ def test_emit_simulation_data(self, temp_dir): emitter.emit(sim_data1) assert emitter.num_emits == 2 emitter.last_batch_future.result() + emitter.finalized = True # Check output t = pl.read_parquet( @@ -561,9 +555,9 @@ def test_emit_simulation_data(self, temp_dir): assert all(t["nested__value"] == [100] * 2) assert emitter.buffered_emits == {} - def test_variable_length_arrays(self, temp_dir): + def test_variable_length_arrays(self, tmp_path): """Test handling arrays with changing dimensions.""" - emitter = ParquetEmitter({"out_dir": temp_dir, "batch_size": 3}) + emitter = ParquetEmitter({"out_dir": tmp_path, "batch_size": 3}) # Configuration emit to initialize variables config_data = { "table": "configuration", @@ -629,6 +623,7 @@ def test_variable_length_arrays(self, temp_dir): # Write to Parquet and check output emitter.emit(sim_data2) emitter.last_batch_future.result() + emitter.finalized = True t = pl.read_parquet( os.path.join( @@ -650,9 +645,9 @@ def test_variable_length_arrays(self, temp_dir): [[1], [1, 2], [1, 2, 3]], ] - def test_extreme_data_types(self, temp_dir): + def test_extreme_data_types(self, tmp_path): """Test with extreme data types and edge cases.""" - emitter = ParquetEmitter({"out_dir": temp_dir, "batch_size": 2}) + emitter = ParquetEmitter({"out_dir": tmp_path, "batch_size": 2}) # Create test data with extreme values and special cases sim_data = { "table": "configuration", @@ -790,6 +785,7 @@ def test_extreme_data_types(self, temp_dir): emitter.emit(sim_data_2) emitter.last_batch_future.result() assert emitter.buffered_emits == {} + emitter.finalized = True out_path = os.path.join( emitter.out_uri, @@ -852,9 +848,9 @@ def test_extreme_data_types(self, temp_dir): f"Mismatch in field {key}" ) - def test_finalize(self, temp_dir): + def test_finalize(self, tmp_path): """Test finalize method that handles remaining data.""" - emitter = ParquetEmitter({"out_dir": temp_dir}) + emitter = ParquetEmitter({"out_dir": tmp_path}) emitter.experiment_id = "test_exp" emitter.partitioning_path = "path/to/output" @@ -886,8 +882,8 @@ def test_finalize(self, temp_dir): assert args[0]["field2"][0] == 20.5 # Test success flag - emitter.success = True - emitter.finalize() + emitter.finalized = False + emitter.finalize(success=True) assert os.path.exists( os.path.join( emitter.out_uri, @@ -898,8 +894,8 @@ def test_finalize(self, temp_dir): ) ) - def test_multiple_agents(self, temp_dir): - emitter = ParquetEmitter({"out_dir": temp_dir}) + def test_multiple_agents(self, tmp_path): + emitter = ParquetEmitter({"out_dir": tmp_path}) emitter.experiment_id = "test_exp" emitter.partitioning_path = "path/to/output" @@ -916,11 +912,12 @@ def test_multiple_agents(self, temp_dir): emitter.emit(sim_data) assert emitter.num_emits == 0 assert emitter.buffered_emits == {} + emitter.finalized = True - def test_batch_processing(self, temp_dir): + def test_batch_processing(self, tmp_path): """Test multiple emits and batch processing.""" # Small batch size for testing - emitter = ParquetEmitter({"out_dir": temp_dir, "batch_size": 3}) + emitter = ParquetEmitter({"out_dir": tmp_path, "batch_size": 3}) # Configuration emit to initialize variables config_data = { @@ -947,6 +944,7 @@ def test_batch_processing(self, temp_dir): sim_data["data"]["agents"]["agent1"]["value"] = i * 10 emitter.emit(sim_data) emitter.last_batch_future.result() + emitter.finalized = True # Verify batch was processed assert emitter.num_emits == 4 @@ -957,15 +955,9 @@ def test_batch_processing(self, temp_dir): class TestParquetEmitterEdgeCases: - @pytest.fixture - def temp_dir(self): - """Create a temporary directory for testing.""" - tmp = tempfile.mkdtemp() - yield tmp - shutil.rmtree(tmp) @patch("ecoli.library.parquet_emitter.ThreadPoolExecutor") - def test_multithreaded_buffer_clearing(self, mock_executor_class, temp_dir): + def test_multithreaded_buffer_clearing(self, mock_executor_class, tmp_path): """ Test to verify that clearing buffers after submitting to ThreadPoolExecutor doesn't cause race conditions with the worker thread. @@ -1006,7 +998,7 @@ def delayed_execution(): mock_executor_class.return_value = mock_executor # Initialize the emitter with a small batch size - emitter = ParquetEmitter({"out_dir": temp_dir, "batch_size": 2}) + emitter = ParquetEmitter({"out_dir": tmp_path, "batch_size": 2}) # Configuration emit to initialize variables config_data = { "table": "configuration", @@ -1086,16 +1078,17 @@ def delayed_execution(): # Changed type for field2 to list so should fail with pytest.raises(pl.exceptions.InvalidOperationError): emitter.finalize() + emitter.finalized = True # Cleanup the real executor real_executor.shutdown() - def test_variable_shape_detection_at_boundaries(self, temp_dir): + def test_variable_shape_detection_at_boundaries(self, tmp_path): """ Test the fixed vs variable shape field detection logic specifically at the boundary points (start of sim, after disk write). """ # Use a small batch size to quickly hit the boundary - emitter = ParquetEmitter({"out_dir": temp_dir, "batch_size": 3}) + emitter = ParquetEmitter({"out_dir": tmp_path, "batch_size": 3}) # Setup: Emit configuration data to intitialize variables config_data = { @@ -1198,6 +1191,7 @@ def test_variable_shape_detection_at_boundaries(self, temp_dir): emitter.emit(sim_data4) emitter.last_batch_future.result() + emitter.finalized = True t = pl.read_parquet( os.path.join( emitter.out_uri, @@ -1224,12 +1218,12 @@ def test_variable_shape_detection_at_boundaries(self, temp_dir): [[1], [2], [3], [4], [5]], ] - def test_expected_failures(self, temp_dir): + def test_expected_failures(self, tmp_path): """ Test a few cases that are expected to fail. """ # Use a small batch size to quickly hit the boundary - emitter = ParquetEmitter({"out_dir": temp_dir, "batch_size": 3}) + emitter = ParquetEmitter({"out_dir": tmp_path, "batch_size": 3}) # Setup: Emit configuration data to intitialize variables config_data = { @@ -1394,10 +1388,11 @@ def test_expected_failures(self, temp_dir): match=re.escape("more than 1 dimension"), ): emitter.emit(sim_data7) + emitter.finalized = True - def test_nested_nullable(self, temp_dir): + def test_nested_nullable(self, tmp_path): """Test handling nullable nested types that increase in depth.""" - emitter = ParquetEmitter({"out_dir": temp_dir, "batch_size": 4}) + emitter = ParquetEmitter({"out_dir": tmp_path, "batch_size": 4}) # Configuration emit to initialize variables config_data = { "table": "configuration", @@ -1518,6 +1513,7 @@ def test_nested_nullable(self, temp_dir): for _ in range(3): emitter.emit(sim_data1) emitter.last_batch_future.result() + emitter.finalized = True # Check output t = pl.read_parquet( @@ -1536,12 +1532,6 @@ def test_nested_nullable(self, temp_dir): class TestEmitPaths: """Tests for the emit_paths filtering feature.""" - @pytest.fixture - def temp_dir(self): - tmp = tempfile.mkdtemp() - yield tmp - shutil.rmtree(tmp) - def _config_emit(self): return { "table": "configuration", @@ -1562,9 +1552,9 @@ def _sim_emit(self, time=1.0, **agent_fields): }, } - def test_no_emit_paths_keeps_all_fields(self, temp_dir): + def test_no_emit_paths_keeps_all_fields(self, tmp_path): """Without emit_paths, all fields are emitted (existing behavior).""" - emitter = ParquetEmitter({"out_dir": temp_dir, "batch_size": 2}) + emitter = ParquetEmitter({"out_dir": tmp_path, "batch_size": 2}) assert emitter.emit_prefixes is None emitter.emit(self._config_emit()) @@ -1585,18 +1575,19 @@ def test_no_emit_paths_keeps_all_fields(self, temp_dir): emitter.last_batch_future.result() t = pl.read_parquet( - os.path.join(temp_dir, "test_exp", "history", "**", "*.pq"), + os.path.join(tmp_path, "test_exp", "history", "**", "*.pq"), hive_partitioning=True, ) assert "listeners__mass__cell_mass" in t.columns assert "listeners__growth__rate" in t.columns assert "bulk" in t.columns + emitter.finalized = True - def test_emit_paths_filters_fields(self, temp_dir): + def test_emit_paths_filters_fields(self, tmp_path): """With emit_paths, only matching fields (plus time) are buffered.""" emitter = ParquetEmitter( { - "out_dir": temp_dir, + "out_dir": tmp_path, "batch_size": 2, "emit_paths": [("listeners", "mass")], } @@ -1621,19 +1612,20 @@ def test_emit_paths_filters_fields(self, temp_dir): emitter.last_batch_future.result() t = pl.read_parquet( - os.path.join(temp_dir, "test_exp", "history", "**", "*.pq"), + os.path.join(tmp_path, "test_exp", "history", "**", "*.pq"), hive_partitioning=True, ) assert "listeners__mass__cell_mass" in t.columns assert "time" in t.columns assert "listeners__growth__rate" not in t.columns assert "bulk" not in t.columns + emitter.finalized = True - def test_emit_paths_exact_key_match(self, temp_dir): + def test_emit_paths_exact_key_match(self, tmp_path): """emit_paths should include fields whose flat key exactly equals a prefix.""" emitter = ParquetEmitter( { - "out_dir": temp_dir, + "out_dir": tmp_path, "batch_size": 2, "emit_paths": [("bulk",)], } @@ -1651,17 +1643,18 @@ def test_emit_paths_exact_key_match(self, temp_dir): emitter.last_batch_future.result() t = pl.read_parquet( - os.path.join(temp_dir, "test_exp", "history", "**", "*.pq"), + os.path.join(tmp_path, "test_exp", "history", "**", "*.pq"), hive_partitioning=True, ) assert "bulk" in t.columns assert "listeners__mass__cell_mass" not in t.columns + emitter.finalized = True - def test_emit_paths_multiple_paths(self, temp_dir): + def test_emit_paths_multiple_paths(self, tmp_path): """Multiple paths in emit_paths each filter correctly.""" emitter = ParquetEmitter( { - "out_dir": temp_dir, + "out_dir": tmp_path, "batch_size": 2, "emit_paths": [ ("listeners", "mass"), @@ -1683,18 +1676,19 @@ def test_emit_paths_multiple_paths(self, temp_dir): emitter.last_batch_future.result() t = pl.read_parquet( - os.path.join(temp_dir, "test_exp", "history", "**", "*.pq"), + os.path.join(tmp_path, "test_exp", "history", "**", "*.pq"), hive_partitioning=True, ) assert "bulk" in t.columns assert "listeners__mass__cell_mass" in t.columns assert "listeners__growth__rate" not in t.columns + emitter.finalized = True - def test_emit_paths_time_always_included(self, temp_dir): + def test_emit_paths_time_always_included(self, tmp_path): """The 'time' field is always included regardless of emit_paths.""" emitter = ParquetEmitter( { - "out_dir": temp_dir, + "out_dir": tmp_path, "batch_size": 2, "emit_paths": [("bulk",)], } @@ -1706,17 +1700,18 @@ def test_emit_paths_time_always_included(self, temp_dir): emitter.last_batch_future.result() t = pl.read_parquet( - os.path.join(temp_dir, "test_exp", "history", "**", "*.pq"), + os.path.join(tmp_path, "test_exp", "history", "**", "*.pq"), hive_partitioning=True, ) assert "time" in t.columns assert t["time"].to_list() == [1.5, 1.5] + emitter.finalized = True - def test_emit_paths_prefix_does_not_match_sibling(self, temp_dir): + def test_emit_paths_prefix_does_not_match_sibling(self, tmp_path): """A prefix like 'listeners__mass' must not match 'listeners__mass_extra'.""" emitter = ParquetEmitter( { - "out_dir": temp_dir, + "out_dir": tmp_path, "batch_size": 2, "emit_paths": [("listeners", "mass")], } @@ -1733,17 +1728,18 @@ def test_emit_paths_prefix_does_not_match_sibling(self, temp_dir): emitter.last_batch_future.result() t = pl.read_parquet( - os.path.join(temp_dir, "test_exp", "history", "**", "*.pq"), + os.path.join(tmp_path, "test_exp", "history", "**", "*.pq"), hive_partitioning=True, ) assert "listeners__mass__cell_mass" in t.columns assert "listeners__mass_extra__value" not in t.columns + emitter.finalized = True - def test_emit_paths_stores_prefixes_correctly(self, temp_dir): + def test_emit_paths_stores_prefixes_correctly(self, tmp_path): """emit_prefixes are stored as double-underscore joined strings.""" emitter = ParquetEmitter( { - "out_dir": temp_dir, + "out_dir": tmp_path, "emit_paths": [ ("a", "b", "c"), ("x",), @@ -1751,3 +1747,4 @@ def test_emit_paths_stores_prefixes_correctly(self, temp_dir): } ) assert emitter.emit_prefixes == {"a__b__c", "x"} + emitter.finalized = True diff --git a/ecoli/library/test_utils.py b/ecoli/library/test_utils.py new file mode 100644 index 000000000..e95ed85da --- /dev/null +++ b/ecoli/library/test_utils.py @@ -0,0 +1,100 @@ + +""" +Utilities for patching execution environments, configurations and functions. +""" + + +from abc import ABC, abstractmethod +from collections.abc import Callable +from functools import reduce +from inspect import ismethod +from unittest.mock import Mock, DEFAULT, _patch, patch +from typing import Any + +import pytest + +from ecoli.library.xarray_emitter.utils import WarningFilter + + +# ============================================================================== +# warnings +# ============================================================================== + + +def filter_warnings(filters: list[WarningFilter]) -> Callable[[Callable], Callable]: + """ + Analogue of :py:func:`ecoli.library.xarray_emitter.utils.filter_warnings`, + but with the effect of applying :py:func:`pytest.mark.filterwarnings` + decorators, instead of :py:func:`warnings.filterwarnings` context modifiers. + """ + return (lambda func: reduce( + lambda fun, wf: pytest.mark.filterwarnings(str(wf))(fun), + filters, func)) + + + +# ============================================================================== +# config patching +# ============================================================================== + + +class PatchConfig(ABC): + """ + Test parameter for modifying an already loaded baseline JSON configuration. + """ + + @abstractmethod + def to_dict(self) -> dict[str, Any]: + """ + Materialise changes to the JSON configuration. + """ + ... + + +# ============================================================================== +# code patching +# ============================================================================== + + +def patch_func(func: str, *, cb: Callable | None = None) -> _patch: + """ + Create a context manager which patches a module-level function, in order to + trace its calls, and to optionally pre-apply a callback. + + .. note:: + ``func`` is passed as the argument ``target`` to + :py:func:`unittest.mock.patch`. + """ + mocked = None + def side_effect(*args, **kwargs) -> Any: + nonlocal cb, mocked + if cb is not None: + cb(*args, **kwargs) + return mocked.temp_original(*args, **kwargs) # type: ignore[attr-defined] + mocked = patch(func, side_effect=side_effect) + return mocked + + +# ------------------------------------------------------------------------------ + + +def patch_meth( + obj: object, meth: str, *, + cb: Callable[..., None] | None = None, + modargs: Callable[..., tuple[tuple, dict]] | None = None +) -> None: + """ + Patch an object instance method, in order to trace its calls, and to + optionally pre-apply a callback or argument modification. + """ + assert ismethod(getattr(obj, meth)) + assert cb is None or modargs is None + def side_effect(*args, **kwargs): + nonlocal obj, meth, cb, modargs + if modargs is not None: + _args, _kwargs = modargs(obj, *args, **kwargs) + return getattr(obj, meth)._mock_wraps(*_args, **_kwargs) + elif cb is not None: + cb(obj, *args, **kwargs) + return DEFAULT + setattr(obj, meth, Mock(wraps=getattr(obj, meth), side_effect=side_effect)) diff --git a/ecoli/library/xarray_emitter/__init__.py b/ecoli/library/xarray_emitter/__init__.py new file mode 100644 index 000000000..e9b042764 --- /dev/null +++ b/ecoli/library/xarray_emitter/__init__.py @@ -0,0 +1,296 @@ + +r""" + +Introduction +============ + +:py:class:`.XarrayEmitter` is an :py:class:`~vivarium.core.emitter.Emitter` +similar to :py:class:`~.ParquetEmitter`, but with a design optimized towards a +different flavour of downstream applications: + +- :py:class:`~.ParquetEmitter` is geared towards emitting a significant fraction + of the simulator state, in a format that supports flexible sparse selections, + `data reductions`_ and time series visualizations, as used in :ref:`analysis + scripts `. +- :py:class:`.XarrayEmitter` is intended for emitting only a pre-selected subset + of statically shaped tensor variables, in a format that supports numerical + algorithms in the high-dimensional and large-sample regime. + +The former type of computations is naturally expressed using `relational query`_ +engines (e.g., :ref:`DuckDB `), whereas the latter type is +naturally expressed using `array programming`_ libraries (e.g., `Cubed`_). Due +to the sheer size of the simulator state, both types may in general require +`out-of-core processing`_ algorithms. + +.. _data reductions: https://en.wikipedia.org/wiki/Data_reduction +.. _relational query: https://en.wikipedia.org/wiki/Relational_database +.. _array programming: https://en.wikipedia.org/wiki/Array_programming +.. _Cubed: https://cubed-dev.github.io/cubed/why-cubed.html +.. _out-of-core processing: https://en.wikipedia.org/wiki/External_memory_algorithm + +In order to facilitate downstream applications based on chunked array +processing, :py:class:`.XarrayEmitter` writes out to any persistent storage +supporting the `Zarr`_ specification, using an in-memory buffer comprised of +`Xarray`_ objects. For optimized throughput, the buffer implements temporal +subsampling, numerical type casting and compression codecs at emission time. +Furthermore, in order to simplify the export of simulation data into external +libraries, the hierarchy of the output `DataTree`_ is decoupled from the +hierarchy of simulation :py:class:`~vivarium.core.store.Store`\ s, using an +output :ref:`variable layout ` specified in the :ref:`simulator +configuration `. + +.. _Xarray: https://xarray.dev/ +.. _Zarr: https://zarr.dev/ +.. _DataTree: https://docs.xarray.dev/en/stable/user-guide/data-structures.html#datatree + + +Comparison with :py:class:`~.ParquetEmitter` +============================================ + +Similarities +------------ + +- Currently only supports simulations of a *single-cell lineage* per + :py:class:`.BufferedEmitter` instance. +- Executes at every time step. +- Buffers emissions into time chunks. +- Uses concurrent threads for writing buffers to persistent storage. +- Produces a hierarchically structured storage layout that supports selective + reading in downstream applications. + +Differences in usage +-------------------- + +- Supports the configuration of *emission predicates*. +- Currently only supports emitting a *static collection* of *statically shaped + tensor variables*. +- Supports *renaming and rearranging* of output variables. +- Requires the configuration of *output data types*. +- Supports the configuration of backend-specific *compression codecs*. +- Supports :ref:`log_updates`, i.e., the emission of individual + :py:class:`~vivarium.core.process.Process` update requests, before they are + aggregated and reallocated by + :py:func:`~ecoli.processes.allocator.calculatePartition` and then applied to + the global cell state by + :py:meth:`~ecoli.processes.partition.PartitionedProcess.evolve_state`. + +.. note:: + See :py:class:`.XarrayEmitter` for an explanation of the JSON configuration + syntax, and for complete examples, see in ``configs/test_configs/``: + + - ``moving_avg_analysis.json`` (parameters for a **realistic workflow**; + Nextflow invocation via :py:func:`runscripts.workflow.main`). + - ``test_xarray_emitter.json`` (parameters for **fast CI tests**; + CLI invocation via :py:func:`ecoli.experiments.ecoli_master_sim.main`) + +.. hint:: + As data structures, `DataTree`_\ s could support changes of variable names and + dimensions across time steps. The constraints currently imposed by + :py:class:`.XarrayEmitter` rather serve to enable I/O optimizations for the + intended use cases. When access to variably sized simulation variables is + desired, users have the choice either of implementing custom :ref:`listeners + ` with static output coordinates, or otherwise of defaulting to the + :py:class:`~.ParquetEmitter`. + +Differences in implementation +----------------------------- + +- Uses the `Xarray`_ API for serialization, buffering, and `metadata + organization`_, including unit annotations (see :py:class:`.VariableSpec` and + :py:class:`.XarrayTransducer` for details). +- Applies a "*process*-major" rather than a "*generation*-major" output layout, + reflecting array variables directly in the output directory tree; this + produces one file per *variable* time chunk, rather than one file per + *simulation* time chunk (compare the :py:class:`.XarrayEmitter` :ref:`storage + layout ` with :py:meth:`.ParquetEmitter.emit`; see + :py:class:`.XarrayStoragePartition` for details). +- Defines the abstract interface :py:class:`.AsyncBufferWriter` for storage + backends with *asynchronous* APIs (currently supported: `Zarr`_), realizing + the opportunity for :ref:`concurrency ` among multiple + `DataArray`_\ s within an output buffer. +- Decouples the *in-memory buffer size* from the *persistent chunk size*, in + order to simplify performance tuning of large-scale simulations (see + :py:class:`.XarrayTransducer` and :py:class:`.AsyncBufferWriter` for details). +- Maintains `consolidated metadata`_ and updates it at the end of each simulated + cell generation, in order to reduce the metadata loading latency for + subsequent storage reads (see :py:class:`.AsyncZarrBufferWriter` for details). + +.. _metadata organization: https://docs.xarray.dev/en/stable/get-help/faq.html#approach-to-metadata +.. _DataArray: https://docs.xarray.dev/en/stable/user-guide/data-structures.html#dataarray +.. _consolidated metadata: https://docs.xarray.dev/en/stable/user-guide/io.html#io-zarr-consolidated-metadata + + +.. _storage_layout: + +Storage layout +============== + +The workflow storage layout, which comprises many individual simulations, is +currently organized as follows --- where file paths in this example are specific +to the Zarr v3 storage backend:: + + {store} ; + ├─ zarr.json ; metadata + └─ experiment_id={}/variant={}/lineage_seed={} ; + ├─ zarr.json ; consolidated metadata + ├─ emitstep_gen={} ;