diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 147afb7c..f9e62d9d 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -7,6 +7,30 @@ Release Notes 0.12 Series ........... +0.12.19 (2026-07-27) +-------------------- + +General: + +* Add execution time profiling to the ABC-SMC run, reporting pure simulation + time, parallel-pipeline setup time, and in-between-iteration time, + including population-size calculation and distance-function adaptation. + The timings are also returned by ``run_generation``. Resolves #325. +* Minor bug fixes and improved typesetting. + +Visualization: + +* ``plot_walltime`` and ``plot_total_walltime`` now report the actual + per-generation walltimes and no longer include the idle time that passed + between a stored analysis and a later resumed run. Resolves #667. + +Storage: + +* Store the per-generation walltime in the database (new ``wall_time`` column, + database version 2). Databases created with older pyABC versions must be + migrated via ``abc-migrate`` before they can be resumed; for such databases + the walltime plots fall back to the previous, end-time-based behavior. + 0.12.18 (2026-04-14) -------------------- diff --git a/pyabc/distance/base.py b/pyabc/distance/base.py index 708d0f43..7178c7db 100644 --- a/pyabc/distance/base.py +++ b/pyabc/distance/base.py @@ -224,7 +224,7 @@ class FunctionDistance(Distance): statistics x and x_0. Returns the distance between both. """ - def __init__(self, fun): + def __init__(self, fun: Callable): super().__init__() self.fun = fun @@ -250,7 +250,7 @@ def get_config(self): return conf @staticmethod - def to_distance(maybe_distance: Callable | Distance) -> Distance: + def to_distance(maybe_distance: Callable | Distance | None) -> Distance: """ Parameters ---------- diff --git a/pyabc/distance/pnorm.py b/pyabc/distance/pnorm.py index 5806cff2..9276d50d 100644 --- a/pyabc/distance/pnorm.py +++ b/pyabc/distance/pnorm.py @@ -597,7 +597,9 @@ def __init__( self.predictor = predictor - self.initial_info_weights: dict[str, float] = initial_info_weights + self.initial_info_weights: dict[str, float] | None = ( + initial_info_weights + ) self.info_weights: dict[int, np.ndarray] = {} if fit_info_ixs is None: diff --git a/pyabc/inference/smc.py b/pyabc/inference/smc.py index b87c408c..2bd61b1e 100644 --- a/pyabc/inference/smc.py +++ b/pyabc/inference/smc.py @@ -4,6 +4,7 @@ import logging from collections.abc import Callable from datetime import datetime, timedelta +from time import perf_counter from typing import TypeVar import numpy as np @@ -54,6 +55,23 @@ def identity(x): return x +class _Timer: + """Context manager measuring wall-clock execution time in seconds. + + The elapsed time is available via the ``elapsed`` attribute after the + ``with`` block has finished, and is used for execution time profiling. + """ + + def __enter__(self) -> '_Timer': + self.elapsed = 0.0 + self._start = perf_counter() + return self + + def __exit__(self, *exc_info) -> bool: + self.elapsed = perf_counter() - self._start + return False + + def run_cleanup(run): """Wrapper: Run and in any case clean up afterwards.""" @@ -172,7 +190,7 @@ class ABCSMC: def __init__( self, models: list[Model] | Model | Callable, - parameter_priors: list[Distribution] | Distribution | Callable, + parameter_priors: list[Distribution] | Distribution, distance_function: Distance | Callable | None = None, population_size: PopulationStrategy | int = 100, summary_statistics: Callable[[model_output], dict] = identity, @@ -545,6 +563,9 @@ def _sample_from_prior(self, t: int) -> Population: Only sample from prior and return results without changing the history of the distance function or the epsilon. """ + # record start time + calibration_start_time = datetime.now() + # create simulate function simulate_one = self._create_simulate_from_prior_function() @@ -567,8 +588,13 @@ def _sample_from_prior(self, t: int) -> Population: population = sample.get_accepted_population() # update information saved in history about calibration + calibration_end_time = datetime.now() self.history.update_after_calibration( - nr_samples=self.sampler.nr_evaluations_, end_time=datetime.now() + nr_samples=self.sampler.nr_evaluations_, + end_time=calibration_end_time, + wall_time=( + calibration_end_time - calibration_start_time + ).total_seconds(), ) return population @@ -792,6 +818,9 @@ def run_generation( generation terminated successfully, and potentially "acceptance_rate". """ + # start execution time profiling for this generation + generation_perf_start = perf_counter() + # get epsilon for generation t current_eps = self.eps(t) if current_eps is None or np.isnan(current_eps): @@ -800,26 +829,29 @@ def run_generation( ) logger.info(f't: {t}, eps: {current_eps:.8e}.') - # create simulate function - simulate_one = self._create_simulate_function(t) - - # population size and maximum number of evaluations - pop_size = self.population_size(t) - max_eval = ( - np.inf - if self.min_acceptance_rate == 0.0 - else pop_size / self.min_acceptance_rate - ) + # set up the simulation pipeline for this generation + with _Timer() as setup_timer: + # create simulate function + simulate_one = self._create_simulate_function(t) + + # population size and maximum number of evaluations + pop_size = self.population_size(t) + max_eval = ( + np.inf + if self.min_acceptance_rate == 0.0 + else pop_size / self.min_acceptance_rate + ) # perform the sampling logger.debug(f'Submitting population {t}.') - sample = self.sampler.sample_until_n_accepted( - n=pop_size, - simulate_one=simulate_one, - t=t, - max_eval=max_eval, - ana_vars=self._vars(t=t), - ) + with _Timer() as simulation_timer: + sample = self.sampler.sample_until_n_accepted( + n=pop_size, + simulate_one=simulate_one, + t=t, + max_eval=max_eval, + ana_vars=self._vars(t=t), + ) # check sample health if not sample.ok: @@ -838,8 +870,9 @@ def run_generation( # save to database n_sim = self.sampler.nr_evaluations_ model_names = [model.name for model in self.models] + wall_time = perf_counter() - generation_perf_start self.history.append_population( - t, current_eps, population, n_sim, model_names + t, current_eps, population, n_sim, model_names, wall_time=wall_time ) logger.debug( f'Total samples up to t = {t}: ' @@ -856,16 +889,42 @@ def run_generation( ) # prepare next iteration - self._prepare_next_iteration( - t=t + 1, - sample=sample, - population=population, - acceptance_rate=acceptance_rate, + with _Timer() as prepare_next_timer: + prepare_next_timings = self._prepare_next_iteration( + t=t + 1, + sample=sample, + population=population, + acceptance_rate=acceptance_rate, + ) + + # execution time profiling + total_time = perf_counter() - generation_perf_start + sim_fraction = ( + 100 * simulation_timer.elapsed / total_time + if total_time > 0 + else 0.0 + ) + timings = { + 'total': total_time, + 'simulation': simulation_timer.elapsed, + 'pipeline_setup': setup_timer.elapsed, + 'prepare_next': prepare_next_timer.elapsed, + **prepare_next_timings, + } + logger.info( + f'Timing t={t} [s]: total={total_time:.3g}, ' + f'simulation={simulation_timer.elapsed:.3g} ' + f'({sim_fraction:.0f}%), ' + f'pipeline-setup={setup_timer.elapsed:.3g}, ' + f'prepare-next={prepare_next_timer.elapsed:.3g} ' + f'(population-size={prepare_next_timings["population_size"]:.3g}, ' + f'distance={prepare_next_timings["distance"]:.3g}).' ) return { 'successful': True, 'acceptance_rate': acceptance_rate, + 'timings': timings, } def check_terminate( @@ -909,7 +968,7 @@ def _prepare_next_iteration( sample: Sample, population: Population, acceptance_rate: float, - ): + ) -> dict: """Update actors for the upcoming iteration. Be aware: The current (finished) iteration is t-1, the next t. @@ -924,6 +983,12 @@ def _prepare_next_iteration( The current iteration's population object. acceptance_rate: float The current iteration's acceptance rate. + + Returns + ------- + timings: + Execution times in seconds of the profiled sub-steps, with keys + ``"population_size"`` and ``"distance"``. """ # make a copy prev_transitions = copy.deepcopy(self.transitions) @@ -932,17 +997,19 @@ def _prepare_next_iteration( self._fit_transitions(t) # update population size - self._adapt_population_size(t) + with _Timer() as population_size_timer: + self._adapt_population_size(t) def get_sample(): return sample # update distance - df_updated = self.distance_function.update( - t=t, - get_sample=get_sample, - total_sims=self.history.total_nr_simulations, - ) + with _Timer() as distance_timer: + df_updated = self.distance_function.update( + t=t, + get_sample=get_sample, + total_sims=self.history.total_nr_simulations, + ) # compute distances with the new distance measure def get_weighted_distances(): @@ -999,6 +1066,11 @@ def get_all_records(): acceptor_config=self.acceptor.get_epsilon_config(t), ) + return { + 'population_size': population_size_timer.elapsed, + 'distance': distance_timer.elapsed, + } + def _adapt_population_size(self, t): """ Adapt population size based on the employed population strategy. @@ -1016,11 +1088,16 @@ def _adapt_population_size(self, t): 'p' ].values + # restrict to models still alive: dead models are never fitted + alive = self.history.alive_models(self.history.max_t) + # make a copy in case the population strategy messes with # the transitions # WARNING: the deepcopy also copies the random states of scipy.stats # distributions - copied_transitions = copy.deepcopy(self.transitions) + copied_transitions = copy.deepcopy( + [self.transitions[m] for m in alive] + ) # update the population size self.population_size.update( diff --git a/pyabc/populationstrategy/populationstrategy.py b/pyabc/populationstrategy/populationstrategy.py index a90a6d60..f29b220a 100644 --- a/pyabc/populationstrategy/populationstrategy.py +++ b/pyabc/populationstrategy/populationstrategy.py @@ -220,7 +220,7 @@ def update( )[0], ) - if not np.isnan(cv_estimate.n_estimated): + if np.isfinite(cv_estimate.n_estimated): self.nr_particles = max( min(int(cv_estimate.n_estimated), self.max_population_size), self.min_population_size, diff --git a/pyabc/storage/db_model.py b/pyabc/storage/db_model.py index 1b7efbf3..1b58d324 100644 --- a/pyabc/storage/db_model.py +++ b/pyabc/storage/db_model.py @@ -98,6 +98,7 @@ class Population(Base): abc_smc_id = Column(Integer, ForeignKey('abc_smc.id')) t = Column(Integer) population_end_time = Column(DateTime) + wall_time = Column(Float) nr_samples = Column(Integer) epsilon = Column(Float) models = relationship('Model') diff --git a/pyabc/storage/history.py b/pyabc/storage/history.py index 256bc365..d08d95a1 100644 --- a/pyabc/storage/history.py +++ b/pyabc/storage/history.py @@ -389,6 +389,8 @@ def get_all_populations(self): * `t`: Population number * `population_end_time`: The end time of the population + * `wall_time`: The wall time in seconds spent on the population, + excluding idle time between resumed runs. * `samples`: The number of sample attempts performed for a population * `epsilon`: The acceptance threshold for the population. @@ -401,6 +403,7 @@ def get_all_populations(self): query = self._session.query( Population.t, Population.population_end_time, + Population.wall_time, Population.nr_samples, Population.epsilon, ).filter(Population.abc_smc_id == self.id) @@ -415,7 +418,7 @@ def get_all_populations(self): @internal_docstring_warning def store_initial_data( self, - ground_truth_model: int, + ground_truth_model: int | None, options: dict, observed_summary_statistics: dict, ground_truth_parameter: dict, @@ -484,7 +487,7 @@ def store_initial_data( @internal_docstring_warning def store_pre_population( self, - ground_truth_model: int, + ground_truth_model: int | None, observed_summary_statistics: dict, ground_truth_parameter: dict, model_names: list[str], @@ -547,7 +550,10 @@ def store_pre_population( @with_session @internal_docstring_warning def update_after_calibration( - self, nr_samples: int, end_time: datetime.datetime + self, + nr_samples: int, + end_time: datetime.datetime, + wall_time: float | None = None, ): """Update after the calibration iteration. In particular set time and number of samples. @@ -559,6 +565,8 @@ def update_after_calibration( Number of samples reported. end_time: End time of the calibration iteration. + wall_time: + Wall time in seconds spent on the calibration iteration. """ # extract population population = ( @@ -572,6 +580,7 @@ def update_after_calibration( # update samples number population.nr_samples = nr_samples population.population_end_time = end_time + population.wall_time = wall_time # commit changes self._session.commit() @@ -697,6 +706,7 @@ def _save_to_population_db( particles_by_model: dict, model_probabilities: pd.DataFrame, model_names, + wall_time: float | None = None, ): # sqlalchemy experimental stuff and highly inefficient implementation # here but that is ok for testing purposes for the moment @@ -706,7 +716,10 @@ def _save_to_population_db( # store the population population = Population( - t=t, nr_samples=nr_simulations, epsilon=current_epsilon + t=t, + nr_samples=nr_simulations, + epsilon=current_epsilon, + wall_time=wall_time, ) abcsmc.populations.append(population) @@ -783,6 +796,7 @@ def append_population( population: PyPopulation, nr_simulations: int, model_names, + wall_time: float | None = None, ): """ Append population to database. @@ -799,6 +813,8 @@ def append_population( The number of model evaluations for this population. model_names: list The model names. + wall_time: float + Wall time in seconds spent on sampling this population. """ particles_by_model = population.get_particles_by_model() model_probabilities = population.get_model_probabilities() @@ -810,6 +826,7 @@ def append_population( particles_by_model, model_probabilities, model_names, + wall_time=wall_time, ) @with_session @@ -921,16 +938,6 @@ def get_weighted_distances(self, t: int | None = None) -> pd.DataFrame: weights.append(weight) distances.append(sample.distance) - # query = (self._session.query(Sample.distance, Particle.w, Model.m) - # .join(Particle) - # .join(Model).join(Population).join(ABCSMC) - # .filter(ABCSMC.id == self.id) - # .filter(Population.t == t)) - # df = pd.read_sql_query(query.statement, self._engine) - # model_probabilities = self.get_model_probabilities(t).reset_index() - # df_weighted = df.merge(model_probabilities) - # df_weighted["w"] *= df_weighted["p"] - return pd.DataFrame({'distance': distances, 'w': weights}) @with_session diff --git a/pyabc/storage/migrations/versions/2_20260724_add_populations_wall_time.py b/pyabc/storage/migrations/versions/2_20260724_add_populations_wall_time.py new file mode 100644 index 00000000..0f3d5af6 --- /dev/null +++ b/pyabc/storage/migrations/versions/2_20260724_add_populations_wall_time.py @@ -0,0 +1,27 @@ +"""add populations.wall_time + +Revision ID: 2 +Revises: 1 +Create Date: 2026-07-24 00:00:00.000000 + +""" + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision = '2' +down_revision = '1' +branch_labels = None +depends_on = None + + +def upgrade(): + op.add_column( + table_name='populations', + column=sa.Column('wall_time', sa.FLOAT, nullable=True), + ) + + +def downgrade(): + pass diff --git a/pyabc/storage/version.py b/pyabc/storage/version.py index 8a97ebc5..a19e08fb 100644 --- a/pyabc/storage/version.py +++ b/pyabc/storage/version.py @@ -1 +1 @@ -__db_version__ = '1' +__db_version__ = '2' diff --git a/pyabc/version.py b/pyabc/version.py index 8327b506..1975b2ee 100644 --- a/pyabc/version.py +++ b/pyabc/version.py @@ -1 +1 @@ -__version__ = '0.12.18' +__version__ = '0.12.19' diff --git a/pyabc/visualization/walltime.py b/pyabc/visualization/walltime.py index 2826720d..7cfd62b0 100644 --- a/pyabc/visualization/walltime.py +++ b/pyabc/visualization/walltime.py @@ -38,8 +38,16 @@ def _prepare_plot_total_walltime( # extract total walltimes walltimes = [] for h in histories: - abc = h.get_abc() - walltimes.append((abc.end_time - abc.start_time).total_seconds()) + wall_time = h.get_all_populations().wall_time + if len(wall_time) > 0 and wall_time.notna().all(): + # sum of the actual per-generation walltimes, excluding idle time + # between resumed runs + walltimes.append(float(wall_time.sum())) + else: + # fall back to the wall-clock duration for runs stored by pyABC + # versions that did not record per-generation walltimes + abc = h.get_abc() + walltimes.append((abc.end_time - abc.start_time).total_seconds()) walltimes = np.asarray(walltimes) # apply time unit @@ -156,7 +164,7 @@ def plot_total_walltime_plotly( def _prepare_walltime( histories: list[History] | History, - show_calibration: bool, + show_calibration: bool | None, ): # preprocess input histories = to_lists(histories) @@ -167,16 +175,19 @@ def _prepare_walltime( h.get_all_populations().samples[0] > 0 for h in histories ) - # extract start times and end times + # extract start times, end times and per-generation walltimes start_times = [] end_times = [] + wall_times = [] for h in histories: # start time start_times.append(h.get_abc().start_time) - # end times - end_times.append(h.get_all_populations().population_end_time) + # end times and walltimes per population + pops = h.get_all_populations() + end_times.append(pops.population_end_time) + wall_times.append(pops.wall_time) - return start_times, end_times, show_calibration + return start_times, end_times, wall_times, show_calibration def plot_walltime( @@ -219,7 +230,7 @@ def plot_walltime( A reference to the axis of the generated plot. """ # preprocess input - start_times, end_times, show_calibration = _prepare_walltime( + start_times, end_times, wall_times, show_calibration = _prepare_walltime( histories=histories, show_calibration=show_calibration ) @@ -233,6 +244,7 @@ def plot_walltime( title=title, size=size, ax=ax, + wall_times=wall_times, ) @@ -248,7 +260,7 @@ def plot_walltime_plotly( ) -> 'go.Figure': """Plot walltimes using plotly.""" # preprocess input - start_times, end_times, show_calibration = _prepare_walltime( + start_times, end_times, wall_times, show_calibration = _prepare_walltime( histories=histories, show_calibration=show_calibration ) @@ -262,6 +274,7 @@ def plot_walltime_plotly( title=title, size=size, fig=fig, + wall_times=wall_times, ) @@ -271,6 +284,7 @@ def _prepare_plot_walltime_lowlevel( labels: list | str | None = None, show_calibration: bool | None = None, unit: str = 's', + wall_times: list | None = None, ): # preprocess input end_times = to_lists(end_times) @@ -290,14 +304,28 @@ def _prepare_plot_walltime_lowlevel( if unit not in TIME_UNITS: raise AssertionError(f'`unit` must be in {TIME_UNITS}') + # per-generation walltimes may be unavailable (e.g. for old databases) + if wall_times is None: + wall_times = [None] * n_run + # extract relative walltimes walltimes = [] - for start_t, end_ts in zip(start_times, end_times): + for start_t, end_ts, wall_t in zip(start_times, end_times, wall_times): times = [start_t, *end_ts] - # compute stacked differences - diffs = [end - start for start, end in zip(times[:-1], times[1:])] - # as seconds - diffs = [diff.total_seconds() for diff in diffs] + # compute stacked differences of the population end times, as seconds + diffs = [ + (end - start).total_seconds() + for start, end in zip(times[:-1], times[1:]) + ] + # prefer the actually measured per-generation walltime where + # available. This excludes idle time between resumed runs + if wall_t is not None: + wall_t = list(wall_t) + for i in range(min(len(diffs), len(wall_t))): + w = wall_t[i] + # skip missing values (NaN) stored by older pyABC versions + if w is not None and not np.isnan(w): + diffs[i] = float(w) # append walltimes.append(diffs) walltimes = np.asarray(walltimes) @@ -332,6 +360,7 @@ def plot_walltime_lowlevel( title: str = 'Walltime by generation', size: tuple | None = None, ax: mpl.axes.Axes | None = None, + wall_times: list | None = None, ) -> mpl.axes.Axes: """Low-level access to `plot_walltime`. @@ -344,6 +373,7 @@ def plot_walltime_lowlevel( labels=labels, show_calibration=show_calibration, unit=unit, + wall_times=wall_times, ) # create figure @@ -388,6 +418,7 @@ def plot_walltime_lowlevel_plotly( title: str = 'Walltime by generation', size: tuple | None = None, fig: 'go.Figure | None' = None, + wall_times: list | None = None, ) -> 'go.Figure': """Low-level access to `plot_walltime_plotly`.""" import plotly.graph_objects as go @@ -399,6 +430,7 @@ def plot_walltime_lowlevel_plotly( labels=labels, show_calibration=show_calibration, unit=unit, + wall_times=wall_times, ) # create figure @@ -557,7 +589,7 @@ def _prepare_plot_eps_walltime_lowlevel( end_times: list, eps: list, labels: list | str, - colors: list[Any], + colors: list[Any] | None, group_by_label: bool, unit: str, ): diff --git a/test/visualization/test_base_viz.py b/test/visualization/test_base_viz.py index 94c38d73..73c7e1df 100644 --- a/test/visualization/test_base_viz.py +++ b/test/visualization/test_base_viz.py @@ -1,3 +1,4 @@ +import datetime import os import tempfile @@ -7,6 +8,7 @@ import pytest import pyabc +from pyabc.visualization.walltime import _prepare_plot_walltime_lowlevel db_path = 'sqlite:///' + tempfile.mkstemp(suffix='.db')[1] log_files = [] @@ -331,6 +333,58 @@ def test_walltime(): plt.close() +def test_walltime_ignores_resume_gap(): + """`plot_walltime` uses the recorded per-generation walltimes instead of + differences of population end times, so that idle time between a stored + and a later resumed run is not attributed to any iteration.""" + base = datetime.datetime(2020, 1, 1, 0, 0, 0) + start_times = [base] + # calibration + 2 generations, with a one-day pause (resume) between + # generation 0 and generation 1 + end_times = [ + [ + base + datetime.timedelta(seconds=10), # calibration end + base + datetime.timedelta(seconds=30), # generation 0 end + base + datetime.timedelta(days=1, seconds=45), # gen 1 (resumed) + ] + ] + # actually measured per-generation walltimes in seconds + wall_times = [[10.0, 20.0, 15.0]] + # generation-1 end-time diff spans the one-day resume gap: + # (1 day + 45 s) - 30 s = 86415 s + gap_seconds = datetime.timedelta(days=1, seconds=45).total_seconds() - 30 + + # without recorded walltimes, the resume gap leaks into generation 1 + matrix_gap, _, _ = _prepare_plot_walltime_lowlevel( + end_times=end_times, + start_times=start_times, + show_calibration=True, + unit='s', + ) + np.testing.assert_allclose(matrix_gap[:, 0], [10.0, 20.0, gap_seconds]) + + # with recorded walltimes, the actual per-generation walltimes are used + matrix, _, _ = _prepare_plot_walltime_lowlevel( + end_times=end_times, + start_times=start_times, + show_calibration=True, + unit='s', + wall_times=wall_times, + ) + np.testing.assert_allclose(matrix[:, 0], [10.0, 20.0, 15.0]) + + # missing (NaN) per-generation walltimes fall back to end-time diffs + wall_times_partial = [[10.0, 20.0, float('nan')]] + matrix_partial, _, _ = _prepare_plot_walltime_lowlevel( + end_times=end_times, + start_times=start_times, + show_calibration=True, + unit='s', + wall_times=wall_times_partial, + ) + np.testing.assert_allclose(matrix_partial[:, 0], [10.0, 20.0, gap_seconds]) + + def test_eps_walltime(): """Test `pyabc.visualization.plot_eps_walltime`""" for group_by_label in [True, False]: