From 1c1f631e0759c9c30641074ce29c256f33379737 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Wed, 9 Aug 2023 14:49:59 +0100 Subject: [PATCH] Autoformatting with black and isort --- src/mici/adapters.py | 20 ++++----- src/mici/autodiff.py | 1 + src/mici/autograd_wrapper.py | 15 +++---- src/mici/integrators.py | 18 ++++---- src/mici/interop.py | 9 ++-- src/mici/matrices.py | 16 ++++--- src/mici/progressbars.py | 2 +- src/mici/samplers.py | 86 ++++++++++++++++++------------------ src/mici/solvers.py | 14 +++--- src/mici/stagers.py | 3 +- src/mici/states.py | 5 ++- src/mici/systems.py | 30 +++++-------- src/mici/transitions.py | 25 ++++++----- src/mici/types.py | 15 ++----- src/mici/utils.py | 4 +- tests/test_adapters.py | 17 +++---- tests/test_integrators.py | 69 ++++++++++------------------- tests/test_interop.py | 3 +- tests/test_matrices.py | 21 +++------ tests/test_progressbars.py | 7 +-- tests/test_samplers.py | 9 ++-- tests/test_solvers.py | 15 +++---- tests/test_stagers.py | 1 + tests/test_states.py | 8 ++-- tests/test_utils.py | 5 +-- 25 files changed, 194 insertions(+), 224 deletions(-) diff --git a/src/mici/adapters.py b/src/mici/adapters.py index a099016..372afed 100644 --- a/src/mici/adapters.py +++ b/src/mici/adapters.py @@ -5,24 +5,24 @@ from abc import ABC, abstractmethod from math import exp, log from typing import TYPE_CHECKING + import numpy as np -from mici.errors import IntegratorError, AdaptationError -from mici.matrices import PositiveDiagonalMatrix, DensePositiveDefiniteMatrix + +from mici.errors import AdaptationError, IntegratorError +from mici.matrices import DensePositiveDefiniteMatrix, PositiveDiagonalMatrix if TYPE_CHECKING: - from typing import Collection, Optional, Iterable, Union + from typing import Collection, Iterable, Optional, Union + from numpy.random import Generator from numpy.typing import ArrayLike + from mici.integrators import Integrator from mici.states import ChainState from mici.systems import System from mici.transitions import Transition - from mici.types import ( - AdaptationStatisticFunction, - AdapterState, - ReducerFunction, - TransitionStatistics, - ) + from mici.types import (AdaptationStatisticFunction, AdapterState, + ReducerFunction, TransitionStatistics) class Adapter(ABC): @@ -486,7 +486,7 @@ def finalize( mean_est /= n_iter var_est += adapt_state["sum_diff_sq"] var_est += ( - mean_diff ** 2 * (adapt_state["iter"] * n_iter_prev) / n_iter + mean_diff**2 * (adapt_state["iter"] * n_iter_prev) / n_iter ) if n_iter < 2: raise AdaptationError( diff --git a/src/mici/autodiff.py b/src/mici/autodiff.py index c77a971..46d050f 100644 --- a/src/mici/autodiff.py +++ b/src/mici/autodiff.py @@ -3,6 +3,7 @@ from __future__ import annotations from typing import TYPE_CHECKING + import mici.autograd_wrapper as autograd_wrapper if TYPE_CHECKING: diff --git a/src/mici/autograd_wrapper.py b/src/mici/autograd_wrapper.py index 2bbb471..99199f3 100644 --- a/src/mici/autograd_wrapper.py +++ b/src/mici/autograd_wrapper.py @@ -7,24 +7,19 @@ AUTOGRAD_AVAILABLE = True try: - from autograd.wrap_util import unary_to_nary + import autograd.numpy as np from autograd.builtins import tuple as atuple from autograd.core import make_vjp from autograd.extend import vspace - import autograd.numpy as np + from autograd.wrap_util import unary_to_nary except ImportError: AUTOGRAD_AVAILABLE = False if TYPE_CHECKING: from typing import Callable - from mici.types import ( - ScalarLike, - ArrayLike, - ScalarFunction, - ArrayFunction, - MatrixHessianProduct, - MatrixTressianProduct, - ) + + from mici.types import (ArrayFunction, ArrayLike, MatrixHessianProduct, + MatrixTressianProduct, ScalarFunction, ScalarLike) def _wrapped_unary_to_nary(func: Callable) -> Callable: diff --git a/src/mici/integrators.py b/src/mici/integrators.py index ca725ab..91fac2e 100644 --- a/src/mici/integrators.py +++ b/src/mici/integrators.py @@ -4,21 +4,21 @@ from abc import ABC, abstractmethod from typing import TYPE_CHECKING + import numpy as np from numpy.typing import ArrayLike -from mici.errors import NonReversibleStepError, AdaptationError -from mici.solvers import ( - maximum_norm, - solve_fixed_point_direct, - solve_projection_onto_manifold_newton, - FixedPointSolver, - ProjectionSolver, -) + +from mici.errors import AdaptationError, NonReversibleStepError +from mici.solvers import (FixedPointSolver, ProjectionSolver, maximum_norm, + solve_fixed_point_direct, + solve_projection_onto_manifold_newton) if TYPE_CHECKING: from typing import Any, Callable, Optional, Sequence + from mici.states import ChainState - from mici.systems import ConstrainedTractableFlowSystem, System, TractableFlowSystem + from mici.systems import (ConstrainedTractableFlowSystem, System, + TractableFlowSystem) from mici.types import NormFunction diff --git a/src/mici/interop.py b/src/mici/interop.py index 29cd50d..89f6ecc 100644 --- a/src/mici/interop.py +++ b/src/mici/interop.py @@ -4,15 +4,18 @@ import importlib import os -import mici -import numpy as np from typing import TYPE_CHECKING +import numpy as np + +import mici + if TYPE_CHECKING: from typing import Literal, Optional, Union - from numpy.typing import ArrayLike + import arviz import pymc + from numpy.typing import ArrayLike def convert_to_inference_data( diff --git a/src/mici/matrices.py b/src/mici/matrices.py index b620e87..a2e397b 100644 --- a/src/mici/matrices.py +++ b/src/mici/matrices.py @@ -5,15 +5,19 @@ import abc import numbers from typing import TYPE_CHECKING + import numpy as np -from mici.errors import LinAlgError import numpy.linalg as nla import scipy.linalg as sla + +from mici.errors import LinAlgError from mici.utils import hash_array if TYPE_CHECKING: from typing import Iterable, Literal, Optional, Tuple, Union + from numpy.typing import NDArray + from mici.types import MatrixLike, ScalarLike @@ -648,7 +652,7 @@ def grad_log_abs_det(self): return self.shape[0] / self._scalar def grad_quadratic_form_inv(self, vector: NDArray) -> float: - return -np.sum(vector ** 2) / self._scalar ** 2 + return -np.sum(vector**2) / self._scalar**2 def __str__(self) -> str: return f"(shape={self.shape}, scalar={self._scalar})" @@ -681,7 +685,7 @@ def _construct_inv(self) -> PositiveScaledIdentityMatrix: return PositiveScaledIdentityMatrix(1 / self._scalar, self.shape[0]) def _construct_sqrt(self) -> PositiveScaledIdentityMatrix: - return PositiveScaledIdentityMatrix(self._scalar ** 0.5, self.shape[0]) + return PositiveScaledIdentityMatrix(self._scalar**0.5, self.shape[0]) class DiagonalMatrix(SymmetricMatrix, DifferentiableMatrix, ImplicitArrayMatrix): @@ -766,7 +770,7 @@ def _construct_inv(self) -> PositiveDiagonalMatrix: return PositiveDiagonalMatrix(1.0 / self.diagonal) def _construct_sqrt(self) -> PositiveDiagonalMatrix: - return PositiveDiagonalMatrix(self.diagonal ** 0.5) + return PositiveDiagonalMatrix(self.diagonal**0.5) def _make_array_triangular(array: NDArray, lower: bool) -> NDArray: @@ -1042,7 +1046,7 @@ def __init__( def _scalar_multiply(self, scalar: ScalarLike) -> TriangularFactoredDefiniteMatrix: if scalar > 0: return TriangularFactoredPositiveDefiniteMatrix( - factor=scalar ** 0.5 * self.factor + factor=scalar**0.5 * self.factor ) else: return super()._scalar_multiply(scalar) @@ -1542,7 +1546,7 @@ def _construct_inv(self) -> EigendecomposedPositiveDefiniteMatrix: return EigendecomposedPositiveDefiniteMatrix(self.eigvec, 1 / self.eigval) def _construct_sqrt(self) -> EigendecomposedPositiveDefiniteMatrix: - return EigendecomposedPositiveDefiniteMatrix(self.eigvec, self.eigval ** 0.5) + return EigendecomposedPositiveDefiniteMatrix(self.eigvec, self.eigval**0.5) class SoftAbsRegularizedPositiveDefiniteMatrix( diff --git a/src/mici/progressbars.py b/src/mici/progressbars.py index 817c305..1281fa5 100644 --- a/src/mici/progressbars.py +++ b/src/mici/progressbars.py @@ -225,7 +225,7 @@ def __init__( @property def description(self) -> str: - """"Description of task being tracked.""" + """Description of task being tracked.""" return self._description @property diff --git a/src/mici/samplers.py b/src/mici/samplers.py index a15145a..05ac1db 100644 --- a/src/mici/samplers.py +++ b/src/mici/samplers.py @@ -2,62 +2,49 @@ from __future__ import annotations +import logging import os import queue +import signal +import tempfile from contextlib import ExitStack, contextmanager, nullcontext from pathlib import Path from pickle import PicklingError -import logging -import tempfile -import signal -from warnings import warn from typing import TYPE_CHECKING, NamedTuple +from warnings import warn + import numpy as np from numpy.random import default_rng -from mici.transitions import ( - IndependentMomentumTransition, - MetropolisRandomIntegrationTransition, - MetropolisStaticIntegrationTransition, - MultinomialDynamicIntegrationTransition, - SliceDynamicIntegrationTransition, - euclidean_no_u_turn_criterion, - riemannian_no_u_turn_criterion, -) -from mici.states import ChainState -from mici.progressbars import ( - SequenceProgressBar, - LabelledSequenceProgressBar, - DummyProgressBar, - _ProxySequenceProgressBar, -) -from mici.errors import AdaptationError + from mici.adapters import DualAveragingStepSizeAdapter +from mici.errors import AdaptationError +from mici.progressbars import (DummyProgressBar, LabelledSequenceProgressBar, + SequenceProgressBar, _ProxySequenceProgressBar) from mici.stagers import WarmUpStager, WindowedWarmUpStager +from mici.states import ChainState +from mici.transitions import (IndependentMomentumTransition, + MetropolisRandomIntegrationTransition, + MetropolisStaticIntegrationTransition, + MultinomialDynamicIntegrationTransition, + SliceDynamicIntegrationTransition, + euclidean_no_u_turn_criterion, + riemannian_no_u_turn_criterion) if TYPE_CHECKING: - from typing import ( - Container, - Generator, - Iterable, - Optional, - Sequence, - Union, - ) + from typing import (Container, Generator, Iterable, Optional, Sequence, + Union) + from numpy.typing import ArrayLike, DTypeLike, NDArray + from mici.adapters import Adapter from mici.integrators import Integrator from mici.progressbars import ProgressBar from mici.stagers import Stager from mici.systems import System - from mici.transitions import IntegrationTransition, MomentumTransition, Transition - from mici.types import ( - AdapterState, - ChainIterator, - ScalarLike, - PyTree, - TraceFunction, - TerminationCriterion, - ) + from mici.transitions import (IntegrationTransition, MomentumTransition, + Transition) + from mici.types import (AdapterState, ChainIterator, PyTree, ScalarLike, + TerminationCriterion, TraceFunction) # Preferentially import from multiprocess library if available as able to # serialize much wider range of types including autograd functions @@ -145,7 +132,10 @@ def _generate_memmap_filenames( def _open_new_memmap( - file_path: str, shape: tuple[int, ...], default_val: ScalarLike, dtype: DTypeLike, + file_path: str, + shape: tuple[int, ...], + default_val: ScalarLike, + dtype: DTypeLike, ) -> np.memmap: """Open a new memory-mapped array object and fill with a default-value. @@ -298,7 +288,15 @@ def _init_traces( ] else: traces[key] = list( - np.full((n_chain, n_iter,) + val.shape, init, val.dtype) + np.full( + ( + n_chain, + n_iter, + ) + + val.shape, + init, + val.dtype, + ) ) return traces @@ -961,7 +959,11 @@ def sample_chains( ) ) stats = _init_stats( - self.transitions, n_chain, n_trace_iter, use_memmap, memmap_path, + self.transitions, + n_chain, + n_trace_iter, + use_memmap, + memmap_path, ) per_chain_rngs = _get_per_chain_rngs(self.rng, n_chain) per_chain_traces = ( @@ -1051,7 +1053,7 @@ class HMCSampleChainsOutputs(NamedTuple): and main sampling stages otherwise. The key for each value is the corresponding key in the dictionary returned by the trace function which computed the traced value. - statistics: Dictionary of chain transition statistic dictionaries. Values in + statistics: Dictionary of chain transition statistic dictionaries. Values in dictionary are lists of arrays of chain statistic values with each array in the list corresponding to a single chain and the leading dimension of each array corresponding to the iteration (draw) index, within the main diff --git a/src/mici/solvers.py b/src/mici/solvers.py index ce1f84a..584c5f7 100644 --- a/src/mici/solvers.py +++ b/src/mici/solvers.py @@ -2,17 +2,17 @@ from __future__ import annotations -from typing import Protocol, TYPE_CHECKING -from mici.errors import ConvergenceError, LinAlgError +from typing import TYPE_CHECKING, Protocol + import numpy as np +from mici.errors import ConvergenceError, LinAlgError + if TYPE_CHECKING: from mici.states import ChainState - from mici.systems import ( - ConstrainedEuclideanMetricSystem, - ConstrainedTractableFlowSystem, - ) - from mici.types import ScalarFunction, ArrayFunction, ArrayLike + from mici.systems import (ConstrainedEuclideanMetricSystem, + ConstrainedTractableFlowSystem) + from mici.types import ArrayFunction, ArrayLike, ScalarFunction def euclidean_norm(vct): diff --git a/src/mici/stagers.py b/src/mici/stagers.py index 8ee295f..0cfe018 100644 --- a/src/mici/stagers.py +++ b/src/mici/stagers.py @@ -3,10 +3,11 @@ from __future__ import annotations import abc -from typing import NamedTuple, TYPE_CHECKING +from typing import TYPE_CHECKING, NamedTuple if TYPE_CHECKING: from typing import Iterable, Optional + from mici.adapters import Adapter from mici.types import TraceFunction diff --git a/src/mici/states.py b/src/mici/states.py index 7dfd18d..255436d 100644 --- a/src/mici/states.py +++ b/src/mici/states.py @@ -3,14 +3,17 @@ from __future__ import annotations import copy -from functools import wraps from collections import Counter +from functools import wraps from typing import TYPE_CHECKING + from mici.errors import ReadOnlyStateError if TYPE_CHECKING: from typing import Any, Callable, Iterable, Optional + from numpy.typing import ArrayLike + from mici.systems import System diff --git a/src/mici/systems.py b/src/mici/systems.py index 5ed8b07..827c32e 100644 --- a/src/mici/systems.py +++ b/src/mici/systems.py @@ -4,32 +4,26 @@ from abc import ABC, abstractmethod from typing import TYPE_CHECKING + import numpy as np -from mici.states import cache_in_state, cache_in_state_with_aux + import mici.matrices as matrices from mici.autodiff import autodiff_fallback +from mici.states import cache_in_state, cache_in_state_with_aux if TYPE_CHECKING: from typing import Any, Optional, Type + from numpy.random import Generator + from mici.states import ChainState - from mici.types import ( - ScalarLike, - ArrayLike, - MatrixLike, - MetricLike, - ArrayFunction, - ScalarFunction, - GradientFunction, - HessianFunction, - MatrixTressianProduct, - MatrixTressianProductFunction, - JacobianFunction, - MatrixHessianProduct, - MatrixHessianProductFunction, - VectorJacobianProduct, - VectorJacobianProductFunction, - ) + from mici.types import (ArrayFunction, ArrayLike, GradientFunction, + HessianFunction, JacobianFunction, + MatrixHessianProduct, MatrixHessianProductFunction, + MatrixLike, MatrixTressianProduct, + MatrixTressianProductFunction, MetricLike, + ScalarFunction, ScalarLike, VectorJacobianProduct, + VectorJacobianProductFunction) class System(ABC): diff --git a/src/mici/transitions.py b/src/mici/transitions.py index 7441a46..4841704 100644 --- a/src/mici/transitions.py +++ b/src/mici/transitions.py @@ -2,23 +2,22 @@ from __future__ import annotations -from abc import ABC, abstractmethod, abstractproperty import logging -from typing import NamedTuple, TYPE_CHECKING +from abc import ABC, abstractmethod, abstractproperty +from typing import TYPE_CHECKING, NamedTuple + import numpy as np + +from mici.errors import (ConvergenceError, Error, HamiltonianDivergenceError, + IntegratorError, NonReversibleStepError) from mici.utils import LogRepFloat -from mici.errors import ( - Error, - IntegratorError, - NonReversibleStepError, - ConvergenceError, - HamiltonianDivergenceError, -) if TYPE_CHECKING: from typing import Optional + from numpy.random import Generator from numpy.typing import ArrayLike, DTypeLike + from mici.integrators import Integrator from mici.states import ChainState from mici.systems import System @@ -184,7 +183,7 @@ def sample( state.mom = self.system.sample_momentum(state, rng) elif self.mom_resample_coeff != 0: mom_ind = self.system.sample_momentum(state, rng) - state.mom *= (1.0 - self.mom_resample_coeff ** 2) ** 0.5 + state.mom *= (1.0 - self.mom_resample_coeff**2) ** 0.5 state.mom += self.mom_resample_coeff * mom_ind return state, None @@ -198,7 +197,7 @@ class IntegrationTransition(Transition): """ @property - def state_variables(self) -> Set[str]: + def state_variables(self) -> set[str]: return {"pos", "mom", "dir"} @property @@ -588,7 +587,9 @@ def _weight_function( @abstractmethod def _weight_ratio( - self, numerator: ScalarLike, denominator: ScalarLike, + self, + numerator: ScalarLike, + denominator: ScalarLike, ) -> ScalarLike: pass diff --git a/src/mici/types.py b/src/mici/types.py index 70270ab..a16b8fc 100644 --- a/src/mici/types.py +++ b/src/mici/types.py @@ -2,23 +2,16 @@ from __future__ import annotations -from typing import ( - Any, - Callable, - Collection, - Iterable, - TypeVar, - Union, -) -from numpy.typing import ArrayLike +from typing import Any, Callable, Collection, Iterable, TypeVar, Union + from numpy import number +from numpy.typing import ArrayLike -from mici.states import ChainState from mici.matrices import Matrix, PositiveDefiniteMatrix +from mici.states import ChainState from mici.systems import System from mici.utils import LogRepFloat - ScalarLike = Union[bool, int, float, LogRepFloat, number] MatrixLike = Union[ArrayLike, Matrix] MetricLike = Union[ArrayLike, PositiveDefiniteMatrix] diff --git a/src/mici/utils.py b/src/mici/utils.py index 14c2570..fb04830 100644 --- a/src/mici/utils.py +++ b/src/mici/utils.py @@ -2,9 +2,10 @@ from __future__ import annotations +from math import exp, expm1, inf, log, log1p, nan from typing import TYPE_CHECKING + import numpy as np -from math import log, exp, log1p, expm1, inf, nan try: import xxhash @@ -15,6 +16,7 @@ if TYPE_CHECKING: from typing import Optional + from mici.types import ScalarLike diff --git a/tests/test_adapters.py b/tests/test_adapters.py index a5bba6d..4c148e1 100644 --- a/tests/test_adapters.py +++ b/tests/test_adapters.py @@ -1,11 +1,12 @@ -import pytest import copy from collections.abc import Mapping -from types import MappingProxyType from math import exp +from types import MappingProxyType + import numpy as np -import mici +import pytest +import mici SEED = 3046987125 STATE_DIM = 10 @@ -126,7 +127,7 @@ def chain_state(self, rng): @pytest.fixture def system(self): return mici.systems.EuclideanMetricSystem( - neg_log_dens=lambda pos: np.sum(pos ** 2) / 2, + neg_log_dens=lambda pos: np.sum(pos**2) / 2, grad_neg_log_dens=lambda pos: pos, ) @@ -151,16 +152,16 @@ def adapter(self): @pytest.fixture def chain_state(self, rng): pos, mom = rng.standard_normal((2, STATE_DIM)) - pos /= np.sum(pos ** 2) ** 0.5 + pos /= np.sum(pos**2) ** 0.5 mom -= (mom @ pos) * pos return mici.states.ChainState(pos=pos, mom=mom, dir=1) @pytest.fixture def system(self): return mici.systems.DenseConstrainedEuclideanMetricSystem( - neg_log_dens=lambda pos: np.sum(pos ** 2) / 2, + neg_log_dens=lambda pos: np.sum(pos**2) / 2, grad_neg_log_dens=lambda pos: pos, - constr=lambda pos: np.sum(pos ** 2)[None] - 1, + constr=lambda pos: np.sum(pos**2)[None] - 1, jacob_constr=lambda pos: 2 * pos[None], ) @@ -182,7 +183,7 @@ def adapter(self): def system(self, rng): var = np.exp(rng.standard_normal(STATE_DIM)) return mici.systems.EuclideanMetricSystem( - neg_log_dens=lambda pos: np.sum(pos ** 2 / var) / 2, + neg_log_dens=lambda pos: np.sum(pos**2 / var) / 2, grad_neg_log_dens=lambda pos: pos / var, ) diff --git a/tests/test_integrators.py b/tests/test_integrators.py index 705c2e0..14f7792 100644 --- a/tests/test_integrators.py +++ b/tests/test_integrators.py @@ -1,11 +1,12 @@ -import pytest import numpy as np +import pytest + import mici.integrators as integrators -import mici.systems as systems import mici.matrices as matrices import mici.solvers as solvers -from mici.states import ChainState +import mici.systems as systems from mici.errors import IntegratorError +from mici.states import ChainState SEED = 3046987125 N_STEPS = {1, 5, 20} @@ -72,7 +73,6 @@ def _integrate_with_reversal(integrator, init_state, n_step): class IntegratorTests: - h_diff_tol = 5e-3 @pytest.mark.parametrize("n_step", N_STEPS) @@ -234,7 +234,7 @@ class LinearEuclideanMetricSystemTests(LinearSystemIntegratorTests): @pytest.fixture def system(self, metric): return systems.EuclideanMetricSystem( - neg_log_dens=lambda q: 0.5 * np.sum(q ** 2), + neg_log_dens=lambda q: 0.5 * np.sum(q**2), metric=metric, grad_neg_log_dens=lambda q: q, ) @@ -244,9 +244,9 @@ class NonLinearEuclideanMetricSystemTests(IntegratorTests): @pytest.fixture def system(self, metric): return systems.EuclideanMetricSystem( - neg_log_dens=lambda q: 0.25 * np.sum(q ** 4), + neg_log_dens=lambda q: 0.25 * np.sum(q**4), metric=metric, - grad_neg_log_dens=lambda q: q ** 3, + grad_neg_log_dens=lambda q: q**3, ) @@ -262,9 +262,9 @@ class NonLinearGaussianEuclideanMetricSystem(IntegratorTests): @pytest.fixture def system(self, metric): return systems.GaussianEuclideanMetricSystem( - neg_log_dens=lambda q: 0.125 * np.sum(q ** 4), + neg_log_dens=lambda q: 0.125 * np.sum(q**4), metric=metric, - grad_neg_log_dens=lambda q: 0.5 * q ** 3, + grad_neg_log_dens=lambda q: 0.5 * q**3, ) @@ -277,7 +277,6 @@ def integrator(self, system): class TestLeapfrogIntegratorLinearEuclideanMetricSystem( LeapfrogIntegratorTests, LinearEuclideanMetricSystemTests ): - h_diff_tol = 2e-3 step_size = 0.25 @@ -285,7 +284,6 @@ class TestLeapfrogIntegratorLinearEuclideanMetricSystem( class TestLeapfrogIntegratorNonLinearEuclideanMetricSystem( LeapfrogIntegratorTests, NonLinearEuclideanMetricSystemTests ): - h_diff_tol = 1e-3 step_size = 0.05 @@ -293,7 +291,6 @@ class TestLeapfrogIntegratorNonLinearEuclideanMetricSystem( class TestLeapfrogIntegratorLinearGaussianEuclideanMetricSystem( LeapfrogIntegratorTests, LinearGaussianEuclideanMetricSystem ): - h_diff_tol = 1e-10 step_size = 0.5 @@ -301,7 +298,6 @@ class TestLeapfrogIntegratorLinearGaussianEuclideanMetricSystem( class TestLeapfrogIntegratorNonLinearGaussianEuclideanMetricSystem( LeapfrogIntegratorTests, NonLinearGaussianEuclideanMetricSystem ): - h_diff_tol = 2e-3 step_size = 0.1 @@ -315,7 +311,6 @@ def integrator(self, system): class TestBCSSTwoStageIntegratorLinearEuclideanMetricSystem( BCSSTwoStageIntegratorTests, LinearEuclideanMetricSystemTests ): - h_diff_tol = 2e-4 step_size = 0.25 @@ -323,7 +318,6 @@ class TestBCSSTwoStageIntegratorLinearEuclideanMetricSystem( class TestBCSSTwoStageIntegratorNonLinearEuclideanMetricSystem( BCSSTwoStageIntegratorTests, NonLinearEuclideanMetricSystemTests ): - h_diff_tol = 1e-3 step_size = 0.05 @@ -331,15 +325,14 @@ class TestBCSSTwoStageIntegratorNonLinearEuclideanMetricSystem( class TestBCSSTwoStageIntegratorLinearGaussianEuclideanMetricSystem( BCSSTwoStageIntegratorTests, LinearGaussianEuclideanMetricSystem ): - h_diff_tol = 1e-10 step_size = 0.5 class TestBCSSTwoStageIntegratorNonLinearGaussianEuclideanMetricSystem( - BCSSTwoStageIntegratorTests, NonLinearGaussianEuclideanMetricSystem, + BCSSTwoStageIntegratorTests, + NonLinearGaussianEuclideanMetricSystem, ): - h_diff_tol = 2e-3 step_size = 0.1 @@ -353,7 +346,6 @@ def integrator(self, system): class TestBCSSThreeStageIntegratorLinearEuclideanMetricSystem( BCSSThreeStageIntegratorTests, LinearEuclideanMetricSystemTests ): - h_diff_tol = 5e-5 step_size = 0.25 @@ -361,7 +353,6 @@ class TestBCSSThreeStageIntegratorLinearEuclideanMetricSystem( class TestBCSSThreeStageIntegratorNonLinearEuclideanMetricSystem( BCSSThreeStageIntegratorTests, NonLinearEuclideanMetricSystemTests ): - h_diff_tol = 5e-4 step_size = 0.25 @@ -369,15 +360,14 @@ class TestBCSSThreeStageIntegratorNonLinearEuclideanMetricSystem( class TestBCSSThreeStageIntegratorLinearGaussianEuclideanMetricSystem( BCSSThreeStageIntegratorTests, LinearGaussianEuclideanMetricSystem ): - h_diff_tol = 1e-10 step_size = 0.5 class TestBCSSThreeStageIntegratorNonLinearGaussianEuclideanMetricSystem( - BCSSThreeStageIntegratorTests, NonLinearGaussianEuclideanMetricSystem, + BCSSThreeStageIntegratorTests, + NonLinearGaussianEuclideanMetricSystem, ): - h_diff_tol = 5e-4 step_size = 0.5 @@ -391,7 +381,6 @@ def integrator(self, system): class TestBCSSFourStageIntegratorLinearEuclideanMetricSystem( BCSSFourStageIntegratorTests, LinearEuclideanMetricSystemTests ): - h_diff_tol = 2e-5 step_size = 1.0 @@ -399,7 +388,6 @@ class TestBCSSFourStageIntegratorLinearEuclideanMetricSystem( class TestBCSSFourStageIntegratorNonLinearEuclideanMetricSystem( BCSSFourStageIntegratorTests, NonLinearEuclideanMetricSystemTests ): - h_diff_tol = 1e-3 step_size = 0.25 @@ -407,15 +395,14 @@ class TestBCSSFourStageIntegratorNonLinearEuclideanMetricSystem( class TestBCSSFourStageIntegratorLinearGaussianEuclideanMetricSystem( BCSSFourStageIntegratorTests, LinearGaussianEuclideanMetricSystem ): - h_diff_tol = 1e-10 step_size = 1.0 class TestBCSSFourStageIntegratorNonLinearGaussianEuclideanMetricSystem( - BCSSFourStageIntegratorTests, NonLinearGaussianEuclideanMetricSystem, + BCSSFourStageIntegratorTests, + NonLinearGaussianEuclideanMetricSystem, ): - h_diff_tol = 5e-4 step_size = 0.5 @@ -455,7 +442,6 @@ def integrator(self, system, fixed_point_solver, reverse_check_norm): class TestImplicitLeapfrogIntegratorLinearEuclideanMetricSystem( ImplicitLeapfrogIntegratorTests, LinearEuclideanMetricSystemTests ): - h_diff_tol = 5e-3 step_size = 0.25 @@ -463,7 +449,6 @@ class TestImplicitLeapfrogIntegratorLinearEuclideanMetricSystem( class TestImplicitMidpointIntegratorLinearEuclideanMetricSystem( ImplicitMidpointIntegratorTests, LinearEuclideanMetricSystemTests ): - h_diff_tol = 1e-7 step_size = 0.25 @@ -471,7 +456,6 @@ class TestImplicitMidpointIntegratorLinearEuclideanMetricSystem( class TestImplicitLeapfrogIntegratorNonLinearEuclideanMetricSystem( ImplicitLeapfrogIntegratorTests, NonLinearEuclideanMetricSystemTests ): - h_diff_tol = 6e-3 step_size = 0.1 @@ -479,7 +463,6 @@ class TestImplicitLeapfrogIntegratorNonLinearEuclideanMetricSystem( class TestImplicitMidpointIntegratorNonLinearEuclideanMetricSystem( ImplicitMidpointIntegratorTests, NonLinearEuclideanMetricSystemTests ): - h_diff_tol = 5e-3 step_size = 0.1 @@ -488,9 +471,9 @@ class NonLinearDiagonalRiemannianMetricSystemTests(IntegratorTests): @pytest.fixture def system(self): return systems.DiagonalRiemannianMetricSystem( - lambda q: np.sum(q ** 2) / 2 + np.sum(q ** 4) / 12, - grad_neg_log_dens=lambda q: q + q ** 3 / 3, - metric_diagonal_func=lambda q: 1 + q ** 2, + lambda q: np.sum(q**2) / 2 + np.sum(q**4) / 12, + grad_neg_log_dens=lambda q: q + q**3 / 3, + metric_diagonal_func=lambda q: 1 + q**2, vjp_metric_diagonal_func=lambda q: lambda m: 2 * m * q, ) @@ -498,7 +481,6 @@ def system(self): class TestImplicitLeapfrogIntegratorNonLinearDiagonalRiemannianMetricSystem( ImplicitLeapfrogIntegratorTests, NonLinearDiagonalRiemannianMetricSystemTests ): - h_diff_tol = 1e-3 step_size = 0.1 @@ -506,7 +488,6 @@ class TestImplicitLeapfrogIntegratorNonLinearDiagonalRiemannianMetricSystem( class TestImplicitMidpointIntegratorNonLinearDiagonalRiemannianMetricSystem( ImplicitMidpointIntegratorTests, NonLinearDiagonalRiemannianMetricSystemTests ): - h_diff_tol = 2e-4 step_size = 0.1 @@ -514,13 +495,12 @@ class TestImplicitMidpointIntegratorNonLinearDiagonalRiemannianMetricSystem( class TestConstrainedLeapfrogIntegratorLinearSystem( ConstrainedLinearSystemIntegratorTests ): - h_diff_tol = 1e-2 @pytest.fixture def integrator(self, metric): system = systems.DenseConstrainedEuclideanMetricSystem( - neg_log_dens=lambda q: 0.5 * np.sum(q ** 2), + neg_log_dens=lambda q: 0.5 * np.sum(q**2), metric=metric, grad_neg_log_dens=lambda q: q, constr=lambda q: q[:1], @@ -532,15 +512,14 @@ def integrator(self, metric): class TestConstrainedLeapfrogIntegratorNonLinearSystem( ConstrainedNonLinearSystemIntegratorTests ): - h_diff_tol = 1e-2 @pytest.fixture def system(self, metric): return systems.DenseConstrainedEuclideanMetricSystem( - neg_log_dens=lambda q: 0.125 * np.sum(q ** 4), + neg_log_dens=lambda q: 0.125 * np.sum(q**4), metric=metric, - grad_neg_log_dens=lambda q: 0.5 * q ** 3, + grad_neg_log_dens=lambda q: 0.5 * q**3, constr=lambda q: q[0:1] ** 2 + q[1:2] ** 2 - 1.0, jacob_constr=lambda q: np.concatenate( [2 * q[0:1], 2 * q[1:2], np.zeros(q.shape[0] - 2)] @@ -563,7 +542,6 @@ def integrator(self, system, request): class TestConstrainedLeapfrogIntegratorGaussianLinearSystem( ConstrainedLinearSystemIntegratorTests ): - h_diff_tol = 1e-4 @pytest.fixture @@ -582,15 +560,14 @@ def integrator(self, metric): class TestConstrainedLeapfrogIntegratorGaussianNonLinearSystem( ConstrainedNonLinearSystemIntegratorTests ): - h_diff_tol = 5e-2 @pytest.fixture def system(self, metric): return systems.GaussianDenseConstrainedEuclideanMetricSystem( - neg_log_dens=lambda q: 0.125 * np.sum(q ** 4), + neg_log_dens=lambda q: 0.125 * np.sum(q**4), metric=metric, - grad_neg_log_dens=lambda q: 0.5 * q ** 3, + grad_neg_log_dens=lambda q: 0.5 * q**3, constr=lambda q: q[0:1] ** 2 + q[1:2] ** 2 - 1.0, jacob_constr=lambda q: np.concatenate( [2 * q[0:1], 2 * q[1:2], np.zeros(q.shape[0] - 2)] diff --git a/tests/test_interop.py b/tests/test_interop.py index f1b33e0..7680530 100644 --- a/tests/test_interop.py +++ b/tests/test_interop.py @@ -1,5 +1,6 @@ -import pytest import numpy as np +import pytest + import mici try: diff --git a/tests/test_matrices.py b/tests/test_matrices.py index 9dba823..6c7448f 100644 --- a/tests/test_matrices.py +++ b/tests/test_matrices.py @@ -1,18 +1,20 @@ -import pytest -from itertools import product -from functools import partial, wraps, reduce from copy import copy, deepcopy +from functools import partial, reduce, wraps +from itertools import product + import numpy as np import numpy.linalg as nla -import scipy.linalg as sla import numpy.testing as npt +import pytest +import scipy.linalg as sla + import mici.matrices as matrices AUTOGRAD_AVAILABLE = True try: import autograd.numpy as anp from autograd import grad - from autograd.core import primitive, defvjp + from autograd.core import defvjp, primitive except ImportError: AUTOGRAD_AVAILABLE = False import warnings @@ -312,7 +314,6 @@ def test_sqrt_array(self, matrix, np_matrix): class DifferentiableMatrixTests(MatrixTests): - if AUTOGRAD_AVAILABLE: @pytest.fixture @@ -383,7 +384,6 @@ class TestScaledIdentityMatrix( ExplicitShapeSymmetricMatrixTests, ExplicitShapeInvertibleMatrixTests, ): - matrix_class = matrices.ScaledIdentityMatrix @staticmethod @@ -394,7 +394,6 @@ def generate_scalar(rng): class TestPositiveScaledIdentityMatrix( DifferentiableScaledIdentityMatrixTests, ExplicitShapePositiveDefiniteMatrixTests ): - matrix_class = matrices.PositiveScaledIdentityMatrix @staticmethod @@ -422,7 +421,6 @@ class TestDiagonalMatrix( ExplicitShapeSymmetricMatrixTests, ExplicitShapeInvertibleMatrixTests, ): - matrix_class = matrices.DiagonalMatrix @staticmethod @@ -433,7 +431,6 @@ def generate_diagonal(size, rng): class TestPositiveDiagonalMatrix( DifferentiableDiagonalMatrixTests, ExplicitShapePositiveDefiniteMatrixTests ): - matrix_class = matrices.PositiveDiagonalMatrix @staticmethod @@ -488,7 +485,6 @@ class TestTriangularFactoredDefiniteMatrix( ExplicitShapeSymmetricMatrixTests, ExplicitShapeInvertibleMatrixTests, ): - matrix_class = matrices.TriangularFactoredDefiniteMatrix @pytest.fixture(params=(+1, -1)) @@ -532,7 +528,6 @@ class TestDenseDefiniteMatrix( ExplicitShapeSymmetricMatrixTests, ExplicitShapeInvertibleMatrixTests, ): - matrix_class = matrices.DenseDefiniteMatrix @pytest.fixture(params=(+1, -1)) @@ -785,7 +780,6 @@ def matrix_pair(self, rng, size, n_block): if AUTOGRAD_AVAILABLE: - # Define new block_diag primitive and corresponding vector-Jacobian-product @primitive @@ -793,7 +787,6 @@ def block_diag(blocks): return sla.block_diag(*blocks) def vjp_block_diag(ans, blocks): - blocks = tuple(blocks) def vjp(g): diff --git a/tests/test_progressbars.py b/tests/test_progressbars.py index e8384e0..6512a0d 100644 --- a/tests/test_progressbars.py +++ b/tests/test_progressbars.py @@ -1,8 +1,10 @@ -import pytest -import mici from collections import OrderedDict from queue import SimpleQueue +import pytest + +import mici + def test_format_time(): assert mici.progressbars._format_time(100) == "01:40" @@ -60,4 +62,3 @@ def test_progress_bar_iter(progress_bar_and_sequence): for (val, iter_dict), val_orig in zip(progress_bar, sequence): assert val == val_orig assert isinstance(iter_dict, dict) - diff --git a/tests/test_samplers.py b/tests/test_samplers.py index aabb4c2..65d6717 100644 --- a/tests/test_samplers.py +++ b/tests/test_samplers.py @@ -1,7 +1,7 @@ -import pytest import numpy as np -import mici +import pytest +import mici SEED = 3046987125 STATE_DIM = 2 @@ -14,7 +14,7 @@ def rng(): def neg_log_dens(pos): - return np.sum(pos ** 2) / 2 + return np.sum(pos**2) / 2 def grad_neg_log_dens(pos): @@ -333,7 +333,6 @@ def check_stats_dict(self, stats, n_iter, n_chain, transitions): class TestStaticMetropolisHMC(HamiltonianMCMCTests): - n_step = 2 def test_max_tree_depth(self, sampler): @@ -351,7 +350,6 @@ def sampler(self, integrator, system, rng): class TestRandomMetropolisHMC(HamiltonianMCMCTests): - n_step_range = (1, 3) def test_max_tree_depth(self, sampler): @@ -374,7 +372,6 @@ def sampler(self, integrator, system, rng): class DynamicHMCTests(HamiltonianMCMCTests): - max_tree_depth = 2 max_delta_h = 1000 diff --git a/tests/test_solvers.py b/tests/test_solvers.py index 758cf9b..ab11333 100644 --- a/tests/test_solvers.py +++ b/tests/test_solvers.py @@ -1,7 +1,9 @@ +from collections import namedtuple + +import numpy as np import pytest + import mici -import numpy as np -from collections import namedtuple @pytest.fixture(params=("direct", "steffensen")) @@ -21,10 +23,10 @@ def fixed_point_solver(request): def convergent_fixed_point_problem(request): if request.param == "babylonian": y = np.array([3.0, 5.0, 7.0]) - return FixedPointProblem(lambda x: (y / x + x) / 2, y ** 0.5, np.ones_like(y)) + return FixedPointProblem(lambda x: (y / x + x) / 2, y**0.5, np.ones_like(y)) elif request.param == "ratio": y = np.array([3.0, 5.0, 7.0]) - return FixedPointProblem(lambda x: (x + y) / (x + 1), y ** 0.5, np.ones_like(y)) + return FixedPointProblem(lambda x: (x + y) / (x + 1), y**0.5, np.ones_like(y)) elif request.param == "cosine": return FixedPointProblem( lambda x: np.cos(x), np.array([0.7390851332151607]), np.array([1.0]) @@ -36,9 +38,7 @@ def divergent_fixed_point_problem(request): if request.param == "doubling": return FixedPointProblem(lambda x: 2 * x, None, np.arange(3)) elif request.param == "quadratic": - return FixedPointProblem( - lambda x: 1 + x**2, None, np.arange(3) - ) + return FixedPointProblem(lambda x: 1 + x**2, None, np.arange(3)) @pytest.fixture(params=(1e-6, 1e-8, 1e-10)) @@ -106,4 +106,3 @@ def func(x): with pytest.raises(mici.errors.ConvergenceError): fixed_point_solver(func=func, x0=np.array([1.0])) - diff --git a/tests/test_stagers.py b/tests/test_stagers.py index 838cbf3..0262dd9 100644 --- a/tests/test_stagers.py +++ b/tests/test_stagers.py @@ -1,4 +1,5 @@ import pytest + import mici N_MAIN_ITER = 1 diff --git a/tests/test_states.py b/tests/test_states.py index 9c5d837..832a06a 100644 --- a/tests/test_states.py +++ b/tests/test_states.py @@ -1,8 +1,10 @@ -import pytest -from unittest.mock import Mock -from itertools import combinations import pickle +from itertools import combinations +from unittest.mock import Mock + import numpy as np +import pytest + import mici diff --git a/tests/test_utils.py b/tests/test_utils.py index c628064..54a21b2 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,7 +1,7 @@ -import pytest import numpy as np -import mici +import pytest +import mici SEED = 3046987125 @@ -41,7 +41,6 @@ def get_val(obj): class TestLogRepFloat: - VALS = sorted((0.0, 0.9, 1, 1.1, 2.1, 120.0)) FIXED = VALS[1]