From e9900520c4e8d2d016b909e8dacfb9c460447a42 Mon Sep 17 00:00:00 2001 From: Valentin Boettcher Date: Fri, 9 Dec 2022 15:06:46 -0500 Subject: [PATCH] allow a callable to decide when to save a snapshot --- hopsflow/util.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/hopsflow/util.py b/hopsflow/util.py index e751e69..a4b8492 100644 --- a/hopsflow/util.py +++ b/hopsflow/util.py @@ -840,7 +840,7 @@ def ensemble_mean_online( save: str, function: Callable[..., np.ndarray], i: Optional[int] = None, - every: Optional[int] = None, + every: Optional[Union[int, Callable[[int], bool]]] = None, ) -> Optional[EnsembleValue]: path = Path(save) @@ -863,7 +863,9 @@ def ensemble_mean_online( aggregate = WelfordAggregator(result, i) aggregate.dump(str(path)) - if every is not None and aggregate.n % every == 0: + if every is not None and ( + aggregate.n % every == 0 if isintance(every, int) else every(aggregate.n) + ): path.with_stem(f"{path.stem}_{aggregate.n}") aggregate.dump(str(path))