allow a callable to decide when to save a snapshot

This commit is contained in:
Valentin Boettcher 2022-12-09 15:06:46 -05:00
parent cd5f82d044
commit e9900520c4
No known key found for this signature in database
GPG key ID: E034E12B7AF56ACE

View file

@ -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))