diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 368681b3..8c1d5d34 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -16,6 +16,9 @@ jobs: matrix: include: # Fast PR matrix + - os: ubuntu-latest + python: "3.14" + toxenv: base - os: ubuntu-latest python: "3.13" toxenv: base @@ -26,42 +29,42 @@ jobs: python: "3.11" toxenv: base - os: ubuntu-latest - python: "3.11" + python: "3.14" toxenv: visualization # macOS sanity - os: macos-latest - python: "3.11" + python: "3.14" toxenv: mac # Quality - os: ubuntu-latest - python: "3.11" + python: "3.14" toxenv: quality - os: ubuntu-latest - python: "3.11" + python: "3.14" toxenv: project - os: ubuntu-latest - python: "3.11" + python: "3.14" toxenv: doc - os: ubuntu-latest python: "3.11" toxenv: migrate - os: ubuntu-latest - python: "3.11" + python: "3.14" toxenv: external-R - os: ubuntu-latest - python: "3.11" + python: "3.14" toxenv: external-other-simulators - os: ubuntu-latest - python: "3.11" + python: "3.14" toxenv: petab - os: ubuntu-latest - python: "3.11" + python: "3.14" toxenv: base-notebooks - os: ubuntu-latest - python: "3.11" + python: "3.14" toxenv: external-notebooks steps: diff --git a/.gitignore b/.gitignore index 9322016a..11c24a1f 100644 --- a/.gitignore +++ b/.gitignore @@ -69,3 +69,6 @@ dask-worker-space *.lock *tmp* *amici_models* + +# macOS +.DS_Store diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 147afb7c..44e1fe8a 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -4,6 +4,39 @@ Release Notes ============= +0.13 Series +........... + +0.13.0 (2026-07-30) +------------------- + +General: + +* Add support for python3.14 +* 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 (breaking): + +* 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. +* Particle weights are now stored using the global normalization (weights sum + to 1 across all particles of all models, matching the in-memory + ``Population`` representation). This is bundled into database version 2 and + handled by ``abc-migrate`` for existing databases. Resolves #47. + + 0.12 Series ........... diff --git a/README.md b/README.md index a5647e90..5f769d7c 100644 --- a/README.md +++ b/README.md @@ -26,4 +26,5 @@ - 📄 **Cite**: [https://pyabc.rtfd.io/en/latest/cite.html](https://pyabc.rtfd.io/en/latest/cite.html) #### Related Projects -- 🧠 **Neural Posterior Estimation**: [BayesFlow](https://bayesflow.org/main/_examples/From_ABC_to_BayesFlow.html) +- **Parameter Estimation with Likelihoods**: [pyPESTO](https://github.com/ICB-DCM/pyPESTO) +- **Neural Posterior Estimation**: [BayesFlow](https://bayesflow.org/main/_examples/From_ABC_to_BayesFlow.html) diff --git a/doc/conf.py b/doc/conf.py index 95f437b5..93ae8cb0 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -13,6 +13,7 @@ # serve to show the default. import os +import subprocess import sys # If extensions (or modules to document with autodoc) are in another directory, @@ -32,18 +33,20 @@ 'sphinx.ext.autodoc', # generate autodoc summaries 'sphinx.ext.autosummary', - # use mathjax for latex formulas - 'sphinx.ext.mathjax', - # link to code - 'sphinx.ext.viewcode', + # test snippets in the documentation + 'sphinx.ext.doctest', + # link to external urls + 'sphinx.ext.extlinks', # link to other projects' docs 'sphinx.ext.intersphinx', + # use mathjax for latex formulas + 'sphinx.ext.mathjax', # support numpy and google style docstrings 'sphinx.ext.napoleon', # support todo items 'sphinx.ext.todo', - # test snippets in the documentation - 'sphinx.ext.doctest', + # link to code + 'sphinx.ext.viewcode', # source parser for jupyter notebook files 'nbsphinx', # code highlighting in jupyter cells @@ -68,7 +71,7 @@ 'numpy': ('https://numpy.org/devdocs/', None), 'scipy': ('https://docs.scipy.org/doc/scipy/', None), 'pandas': ('https://pandas.pydata.org/pandas-docs/dev', None), - 'petab': ('https://petab.readthedocs.io/en/stable/', None), + 'petab': ('https://petab.readthedocs.io/en/latest/', None), 'amici': ('https://amici.readthedocs.io/en/latest/', None), 'sklearn': ('https://scikit-learn.org/stable/', None), } @@ -99,6 +102,30 @@ # # The short X.Y version. + +# Resolve the current branch name +def get_branch(): + # Read the Docs sets this automatically ("latest", "stable", branch name...) + rtd_version = os.environ.get('READTHEDOCS_VERSION') + if rtd_version and rtd_version not in ('latest', 'stable'): + return rtd_version + try: + return subprocess.check_output( + ['git', 'rev-parse', '--abbrev-ref', 'HEAD'], text=True + ).strip() + except Exception: + return 'main' + + +# Set up links that resolve to the branch of the documentation +branch = get_branch() +extlinks = { + 'repository': ( + f'https://github.com/ICB-DCM/pyABC/blob/{branch}/%s', + '%s', + ), +} + import pyabc # noqa: E402 version = pyabc.__version__ diff --git a/doc/index.rst b/doc/index.rst index 03f210b5..ce70b67d 100644 --- a/doc/index.rst +++ b/doc/index.rst @@ -11,8 +11,8 @@ pyABC - distributed, likelihood-free inference :target: https://badge.fury.io/py/pyabc .. image:: https://zenodo.org/badge/DOI/10.5281/zenodo.3257587.svg :target: https://doi.org/10.5281/zenodo.3257587 -.. image:: https://img.shields.io/badge/code%20style-black-000000.svg - :target: https://github.com/psf/black +.. image:: https://img.shields.io/badge/code%20style-ruff-41B5BE.svg + :target: https://github.com/astral-sh/ruff :Release: |version| :Source code: https://github.com/icb-dcm/pyabc diff --git a/doc/installation.rst b/doc/installation.rst index 1fc670e1..21e12b34 100644 --- a/doc/installation.rst +++ b/doc/installation.rst @@ -6,7 +6,7 @@ Install Preparation ----------- -This package requires Python 3.8 or later. +This package requires Python 3.11 or later. The package is continuously tested on Linux, and in parts on iOS, via GitHub Actions. @@ -77,20 +77,11 @@ home directory with .. _anacondaCluster: -Installing Anaconda on a Cluster environment --------------------------------------------- - -To install `Anaconda `_, run:: - - wget https://repo.anaconda.com/archive/Anaconda3-2021.11-Linux-x86_64.sh - bash Anaconda3-2021.11-Linux-x86_64.sh - -and follow the installation guide. -Replace the "2021.11" by the most recent version of Anaconda, see -https://repo.anaconda.com/archive. +Installing Miniconda on a Cluster environment +--------------------------------------------- `Miniconda `_ -provides an alternative, minimal installer for conda, including +provides a minimal installer for conda, including only conda, Python, and some core and useful packages. Install the latest version via:: @@ -102,7 +93,7 @@ version via:: Optional dependencies --------------------- -pyABC has various optional dependencies, see `setup.cfg`. +pyABC has various optional dependencies, see :repository:`pyproject.toml`. In particular, pyABC has optional interfaces to the :ref:`R ` and :ref:`Julia ` languages, see the API documentation diff --git a/pyabc/acceptor/acceptor.py b/pyabc/acceptor/acceptor.py index 21d2497a..9f4917a7 100644 --- a/pyabc/acceptor/acceptor.py +++ b/pyabc/acceptor/acceptor.py @@ -347,9 +347,9 @@ class StochasticAcceptor(Acceptor): def __init__( self, - pdf_norm_method: Callable = None, + pdf_norm_method: Callable | None = None, apply_importance_weighting: bool = True, - log_file: str = None, + log_file: str | None = None, ): """ Parameters @@ -429,7 +429,7 @@ def _update( self, t: int, get_weighted_distances: Callable[[], pd.DataFrame], - prev_temp: float = None, + prev_temp: float | None = None, acceptance_rate: float = 1.0, ): """ diff --git a/pyabc/acceptor/pdf_norm.py b/pyabc/acceptor/pdf_norm.py index f9afbf29..be3c9607 100644 --- a/pyabc/acceptor/pdf_norm.py +++ b/pyabc/acceptor/pdf_norm.py @@ -32,7 +32,10 @@ def pdf_norm_max_found( prev_pdf_norm = -np.inf # take maximum over all normalizations - pdf_norm = max(prev_pdf_norm, *pdfs) + if len(pdfs) == 0: + pdf_norm = prev_pdf_norm + else: + pdf_norm = max(prev_pdf_norm, float(pdfs.max())) return pdf_norm diff --git a/pyabc/copasi/model.py b/pyabc/copasi/model.py index ba0f52ff..07d08e05 100644 --- a/pyabc/copasi/model.py +++ b/pyabc/copasi/model.py @@ -34,16 +34,16 @@ class BasicoModel(Model): def __init__( self, sbml_file: str, - changes: dict[str, float] = None, + changes: dict[str, float] | None = None, change_unit: bool = True, method: str = 'stochastic', - t0: float = None, - duration: float = None, - num_steps: int = None, + t0: float | None = None, + duration: float | None = None, + num_steps: int | None = None, automatic: bool = True, use_numbers: bool = False, - output: list[str] = None, - model_name: str = None, + output: list[str] | None = None, + model_name: str | None = None, ): """ Parameters diff --git a/pyabc/distance/aggregate.py b/pyabc/distance/aggregate.py index 85372819..2472a9b4 100644 --- a/pyabc/distance/aggregate.py +++ b/pyabc/distance/aggregate.py @@ -26,8 +26,8 @@ class AggregatedDistance(Distance): def __init__( self, distances: list[Distance | Callable], - weights: list | dict = None, - factors: list | dict = None, + weights: list | dict | None = None, + factors: list | dict | None = None, ): """ Parameters @@ -126,8 +126,8 @@ def __call__( self, x: dict, x_0: dict, - t: int = None, - par: dict = None, + t: int | None = None, + par: dict | None = None, ) -> float: """ Applies all distance functions and computes the weighted sum of all @@ -221,21 +221,21 @@ class AdaptiveAggregatedDistance(AggregatedDistance): def __init__( self, distances: list[Distance], - initial_weights: list = None, - factors: list | dict = None, + initial_weights: list | None = None, + factors: list | dict | None = None, adaptive: bool = True, - scale_function: Callable = None, - log_file: str = None, + scale_function: Callable | None = None, + log_file: str | None = None, ): super().__init__(distances=distances) - self.initial_weights: list = initial_weights - self.factors: list | dict = factors + self.initial_weights: list | None = initial_weights + self.factors: list | dict | None = factors self.adaptive: bool = adaptive self.x_0: dict | None = None if scale_function is None: scale_function = span self.scale_function: Callable = scale_function - self.log_file: str = log_file + self.log_file: str | None = log_file def requires_calibration(self) -> bool: return self.initial_weights is None or any( diff --git a/pyabc/distance/base.py b/pyabc/distance/base.py index b7f60f95..7178c7db 100644 --- a/pyabc/distance/base.py +++ b/pyabc/distance/base.py @@ -89,8 +89,8 @@ def __call__( self, x: dict, x_0: dict, - t: int = None, - par: dict = None, + t: int | None = None, + par: dict | None = None, ) -> float: """ Evaluate at time point t the distance of the summary statistics of @@ -183,8 +183,8 @@ def __call__( self, x: dict, # noqa: ARG002 x_0: dict, # noqa: ARG002 - t: int = None, # noqa: ARG002 - par: dict = None, # noqa: ARG002 + t: int | None = None, # noqa: ARG002 + par: dict | None = None, # noqa: ARG002 ) -> float: raise AssertionError( f'Distance {self.__class__.__name__} should not be called.' @@ -203,8 +203,8 @@ def __call__( self, x: dict, # noqa: ARG002 x_0: dict, # noqa: ARG002 - t: int = None, # noqa: ARG002 - par: dict = None, # noqa: ARG002 + t: int | None = None, # noqa: ARG002 + par: dict | None = None, # noqa: ARG002 ) -> float: return -1 @@ -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 @@ -232,8 +232,8 @@ def __call__( self, x: dict, x_0: dict, - t: int = None, # noqa: ARG002 - par: dict = None, # noqa: ARG002 + t: int | None = None, # noqa: ARG002 + par: dict | None = None, # noqa: ARG002 ) -> float: return self.fun(x, x_0) @@ -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/distance.py b/pyabc/distance/distance.py index c6cb12ef..14f55f12 100644 --- a/pyabc/distance/distance.py +++ b/pyabc/distance/distance.py @@ -65,8 +65,8 @@ def __call__( self, x: dict, x_0: dict, - t: int = None, # noqa: ARG002 - par: dict = None, # noqa: ARG002 + t: int | None = None, # noqa: ARG002 + par: dict | None = None, # noqa: ARG002 ) -> float: return sum( abs((x[key] - x_0[key]) / x_0[key]) @@ -105,7 +105,9 @@ def _dict_to_vect(self, x): def _calculate_whitening_transformation_matrix(self, sum_stats): # create data matrix, shape (n_sample, n_y) - x = np.asarray([self._dict_to_vect(x) for x in sum_stats]) + # force float dtype so in-place centering below works also for + # integer-valued summary statistics (e.g. counts) + x = np.asarray([self._dict_to_vect(x) for x in sum_stats], dtype=float) # center mean = np.mean(x, axis=0) x -= mean @@ -142,8 +144,8 @@ def __call__( self, x: dict, x_0: dict, - t: int = None, # noqa: ARG002 - par: dict = None, # noqa: ARG002 + t: int | None = None, # noqa: ARG002 + par: dict | None = None, # noqa: ARG002 ) -> float: x_vec, x_0_vec = self._dict_to_vect(x), self._dict_to_vect(x_0) distance = la.norm( @@ -249,8 +251,8 @@ def __call__( self, x: dict, x_0: dict, - t: int = None, # noqa: ARG002 - par: dict = None, # noqa: ARG002 + t: int | None = None, # noqa: ARG002 + par: dict | None = None, # noqa: ARG002 ) -> float: distance = sum( abs((x[key] - x_0[key]) / self.normalization[key]) diff --git a/pyabc/distance/kernel.py b/pyabc/distance/kernel.py index a01b81be..6386008b 100644 --- a/pyabc/distance/kernel.py +++ b/pyabc/distance/kernel.py @@ -45,8 +45,8 @@ class StochasticKernel(Distance): def __init__( self, ret_scale: str = SCALE_LIN, - keys: list[str] = None, - pdf_max: float = None, + keys: list[str] | None = None, + pdf_max: float | None = None, ): StochasticKernel.check_ret_scale(ret_scale) self.ret_scale = ret_scale @@ -96,8 +96,8 @@ def __init__( self, fun: Callable, ret_scale: str = SCALE_LIN, - keys: list[str] = None, - pdf_max: float = None, + keys: list[str] | None = None, + pdf_max: float | None = None, ): super().__init__(ret_scale=ret_scale, keys=keys, pdf_max=pdf_max) self.fun = fun @@ -106,8 +106,8 @@ def __call__( self, x: dict, x_0: dict, - t: int = None, - par: dict = None, + t: int | None = None, + par: dict | None = None, ) -> float: return self.fun(x=x, x_0=x_0, t=t, par=par) @@ -133,10 +133,10 @@ class NormalKernel(StochasticKernel): def __init__( self, - cov: np.ndarray = None, + cov: np.ndarray | None = None, ret_scale: str = SCALE_LOG, - keys: list[str] = None, - pdf_max: float = None, + keys: list[str] | None = None, + pdf_max: float | None = None, ): super().__init__(ret_scale=ret_scale, keys=keys, pdf_max=pdf_max) self.cov = cov @@ -182,8 +182,8 @@ def __call__( self, x: dict, x_0: dict, - t: int = None, # noqa: ARG002 - par: dict = None, # noqa: ARG002 + t: int | None = None, # noqa: ARG002 + par: dict | None = None, # noqa: ARG002 ) -> float: """ Return the value of the normal distribution at x - x_0, or its @@ -229,9 +229,9 @@ class IndependentNormalKernel(StochasticKernel): def __init__( self, - var: Callable | Sequence[float] | float = None, - keys: list[str] = None, - pdf_max: float = None, + var: Callable | Sequence[float] | float | None = None, + keys: list[str] | None = None, + pdf_max: float | None = None, ): super().__init__(ret_scale=SCALE_LOG, keys=keys, pdf_max=pdf_max) self.var = var @@ -269,8 +269,8 @@ def __call__( self, x: dict, x_0: dict, - t: int = None, # noqa: ARG002 - par: dict = None, + t: int | None = None, # noqa: ARG002 + par: dict | None = None, ): # safety check if self.keys is None: @@ -320,9 +320,9 @@ class IndependentLaplaceKernel(StochasticKernel): def __init__( self, - scale: Callable | Sequence[float] | float = None, - keys: list[str] = None, - pdf_max: float = None, + scale: Callable | Sequence[float] | float | None = None, + keys: list[str] | None = None, + pdf_max: float | None = None, ): super().__init__(ret_scale=SCALE_LOG, keys=keys, pdf_max=pdf_max) self.scale = scale @@ -361,8 +361,8 @@ def __call__( self, x: dict, x_0: dict, - t: int = None, # noqa: ARG002 - par: dict = None, + t: int | None = None, # noqa: ARG002 + par: dict | None = None, ): # safety check if self.keys is None: @@ -404,8 +404,8 @@ def __init__( self, p: float | Callable, ret_scale: str = SCALE_LOG, - keys: list[str] = None, - pdf_max: float = None, + keys: list[str] | None = None, + pdf_max: float | None = None, ): super().__init__(ret_scale=ret_scale, keys=keys, pdf_max=pdf_max) @@ -441,8 +441,8 @@ def __call__( self, x: dict, x_0: dict, - t: int = None, # noqa: ARG002 - par: dict = None, + t: int | None = None, # noqa: ARG002 + par: dict | None = None, ) -> float: x = np.asarray(_arr(x, self.keys), dtype=int) x_0 = np.asarray(_arr(x_0, self.keys), dtype=int) @@ -470,8 +470,8 @@ class PoissonKernel(StochasticKernel): def __init__( self, ret_scale: str = SCALE_LOG, - keys: list[str] = None, - pdf_max: float = None, + keys: list[str] | None = None, + pdf_max: float | None = None, ): super().__init__(ret_scale=ret_scale, keys=keys, pdf_max=pdf_max) @@ -500,8 +500,8 @@ def __call__( self, x: dict, x_0: dict, - t: int = None, # noqa: ARG002 - par: dict = None, # noqa: ARG002 + t: int | None = None, # noqa: ARG002 + par: dict | None = None, # noqa: ARG002 ) -> float: x = np.asarray(_arr(x, self.keys), dtype=int) x_0 = np.asarray(_arr(x_0, self.keys), dtype=int) @@ -530,8 +530,8 @@ def __init__( self, p: float, ret_scale: str = SCALE_LOG, - keys: list[str] = None, - pdf_max: float = None, + keys: list[str] | None = None, + pdf_max: float | None = None, ): super().__init__(ret_scale=ret_scale, keys=keys, pdf_max=pdf_max) @@ -562,8 +562,8 @@ def __call__( self, x: dict, x_0: dict, - t: int = None, # noqa: ARG002 - par: dict = None, + t: int | None = None, # noqa: ARG002 + par: dict | None = None, ) -> float: x = np.asarray(_arr(x, self.keys), dtype=int) x_0 = np.asarray(_arr(x_0, self.keys), dtype=int) diff --git a/pyabc/distance/ot.py b/pyabc/distance/ot.py index c0bfa343..b2be59fc 100644 --- a/pyabc/distance/ot.py +++ b/pyabc/distance/ot.py @@ -68,8 +68,8 @@ def __init__( self, sumstat: Sumstat, p: float = 2.0, - dist: str | Callable = None, - emd_args: dict = None, + dist: str | Callable | None = None, + emd_args: dict | None = None, ): """ Parameters @@ -121,9 +121,9 @@ def __init__( def initialize( self, x_0: dict, - t: int = None, - get_sample: Callable[[], Sample] = None, - total_sims: int = None, + t: int | None = None, + get_sample: Callable[[], Sample] | None = None, + total_sims: int | None = None, ) -> None: # initialize summary statistics self.sumstat.initialize( @@ -157,8 +157,8 @@ def __call__( self, x: dict, x_0: dict, - t: int = None, # noqa: ARG002 - par: dict = None, # noqa: ARG002 + t: int | None = None, # noqa: ARG002 + par: dict | None = None, # noqa: ARG002 ) -> float: # compute summary statistics, shape (n, dim), (n0, dim) s, s0 = self.sumstat(x), self.sumstat(x_0) @@ -225,8 +225,8 @@ def __init__( metric: str = 'sqeuclidean', p: float = 2.0, n_proj: int = 50, - seed: int | np.random.RandomState = None, - emd_1d_args: dict = None, + seed: int | np.random.RandomState | None = None, + emd_1d_args: dict | None = None, ): """ Parameters @@ -270,9 +270,9 @@ def __init__( def initialize( self, x_0: dict, - t: int = None, - get_sample: Callable[[], Sample] = None, - total_sims: int = None, + t: int | None = None, + get_sample: Callable[[], Sample] | None = None, + total_sims: int | None = None, ) -> None: # initialize summary statistics self.sumstat.initialize( @@ -302,8 +302,8 @@ def __call__( self, x: dict, x_0: dict, - t: int = None, # noqa: ARG002 - par: dict = None, # noqa: ARG002 + t: int | None = None, # noqa: ARG002 + par: dict | None = None, # noqa: ARG002 ) -> float: # compute summary statistics, shape (n, dim), (n0, dim) s, s0 = self.sumstat(x), self.sumstat(x_0) @@ -355,7 +355,7 @@ def __call__( def uniform_unit_sphere_samples( n_proj: int, dim: int, - seed: int | np.random.RandomState = None, + seed: int | np.random.RandomState | None = None, ) -> np.ndarray: r""" Generate uniformly distributed samples from the :math:`d-1`-dim. diff --git a/pyabc/distance/pnorm.py b/pyabc/distance/pnorm.py index e77d6298..c761dad3 100644 --- a/pyabc/distance/pnorm.py +++ b/pyabc/distance/pnorm.py @@ -60,8 +60,10 @@ class PNormDistance(Distance): def __init__( self, p: float = 1, - fixed_weights: dict[str, float] | dict[int, dict[str, float]] = None, - sumstat: Sumstat = None, + fixed_weights: dict[str, float] + | dict[int, dict[str, float]] + | None = None, + sumstat: Sumstat | None = None, ): super().__init__() @@ -146,7 +148,7 @@ def get_weights(self, t: int) -> np.ndarray: @staticmethod def format_dict( - vals: dict[str, float] | dict[int, dict[str, float]], + vals: dict[str, float] | dict[int, dict[str, float]] | None, t: int, s_ids: list[str], ) -> dict[int, float | np.ndarray]: @@ -170,8 +172,7 @@ def format_dict( vals = {t: vals} # convert dicts to arrays - for _t, dct in vals.items(): - vals[_t] = dict2arr(dct, keys=s_ids) + vals = {_t: dict2arr(dct, keys=s_ids) for _t, dct in vals.items()} return vals @@ -201,8 +202,8 @@ def __call__( self, x: dict, x_0: dict, - t: int = None, - par: dict = None, # noqa: ARG002 + t: int | None = None, + par: dict | None = None, # noqa: ARG002 ) -> float: # extract weights for given time point weights = self.get_weights(t=t) @@ -308,19 +309,21 @@ class AdaptivePNormDistance(PNormDistance): def __init__( self, p: float = 1, - initial_scale_weights: dict[str, float] = None, - fixed_weights: dict[str, float] = None, + initial_scale_weights: dict[str, float] | None = None, + fixed_weights: dict[str, float] | None = None, fit_scale_ixs: EventIxs | Collection[int] | int = np.inf, - scale_function: Callable = None, - max_scale_weight_ratio: float = None, - scale_log_file: str = None, + scale_function: Callable | None = None, + max_scale_weight_ratio: float | None = None, + scale_log_file: str | None = None, all_particles_for_scale: bool = True, - sumstat: Sumstat = None, + sumstat: Sumstat | None = None, ): # call p-norm constructor super().__init__(p=p, fixed_weights=fixed_weights, sumstat=sumstat) - self.initial_scale_weights: dict[str, float] = initial_scale_weights + self.initial_scale_weights: dict[str, float] | None = ( + initial_scale_weights + ) self.scale_weights: dict[int, np.ndarray] = {} @@ -505,25 +508,25 @@ def __init__( self, predictor: Predictor, p: float = 1, - initial_scale_weights: dict[str, float] = None, - initial_info_weights: dict[str, float] = None, - fixed_weights: dict[str, float] = None, + initial_scale_weights: dict[str, float] | None = None, + initial_info_weights: dict[str, float] | None = None, + fixed_weights: dict[str, float] | None = None, fit_scale_ixs: EventIxs | Collection | int = np.inf, - fit_info_ixs: EventIxs | Collection | int = None, + fit_info_ixs: EventIxs | Collection | int | None = None, normalize_by_par: bool = True, - scale_function: Callable = None, - max_scale_weight_ratio: float = None, - max_info_weight_ratio: float = None, - scale_log_file: str = None, - info_log_file: str = None, - info_sample_log_file: str = None, - sumstat: Sumstat = None, - fd_deltas: list[float] | float = None, - subsetter: Subsetter = None, + scale_function: Callable | None = None, + max_scale_weight_ratio: float | None = None, + max_info_weight_ratio: float | None = None, + scale_log_file: str | None = None, + info_log_file: str | None = None, + info_sample_log_file: str | None = None, + sumstat: Sumstat | None = None, + fd_deltas: list[float] | float | None = None, + subsetter: Subsetter | None = None, all_particles_for_scale: bool = True, all_particles_for_prediction: bool = True, feature_normalization: str = WEIGHTS, - par_trafo: ParTrafoBase = None, + par_trafo: ParTrafoBase | None = None, ): """ Parameters @@ -595,7 +598,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/distance/scale.py b/pyabc/distance/scale.py index ef4b1ca2..13760b19 100644 --- a/pyabc/distance/scale.py +++ b/pyabc/distance/scale.py @@ -141,7 +141,7 @@ def root_mean_square_deviation( rmse = np.sqrt(mse) # debugging - warn_obs_off(off_ixs=np.flatnonzero(bs > 2 * std), s_ids=s_ids) + warn_obs_off(off_ixs=np.flatnonzero(np.abs(bs) > 2 * std), s_ids=s_ids) return rmse @@ -161,7 +161,7 @@ def std_or_rmsd( bs = bias(samples=samples, s0=s0) std = standard_deviation(samples=samples) - if sum(bs > 2 * std) > 1 / 3 * len(std): + if sum(np.abs(bs) > 2 * std) > 1 / 3 * len(std): logger.info('Too many high-bias values, correcting only for scale.') return std @@ -169,7 +169,7 @@ def std_or_rmsd( rmse = np.sqrt(mse) # debugging - warn_obs_off(off_ixs=np.flatnonzero(bs > 2 * std), s_ids=s_ids) + warn_obs_off(off_ixs=np.flatnonzero(np.abs(bs) > 2 * std), s_ids=s_ids) return rmse diff --git a/pyabc/distance/util.py b/pyabc/distance/util.py index 292309de..3471a452 100644 --- a/pyabc/distance/util.py +++ b/pyabc/distance/util.py @@ -84,7 +84,7 @@ def log_weights( def fd_nabla1_multi_delta( x: np.ndarray, fun: Callable, - test_deltas: Sequence[float] = None, + test_deltas: Sequence[float] | None = None, ) -> np.ndarray: """Calculate FD approximation to 1st order derivative (Jacobian/gradient) with automatic step size selection. diff --git a/pyabc/epsilon/temperature.py b/pyabc/epsilon/temperature.py index 65e257b7..7e0e84ba 100644 --- a/pyabc/epsilon/temperature.py +++ b/pyabc/epsilon/temperature.py @@ -80,12 +80,15 @@ class Temperature(TemperatureBase): def __init__( self, - schemes: Callable | list[Callable] = None, - aggregate_fun: Callable[[list[float]], float] = None, - initial_temperature: float = None, + schemes: Callable | list[Callable] | None = None, + aggregate_fun: Callable[[list[float]], float] | None = None, + initial_temperature: float | None = None, enforce_exact_final_temperature: bool = True, - log_file: str = None, + log_file: str | None = None, ): + # normalize a single callable to a list, as all consumers iterate + if schemes is not None and callable(schemes): + schemes = [schemes] self.schemes = schemes if aggregate_fun is None: @@ -317,7 +320,9 @@ class AcceptanceRateScheme(TemperatureScheme): 2) to avoid uneccessary computations. """ - def __init__(self, target_rate: float = 0.3, min_rate: float = None): + def __init__( + self, target_rate: float = 0.3, min_rate: float | None = None + ): self.target_rate = target_rate self.min_rate = min_rate @@ -434,7 +439,6 @@ class ExpDecayFixedIterScheme(TemperatureScheme): Parameters ---------- - alpha: float Factor by which to reduce the temperature, if `max_nr_populations` is infinite. diff --git a/pyabc/external/base.py b/pyabc/external/base.py index c6855250..67ecf51a 100644 --- a/pyabc/external/base.py +++ b/pyabc/external/base.py @@ -31,16 +31,16 @@ class ExternalHandler: def __init__( self, executable: str, - file: str = None, - fixed_args: list = None, + file: str | None = None, + fixed_args: list | None = None, create_folder: bool = False, - suffix: str = None, - prefix: str = None, - dir: str = None, + suffix: str | None = None, + prefix: str | None = None, + dir: str | None = None, show_stdout: bool = False, show_stderr: bool = True, raise_on_error: bool = False, - timeout: float = None, + timeout: float | None = None, ): """ Parameters @@ -110,7 +110,12 @@ def create_executable(self, loc): executable = self.executable.replace('{loc}', loc) return executable - def run(self, args: list[str] = None, cmd: str = None, loc: str = None): + def run( + self, + args: list[str] | None = None, + cmd: str | None = None, + loc: str | None = None, + ): """Run the script for the given arguments. Parameters @@ -196,15 +201,15 @@ def __init__( self, executable: str, file: str, - fixed_args: list = None, + fixed_args: list | None = None, create_folder: bool = False, - suffix: str = None, + suffix: str | None = None, prefix: str = 'modelsim_', - dir: str = None, + dir: str | None = None, show_stdout: bool = False, show_stderr: bool = True, raise_on_error: bool = False, - timeout: float = None, + timeout: float | None = None, name: str = 'ExternalModel', ): """Initialize the model. @@ -336,15 +341,15 @@ def __init__( self, executable: str, file: str, - fixed_args: list = None, + fixed_args: list | None = None, create_folder: bool = False, - suffix: str = None, + suffix: str | None = None, prefix: str = 'sumstat_', - dir: str = None, + dir: str | None = None, show_stdout: bool = False, show_stderr: bool = True, raise_on_error: bool = False, - timeout: float = None, + timeout: float | None = None, ): self.eh = ExternalHandler( executable=executable, @@ -385,14 +390,14 @@ def __init__( self, executable: str, file: str, - fixed_args: list = None, - suffix: str = None, + fixed_args: list | None = None, + suffix: str | None = None, prefix: str = 'dist_', - dir: str = None, + dir: str | None = None, show_stdout: bool = False, show_stderr: bool = True, raise_on_error: bool = False, - timeout: float = None, + timeout: float | None = None, ): self.eh = ExternalHandler( executable=executable, diff --git a/pyabc/external/julia/jl_pyjulia.py b/pyabc/external/julia/jl_pyjulia.py index df01ebf7..ddd484e5 100644 --- a/pyabc/external/julia/jl_pyjulia.py +++ b/pyabc/external/julia/jl_pyjulia.py @@ -155,7 +155,7 @@ class Julia: lazy implementation. """ - def __init__(self, module_name: str, source_file: str = None): + def __init__(self, module_name: str, source_file: str | None = None): if Main is None: raise ImportError( 'Install PyJulia, e.g. via `pip install pyabc[julia]`, ' diff --git a/pyabc/inference/smc.py b/pyabc/inference/smc.py index 7185bbe2..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,16 +190,16 @@ class ABCSMC: def __init__( self, models: list[Model] | Model | Callable, - parameter_priors: list[Distribution] | Distribution | Callable, - distance_function: Distance | Callable = None, + parameter_priors: list[Distribution] | Distribution, + distance_function: Distance | Callable | None = None, population_size: PopulationStrategy | int = 100, summary_statistics: Callable[[model_output], dict] = identity, - model_prior: RV = None, - model_perturbation_kernel: ModelPerturbationKernel = None, - transitions: list[Transition] | Transition = None, - eps: Epsilon = None, - sampler: Sampler = None, - acceptor: Acceptor = None, + model_prior: RV | None = None, + model_perturbation_kernel: ModelPerturbationKernel | None = None, + transitions: list[Transition] | Transition | None = None, + eps: Epsilon | None = None, + sampler: Sampler | None = None, + acceptor: Acceptor | None = None, stop_if_only_single_model_alive: bool = False, max_nr_recorded_particles: int = np.inf, ): @@ -290,10 +308,10 @@ def __getstate__(self): def new( self, db: str, - observed_sum_stat: dict = None, + observed_sum_stat: dict | None = None, *, - gt_model: int = None, - gt_par: dict = None, + gt_model: int | None = None, + gt_par: dict | None = None, meta_info=None, ) -> History: """ @@ -392,7 +410,7 @@ def load( self, db: str, abc_id: int = 1, - observed_sum_stat: dict = None, + observed_sum_stat: dict | None = None, ) -> History: """ Load an ABC-SMC run for continuation. @@ -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 @@ -628,11 +654,11 @@ def _create_transition_pdf(self, t: int, transitions): @run_cleanup def run( self, - minimum_epsilon: float = None, + minimum_epsilon: float | None = None, max_nr_populations: int = np.inf, min_acceptance_rate: float = 0.0, max_total_nr_simulations: int = np.inf, - max_walltime: timedelta = None, + max_walltime: timedelta | None = None, min_eps_diff: float = 0.0, ) -> History: """ @@ -715,11 +741,11 @@ def run( def initialize_components_before_run( self, - minimum_epsilon: float, + minimum_epsilon: float | None, max_nr_populations: int, min_acceptance_rate: float, max_total_nr_simulations: int, - max_walltime: timedelta, + max_walltime: timedelta | None, min_eps_diff: float, ) -> int: """Initialize everything before starting a run. @@ -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/inference_util/inference_util.py b/pyabc/inference_util/inference_util.py index 89b9f28d..f255c980 100644 --- a/pyabc/inference_util/inference_util.py +++ b/pyabc/inference_util/inference_util.py @@ -4,6 +4,7 @@ import uuid from collections.abc import Callable from datetime import datetime, timedelta +from functools import partial from typing import TYPE_CHECKING import numpy as np @@ -23,6 +24,138 @@ logger = logging.getLogger('ABC') +def _simulate_one_from_prior( + model_prior: RV, + parameter_priors: list[Distribution], + models: list[Model], + summary_statistics: Callable, +): + """Sample one particle from the prior.""" + from ..population import Particle + + # sample model + m = int(model_prior.rvs()) + # sample parameter + theta = parameter_priors[m].rvs() + # simulate summary statistics + model_result = models[m].summary_statistics(0, theta, summary_statistics) + # sampled from prior, so all have uniform weight + weight = 1.0 + # distance will be computed after initialization of the + # distance function + distance = np.inf + # all are happy and accepted + accepted = True + + return Particle( + m=m, + parameter=theta, + weight=weight, + sum_stat=model_result.sum_stat, + distance=distance, + accepted=accepted, + proposal_id=0, + preliminary=False, + ) + + +def _simulate_one( + *, + t: int, + m: np.ndarray, + p: np.ndarray, + model_prior: RV, + parameter_priors: list[Distribution], + model_perturbation_kernel: ModelPerturbationKernel, + transitions: list[Transition], + models: list[Model], + summary_statistics: Callable, + x_0: dict, + distance_function: Distance, + eps: Epsilon, + acceptor: Acceptor, + weight_function: Callable, + evaluate: bool, + proposal_id: int, +): + """Sample one parameter and evaluate/simulate one particle.""" + parameter = generate_valid_proposal( + t=t, + m=m, + p=p, + model_prior=model_prior, + parameter_priors=parameter_priors, + model_perturbation_kernel=model_perturbation_kernel, + transitions=transitions, + ) + if evaluate: + particle = evaluate_proposal( + *parameter, + t=t, + models=models, + summary_statistics=summary_statistics, + distance_function=distance_function, + eps=eps, + acceptor=acceptor, + x_0=x_0, + weight_function=weight_function, + proposal_id=proposal_id, + ) + else: + particle = only_simulate_data_for_proposal( + *parameter, + t=t, + models=models, + summary_statistics=summary_statistics, + weight_function=weight_function, + proposal_id=proposal_id, + ) + return particle + + +def _prior_pdf( + m_ss: int, + theta_ss: Parameter, + model_prior: RV, + parameter_priors: list[Distribution], +) -> float: + """Evaluate the prior density for a proposed sample.""" + return model_prior.pmf(m_ss) * parameter_priors[m_ss].pdf(theta_ss) + + +def _transition_pdf( + m_ss: int, + theta_ss: Parameter, + transitions: list[Transition], + model_probabilities: pd.DataFrame, + model_perturbation_kernel: ModelPerturbationKernel, +) -> float: + """Evaluate the transition density for a proposed sample.""" + model_factor = sum( + row.p * model_perturbation_kernel.pmf(m_ss, m) + for m, row in model_probabilities.iterrows() + ) + particle_factor = transitions[m_ss].pdf(theta_ss) + + transition_pd = model_factor * particle_factor + if transition_pd == 0: + logger.debug('Transition density is zero!') + return transition_pd + + +def _weight_function( + m_ss: int, + theta_ss: Parameter, + acceptance_weight: float, + prior_pdf: Callable, + transition_pdf: Callable, +) -> float: + """Calculate total weight from sampling and acceptance weight.""" + prior_pd = prior_pdf(m_ss, theta_ss) + transition_pd = transition_pdf(m_ss, theta_ss) + return acceptance_weight * prior_pd / transition_pd + + class AnalysisVars: """Contract object class for passing analysis variables. @@ -97,38 +230,13 @@ def create_simulate_from_prior_function( simulate_one: A function that returns a sampled particle. """ - # simulation function, simplifying some parts compared to later - from ..population import Particle - - def simulate_one(): - # sample model - m = int(model_prior.rvs()) - # sample parameter - theta = parameter_priors[m].rvs() - # simulate summary statistics - model_result = models[m].summary_statistics( - 0, theta, summary_statistics - ) - # sampled from prior, so all have uniform weight - weight = 1.0 - # distance will be computed after initialization of the - # distance function - distance = np.inf - # all are happy and accepted - accepted = True - - return Particle( - m=m, - parameter=theta, - weight=weight, - sum_stat=model_result.sum_stat, - distance=distance, - accepted=accepted, - proposal_id=0, - preliminary=False, - ) - - return simulate_one + return partial( + _simulate_one_from_prior, + model_prior=model_prior, + parameter_priors=parameter_priors, + models=models, + summary_statistics=summary_statistics, + ) def generate_valid_proposal( @@ -273,11 +381,11 @@ def create_prior_pdf( prior_pdf: The prior density function. """ - def prior_pdf(m_ss, theta_ss): - prior_pd = model_prior.pmf(m_ss) * parameter_priors[m_ss].pdf(theta_ss) - return prior_pd - - return prior_pdf + return partial( + _prior_pdf, + model_prior=model_prior, + parameter_priors=parameter_priors, + ) def create_transition_pdf( @@ -298,20 +406,12 @@ def create_transition_pdf( transition_pdf: The transition density function. """ - def transition_pdf(m_ss, theta_ss): - model_factor = sum( - row.p * model_perturbation_kernel.pmf(m_ss, m) - for m, row in model_probabilities.iterrows() - ) - particle_factor = transitions[m_ss].pdf(theta_ss) - - transition_pd = model_factor * particle_factor - - if transition_pd == 0: - logger.debug('Transition density is zero!') - return transition_pd - - return transition_pdf + return partial( + _transition_pdf, + transitions=transitions, + model_probabilities=model_probabilities, + model_perturbation_kernel=model_perturbation_kernel, + ) def create_weight_function( @@ -332,27 +432,11 @@ def create_weight_function( weight_function: The importance sample weight function. """ - def weight_function(m_ss, theta_ss, acceptance_weight: float): - """Calculate total weight, from sampling and acceptance weight. - - Parameters - ---------- - m_ss: The model sample. - theta_ss: The parameter sample. - acceptance_weight: The acceptance weight sample. In most cases 1. - - Returns - ------- - weight: The total weight. - """ - # prior and transition density (can be equal) - prior_pd = prior_pdf(m_ss, theta_ss) - transition_pd = transition_pdf(m_ss, theta_ss) - # calculate weight - weight = acceptance_weight * prior_pd / transition_pd - return weight - - return weight_function + return partial( + _weight_function, + prior_pdf=prior_pdf, + transition_pdf=transition_pdf, + ) def create_simulate_function( @@ -431,42 +515,25 @@ def create_simulate_function( prior_pdf=prior_pdf, transition_pdf=transition_pdf ) - # simulation function - def simulate_one(): - parameter = generate_valid_proposal( - t=t, - m=m, - p=p, - model_prior=model_prior, - parameter_priors=parameter_priors, - model_perturbation_kernel=model_perturbation_kernel, - transitions=transitions, - ) - if evaluate: - particle = evaluate_proposal( - *parameter, - t=t, - models=models, - summary_statistics=summary_statistics, - distance_function=distance_function, - eps=eps, - acceptor=acceptor, - x_0=x_0, - weight_function=weight_function, - proposal_id=proposal_id, - ) - else: - particle = only_simulate_data_for_proposal( - *parameter, - t=t, - models=models, - summary_statistics=summary_statistics, - weight_function=weight_function, - proposal_id=proposal_id, - ) - return particle - - return simulate_one + return partial( + _simulate_one, + t=t, + m=m, + p=p, + model_prior=model_prior, + parameter_priors=parameter_priors, + model_perturbation_kernel=model_perturbation_kernel, + transitions=transitions, + models=models, + summary_statistics=summary_statistics, + x_0=x_0, + distance_function=distance_function, + eps=eps, + acceptor=acceptor, + weight_function=weight_function, + evaluate=evaluate, + proposal_id=proposal_id, + ) def only_simulate_data_for_proposal( @@ -623,10 +690,10 @@ def create_analysis_id(): return str(uuid.uuid4()) -def eps_from_hist(history: History, t: int = None) -> float: +def eps_from_hist(history: History, t: int | None = None) -> float | None: """Read epsilon value for time `t` from `history`. Defaults to latest.""" pops = history.get_all_populations() - if len(pops) == 0 or (t is not None and t not in pops.t): + if len(pops) == 0 or (t is not None and t not in pops.t.values): return None if t is None: return pops.epsilon.to_numpy()[-1] diff --git a/pyabc/model/model.py b/pyabc/model/model.py index 901c103d..80cdb340 100644 --- a/pyabc/model/model.py +++ b/pyabc/model/model.py @@ -16,9 +16,9 @@ class ModelResult: def __init__( self, - sum_stat: dict = None, - distance: float = None, - accepted: bool = None, + sum_stat: dict | None = None, + distance: float | None = None, + accepted: bool | None = None, weight: float = 1.0, ): self.sum_stat = sum_stat if sum_stat is not None else {} @@ -233,7 +233,9 @@ class FunctionModel(Model): """ def __init__( - self, sample_function: Callable[[Parameter], Any], name: str = None + self, + sample_function: Callable[[Parameter], Any], + name: str | None = None, ): if name is None: # try to get the model name diff --git a/pyabc/petab/amici.py b/pyabc/petab/amici.py index f67c8ab4..07f282c7 100644 --- a/pyabc/petab/amici.py +++ b/pyabc/petab/amici.py @@ -186,8 +186,8 @@ class AmiciPetabImporter(PetabImporter): def __init__( self, petab_problem: petab.Problem, - amici_model: amici.sim.sundials.Model = None, - amici_solver: amici.sim.sundials.Solver = None, + amici_model: amici.sim.sundials.Model | None = None, + amici_solver: amici.sim.sundials.Solver | None = None, ): super().__init__(petab_problem=petab_problem) diff --git a/pyabc/populationstrategy/populationstrategy.py b/pyabc/populationstrategy/populationstrategy.py index 3a528c45..f29b220a 100644 --- a/pyabc/populationstrategy/populationstrategy.py +++ b/pyabc/populationstrategy/populationstrategy.py @@ -38,14 +38,14 @@ class directly. Subclasses must override the `update` method. Number of calibration particles. """ - def __init__(self, nr_calibration_particles: int = None): + def __init__(self, nr_calibration_particles: int | None = None): self.nr_calibration_particles = nr_calibration_particles def update( # noqa: B027 self, transitions: list[Transition], model_weights: np.ndarray, - t: int = None, + t: int | None = None, ): """ Select the population size for the next population. @@ -62,7 +62,7 @@ def update( # noqa: B027 pass @abstractmethod - def __call__(self, t: int = None) -> int: + def __call__(self, t: int | None = None) -> int: raise NotImplementedError() def get_config(self) -> dict: @@ -108,13 +108,13 @@ class ConstantPopulationSize(PopulationStrategy): def __init__( self, nr_particles: int, - nr_calibration_particles: int = None, + nr_calibration_particles: int | None = None, ): super().__init__(nr_calibration_particles=nr_calibration_particles) self.nr_particles = nr_particles @dec_bound_pop_size_from_env - def __call__(self, t: int = None) -> int: + def __call__(self, t: int | None = None) -> int: if t == -1 and self.nr_calibration_particles is not None: return self.nr_calibration_particles return self.nr_particles @@ -173,7 +173,7 @@ def __init__( max_population_size: int = np.inf, min_population_size: int = 10, n_bootstrap: int = 10, - nr_calibration_particles: int = None, + nr_calibration_particles: int | None = None, ): super().__init__( nr_calibration_particles=nr_calibration_particles, @@ -200,7 +200,7 @@ def update( self, transitions: list[Transition], model_weights: np.ndarray, - t: int = None, # noqa: ARG002 + t: int | None = None, # noqa: ARG002 ): test_X = [trans.X for trans in transitions] test_w = [trans.w for trans in transitions] @@ -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, @@ -231,7 +231,7 @@ def update( ) @dec_bound_pop_size_from_env - def __call__(self, t: int = None) -> int: + def __call__(self, t: int | None = None) -> int: if t == -1 and self.nr_calibration_particles is not None: return self.nr_calibration_particles return self.nr_particles @@ -254,7 +254,7 @@ class ListPopulationSize(PopulationStrategy): def __init__( self, values: list[int] | dict[int, int], - nr_calibration_particles: int = None, + nr_calibration_particles: int | None = None, ): super().__init__(nr_calibration_particles=nr_calibration_particles) self.values = values @@ -265,7 +265,7 @@ def get_config(self) -> dict: return config @dec_bound_pop_size_from_env - def __call__(self, t: int = None) -> int: + def __call__(self, t: int | None = None) -> int: if t == -1 and self.nr_calibration_particles is not None: return self.nr_calibration_particles return self.values[t] diff --git a/pyabc/predictor/predictor.py b/pyabc/predictor/predictor.py index 5a29ec00..1e718396 100644 --- a/pyabc/predictor/predictor.py +++ b/pyabc/predictor/predictor.py @@ -37,7 +37,9 @@ class Predictor(ABC): """ @abstractmethod - def fit(self, x: np.ndarray, y: np.ndarray, w: np.ndarray = None) -> None: + def fit( + self, x: np.ndarray, y: np.ndarray, w: np.ndarray | None = None + ) -> None: """Fit the predictor to labeled data. Parameters @@ -147,7 +149,9 @@ def __init__( self.std_y: np.ndarray | None = None @wrap_fit_log - def fit(self, x: np.ndarray, y: np.ndarray, w: np.ndarray = None) -> None: + def fit( + self, x: np.ndarray, y: np.ndarray, w: np.ndarray | None = None + ) -> None: """Fit the predictor to labeled data. Parameters @@ -307,7 +311,9 @@ def __init__( log_pearson=log_pearson, ) - def fit(self, x: np.ndarray, y: np.ndarray, w: np.ndarray = None) -> None: + def fit( + self, x: np.ndarray, y: np.ndarray, w: np.ndarray | None = None + ) -> None: super().fit(x, y, w) # log if self.joint: @@ -372,7 +378,7 @@ class GPPredictor(SimplePredictor): def __init__( self, - kernel: Callable | skl_gp.kernels.Kernel = None, + kernel: Callable | skl_gp.kernels.Kernel | None = None, normalize_features: bool = True, normalize_labels: bool = True, joint: bool = True, @@ -408,7 +414,9 @@ def __init__( log_pearson=log_pearson, ) - def fit(self, x: np.ndarray, y: np.ndarray, w: np.ndarray = None) -> None: + def fit( + self, x: np.ndarray, y: np.ndarray, w: np.ndarray | None = None + ) -> None: # need to recreate the model # set indices to keep @@ -438,8 +446,8 @@ class GPKernelHandle: def __init__( self, - kernels: list[str] = None, - kernel_kwargs: list[dict] = None, + kernels: list[str] | None = None, + kernel_kwargs: list[dict] | None = None, ard: bool = True, ): """ @@ -508,7 +516,7 @@ def __init__( normalize_features: bool = True, normalize_labels: bool = True, joint: bool = True, - hidden_layer_sizes: tuple[int, ...] | Callable = None, + hidden_layer_sizes: tuple[int, ...] | Callable | None = None, log_pearson: bool = True, **kwargs, ): @@ -549,7 +557,9 @@ def __init__( log_pearson=log_pearson, ) - def fit(self, x: np.ndarray, y: np.ndarray, w: np.ndarray = None) -> None: + def fit( + self, x: np.ndarray, y: np.ndarray, w: np.ndarray | None = None + ) -> None: # need to recreate the model # set indices to keep @@ -684,7 +694,7 @@ def __init__( split_method: str = TRAIN_TEST_SPLIT, n_splits: int = 5, test_size: float = 0.2, - f_score: Callable = None, + f_score: Callable | None = None, ): """ Parameters @@ -722,11 +732,15 @@ def __init__( if f_score is None: self.f_score = root_mean_square_error + else: + self.f_score = f_score # holds the chosen predictor model self.chosen_one: Predictor | None = None - def fit(self, x: np.ndarray, y: np.ndarray, w: np.ndarray = None) -> None: + def fit( + self, x: np.ndarray, y: np.ndarray, w: np.ndarray | None = None + ) -> None: # output normalization std_y = np.std(y, axis=0) @@ -844,7 +858,7 @@ def root_mean_square_relative_error( Returns ------- - val: The normalized root mean square relative error value. + val: The normalized root mean squared relative error value. """ return np.sqrt( np.sum( diff --git a/pyabc/sampler/base.py b/pyabc/sampler/base.py index a391ec5d..ec6869c1 100644 --- a/pyabc/sampler/base.py +++ b/pyabc/sampler/base.py @@ -96,7 +96,7 @@ def sample_until_n_accepted( *, max_eval: Real = np.inf, all_accepted: bool = False, - ana_vars: AnalysisVars = None, + ana_vars: AnalysisVars | None = None, ) -> Sample: """ Performs the sampling, i.e. creation of a new generation (i.e. diff --git a/pyabc/sampler/dask_sampler.py b/pyabc/sampler/dask_sampler.py index 9ce6ba9b..29bdfc64 100644 --- a/pyabc/sampler/dask_sampler.py +++ b/pyabc/sampler/dask_sampler.py @@ -41,7 +41,7 @@ class DaskDistributedSampler(EPSMixin, Sampler): def __init__( self, - dask_client: Client = None, + dask_client: Client | None = None, client_max_jobs: int = np.inf, default_pickle: bool = True, batch_size: int = 1, diff --git a/pyabc/sampler/multicore.py b/pyabc/sampler/multicore.py index fe7a7c79..44660299 100644 --- a/pyabc/sampler/multicore.py +++ b/pyabc/sampler/multicore.py @@ -1,6 +1,5 @@ import logging import random -from multiprocessing import Process, Queue import cloudpickle as pickle import numpy as np @@ -8,6 +7,7 @@ from .multicorebase import MultiCoreSampler, get_if_worker_healthy from .singlecore import SingleCoreSampler +from .util import get_mp_context logger = logging.getLogger('ABC.Sampler') @@ -36,7 +36,10 @@ def work(feed_q, result_q, simulate_one, max_eval, single_core_sampler): break res = single_core_sampler.sample_until_n_accepted( - 1, simulate_one, max_eval + 1, + simulate_one, + t=0, # t is not used in this sampler + max_eval=max_eval, ) result_q.put((res, single_core_sampler.nr_evaluations_)) @@ -88,10 +91,12 @@ def sample_until_n_accepted( logger.debug( f'Start sampling on {n_procs} cores ({self.n_procs} requested)' ) - feed_q = Queue() - result_q = Queue() + # queues and processes must share one context + ctx = get_mp_context() + feed_q = ctx.Queue() + result_q = ctx.Queue() - feed_process = Process(target=feed, args=(feed_q, n, n_procs)) + feed_process = ctx.Process(target=feed, args=(feed_q, n, n_procs)) single_core_sampler = SingleCoreSampler( check_max_eval=self.check_max_eval @@ -105,7 +110,7 @@ def sample_until_n_accepted( args = (feed_q, result_q, simulate_one, max_eval, single_core_sampler) worker_processes = [ - Process(target=work, args=args) for _ in range(n_procs) + ctx.Process(target=work, args=args) for _ in range(n_procs) ] for proc in worker_processes: diff --git a/pyabc/sampler/multicore_evaluation_parallel.py b/pyabc/sampler/multicore_evaluation_parallel.py index 7d138afb..4f33f514 100644 --- a/pyabc/sampler/multicore_evaluation_parallel.py +++ b/pyabc/sampler/multicore_evaluation_parallel.py @@ -1,12 +1,12 @@ import random from ctypes import c_longlong -from multiprocessing import Process, Queue, Value import cloudpickle as pickle import numpy as np from jabbar import jabbar from .multicorebase import MultiCoreSampler, get_if_worker_healthy +from .util import get_mp_context DONE = 'Done' @@ -14,8 +14,8 @@ def work( simulate_one, queue, - n_eval: Value, - n_acc: Value, + n_eval, + n_acc, n: int, check_max_eval: bool, max_eval: int, @@ -100,13 +100,16 @@ def sample_until_n_accepted( all_accepted=False, ana_vars=None, # noqa: ARG002 ): - n_eval = Value(c_longlong) + # values, queue and processes must share one context + ctx = get_mp_context() + + n_eval = ctx.Value(c_longlong) n_eval.value = 0 - n_acc = Value(c_longlong) + n_acc = ctx.Value(c_longlong) n_acc.value = 0 - queue = Queue() + queue = ctx.Queue() # wrap arguments if self.pickle: @@ -124,7 +127,7 @@ def sample_until_n_accepted( ) processes = [ - Process(target=work, args=args, daemon=self.daemon) + ctx.Process(target=work, args=args, daemon=self.daemon) for _ in range(self.n_procs) ] diff --git a/pyabc/sampler/multicorebase.py b/pyabc/sampler/multicorebase.py index 6204eb27..946ed5b8 100644 --- a/pyabc/sampler/multicorebase.py +++ b/pyabc/sampler/multicorebase.py @@ -31,9 +31,9 @@ class MultiCoreSampler(Sampler): def __init__( self, - n_procs: int = None, + n_procs: int | None = None, daemon: bool = True, - pickle: bool = None, + pickle: bool | None = None, check_max_eval: bool = False, ): super().__init__() diff --git a/pyabc/sampler/redis_eps/cli.py b/pyabc/sampler/redis_eps/cli.py index 36e60c27..a82cf324 100644 --- a/pyabc/sampler/redis_eps/cli.py +++ b/pyabc/sampler/redis_eps/cli.py @@ -2,13 +2,13 @@ import os import random import socket -from multiprocessing import Process from time import time import click import numpy as np from redis import StrictRedis +from ..util import get_mp_process from .cmd import ( ANALYSIS_ID, DYNAMIC, @@ -89,8 +89,9 @@ def work( return _work(host, port, runtime, password, catch) # define parallel processes + process_cls = get_mp_process() procs = [ - Process( + process_cls( target=_work, args=(host, port, runtime, password, catch), daemon=daemon, @@ -198,7 +199,8 @@ def _work( logger.info('Received stop signal. Shutdown redis worker.') return - # TODO other messages (some integers?) are ignored + # Any other messages on the channel (e.g. redis subscription-count + # notifications) are ignored. # check total time condition elapsed_time = time() - start_time diff --git a/pyabc/sampler/redis_eps/redis_logging.py b/pyabc/sampler/redis_eps/redis_logging.py index 79a23b2d..aa4fb559 100644 --- a/pyabc/sampler/redis_eps/redis_logging.py +++ b/pyabc/sampler/redis_eps/redis_logging.py @@ -9,7 +9,7 @@ class RedisSamplerLogger: """A logger for the redis sampler with enhanced interest-variables output.""" - def __init__(self, log_file: str = None): + def __init__(self, log_file: str | None = None): self.log_file = log_file if log_file: if os.path.exists(log_file) and os.stat(log_file).st_size > 0: diff --git a/pyabc/sampler/redis_eps/sampler.py b/pyabc/sampler/redis_eps/sampler.py index 5138444a..ea2abb75 100644 --- a/pyabc/sampler/redis_eps/sampler.py +++ b/pyabc/sampler/redis_eps/sampler.py @@ -74,8 +74,8 @@ def __init__( self, host: str = 'localhost', port: int = 6379, - password: str = None, - log_file: str = None, + password: str | None = None, + log_file: str | None = None, ): super().__init__() logger.debug(f'Redis sampler: host={host} port={port}') @@ -113,7 +113,7 @@ def sample_until_n_accepted( *, max_eval: int = np.inf, all_accepted: bool = False, - ana_vars: AnalysisVars = None, + ana_vars: AnalysisVars | None = None, ) -> Sample: raise NotImplementedError() @@ -219,14 +219,14 @@ def __init__( self, host: str = 'localhost', port: int = 6379, - password: str = None, + password: str | None = None, batch_size: int = 1, look_ahead: bool = False, look_ahead_delay_evaluation: bool = True, max_n_eval_look_ahead_factor: float = 10.0, wait_for_all_samples: bool = False, adapt_look_ahead_proposal: bool = False, - log_file: str = None, + log_file: str | None = None, ): super().__init__( host=host, port=port, password=password, log_file=log_file diff --git a/pyabc/sampler/redis_eps/server_starter.py b/pyabc/sampler/redis_eps/server_starter.py index 14bac973..29a9d951 100644 --- a/pyabc/sampler/redis_eps/server_starter.py +++ b/pyabc/sampler/redis_eps/server_starter.py @@ -1,11 +1,11 @@ import tempfile import time -from multiprocessing import Process from subprocess import Popen # noqa: S404 from time import sleep import psutil +from ..util import get_mp_process from .cli import _manage, work from .sampler import RedisEvalParallelSampler from .sampler_static import RedisStaticSampler @@ -14,7 +14,7 @@ class RedisServerStarter: def __init__( self, - password: str = None, + password: str | None = None, workers: int = 2, processes_per_worker: int = 1, daemon: bool = True, @@ -47,8 +47,9 @@ def __init__( # initiate worker processes maybe_password = [] if password is None else ['--password', password] maybe_daemon = [] if daemon is None else ['--daemon', str(daemon)] + process_cls = get_mp_process() self.workers = [ - Process( + process_cls( target=work, args=( [ @@ -105,7 +106,7 @@ class RedisEvalParallelSamplerServerStarter(RedisEvalParallelSampler): def __init__( self, - password: str = None, + password: str | None = None, batch_size: int = 1, wait_for_all_samples: bool = False, look_ahead: bool = False, @@ -116,7 +117,7 @@ def __init__( processes_per_worker: int = 1, daemon: bool = True, catch: bool = True, - log_file: str = None, + log_file: str | None = None, ): self.server_starter = RedisServerStarter( password=password, @@ -151,12 +152,12 @@ class RedisStaticSamplerServerStarter(RedisStaticSampler): def __init__( self, - password: str = None, + password: str | None = None, workers: int = 2, processes_per_worker: int = 1, daemon: bool = True, catch: bool = True, - log_file: str = None, + log_file: str | None = None, ): self.server_starter = RedisServerStarter( password=password, diff --git a/pyabc/sampler/util.py b/pyabc/sampler/util.py index d1012f38..74de6a87 100644 --- a/pyabc/sampler/util.py +++ b/pyabc/sampler/util.py @@ -1,8 +1,35 @@ """Sampling util functions.""" +import multiprocessing as mp + from ..population import Sample def any_particle_preliminary(sample: Sample) -> bool: """Determine whether any particle in that sample is preliminary.""" return any(particle.preliminary for particle in sample.all_particles) + + +def get_mp_context(): + """Get a multiprocessing context. + + On POSIX, prefer ``fork`` when available to support non-picklable + callables (e.g. local functions in tests) in nested worker setups. + + Queues, Values and Processes that are shared with each other must all be + created from the *same* context. Mixing contexts (e.g. a default-context + ``spawn`` queue with a ``fork`` process) can crash at runtime, so callers + should obtain a single context here and derive all of them from it. + """ + if 'fork' in mp.get_all_start_methods(): + return mp.get_context('fork') + return mp.get_context() + + +def get_mp_process(): + """Get a multiprocessing Process constructor. + + On POSIX, prefer ``fork`` when available to support non-picklable + callables (e.g. local functions in tests) in nested worker setups. + """ + return get_mp_context().Process diff --git a/pyabc/settings/settings.py b/pyabc/settings/settings.py index 7c0b211b..7a8e62c1 100644 --- a/pyabc/settings/settings.py +++ b/pyabc/settings/settings.py @@ -5,9 +5,9 @@ def set_figure_params( theme: str = 'pyabc', - style: str = None, + style: str | None = None, color_map: str = 'plasma', - color_cycle: list[str] = None, + color_cycle: list[str] | None = None, ) -> None: """Set global figure parameters for a consistent, beautified design. diff --git a/pyabc/sge/db.py b/pyabc/sge/db.py index 1463f608..b8d6db1a 100644 --- a/pyabc/sge/db.py +++ b/pyabc/sge/db.py @@ -49,11 +49,9 @@ def wait_for_job(self, ID, max_run_time_h): Return true if we should still wait for the job. Return false otherwise """ - # TODO Possible SQL injection error should be fixed, e.g. via - # pre-calculated expressions with self.connection: results = self.connection.execute( - 'SELECT status, time from status WHERE ID=' + str(ID) + 'SELECT status, time from status WHERE ID=?', (ID,) ).fetchall() nr_rows = len(results) diff --git a/pyabc/sge/sge.py b/pyabc/sge/sge.py index dac9db6a..15e6d5de 100644 --- a/pyabc/sge/sge.py +++ b/pyabc/sge/sge.py @@ -119,12 +119,12 @@ class SGE: def __init__( self, - tmp_directory: str = None, + tmp_directory: str | None = None, memory: str = '3G', time_h: int = 100, - python_executable_path: str = None, - sge_error_file: str = None, - sge_output_file: str = None, + python_executable_path: str | None = None, + sge_error_file: str | None = None, + sge_output_file: str | None = None, parallel_environment=None, name='map', queue=None, diff --git a/pyabc/storage/bytes_storage.py b/pyabc/storage/bytes_storage.py index e4d006b3..5f95b4b6 100644 --- a/pyabc/storage/bytes_storage.py +++ b/pyabc/storage/bytes_storage.py @@ -16,7 +16,7 @@ def r_to_py(object_): py_object_ = conv.rpy2py(object_) # Ensure factor columns are converted to strings for col in py_object_.columns: - if isinstance(py_object_[col], pd.CategoricalDtype): + if isinstance(py_object_[col].dtype, pd.CategoricalDtype): py_object_[col] = py_object_[col].astype(str) return py_object_ return object_ diff --git a/pyabc/storage/db_model.py b/pyabc/storage/db_model.py index 1b7efbf3..d5eae1bb 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') @@ -168,6 +169,8 @@ class SummaryStatistic(Base): value = Column(BytesStorage) -def datetime2str(datetime: datetime.datetime) -> str: +def datetime2str(datetime: datetime.datetime | None) -> str: """Format print datetime.""" + if datetime is None: + return 'None' return datetime.strftime('%Y-%m-%d %H:%M:%S') diff --git a/pyabc/storage/history.py b/pyabc/storage/history.py index 41177a1c..3a0024a1 100644 --- a/pyabc/storage/history.py +++ b/pyabc/storage/history.py @@ -80,7 +80,7 @@ def git_hash(): return hash_ -def create_sqlite_db_id(dir_: str = None, file_: str = 'pyabc_test.db'): +def create_sqlite_db_id(dir_: str | None = None, file_: str = 'pyabc_test.db'): """ Convenience function to create a sqlite database identifier which can be understood by sqlalchemy. @@ -144,7 +144,7 @@ def __init__( self, db: str, stores_sum_stats: bool = True, - _id: int = None, + _id: int | None = None, create: bool = True, ): """Initialize history object. @@ -273,7 +273,7 @@ def id(self, val): self._id = val @with_session - def alive_models(self, t: int = None) -> list: + def alive_models(self, t: int | None = None) -> list: """ Get the models which are still alive at time `t`. @@ -303,7 +303,7 @@ def alive_models(self, t: int = None) -> list: @with_session def get_distribution( - self, m: int = 0, t: int = None + self, m: int = 0, t: int | None = None ) -> tuple[pd.DataFrame, np.ndarray]: """ Returns the weighted population sample for model m and timepoint t @@ -344,10 +344,10 @@ def get_distribution( ).sort_index() w = df[['id', 'w']].drop_duplicates().set_index('id').sort_index() w_arr = w.w.values - if w_arr.size > 0 and not np.isclose(w_arr.sum(), 1): - raise AssertionError( - f'Weight not close to 1, w.sum()={w_arr.sum()}' - ) + # Stored weights are global (summing to 1 across all models); within a + # single model they sum to the model probability. + if w_arr.size > 0: + w_arr = w_arr / w_arr.sum() return pars, w_arr @with_session @@ -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, @@ -423,7 +426,7 @@ def store_initial_data( distance_function_json_str: str, eps_function_json_str: str, population_strategy_json_str: str, - start_time: datetime.datetime = None, + start_time: datetime.datetime | None = None, ) -> None: """ Store the initial configuration data. @@ -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() @@ -671,7 +680,7 @@ def __getstate__(self): @with_session @internal_docstring_warning - def done(self, end_time: datetime.datetime = None): + def done(self, end_time: datetime.datetime | None = None): """ Close database sessions and store end time of the analysis. @@ -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) @@ -722,18 +735,16 @@ def _save_to_population_db( # append model population.models.append(model) - # TODO This normalization is different than in the in-memory - # population. It would be cleaner to update the db too. - total_model_weight = sum(p.weight for p in model_population) - # iterate over model population of particles for py_particle in model_population: # a store_item is a Particle py_parameter = py_particle.parameter - # create new particle + # create new particle. The stored weight is the global + # particle weight (normalized across all particles of all + # models so that they sum to 1) particle = Particle( - w=py_particle.weight / total_model_weight, + w=py_particle.weight, proposal_id=py_particle.proposal_id, ) # append particle to model @@ -783,6 +794,7 @@ def append_population( population: PyPopulation, nr_simulations: int, model_names, + wall_time: float | None = None, ): """ Append population to database. @@ -799,6 +811,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 +824,7 @@ def append_population( particles_by_model, model_probabilities, model_names, + wall_time=wall_time, ) @with_session @@ -841,14 +856,13 @@ def get_model_probabilities(self, t: int | None = None) -> pd.DataFrame: .order_by(Model.m) .all() ) - # TODO this is a mess + # Two return shapes: for a single t, a frame indexed by model id `m` + # with a `p` column; for t=None, a t-by-model pivot table (below). if t is not None: p_models_df = pd.DataFrame( [p[:2] for p in p_models], columns=['p', 'm'] ).set_index('m') - # TODO the following line is redundant - # only models with no-zero weight are stored for each population - p_models_df = p_models_df[p_models_df.p >= 0] + # Only models with non-zero weight are stored per population return p_models_df else: p_models_df = ( @@ -858,7 +872,7 @@ def get_model_probabilities(self, t: int | None = None) -> pd.DataFrame: ) return p_models_df - def nr_of_models_alive(self, t: int = None) -> int: + def nr_of_models_alive(self, t: int | None = None) -> int: """ Number of models still alive. @@ -880,7 +894,7 @@ def nr_of_models_alive(self, t: int = None) -> int: return int((model_probs.p > 0).sum()) @with_session - def get_weighted_distances(self, t: int = None) -> pd.DataFrame: + def get_weighted_distances(self, t: int | None = None) -> pd.DataFrame: """ Population's weighted distances to the measured sample. These weights do not necessarily sum up to 1. @@ -916,21 +930,11 @@ def get_weighted_distances(self, t: int = None) -> pd.DataFrame: distances = [] for model in models: for particle in model.particles: - weight = particle.w * model.p_model + weight = particle.w for sample in particle.samples: 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 @@ -981,7 +985,7 @@ def n_populations(self): @with_session def get_weighted_sum_stats_for_model( - self, m: int = 0, t: int = None + self, m: int = 0, t: int | None = None ) -> tuple[np.ndarray, list]: """ Summary statistics for model `m`. The weights sum to 1, unless @@ -1011,14 +1015,31 @@ def get_weighted_sum_stats_for_model( .filter(ABCSMC.id == self.id) .filter(Population.t == t) .filter(Model.m == m) + .options( + subqueryload(Particle.samples).subqueryload( + Sample.summary_statistics + ) + ) .all() ) + # model probability, used to renormalize the stored global weights + # back to the within-model posterior (weights summing to 1) + p_model = ( + self._session.query(Model.p_model) + .join(Population) + .join(ABCSMC) + .filter(ABCSMC.id == self.id) + .filter(Population.t == t) + .filter(Model.m == m) + .scalar() + ) + results = [] weights = [] for particle in particles: for sample in particle.samples: - weights.append(particle.w) + weights.append(particle.w / p_model if p_model else particle.w) sum_stats = {} for ss in sample.summary_statistics: sum_stats[ss.name] = ss.value @@ -1027,7 +1048,7 @@ def get_weighted_sum_stats_for_model( @with_session def get_weighted_sum_stats( - self, t: int = None + self, t: int | None = None ) -> tuple[list[float], list[dict]]: """ Population's weighted summary statistics. @@ -1068,7 +1089,7 @@ def get_weighted_sum_stats( for model in models: for particle in model.particles: - weight = particle.w * model.p_model + weight = particle.w for sample in particle.samples: # extract sum stats sum_stats = {} @@ -1080,7 +1101,7 @@ def get_weighted_sum_stats( return all_weights, all_sum_stats @with_session - def get_population(self, t: int = None): + def get_population(self, t: int | None = None): """ Create a pyabc.Population object containing all particles, as far as those can be recreated from the database. In particular, @@ -1118,7 +1139,7 @@ def get_population(self, t: int = None): py_m = model.m for particle in model.particles: # weight - py_weight = particle.w * model.p_model + py_weight = particle.w # parameter py_parameter = {} @@ -1127,7 +1148,9 @@ def get_population(self, t: int = None): py_parameter = PyParameter(**py_parameter) # simulations - # TODO this is legacy from when there were multiple + # NOTE: samples is a one-to-many relationship for legacy + # reasons (a particle could once store multiple samples); + # today exactly one is expected. if len(particle.samples) != 1: raise AssertionError('There should be exactly one sample.') sample = particle.samples[0] diff --git a/pyabc/storage/json.py b/pyabc/storage/json.py index d2f91fdd..e90afcae 100644 --- a/pyabc/storage/json.py +++ b/pyabc/storage/json.py @@ -19,7 +19,7 @@ def save_dict_to_json(dct: dict, file_: str): for key, val in dct.items(): # cannot handle ndarrays if isinstance(val, np.ndarray): - dct[key] = list(val) + dct[key] = val.tolist() with open(file_, 'w') as f: json.dump(dct, f) diff --git a/pyabc/storage/migrate.py b/pyabc/storage/migrate.py index a5ddda71..1a22ced6 100644 --- a/pyabc/storage/migrate.py +++ b/pyabc/storage/migrate.py @@ -14,6 +14,57 @@ SQLITE_STR = 'sqlite:///' +def _to_db_file(db: str) -> str: + """Normalize a database identifier to a file name. + + Parameters + ---------- + db: Database file name, or sqlite URL ``sqlite:///``. + + Returns + ------- + db_file: The database file name. + + Raises + ------ + ValueError: If a URL of a dialect other than sqlite is passed. + """ + if db.startswith(SQLITE_STR): + return db[len(SQLITE_STR) :] + if '://' in db: + raise ValueError( + f'Cannot handle database identifier {db}: migration currently ' + f'only supports sqlite databases, i.e. either a file name, or a ' + f'URL of the form {SQLITE_STR}.' + ) + return db + + +def _alembic_config(db: str) -> 'Config': + """Create the alembic configuration operating on a database. + + Parameters + ---------- + db: Database the migrations are applied to, either as a file name or as a + sqlite URL ``sqlite:///``. + + Returns + ------- + cfg: The alembic configuration. + """ + # config base path + base_path = os.path.dirname(os.path.abspath(__file__)) + # read configuration file + cfg = Config(os.path.join(base_path, 'alembic.ini')) + # set absolute script location path + cfg.set_main_option( + 'script_location', os.path.join(base_path, 'migrations') + ) + # set target database file + cfg.set_main_option('sqlalchemy.url', SQLITE_STR + _to_db_file(db)) + return cfg + + @click.command( help='**Migrate pyABC database**\n\n' "Sometimes, changes to pyABC's storage format are unavoidable. " @@ -36,8 +87,8 @@ def migrate(src: str, dst: str, version: str) -> None: Parameters ---------- - src: Source - dst: Destination + src: Source, either a file name or a sqlite URL + dst: Destination, either a file name or a sqlite URL version: Version to migrate to """ if Config is None or command is None: @@ -48,10 +99,11 @@ def migrate(src: str, dst: str, version: str) -> None: return # to file paths if URLs - if src.startswith(SQLITE_STR): - src = src[len(SQLITE_STR) :] - if dst.startswith(SQLITE_STR): - dst = dst[len(SQLITE_STR) :] + try: + src, dst = _to_db_file(src), _to_db_file(dst) + except ValueError as e: + print(f'Error: {e}') + return # copy file if src != dst: @@ -61,16 +113,5 @@ def migrate(src: str, dst: str, version: str) -> None: # copy source to destination shutil.copyfile(src=src, dst=dst) - # config base path - base_path = os.path.dirname(os.path.abspath(__file__)) - # read configuration file - cfg = Config(os.path.join(base_path, 'alembic.ini')) - # set absolute script location path - cfg.set_main_option( - 'script_location', os.path.join(base_path, 'migrations') - ) - # set target database file - cfg.set_main_option('sqlalchemy.url', SQLITE_STR + dst) - # run the actual upgrade - command.upgrade(cfg, version) + command.upgrade(_alembic_config(dst), version) 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..c618ee19 --- /dev/null +++ b/pyabc/storage/migrations/versions/2_20260724_add_populations_wall_time.py @@ -0,0 +1,48 @@ +"""add populations.wall_time and store global particle weights + +Revision ID: 2 +Revises: 1 +Create Date: 2026-07-24 00:00:00.000000 + +This revision bundles two v2 storage-format changes: + +* adds the ``populations.wall_time`` column (wall-time tracking), and +* converts particle weights from the old per-model normalization + (``w = g_i / p_model``, summing to 1 within each model) to the global + convention (``w = g_i``, summing to 1 across all particles of all models), + matching the in-memory ``Population`` representation. The transform is + ``w := w * p_model``; within-model normalization is recovered on read. +""" + +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(): + # the column may exist already, as the downgrade does not remove it + inspector = sa.inspect(op.get_bind()) + columns = [col['name'] for col in inspector.get_columns('populations')] + if 'wall_time' not in columns: + op.add_column( + table_name='populations', + column=sa.Column('wall_time', sa.FLOAT, nullable=True), + ) + # per-model-normalized weights -> global weights (w := w * p_model) + op.execute( + 'UPDATE particles SET w = w * (' + 'SELECT p_model FROM models WHERE models.id = particles.model_id)' + ) + + +def downgrade(): + # global weights -> per-model-normalized weights (w := w / p_model) + op.execute( + 'UPDATE particles SET w = w / (' + 'SELECT p_model FROM models WHERE models.id = particles.model_id)' + ) diff --git a/pyabc/storage/numpy_bytes_storage.py b/pyabc/storage/numpy_bytes_storage.py index c1553135..7d7c7257 100644 --- a/pyabc/storage/numpy_bytes_storage.py +++ b/pyabc/storage/numpy_bytes_storage.py @@ -52,6 +52,12 @@ def np_from_bytes(arr_bytes): # try to convert to primitive types for type_ in _primitive_types: + if ( + type_ is int + and not np.issubdtype(arr.dtype, np.integer) + and arr.dtype != np.bool_ + ): + continue try: if type_(arr) == arr: return type_(arr) 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/sumstat/base.py b/pyabc/sumstat/base.py index 0aef357a..411ea021 100644 --- a/pyabc/sumstat/base.py +++ b/pyabc/sumstat/base.py @@ -18,7 +18,7 @@ class Sumstat(ABC): concatenated/chained. """ - def __init__(self, pre: 'Sumstat' = None): + def __init__(self, pre: 'Sumstat | None' = None): """ Parameters ---------- @@ -166,8 +166,8 @@ class IdentitySumstat(Sumstat): def __init__( self, - trafos: list[Callable[[np.ndarray], np.ndarray]] = None, - pre: Sumstat = None, + trafos: list[Callable[[np.ndarray], np.ndarray]] | None = None, + pre: Sumstat | None = None, shape_out: tuple[int, ...] = (-1,), ): """ @@ -188,7 +188,7 @@ def __init__( deriving from Sumstat or IdentitySumstat. """ super().__init__(pre=pre) - self.trafos: list[Callable[[np.ndarray], np.ndarray]] = trafos + self.trafos: list[Callable[[np.ndarray], np.ndarray]] | None = trafos self.shape_out: tuple[int, ...] = shape_out @io_dict2arr diff --git a/pyabc/sumstat/learn.py b/pyabc/sumstat/learn.py index 89f8fba1..2d07171d 100644 --- a/pyabc/sumstat/learn.py +++ b/pyabc/sumstat/learn.py @@ -48,14 +48,14 @@ class PredictorSumstat(Sumstat): def __init__( self, predictor: Predictor | Callable, - fit_ixs: EventIxs | Collection[int] | int = None, + fit_ixs: EventIxs | Collection[int] | int | None = None, all_particles: bool = True, normalize_labels: bool = True, fitted: bool = False, - subsetter: Subsetter = None, - pre: Sumstat = None, + subsetter: Subsetter | None = None, + pre: Sumstat | None = None, pre_before_fit: bool = False, - par_trafo: ParTrafoBase = None, + par_trafo: ParTrafoBase | None = None, ): """ Parameters diff --git a/pyabc/sumstat/subset.py b/pyabc/sumstat/subset.py index 60d4b242..24c668a5 100644 --- a/pyabc/sumstat/subset.py +++ b/pyabc/sumstat/subset.py @@ -102,7 +102,7 @@ def __init__( n_components_max: int = 5, min_fraction: float = 0.3, normalize_labels: bool = True, - gmm_args: dict = None, + gmm_args: dict | None = None, ): if skl_mx is None: raise ImportError( @@ -135,7 +135,9 @@ def select( """Select based on GMM clusters.""" # normalize if self.normalize_labels: - y_norm = (y - np.mean(y, axis=0)) / np.std(y, axis=0) + std = np.std(y, axis=0) + # avoid division by zero for constant (zero-variance) columns + y_norm = (y - np.mean(y, axis=0)) / np.where(std == 0, 1.0, std) else: y_norm = y @@ -242,13 +244,18 @@ def get_augmented_subset( # sort remaining values by distance to reference point y_left = y[~in_cluster] distances: np.ndarray = np.linalg.norm(y_left - ref, ord=2, axis=1) - # indices of the required closest parameters - ixs_nearest: np.ndarray = np.argpartition(distances, required)[:required] + # indices of the required closest parameters (np.argpartition requires + # kth < len, so select all remaining when we need at least all of them) + n_left = len(distances) + if required >= n_left: + ixs_nearest = np.arange(n_left) + else: + ixs_nearest = np.argpartition(distances, required)[:required] ixs_not_in_cluster: np.ndarray = np.flatnonzero(~in_cluster) in_cluster[ixs_not_in_cluster[ixs_nearest]] = True - if sum(in_cluster) != desired: + if sum(in_cluster) != min(desired, len(y)): raise AssertionError('Unexpected number of entries.') return in_cluster diff --git a/pyabc/transition/base.py b/pyabc/transition/base.py index 818edf7c..f698b4ec 100644 --- a/pyabc/transition/base.py +++ b/pyabc/transition/base.py @@ -63,7 +63,7 @@ def rvs_single(self) -> Parameter: A sample from the fitted model. """ - def rvs(self, size: int = None) -> Parameter | pd.DataFrame: + def rvs(self, size: int | None = None) -> Parameter | pd.DataFrame: """ Sample from the density. diff --git a/pyabc/transition/jump.py b/pyabc/transition/jump.py index ab9f0d2c..19d679ed 100644 --- a/pyabc/transition/jump.py +++ b/pyabc/transition/jump.py @@ -30,7 +30,12 @@ def __init__(self, domain: np.ndarray, p_stay: float = 0.7): if not 0 <= p_stay <= 1: raise ValueError('p_stay must be in [0, 1].') self.p_stay = p_stay - self.p_move = (1 - p_stay) / (len(self.domain) - 1) + # guard against a single-value domain (no other value to move to) + self.p_move = ( + 0.0 + if len(self.domain) == 1 + else (1 - p_stay) / (len(self.domain) - 1) + ) # cache a random variable (later the start index and 0 must be swapped) indices = np.arange(len(domain)) diff --git a/pyabc/transition/multivariatenormal.py b/pyabc/transition/multivariatenormal.py index 1f830c74..fa3601da 100644 --- a/pyabc/transition/multivariatenormal.py +++ b/pyabc/transition/multivariatenormal.py @@ -91,7 +91,7 @@ def fit(self, X: pd.DataFrame, w: np.ndarray) -> None: # cache range array self._range = np.arange(len(self.X)) - def rvs(self, size: int = None) -> Parameter | pd.DataFrame: + def rvs(self, size: int | None = None) -> Parameter | pd.DataFrame: if size is None: return self.rvs_single() sample_ind = np.random.choice( diff --git a/pyabc/transition/transitionmeta.py b/pyabc/transition/transitionmeta.py index 69ca4a0b..7dfe5b52 100644 --- a/pyabc/transition/transitionmeta.py +++ b/pyabc/transition/transitionmeta.py @@ -9,13 +9,15 @@ def wrap_fit(f): @functools.wraps(f) def fit(self, X: pd.DataFrame, w: np.ndarray): self.X = X - self.w = w if len(X.columns) == 0: + self.w = w self.no_parameters = True return self.no_parameters = False if w.size > 0 and not np.isclose(w.sum(), 1): - w /= w.sum() + # normalize out-of-place so the caller's array is not mutated + w = w / w.sum() + self.w = w f(self, X, w) return fit @@ -33,7 +35,7 @@ def pdf(self, x: pd.Series | pd.DataFrame): def wrap_rvs(f): @functools.wraps(f) - def rvs(self, size: int = None): + def rvs(self, size: int | None = None): if self.no_parameters: return pd.DataFrame(dtype=float) return f(self, size) diff --git a/pyabc/util/dict2arr.py b/pyabc/util/dict2arr.py index 4d0602f7..1bc5205e 100644 --- a/pyabc/util/dict2arr.py +++ b/pyabc/util/dict2arr.py @@ -41,8 +41,7 @@ def dict2arr(dct: dict | np.ndarray, keys: list) -> np.ndarray: if len(arr) == 1: return np.asarray(arr[0]) # flatten - arr = [val for sub_arr in arr for val in sub_arr] - return np.asarray(arr) + return np.concatenate([np.asarray(sub_arr) for sub_arr in arr]) def dict2arrlabels(dct: dict, keys: list) -> list[str]: diff --git a/pyabc/util/event_ixs.py b/pyabc/util/event_ixs.py index e71fcc36..a1c594dc 100644 --- a/pyabc/util/event_ixs.py +++ b/pyabc/util/event_ixs.py @@ -13,10 +13,10 @@ class EventIxs: def __init__( self, - ts: Collection[int] | int = None, - sims: Collection[int] | int = None, - from_t: int = None, - from_sims: int = None, + ts: Collection[int] | int | None = None, + sims: Collection[int] | int | None = None, + from_t: int | None = None, + from_sims: int | None = None, ): """ Parameters diff --git a/pyabc/util/par_trafo.py b/pyabc/util/par_trafo.py index f5984892..daa942dc 100644 --- a/pyabc/util/par_trafo.py +++ b/pyabc/util/par_trafo.py @@ -46,7 +46,7 @@ class ParTrafo(ParTrafoBase): def __init__( self, - trafos: list[Callable[[np.ndarray], np.ndarray]] = None, + trafos: list[Callable[[np.ndarray], np.ndarray]] | None = None, trafo_ids: str | list[str] = '{par_id}_{trafo_ix}', ): self.trafos = trafos diff --git a/pyabc/version.py b/pyabc/version.py index 8327b506..2d7893e3 100644 --- a/pyabc/version.py +++ b/pyabc/version.py @@ -1 +1 @@ -__version__ = '0.12.18' +__version__ = '0.13.0' diff --git a/pyabc/visualization/contour.py b/pyabc/visualization/contour.py index 5c5e5f35..ff35c790 100644 --- a/pyabc/visualization/contour.py +++ b/pyabc/visualization/contour.py @@ -22,24 +22,24 @@ def plot_contour_2d( x: str, y: str, m: int = 0, - t: int = None, - xmin: float = None, - xmax: float = None, - ymin: float = None, - ymax: float = None, + t: int | None = None, + xmin: float | None = None, + xmax: float | None = None, + ymin: float | None = None, + ymax: float | None = None, numx: int = 50, numy: int = 50, ax=None, size=None, - title: str = None, + title: str | None = None, refval=None, refval_color='C1', kde=None, - xname: str = None, - yname: str = None, + xname: str | None = None, + yname: str | None = None, show_clabel: bool = False, show_legend: bool = False, - clabel_kwargs: dict = None, + clabel_kwargs: dict | None = None, **kwargs, ): """ @@ -151,15 +151,15 @@ def plot_contour_2d_lowlevel( numy=50, ax=None, size=None, - title: str = None, + title: str | None = None, refval=None, refval_color='C1', kde=None, - xname: str = None, - yname: str = None, + xname: str | None = None, + yname: str | None = None, show_clabel: bool = False, show_legend: bool = False, - clabel_kwargs: dict = None, + clabel_kwargs: dict | None = None, **kwargs, ): """ @@ -235,18 +235,18 @@ def plot_contour_2d_lowlevel( def plot_contour_matrix( history, m: int = 0, - t: int = None, - limits: dict = None, + t: int | None = None, + limits: dict | None = None, height: float = 2.5, numx: int = 50, numy: int = 50, - refval: dict = None, + refval: dict | None = None, refval_color='C1', kde=None, - names: dict = None, + names: dict | None = None, show_clabel: bool = False, show_legend: bool = False, - clabel_kwargs: dict = None, + clabel_kwargs: dict | None = None, arr_ax=None, **kwargs, ): @@ -322,17 +322,17 @@ def plot_contour_matrix( def plot_contour_matrix_lowlevel( df: pd.DataFrame, w: np.ndarray, - limits: dict = None, + limits: dict | None = None, height: int = 2.5, numx: int = 50, numy: int = 50, - refval: dict = None, + refval: dict | None = None, refval_color='C1', kde=None, - names: dict = None, + names: dict | None = None, show_clabel: bool = False, show_legend: bool = False, - clabel_kwargs: dict = None, + clabel_kwargs: dict | None = None, arr_ax=None, **kwargs, ): diff --git a/pyabc/visualization/credible.py b/pyabc/visualization/credible.py index 1e76b851..b570685c 100644 --- a/pyabc/visualization/credible.py +++ b/pyabc/visualization/credible.py @@ -14,14 +14,14 @@ def _prepare_credible_intervals( history: History, m: int, - ts: list[int] | int, - par_names: list, - levels: list, + ts: list[int] | int | None, + par_names: list | None, + levels: list | None, show_mean: bool, show_kde_max: bool, show_kde_max_1d: bool, - kde: Transition, - kde_1d: Transition, + kde: Transition | None, + kde_1d: Transition | None, ): if levels is None: levels = [0.95] @@ -98,23 +98,23 @@ def _prepare_credible_intervals( def plot_credible_intervals( history: History, m: int = 0, - ts: list[int] | int = None, - par_names: list = None, - levels: list = None, - colors: list = None, - color_median: str = None, + ts: list[int] | int | None = None, + par_names: list | None = None, + levels: list | None = None, + colors: list | None = None, + color_median: str | None = None, show_mean: bool = False, - color_mean: str = None, + color_mean: str | None = None, show_kde_max: bool = False, - color_kde_max: str = None, + color_kde_max: str | None = None, show_kde_max_1d: bool = False, - color_kde_max_1d: str = None, - size: tuple = None, - refval: dict = None, + color_kde_max_1d: str | None = None, + size: tuple | None = None, + refval: dict | None = None, refval_color: str = 'C1', - kde: Transition = None, - kde_1d: Transition = None, - arr_ax: list[matplotlib.axes.Axes] = None, + kde: Transition | None = None, + kde_1d: Transition | None = None, + arr_ax: list[matplotlib.axes.Axes] | None = None, ): """Plot credible intervals over time. @@ -274,15 +274,15 @@ def plot_credible_intervals( def plot_credible_intervals_plotly( history: History, m: int = 0, - ts: list[int] | int = None, - par_names: list = None, - levels: list = None, + ts: list[int] | int | None = None, + par_names: list | None = None, + levels: list | None = None, colors=None, - size: tuple = None, - refval: dict = None, + size: tuple | None = None, + refval: dict | None = None, refval_color: str = 'gray', - kde: Transition = None, - kde_1d: Transition = None, + kde: Transition | None = None, + kde_1d: Transition | None = None, ): """Plot credible intervals over time using plotly.""" import plotly.graph_objects as go @@ -337,8 +337,8 @@ def plot_credible_intervals_plotly( error_y={ 'type': 'data', 'symmetric': False, - 'array': cis[i_par, :, i_c] - median[i_par], - 'arrayminus': median[i_par] - cis[i_par, :, -1 - i_c], + 'array': cis[i_par, :, -1 - i_c] - median[i_par], + 'arrayminus': median[i_par] - cis[i_par, :, i_c], }, mode='lines+markers', marker={'color': colors[i_c]}, @@ -375,19 +375,19 @@ def plot_credible_intervals_plotly( def plot_credible_intervals_for_time( histories: list[History] | History, - labels: list[str] | str = None, - ms: list[int] | int = None, - ts: list[int] | int = None, - par_names: list[str] = None, - levels: list[float] = None, + labels: list[str] | str | None = None, + ms: list[int] | int | None = None, + ts: list[int] | int | None = None, + par_names: list[str] | None = None, + levels: list[float] | None = None, show_mean: bool = False, show_kde_max: bool = False, show_kde_max_1d: bool = False, - size: tuple = None, + size: tuple | None = None, rotation: int = 0, - refvals: list[dict] | dict = None, - kde: Transition = None, - kde_1d: Transition = None, + refvals: list[dict] | dict | None = None, + kde: Transition | None = None, + kde_1d: Transition | None = None, ): """ Plot credible intervals over time. @@ -427,7 +427,8 @@ def plot_credible_intervals_for_time( if ms is None: ms = [0] * n_run elif not isinstance(ms, list) or len(ms) == 1: - ms = [ms] * n_run + # broadcast a single model id (int, or length-1 list) across runs + ms = [ms[0] if isinstance(ms, list) else ms] * n_run if levels is None: levels = [0.95] levels = sorted(levels) @@ -512,7 +513,7 @@ def plot_credible_intervals_for_time( color=f'C{i_c}', ) # reference value - if refvals[i_run] is not None: + if refvals is not None and refvals[i_run] is not None: ax.plot([i_run], [refvals[i_run][par]], 'x', color='black') ax.set_title(f'Parameter {par}') # mean diff --git a/pyabc/visualization/data.py b/pyabc/visualization/data.py index cc60e070..2a94c0af 100644 --- a/pyabc/visualization/data.py +++ b/pyabc/visualization/data.py @@ -15,11 +15,11 @@ def plot_data_callback( history: History, - f_plot: Callable = None, - f_plot_aggregated: Callable = None, - t: int = None, - n_sample: int = None, - ax: matplotlib.axes.Axes = None, + f_plot: Callable | None = None, + f_plot_aggregated: Callable | None = None, + t: int | None = None, + n_sample: int | None = None, + ax: matplotlib.axes.Axes | None = None, **kwargs, ): """ @@ -66,8 +66,8 @@ def plot_data_callback_lowlevel( sum_stats: list, weights: list, f_plot: Callable, - f_plot_aggregated: Callable = None, - n_sample: int = None, + f_plot_aggregated: Callable | None = None, + n_sample: int | None = None, ax=None, **kwargs, ): @@ -100,7 +100,7 @@ def plot_data_callback_lowlevel( def plot_data_default( - obs_data: dict, sim_data: dict, keys: list[str] | str = None + obs_data: dict, sim_data: dict, keys: list[str] | str | None = None ): """ Plot summary statistic data. diff --git a/pyabc/visualization/distance.py b/pyabc/visualization/distance.py index c2023763..5845aa1b 100644 --- a/pyabc/visualization/distance.py +++ b/pyabc/visualization/distance.py @@ -15,18 +15,18 @@ def plot_distance_weights( log_files: list[str] | str, ts: list[int] | list[str] | int | str = 'last', - labels: list[str] | str = None, - colors: list[Any] | Any = None, - linestyles: list[str] | str = None, + labels: list[str] | str | None = None, + colors: list[Any] | Any | None = None, + linestyles: list[str] | str | None = None, keys_as_labels: bool = True, - keys: list[str] = None, + keys: list[str] | None = None, xticklabel_rotation: float = 0, normalize: bool = True, - size: tuple[float, float] = None, + size: tuple[float, float] | None = None, xlabel: str = 'Summary statistic', ylabel: str = 'Weight', - title: str = None, - ax: mpl.axes.Axes = None, + title: str | None = None, + ax: mpl.axes.Axes | None = None, **kwargs, ) -> mpl.axes.Axes: """Plot distance weights, one curve per argument. diff --git a/pyabc/visualization/effective_sample_size.py b/pyabc/visualization/effective_sample_size.py index 6a535915..66dee092 100644 --- a/pyabc/visualization/effective_sample_size.py +++ b/pyabc/visualization/effective_sample_size.py @@ -44,13 +44,13 @@ def _prepare_plot_effective_sample_sizes( def plot_effective_sample_sizes( histories: list | History, - labels: list | str = None, + labels: list | str | None = None, rotation: int = 0, title: str = 'Effective sample size', relative: bool = False, - colors: list = None, - size: tuple = None, - ax: mpl.axes.Axes = None, + colors: list | None = None, + size: tuple | None = None, + ax: mpl.axes.Axes | None = None, ) -> mpl.axes.Axes: """ Plot effective sample sizes over all iterations. @@ -121,13 +121,13 @@ def plot_effective_sample_sizes( def plot_effective_sample_sizes_plotly( histories: list | History, - labels: list | str = None, + labels: list | str | None = None, rotation: int = 0, title: str = 'Effective sample size', relative: bool = False, - colors: list = None, - size: tuple = None, - fig: 'go.Figure' = None, + colors: list | None = None, + size: tuple | None = None, + fig: 'go.Figure | None' = None, ) -> 'go.Figure': """Plot effective sample sizes using plotly.""" import plotly.graph_objects as go diff --git a/pyabc/visualization/epsilon.py b/pyabc/visualization/epsilon.py index 4dc67d27..a5a8d123 100644 --- a/pyabc/visualization/epsilon.py +++ b/pyabc/visualization/epsilon.py @@ -37,12 +37,12 @@ def _prepare( def plot_epsilons( histories: list | History, - labels: list | str = None, - colors: list = None, + labels: list | str | None = None, + colors: list | None = None, yscale: str = 'log', title: str = 'Epsilon values', - size: tuple = None, - ax: mpl.axes.Axes = None, + size: tuple | None = None, + ax: mpl.axes.Axes | None = None, ) -> mpl.axes.Axes: """ Plot epsilon trajectory. @@ -103,12 +103,12 @@ def plot_epsilons( def plot_epsilons_plotly( histories: list | History, - labels: list | str = None, - colors: list = None, + labels: list | str | None = None, + colors: list | None = None, yscale: str = 'log', title: str = 'Epsilon values', - size: tuple = None, - fig: 'go.Figure' = None, + size: tuple | None = None, + fig: 'go.Figure | None' = None, ) -> 'go.Figure': """Plot epsilon trajectory using plotly.""" import plotly.graph_objects as go diff --git a/pyabc/visualization/histogram.py b/pyabc/visualization/histogram.py index 2183d615..7fa06b60 100644 --- a/pyabc/visualization/histogram.py +++ b/pyabc/visualization/histogram.py @@ -11,14 +11,14 @@ def plot_histogram_1d( history: History, x: str, m: int = 0, - t: int = None, + t: int | None = None, xmin=None, xmax=None, ax=None, size=None, refval=None, refval_color='C1', - xname: str = None, + xname: str | None = None, **kwargs, ): """ @@ -80,7 +80,7 @@ def plot_histogram_1d_lowlevel( size=None, refval=None, refval_color='C1', - xname: str = None, + xname: str | None = None, **kwargs, ): """ @@ -121,7 +121,7 @@ def plot_histogram_2d( x: str, y: str, m: int = 0, - t: int = None, + t: int | None = None, xmin=None, xmax=None, ymin=None, @@ -130,8 +130,8 @@ def plot_histogram_2d( size=None, refval=None, refval_color='C1', - xname: str = None, - yname: str = None, + xname: str | None = None, + yname: str | None = None, **kwargs, ): """ @@ -202,8 +202,8 @@ def plot_histogram_2d_lowlevel( size=None, refval=None, refval_color='C1', - xname: str = None, - yname: str = None, + xname: str | None = None, + yname: str | None = None, **kwargs, ): """ @@ -251,11 +251,11 @@ def plot_histogram_2d_lowlevel( def plot_histogram_matrix( history: History, m: int = 0, - t: int = None, + t: int | None = None, size=None, refval=None, refval_color='C1', - names: dict = None, + names: dict | None = None, **kwargs, ): """ @@ -298,7 +298,7 @@ def plot_histogram_matrix_lowlevel( size=None, refval=None, refval_color='C1', - names: dict = None, + names: dict | None = None, **kwargs, ): """ diff --git a/pyabc/visualization/kde.py b/pyabc/visualization/kde.py index 82f997b9..c0ea0b0b 100644 --- a/pyabc/visualization/kde.py +++ b/pyabc/visualization/kde.py @@ -85,17 +85,17 @@ def plot_kde_1d_highlevel( history: History, x: str, m: int = 0, - t: int = None, + t: int | None = None, xmin=None, xmax=None, numx=50, - ax: mpl.axes.Axes = None, + ax: mpl.axes.Axes | None = None, size=None, - title: str = None, + title: str | None = None, refval=None, refval_color='C1', kde=None, - xname: str = None, + xname: str | None = None, **kwargs, ) -> mpl.axes.Axes: """ @@ -170,20 +170,20 @@ def plot_kde_1d_highlevel_plotly( history: History, x: str, m: int = 0, - t: int = None, + t: int | None = None, xmin=None, xmax=None, numx: int = 50, - fig: 'go.Figure' = None, + fig: 'go.Figure | None' = None, row: int = 1, col: int = 1, size=None, - title: str = None, + title: str | None = None, refval=None, refval_color='gray', marker_color=None, kde=None, - xname: str = None, + xname: str | None = None, **kwargs, ): df, w = history.get_distribution(m=m, t=t) @@ -216,13 +216,13 @@ def plot_kde_1d( xmin=None, xmax=None, numx=50, - ax: mpl.axes.Axes = None, + ax: mpl.axes.Axes | None = None, size=None, - title: str = None, + title: str | None = None, refval=None, refval_color='C1', kde=None, - xname: str = None, + xname: str | None = None, **kwargs, ) -> mpl.axes.Axes: """ @@ -246,9 +246,11 @@ def plot_kde_1d( xname = x if ax is None: _, ax = plt.subplots() - ax.plot(x_vals, pdf, **kwargs) - # TODO This fixes the upper bound inadequately - # ax.set_ylim(bottom=min(ax.get_ylim()[0], 0)) + (line,) = ax.plot(x_vals, pdf, **kwargs) + # a density is non-negative, but `set_ylim` would switch off autoscaling + line.sticky_edges.y.append(0.0) + ax.update_datalim([(x_vals[0], 0.0)]) + ax.autoscale_view() ax.set_xlabel(xname) ax.set_ylabel('Posterior') ax.set_xlim(xmin, xmax) @@ -271,16 +273,16 @@ def plot_kde_1d_plotly( xmin=None, xmax=None, numx=50, - fig: 'go.Figure' = None, + fig: 'go.Figure | None' = None, row: int = 1, col: int = 1, size=None, - title: str = None, + title: str | None = None, refval=None, refval_color='gray', marker_color=None, kde=None, - xname: str = None, + xname: str | None = None, **kwargs, ) -> 'go.Figure': """Plot 1d kde using plotly.""" @@ -307,14 +309,11 @@ def plot_kde_1d_plotly( row=row, col=col, ) - # set trace color to blue - # fig.update_traces(marker_color="blue", row=row, col=col) - # fig.add_trace( - # go.Scatter(x=x_vals, y=pdf, name=xname, **kwargs), - # row=row, - # col=col, - # ) fig.update_xaxes(title_text=xname, range=[xmin, xmax], row=row, col=col) + # a density is non-negative, plotly's default rangemode ignores zero + fig.update_yaxes( + title_text='Posterior', rangemode='tozero', row=row, col=col + ) # add vertical line for reference value if refval is not None: @@ -427,22 +426,22 @@ def plot_kde_2d_highlevel( x: str, y: str, m: int = 0, - t: int = None, - xmin: float = None, - xmax: float = None, - ymin: float = None, - ymax: float = None, + t: int | None = None, + xmin: float | None = None, + xmax: float | None = None, + ymin: float | None = None, + ymax: float | None = None, numx: int = 50, numy: int = 50, - ax: mpl.axes.Axes = None, + ax: mpl.axes.Axes | None = None, size=None, colorbar=True, - title: str = None, + title: str | None = None, refval=None, refval_color='C1', kde=None, - xname: str = None, - yname: str = None, + xname: str | None = None, + yname: str | None = None, **kwargs, ) -> mpl.axes.Axes: """ @@ -539,25 +538,25 @@ def plot_kde_2d_highlevel_plotly( x: str, y: str, m: int = 0, - t: int = None, - xmin: float = None, - xmax: float = None, - ymin: float = None, - ymax: float = None, + t: int | None = None, + xmin: float | None = None, + xmax: float | None = None, + ymin: float | None = None, + ymax: float | None = None, numx: int = 50, numy: int = 50, - fig: 'go.Figure' = None, + fig: 'go.Figure | None' = None, row: int = 1, col: int = 1, size=None, showscale=True, showlegend=True, - title: str = None, + title: str | None = None, refval=None, refval_color='gray', kde=None, - xname: str = None, - yname: str = None, + xname: str | None = None, + yname: str | None = None, **kwargs, ) -> 'go.Figure': """ @@ -603,15 +602,15 @@ def plot_kde_2d( ymax=None, numx=50, numy=50, - ax: mpl.axes.Axes = None, + ax: mpl.axes.Axes | None = None, size=None, colorbar=True, - title: str = None, + title: str | None = None, refval=None, refval_color='C1', kde=None, - xname: str = None, - yname: str = None, + xname: str | None = None, + yname: str | None = None, **kwargs, ) -> mpl.axes.Axes: """ @@ -679,18 +678,18 @@ def plot_kde_2d_plotly( ymax=None, numx=50, numy=50, - fig: 'go.Figure' = None, + fig: 'go.Figure | None' = None, row: int = 1, col: int = 1, size=None, showscale=True, showlegend=True, - title: str = None, + title: str | None = None, refval=None, refval_color='gray', kde=None, - xname: str = None, - yname: str = None, + xname: str | None = None, + yname: str | None = None, **kwargs, ): """ @@ -760,7 +759,7 @@ def plot_kde_2d_plotly( def plot_kde_matrix_highlevel( history, m: int = 0, - t: int = None, + t: int | None = None, limits=None, colorbar: bool = True, height: float = 2.5, @@ -769,7 +768,7 @@ def plot_kde_matrix_highlevel( refval=None, refval_color='C1', kde=None, - names: dict = None, + names: dict | None = None, arr_ax=None, ): """ @@ -836,7 +835,7 @@ def plot_kde_matrix_highlevel( def plot_kde_matrix_highlevel_plotly( history, m: int = 0, - t: int = None, + t: int | None = None, limits=None, height: int = 30, numx: int = 50, @@ -844,7 +843,7 @@ def plot_kde_matrix_highlevel_plotly( refval=None, refval_color='gray', kde=None, - names: dict = None, + names: dict | None = None, title: str = 'Univariate and bivariate distributions using KDE', ) -> 'go.Figure': """ @@ -879,7 +878,7 @@ def plot_kde_matrix( refval=None, refval_color='C1', kde=None, - names: dict = None, + names: dict | None = None, arr_ax=None, ): """ @@ -1009,7 +1008,7 @@ def plot_kde_matrix_plotly( refval_color='gray', marker_color=None, kde=None, - names: dict = None, + names: dict | None = None, title: str = 'Univariate and bivariate distributions using KDE', ) -> 'go.Figure': """ diff --git a/pyabc/visualization/model_probabilities.py b/pyabc/visualization/model_probabilities.py index c14410e3..3d5ce8bc 100644 --- a/pyabc/visualization/model_probabilities.py +++ b/pyabc/visualization/model_probabilities.py @@ -15,8 +15,8 @@ def plot_model_probabilities( history: History, rotation: int = 0, title: str = 'Model probabilities', - size: tuple = None, - ax: mpl.axes.Axes = None, + size: tuple | None = None, + ax: mpl.axes.Axes | None = None, ) -> mpl.axes.Axes: """ Plot the probabilities of models over time. @@ -68,8 +68,8 @@ def plot_model_probabilities_plotly( history: History, rotation: int = 0, title: str = 'Model probabilities', - size: tuple = None, - fig: 'go.Figure' = None, + size: tuple | None = None, + fig: 'go.Figure | None' = None, ) -> 'go.Figure': """Plot model probabilities using plotly.""" import plotly.graph_objects as go diff --git a/pyabc/visualization/sample.py b/pyabc/visualization/sample.py index 36fcec77..e43a1a99 100644 --- a/pyabc/visualization/sample.py +++ b/pyabc/visualization/sample.py @@ -44,11 +44,11 @@ def _prepare_plot_sample_numbers( def plot_sample_numbers( histories: list[History] | History, - labels: list[str] | str = None, + labels: list[str] | str | None = None, rotation: int = 0, title: str = 'Required samples', - size: tuple[float, float] = None, - ax: mpl.axes.Axes = None, + size: tuple[float, float] | None = None, + ax: mpl.axes.Axes | None = None, ) -> mpl.axes.Axes: """ Stacked bar plot of required numbers of samples over all iterations. @@ -112,11 +112,11 @@ def plot_sample_numbers( def plot_sample_numbers_plotly( histories: list[History] | History, - labels: list[str] | str = None, + labels: list[str] | str | None = None, rotation: int = 0, title: str = 'Required samples', - size: tuple[float, float] = None, - fig: 'go.Figure' = None, + size: tuple[float, float] | None = None, + fig: 'go.Figure | None' = None, ) -> 'go.Figure': """Plot sample numbers using plotly.""" import plotly.graph_objects as go @@ -199,12 +199,12 @@ def _prepare_plot_total_sample_numbers( def plot_total_sample_numbers( histories: list | History, - labels: list | str = None, + labels: list | str | None = None, rotation: int = 0, title: str = 'Total required samples', yscale: str = 'lin', - size: tuple = None, - ax: mpl.axes.Axes = None, + size: tuple | None = None, + ax: mpl.axes.Axes | None = None, ) -> mpl.axes.Axes: """ Bar plot of total required sample number over all iterations, i.e. @@ -266,12 +266,12 @@ def plot_total_sample_numbers( def plot_total_sample_numbers_plotly( histories: list | History, - labels: list | str = None, + labels: list | str | None = None, rotation: int = 0, title: str = 'Total required samples', yscale: str = 'lin', - size: tuple = None, - fig: 'go.Figure' = None, + size: tuple | None = None, + fig: 'go.Figure | None' = None, ) -> 'go.Figure': """Plot total sample numbers using plotly.""" import plotly.graph_objects as go @@ -349,11 +349,11 @@ def _prepare_plot_sample_numbers_trajectory( def plot_sample_numbers_trajectory( histories: list | History, - labels: list | str = None, + labels: list | str | None = None, title: str = 'Required samples', yscale: str = 'lin', - size: tuple = None, - ax: mpl.axes.Axes = None, + size: tuple | None = None, + ax: mpl.axes.Axes | None = None, ) -> mpl.axes.Axes: """ Plot of required sample number over all iterations, i.e. one trajectory @@ -412,11 +412,11 @@ def plot_sample_numbers_trajectory( def plot_sample_numbers_trajectory_plotly( histories: list | History, - labels: list | str = None, + labels: list | str | None = None, title: str = 'Required samples', yscale: str = 'lin', - size: tuple = None, - fig: 'go.Figure' = None, + size: tuple | None = None, + fig: 'go.Figure | None' = None, ) -> 'go.Figure': """Plot sample number trajectories using plotly.""" import plotly.graph_objects as go @@ -506,13 +506,13 @@ def _prepare_plot_acceptance_rates_trajectory( def plot_acceptance_rates_trajectory( histories: list | History, - labels: list | str = None, + labels: list | str | None = None, title: str = 'Acceptance rates', yscale: str = 'lin', - size: tuple = None, - colors: list[str] = None, + size: tuple | None = None, + colors: list[str] | None = None, normalize_by_ess: bool = False, - ax: mpl.axes.Axes = None, + ax: mpl.axes.Axes | None = None, ) -> mpl.axes.Axes: """ Plot of acceptance rates over all iterations, i.e. one trajectory @@ -583,13 +583,13 @@ def plot_acceptance_rates_trajectory( def plot_acceptance_rates_trajectory_plotly( histories: list | History, - labels: list | str = None, + labels: list | str | None = None, title: str = 'Acceptance rates', yscale: str = 'lin', - size: tuple = None, - colors: list[str] = None, + size: tuple | None = None, + colors: list[str] | None = None, normalize_by_ess: bool = False, - fig: 'go.Figure' = None, + fig: 'go.Figure | None' = None, ) -> 'go.Figure': """Plot acceptance rates trajectories using plotly.""" import plotly.graph_objects as go @@ -638,11 +638,11 @@ def plot_lookahead_evaluations( sampler_df: pd.DataFrame | str, relative: bool = False, fill: bool = False, - alpha: float = None, + alpha: float | None = None, t_min: int = 0, title: str = 'Total evaluations', - size: tuple = None, - ax: mpl.axes.Axes = None, + size: tuple | None = None, + ax: mpl.axes.Axes | None = None, ): """Plot total vs look-ahead evaluations over the generations. @@ -735,11 +735,11 @@ def plot_lookahead_final_acceptance_fractions( population_sizes: np.ndarray | History, relative: bool = False, fill: bool = False, - alpha: float = None, + alpha: float | None = None, t_min: int = 0, title: str = 'Composition of final acceptances', - size: tuple = None, - ax: mpl.axes.Axes = None, + size: tuple | None = None, + ax: mpl.axes.Axes | None = None, ): """Plot fraction of look-ahead samples in final acceptances, over generations. @@ -852,8 +852,8 @@ def plot_lookahead_acceptance_rates( sampler_df: pd.DataFrame | str, t_min: int = 0, title: str = 'Acceptance rates', - size: tuple = None, - ax: mpl.axes.Axes = None, + size: tuple | None = None, + ax: mpl.axes.Axes | None = None, ): """Plot acceptance rates for look-ahead vs ordinary samples. The ratios are relative to all accepted particles, including eventually diff --git a/pyabc/visualization/sankey.py b/pyabc/visualization/sankey.py index fe971992..285789a6 100644 --- a/pyabc/visualization/sankey.py +++ b/pyabc/visualization/sankey.py @@ -23,20 +23,20 @@ def plot_sensitivity_sankey( t: int | str, h: pyabc.storage.History, predictor: pyabc.predictor.Predictor, - par_trafo: pyabc.util.ParTrafoBase = None, - sumstat: pyabc.sumstat.Sumstat = None, - subsetter: pyabc.sumstat.Subsetter = None, + par_trafo: pyabc.util.ParTrafoBase | None = None, + sumstat: pyabc.sumstat.Sumstat | None = None, + subsetter: pyabc.sumstat.Subsetter | None = None, feature_normalization: str = pyabc.distance.InfoWeightedPNormDistance.MAD, normalize_by_par: bool = True, - fd_deltas: list[float] | float = None, - scale_weights: dict[int, np.ndarray] = None, + fd_deltas: list[float] | float | None = None, + scale_weights: dict[int, np.ndarray] | None = None, title: str = 'Data-parameter sensitivities', - width: float = None, - height: float = None, - sumstat_color: Callable[[str], str] = None, - par_color: Callable[[str], str] = None, - node_kwargs: dict = None, - layout_kwargs: dict = None, + width: float | None = None, + height: float | None = None, + sumstat_color: Callable[[str], str] | None = None, + par_color: Callable[[str], str] | None = None, + node_kwargs: dict | None = None, + layout_kwargs: dict | None = None, ): """Plot sensitivity matrix as a Sankey flow plot. diff --git a/pyabc/visualization/walltime.py b/pyabc/visualization/walltime.py index 52b9c5c2..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 @@ -55,12 +63,12 @@ def _prepare_plot_total_walltime( def plot_total_walltime( histories: list[History] | History, - labels: list | str = None, + labels: list | str | None = None, unit: str = 's', rotation: int = 0, title: str = 'Total walltimes', - size: tuple = None, - ax: mpl.axes.Axes = None, + size: tuple | None = None, + ax: mpl.axes.Axes | None = None, ) -> mpl.axes.Axes: """Plot total walltimes, for each history one single-color bar. @@ -114,12 +122,12 @@ def plot_total_walltime( def plot_total_walltime_plotly( histories: list[History] | History, - labels: list | str = None, + labels: list | str | None = None, unit: str = 's', rotation: int = 0, title: str = 'Total walltimes', - size: tuple = None, - fig: 'go.Figure' = None, + size: tuple | None = None, + fig: 'go.Figure | None' = None, ) -> 'go.Figure': """Plot total walltimes using plotly.""" import plotly.graph_objects as go @@ -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,27 +175,30 @@ 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( histories: list[History] | History, - labels: list | str = None, - show_calibration: bool = None, + labels: list | str | None = None, + show_calibration: bool | None = None, unit: str = 's', rotation: int = 0, title: str = 'Walltime by generation', - size: tuple = None, - ax: mpl.axes.Axes = None, + size: tuple | None = None, + ax: mpl.axes.Axes | None = None, ) -> mpl.axes.Axes: """Plot walltimes, with different colors indicating different iterations. @@ -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,22 +244,23 @@ def plot_walltime( title=title, size=size, ax=ax, + wall_times=wall_times, ) def plot_walltime_plotly( histories: list[History] | History, - labels: list | str = None, - show_calibration: bool = None, + labels: list | str | None = None, + show_calibration: bool | None = None, unit: str = 's', rotation: int = 0, title: str = 'Walltime by generation', - size: tuple = None, - fig: 'go.Figure' = None, + size: tuple | None = None, + fig: 'go.Figure | None' = None, ) -> '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,15 +274,17 @@ def plot_walltime_plotly( title=title, size=size, fig=fig, + wall_times=wall_times, ) def _prepare_plot_walltime_lowlevel( end_times: list, start_times: list | None = None, - labels: list | str = None, - show_calibration: bool = None, + 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) @@ -325,13 +353,14 @@ def _prepare_plot_walltime_lowlevel( def plot_walltime_lowlevel( end_times: list, start_times: list | None = None, - labels: list | str = None, - show_calibration: bool = None, + labels: list | str | None = None, + show_calibration: bool | None = None, unit: str = 's', rotation: int = 0, title: str = 'Walltime by generation', - size: tuple = None, - ax: mpl.axes.Axes = None, + 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 @@ -381,13 +411,14 @@ def plot_walltime_lowlevel( def plot_walltime_lowlevel_plotly( end_times: list, start_times: list | None = None, - labels: list | str = None, - show_calibration: bool = None, + labels: list | str | None = None, + show_calibration: bool | None = None, unit: str = 's', rotation: int = 0, title: str = 'Walltime by generation', - size: tuple = None, - fig: 'go.Figure' = None, + 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 @@ -456,16 +488,16 @@ def _prepare_plot_eps_walltime( def plot_eps_walltime( histories: list[History] | History, - labels: list | str = None, - colors: list[Any] = None, + labels: list | str | None = None, + colors: list[Any] | None = None, group_by_label: bool = True, indicate_end: bool = True, unit: str = 's', xscale: str = 'linear', yscale: str = 'log', title: str = 'Epsilon over walltime', - size: tuple = None, - ax: mpl.axes.Axes = None, + size: tuple | None = None, + ax: mpl.axes.Axes | None = None, ) -> mpl.axes.Axes: """Plot epsilon values (y-axis) over the walltime (x-axis), iterating over the generations. @@ -522,16 +554,16 @@ def plot_eps_walltime( def plot_eps_walltime_plotly( histories: list[History] | History, - labels: list | str = None, - colors: list[Any] = None, + labels: list | str | None = None, + colors: list[Any] | None = None, group_by_label: bool = True, indicate_end: bool = True, unit: str = 's', xscale: str = 'linear', yscale: str = 'log', title: str = 'Epsilon over walltime', - size: tuple = None, - fig: 'go.Figure' = None, + size: tuple | None = None, + fig: 'go.Figure | None' = None, ) -> 'go.Figure': """Plot epsilon values over walltime using plotly.""" # preprocess input @@ -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, ): @@ -604,16 +636,16 @@ def _prepare_plot_eps_walltime_lowlevel( def plot_eps_walltime_lowlevel( end_times: list, eps: list, - labels: list | str = None, - colors: list[Any] = None, + labels: list | str | None = None, + colors: list[Any] | None = None, group_by_label: bool = True, indicate_end: bool = True, unit: str = 's', xscale: str = 'linear', yscale: str = 'log', title: str = 'Epsilon over walltime', - size: tuple = None, - ax: mpl.axes.Axes = None, + size: tuple | None = None, + ax: mpl.axes.Axes | None = None, ) -> mpl.axes.Axes: """Low-level access to `plot_eps_walltime`. Directly define `end_times` and `eps`. Note that both should be arrays of @@ -679,16 +711,16 @@ def plot_eps_walltime_lowlevel( def plot_eps_walltime_lowlevel_plotly( end_times: list, eps: list, - labels: list | str = None, - colors: list[Any] = None, + labels: list | str | None = None, + colors: list[Any] | None = None, group_by_label: bool = True, indicate_end: bool = True, unit: str = 's', xscale: str = 'linear', yscale: str = 'log', title: str = 'Epsilon over walltime', - size: tuple = None, - fig: 'go.Figure' = None, + size: tuple | None = None, + fig: 'go.Figure | None' = None, ) -> 'go.Figure': """Plot epsilon values over walltime using plotly.""" import plotly.graph_objects as go diff --git a/pyabc/weighted_statistics/weighted_statistics.py b/pyabc/weighted_statistics/weighted_statistics.py index 4a6e1dc5..42b64d10 100644 --- a/pyabc/weighted_statistics/weighted_statistics.py +++ b/pyabc/weighted_statistics/weighted_statistics.py @@ -127,8 +127,8 @@ def resample(points, weights, n): A total of `n` points sampled from `points` with putting back according to `weights`. """ - weights = np.asarray(weights) - weights /= np.sum(weights) + weights = np.asarray(weights, dtype=float) + weights = weights / np.sum(weights) indices = np.random.choice( points.shape[0], size=n, p=weights ) # sample index from multi-dimensional sample diff --git a/pyproject.toml b/pyproject.toml index 7b274132..5e16f673 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,6 +36,7 @@ classifiers = [ "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", ] dependencies = [ @@ -149,6 +150,7 @@ select = [ "UP", # pyupgrade "ARG", # flake8-unused-arguments "SIM", # flake8-simplify + "RUF013", # implicit-optional ] ignore = [ "E501", # line too long (handled by formatter) diff --git a/test/base/test_epsilon.py b/test/base/test_epsilon.py index 4b4faacf..80c8adf6 100644 --- a/test/base/test_epsilon.py +++ b/test/base/test_epsilon.py @@ -253,3 +253,23 @@ def model(p): < 3 < pyabc.weighted_quantile(df.theta.to_numpy(), w, alpha=0.75) ) + + +def test_temperature_single_scheme_normalized(): + """Regression: a single callable ``schemes`` argument is normalized to a + list, so all consumers (which iterate over ``self.schemes``) work. A bare + callable previously broke iteration (e.g. in ``is_adaptive``).""" + from pyabc.epsilon.temperature import PolynomialDecayFixedIterScheme + + scheme = PolynomialDecayFixedIterScheme() + temp = pyabc.Temperature(schemes=scheme) + assert temp.schemes == [scheme] + # iterating over the schemes must not raise on a scalar callable + temp.is_adaptive() + + # a list of schemes is preserved unchanged + schemes = [ + PolynomialDecayFixedIterScheme(), + PolynomialDecayFixedIterScheme(), + ] + assert pyabc.Temperature(schemes=schemes).schemes == schemes diff --git a/test/base/test_populationstrategy.py b/test/base/test_populationstrategy.py index 0a9689ca..25058f88 100644 --- a/test/base/test_populationstrategy.py +++ b/test/base/test_populationstrategy.py @@ -11,7 +11,7 @@ from pyabc.transition import MultivariateNormalTransition -def Adaptive(nr_calibration_particles: int = None): +def Adaptive(nr_calibration_particles: int | None = None): # Only 4 bootstraps for faster testing ada = AdaptivePopulationSize( 100, @@ -22,13 +22,13 @@ def Adaptive(nr_calibration_particles: int = None): return ada -def Constant(nr_calibration_particles: int = None): +def Constant(nr_calibration_particles: int | None = None): return ConstantPopulationSize( 100, nr_calibration_particles=nr_calibration_particles ) -def List(nr_calibration_particles: int = None): +def List(nr_calibration_particles: int | None = None): return ListPopulationSize( [100] * 10, nr_calibration_particles=nr_calibration_particles ) diff --git a/test/base/test_predictor.py b/test/base/test_predictor.py index b5c13531..24423428 100644 --- a/test/base/test_predictor.py +++ b/test/base/test_predictor.py @@ -174,3 +174,24 @@ def test_wrong_input(): """Test all kinds of wrong inputs.""" with pytest.raises(ValueError): HiddenLayerHandle(method='potato')(n_in=10, n_out=10, n_sample=100) + + +def test_model_selection_custom_f_score(): + """Regression: a custom ``f_score`` passed to ``ModelSelectionPredictor`` + must be stored and used. Previously only the default branch set + ``self.f_score``, so a custom scorer was silently dropped and ``fit()`` + raised ``AttributeError``.""" + + def my_score(y1, y2, sigma): + return float(np.mean(np.abs(y1 - y2))) + + msp = ModelSelectionPredictor( + predictors=[LinearPredictor()], f_score=my_score + ) + assert msp.f_score is my_score + + rng = np.random.RandomState(0) + x = rng.normal(size=(40, 2)) + y = x @ np.array([[1.0], [2.0]]) + 0.01 * rng.normal(size=(40, 1)) + msp.fit(x, y) # would previously raise AttributeError + assert msp.chosen_one is not None diff --git a/test/base/test_storage.py b/test/base/test_storage.py index b0f9259f..04e98bf5 100644 --- a/test/base/test_storage.py +++ b/test/base/test_storage.py @@ -72,7 +72,9 @@ def history_uninitialized(): pass -def rand_pop_list(m: int = 0, normalized: bool = True, n_sample: int = None): +def rand_pop_list( + m: int = 0, normalized: bool = True, n_sample: int | None = None +): """ Create a population for model m, of random size >= 3. @@ -280,6 +282,86 @@ def test_sum_stats_save_load(history: History): assert (sum_stats[1]['ss33'] == example_df()).all().all() +def test_global_particle_weight_convention(history: History): + """Regression: particle weights are stored with the global convention + (summing to 1 across all particles of all models, matching the in-memory + ``Population``), while within-model read paths renormalize via the model + probability. Previously the DB stored per-model-normalized weights. + """ + # two models: p_model(0)=0.4, p_model(1)=0.6; global weights sum to 1 + particle_list = [ + Particle( + m=0, + parameter=Parameter({'a': 1.0}), + weight=0.3, + sum_stat={'ss': 1.0}, + distance=0.1, + ), + Particle( + m=0, + parameter=Parameter({'a': 2.0}), + weight=0.1, + sum_stat={'ss': 2.0}, + distance=0.2, + ), + Particle( + m=1, + parameter=Parameter({'a': 3.0}), + weight=0.2, + sum_stat={'ss': 3.0}, + distance=0.3, + ), + Particle( + m=1, + parameter=Parameter({'a': 4.0}), + weight=0.4, + sum_stat={'ss': 4.0}, + distance=0.4, + ), + ] + history.append_population( + 0, 0.5, Population(particle_list), 4, ['m0', 'm1'] + ) + + # model probabilities are the per-model sums of the global weights + mp = history.get_model_probabilities(t=0) + assert np.isclose(mp.loc[0, 'p'], 0.4) + assert np.isclose(mp.loc[1, 'p'], 0.6) + + # the RAW stored weights are global and sum to 1 (per-model storage would + # have summed to the number of models, i.e. 2) + ext = history.get_population_extended(t=0, tidy=False) + raw_w = ext[['particle_id', 'w']].drop_duplicates().w.values + assert np.isclose(raw_w.sum(), 1.0) + assert np.allclose(sorted(raw_w), [0.1, 0.2, 0.3, 0.4]) + + # within-model posterior weights sum to 1 (= g_i / p_model) + _, w0 = history.get_distribution(m=0, t=0) + _, w1 = history.get_distribution(m=1, t=0) + assert np.isclose(w0.sum(), 1.0) and np.isclose(w1.sum(), 1.0) + assert np.allclose(sorted(w0), sorted([0.3 / 0.4, 0.1 / 0.4])) + assert np.allclose(sorted(w1), sorted([0.2 / 0.6, 0.4 / 0.6])) + + # global weighted distances / summary statistics sum to 1 + wd = history.get_weighted_distances(t=0) + assert np.isclose(wd.w.sum(), 1.0) + assert np.allclose(sorted(wd.w.values), [0.1, 0.2, 0.3, 0.4]) + w_ss, _ = history.get_weighted_sum_stats(t=0) + assert np.isclose(sum(w_ss), 1.0) + + # within-model summary statistics for a single model sum to 1 + w_m0, _ = history.get_weighted_sum_stats_for_model(m=0, t=0) + assert np.isclose(w_m0.sum(), 1.0) + assert np.allclose(sorted(w_m0), sorted([0.3 / 0.4, 0.1 / 0.4])) + + # get_population reconstructs the global weights, summing to 1 + pop = history.get_population(t=0) + assert np.isclose(sum(p.weight for p in pop.particles), 1.0) + assert np.allclose( + sorted(p.weight for p in pop.particles), [0.1, 0.2, 0.3, 0.4] + ) + + def test_total_nr_samples(history: History): particle_list = [ Particle( @@ -593,3 +675,35 @@ def model(p): finally: if os.path.exists(db_file): os.remove(db_file) + + +def test_save_dict_to_json_numpy_arrays(): + """Regression: `save_dict_to_json` handles integer and multi-dimensional + numpy arrays.""" + from pyabc.storage.json import load_dict_from_json, save_dict_to_json + + f = tempfile.mkstemp(suffix='.json')[1] + try: + save_dict_to_json( + {1: np.array([1, 2, 3]), 2: np.array([[1.0, 2.0], [3.0, 4.0]])}, f + ) + got = load_dict_from_json(f) + finally: + if os.path.exists(f): + os.remove(f) + assert got[1] == [1, 2, 3] + assert got[2] == [[1.0, 2.0], [3.0, 4.0]] + + +def test_abcsmc_repr_with_unfinished_run(): + """Regression: `ABCSMC.__repr__` must not crash for a run that has been + started but not finished (`end_time is None`).""" + from pyabc.storage.db_model import ABCSMC, datetime2str + + row = ABCSMC( + id=1, + start_time=datetime.datetime(2024, 1, 1, 12, 0, 0), + end_time=None, + ) + assert 'end_time=None' in repr(row) + assert datetime2str(None) == 'None' diff --git a/test/base/test_sumstat.py b/test/base/test_sumstat.py index 9acd6f9a..048d25ae 100644 --- a/test/base/test_sumstat.py +++ b/test/base/test_sumstat.py @@ -297,3 +297,53 @@ def model(p): df_info, w_info = h.get_distribution() off_info = abs(pyabc.weighted_mean(df_info.p0, w_info) - 0.1) assert off_comp > off_info + + +def test_dict2arr_multi_key_concatenation(): + """Regression/efficiency: `dict2arr` concatenates the values of several + keys into a single flat 1d array.""" + dct = { + 'a': np.array([1.0, 2.0]), + 'b': 3.0, + 'c': np.array([4.0, 5.0, 6.0]), + } + out = dict2arr(dct, keys=['a', 'b', 'c']) + assert np.array_equal(out, np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0])) + + +def test_gmm_subsetter_zero_variance_labels(): + """Regression: GMMSubsetter z-score normalization must guard against a + constant (zero-variance) label column instead of dividing by zero (which + produced NaNs that broke the subsequent GMM fit).""" + rng = np.random.RandomState(0) + n = 200 + x = rng.normal(size=(n, 2)) + # second label column is constant -> zero standard deviation + y = np.column_stack([rng.normal(size=n), np.full(n, 7.0)]) + w = np.full( + (n, 1), 1.0 / n + ) # weights shaped (n_sample, 1), as read_sample + x_new, y_new, w_new = GMMSubsetter().select(x, y, w) + assert np.all(np.isfinite(x_new)) + assert len(w_new) == len(x_new) + + +def test_get_augmented_subset_full_fraction(): + """Regression: `get_augmented_subset` must not pass an out-of-bounds `kth` + to `np.argpartition` when the required count equals the number of + remaining samples (`min_fraction=1.0`).""" + from pyabc.sumstat.subset import get_augmented_subset + + y = np.arange(10, dtype=float).reshape(-1, 1) + ref = np.array([0.0]) + + in_cluster = np.zeros(10, dtype=bool) + in_cluster[:3] = True + res = get_augmented_subset(y, ref, in_cluster, min_fraction=1.0) + assert res.sum() == 10 + + # partial augmentation still selects exactly the desired count + ic2 = np.zeros(10, dtype=bool) + ic2[:2] = True + res2 = get_augmented_subset(y, ref, ic2, min_fraction=0.5) + assert res2.sum() == 5 diff --git a/test/migrate/test_migrate.py b/test/migrate/test_migrate.py index 752a405c..94903cda 100644 --- a/test/migrate/test_migrate.py +++ b/test/migrate/test_migrate.py @@ -1,11 +1,37 @@ """Migration tests.""" import os +import sqlite3 import tempfile +import numpy as np +import pandas as pd import pytest import pyabc +from pyabc.parameters import Parameter +from pyabc.population import Particle, Population +from pyabc.storage.version import __db_version__ + +# model names of the test database created below +MODEL_NAMES = ['m0', 'm1'] + +# global particle weights per population index and model, i.e. as stored by +# the current format: they sum to 1 across all particles of all models, +# and within a model to that model's probability +WEIGHTS = { + 0: {0: [0.3, 0.1], 1: [0.2, 0.4]}, + 1: {0: [0.1, 0.15], 1: [0.5, 0.25]}, +} + +# wall times per population index +WALL_TIMES = {0: 12.5, 1: 7.25} + +# SQLite can only drop table columns from version 3.35 on +requires_drop_column = pytest.mark.skipif( + sqlite3.sqlite_version_info < (3, 35), + reason='Dropping a table column requires SQLite>=3.35', +) def test_db_import(script_runner): @@ -29,3 +55,253 @@ def test_db_import(script_runner): # remove file os.remove(db_file) + + +def create_current_db(db_file: str) -> None: + """Create a database in the current format, holding two models.""" + h = pyabc.History('sqlite:///' + db_file) + h.store_initial_data(None, {}, {'ss': 0.0}, {}, MODEL_NAMES, '', '', '{}') + for t, weights in WEIGHTS.items(): + particles = [ + Particle( + m=m, + parameter=Parameter({'a': float(m), 'b': float(ix)}), + weight=w, + sum_stat={'ss': float(ix)}, + distance=0.1 * (ix + 1), + ) + for m, ws in weights.items() + for ix, w in enumerate(ws) + ] + h.append_population( + t, + 0.5 / (t + 1), + Population(particles), + 10, + MODEL_NAMES, + wall_time=WALL_TIMES[t], + ) + + +def query(db_file: str, sql: str) -> list: + """Run a raw SQL query on the database file.""" + con = sqlite3.connect(db_file) + try: + return con.execute(sql).fetchall() + finally: + con.close() + + +def db_version(db_file: str) -> str: + """Storage format version of the database.""" + return str(query(db_file, 'SELECT version_num FROM version')[0][0]) + + +def columns(db_file: str, table: str) -> list[str]: + """Column names of a database table.""" + return [row[1] for row in query(db_file, f'PRAGMA table_info({table})')] + + +def stored_weights(db_file: str) -> dict: + """Stored particle weights as ``{t: {m: sorted weights}}``. + + The pre-population is not included. + """ + rows = query( + db_file, + 'SELECT populations.t, models.m, particles.w FROM particles ' + 'JOIN models ON models.id = particles.model_id ' + 'JOIN populations ON populations.id = models.population_id ' + 'WHERE populations.t >= 0', + ) + weights = {} + for t, m, w in rows: + weights.setdefault(t, {}).setdefault(m, []).append(w) + return { + t: {m: sorted(ws) for m, ws in per_model.items()} + for t, per_model in weights.items() + } + + +def global_weights() -> dict: + """Expected weights in the current format, summing to 1 over all models.""" + return { + t: {m: sorted(ws) for m, ws in per_model.items()} + for t, per_model in WEIGHTS.items() + } + + +def per_model_weights() -> dict: + """Expected weights in version 1, summing to 1 within each model.""" + return { + t: {m: sorted(np.asarray(ws) / sum(ws)) for m, ws in per_model.items()} + for t, per_model in WEIGHTS.items() + } + + +def assert_weights_close(actual: dict, expected: dict) -> None: + """Assert that two ``{t: {m: weights}}`` dictionaries match.""" + assert actual.keys() == expected.keys() + for t, per_model in expected.items(): + assert actual[t].keys() == per_model.keys() + for m, ws in per_model.items(): + assert np.allclose(actual[t][m], ws) + + +def alembic_config(db: str): + """Alembic configuration, skipping the test if alembic is missing.""" + pytest.importorskip('alembic') + from pyabc.storage.migrate import _alembic_config + + return _alembic_config(db) + + +def to_v1(db_file: str) -> None: + """Turn a current-format database into a version 1 database. + + Applies the version 2 downgrade, which reverts the weight normalization + (from global back to within-model, ``w = g_i / p_model``), and in addition + drops the wall time column, which did not exist in version 1 but is kept + by the downgrade. + """ + command = pytest.importorskip('alembic.command') + command.downgrade(alembic_config(db_file), '1') + + con = sqlite3.connect(db_file) + with con: + con.execute('ALTER TABLE populations DROP COLUMN wall_time') + con.close() + + +@requires_drop_column +def test_migrate_v1_to_v2(script_runner, tmp_path): + """Migrating a version 1 database to the current format. + + Checks that the wall time column is added and that particle weights are + converted from the per-model to the global normalization, on a database + with two models, i.e. with model probabilities != 1. + """ + src = str(tmp_path / 'v1.db') + dst = str(tmp_path / 'v2.db') + + # create a database in the current format and record reference values + create_current_db(src) + h = pyabc.History('sqlite:///' + src) + p_models = {t: h.get_model_probabilities(t=t) for t in WEIGHTS} + distributions = { + (t, m): h.get_distribution(m=m, t=t) for t in WEIGHTS for m in [0, 1] + } + assert_weights_close(stored_weights(src), global_weights()) + + # turn it into a version 1 database + to_v1(src) + assert db_version(src) == '1' + assert 'wall_time' not in columns(src, 'populations') + assert_weights_close(stored_weights(src), per_model_weights()) + + # an outdated database cannot be imported + with pytest.raises(AssertionError, match='Database has version 1'): + pyabc.History('sqlite:///' + src) + + # call the migration script + ret = script_runner.run(['abc-migrate', '--src', src, '--dst', dst]) + assert ret.success + + # the source database is left untouched + assert db_version(src) == '1' + assert 'wall_time' not in columns(src, 'populations') + assert_weights_close(stored_weights(src), per_model_weights()) + + # the destination database is up-to-date and has the new column + assert db_version(dst) == __db_version__ == '2' + assert 'wall_time' in columns(dst, 'populations') + + # weights are back to the global normalization + assert_weights_close(stored_weights(dst), global_weights()) + + # the pre-population's dummy particle is not affected + assert query( + dst, + 'SELECT particles.w FROM particles ' + 'JOIN models ON models.id = particles.model_id ' + 'JOIN populations ON populations.id = models.population_id ' + 'WHERE populations.t = -1', + ) == [(1.0,)] + + # the migrated database gives the same results as the original one + h = pyabc.History('sqlite:///' + dst) + for t in WEIGHTS: + assert np.allclose( + h.get_model_probabilities(t=t).p.values, p_models[t].p.values + ) + for m in [0, 1]: + df, w = h.get_distribution(m=m, t=t) + df_expected, w_expected = distributions[(t, m)] + pd.testing.assert_frame_equal(df, df_expected) + # within-model weights sum to 1 + assert np.allclose(w, w_expected) + assert np.isclose(w.sum(), 1.0) + + # wall times are unknown for migrated populations + populations = h.get_all_populations() + assert 'wall_time' in populations.columns + assert populations.wall_time.isna().all() + + +def test_migrate_v2_downgrade(tmp_path): + """The version 2 revision can be reverted and then applied again. + + The downgrade inverts the weight conversion, but keeps the wall time + column, so that the upgrade must tolerate an existing column. + """ + command = pytest.importorskip('alembic.command') + + db_file = str(tmp_path / 'db.db') + create_current_db(db_file) + cfg = alembic_config(db_file) + + # revert to version 1 + command.downgrade(cfg, '1') + assert db_version(db_file) == '1' + assert 'wall_time' in columns(db_file, 'populations') + assert_weights_close(stored_weights(db_file), per_model_weights()) + + # and migrate back to the current version + command.upgrade(cfg, 'head') + assert db_version(db_file) == __db_version__ + assert 'wall_time' in columns(db_file, 'populations') + assert_weights_close(stored_weights(db_file), global_weights()) + + +def test_db_identifier(tmp_path): + """Databases can be specified as file names or as sqlite URLs.""" + pytest.importorskip('alembic') + from pyabc.storage.migrate import _alembic_config, _to_db_file + + db_file = str(tmp_path / 'db.db') + + # file names and URLs are equivalent + assert _to_db_file(db_file) == _to_db_file('sqlite:///' + db_file) + assert _alembic_config(db_file).get_main_option( + 'sqlalchemy.url' + ) == _alembic_config('sqlite:///' + db_file).get_main_option( + 'sqlalchemy.url' + ) + + # other dialects are not supported + with pytest.raises(ValueError, match='only supports sqlite'): + _to_db_file('postgresql://user@localhost/db') + + +def test_migrate_unsupported_dialect(script_runner, tmp_path): + """Migrating a non-sqlite database gives an error message.""" + ret = script_runner.run( + [ + 'abc-migrate', + '--src', + 'postgresql://user@localhost/db', + '--dst', + str(tmp_path / 'db.db'), + ] + ) + assert 'only supports sqlite' in ret.stdout 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]: