From a41d5244b8ec78372751924ba0484ed02d57b2cb Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Tue, 7 Mar 2023 09:36:47 +0100 Subject: [PATCH] Bring back testing utilities used in downstream packages Follow up to * 534a9ae28a6bfea3bc31921cd92e4e4d2875b6c9 * e1d36cacb2cfb946f56872ee03e8678dc6ccf7f4 --- .github/workflows/tests.yml | 1 + .../contributing/implementing_distribution.md | 58 +++---- .../distributions/util.py => pymc/testing.py | 161 ++++++++++-------- tests/distributions/test_continuous.py | 8 +- tests/distributions/test_discrete.py | 4 +- tests/distributions/test_distribution.py | 2 +- tests/distributions/test_mixture.py | 10 +- tests/distributions/test_multivariate.py | 8 +- tests/distributions/test_simulator.py | 2 +- tests/distributions/test_timeseries.py | 3 +- tests/distributions/test_transform.py | 6 +- tests/distributions/test_truncated.py | 2 +- tests/helpers.py | 48 +----- tests/logprob/test_censoring.py | 2 +- tests/logprob/test_composite_logprob.py | 2 +- tests/logprob/test_cumsum.py | 2 +- tests/logprob/test_joint_logprob.py | 2 +- tests/logprob/test_mixture.py | 2 +- tests/logprob/test_scan.py | 2 +- tests/logprob/test_tensor.py | 2 +- tests/logprob/test_transforms.py | 2 +- tests/logprob/test_utils.py | 2 +- tests/ode/test_ode.py | 2 +- tests/sampler_fixtures.py | 2 +- tests/sampling/test_forward.py | 2 +- tests/sampling/test_mcmc.py | 2 +- tests/smc/test_smc.py | 3 +- tests/step_methods/test_compound.py | 3 +- tests/step_methods/test_metropolis.py | 7 +- tests/test_data.py | 2 +- tests/test_math.py | 3 +- tests/test_model.py | 2 +- tests/test_model_graph.py | 2 +- tests/test_pytensorf.py | 2 +- tests/test_testing.py | 34 ++++ tests/tuning/test_starting.py | 2 +- tests/variational/test_minibatch_rv.py | 2 +- 37 files changed, 207 insertions(+), 194 deletions(-) rename tests/distributions/util.py => pymc/testing.py (90%) create mode 100644 tests/test_testing.py diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 670f5895c23..306ea93232e 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -47,6 +47,7 @@ jobs: tests/test_func_utils.py tests/distributions/test_shape_utils.py tests/distributions/test_mixture.py + tests/test_testing.py - | tests/distributions/test_continuous.py diff --git a/docs/source/contributing/implementing_distribution.md b/docs/source/contributing/implementing_distribution.md index 78eec6b55fa..b48584f0969 100644 --- a/docs/source/contributing/implementing_distribution.md +++ b/docs/source/contributing/implementing_distribution.md @@ -240,10 +240,11 @@ Most tests can be accommodated by the default `BaseTestDistributionRandom` class 1. Shape variable inference is correct, via `check_rv_size` ```python -from tests.distributions.util import BaseTestDistributionRandom, seeded_scipy_distribution_builder -class TestBlah(BaseTestDistributionRandom): +from pymc.testing import BaseTestDistributionRandom, seeded_scipy_distribution_builder + +class TestBlah(BaseTestDistributionRandom): pymc_dist = pm.Blah # Parameters with which to test the blah pymc Distribution pymc_dist_params = {"param1": 0.25, "param2": 2.0} @@ -311,38 +312,36 @@ Tests for the `logp` and `logcdf` mostly make use of the helpers `check_logp`, ` `check_selfconsistency_discrete_logcdf` implemented in `~tests.distributions.util` ```python -from tests.helpers import select_by_precision -from tests.distributions.util import check_logp, check_logcdf, Domain + +from pymc.testing import Domain, check_logp, check_logcdf, select_by_precision R = Domain([-np.inf, -2.1, -1, -0.01, 0.0, 0.01, 1, 2.1, np.inf]) Rplus = Domain([0, 0.01, 0.1, 0.9, 0.99, 1, 1.5, 2, 100, np.inf]) - def test_blah(): - - check_logp( - pymc_dist=pm.Blah, - # Domain of the distribution values - domain=R, - # Domains of the distribution parameters - paramdomains={"mu": R, "sigma": Rplus}, - # Reference scipy (or other) logp function - scipy_logp = lambda value, mu, sigma: sp.norm.logpdf(value, mu, sigma), - # Number of decimal points expected to match between the pymc and reference functions - decimal=select_by_precision(float64=6, float32=3), - # Maximum number of combinations of domain * paramdomains to test - n_samples=100, - ) - - check_logcdf( - pymc_dist=pm.Blah, - domain=R, - paramdomains={"mu": R, "sigma": Rplus}, - scipy_logcdf=lambda value, mu, sigma: sp.norm.logcdf(value, mu, sigma), - decimal=select_by_precision(float64=6, float32=1), - n_samples=-1, - ) + check_logp( + pymc_dist=pm.Blah, + # Domain of the distribution values + domain=R, + # Domains of the distribution parameters + paramdomains={"mu": R, "sigma": Rplus}, + # Reference scipy (or other) logp function + scipy_logp=lambda value, mu, sigma: sp.norm.logpdf(value, mu, sigma), + # Number of decimal points expected to match between the pymc and reference functions + decimal=select_by_precision(float64=6, float32=3), + # Maximum number of combinations of domain * paramdomains to test + n_samples=100, + ) + + check_logcdf( + pymc_dist=pm.Blah, + domain=R, + paramdomains={"mu": R, "sigma": Rplus}, + scipy_logcdf=lambda value, mu, sigma: sp.norm.logcdf(value, mu, sigma), + decimal=select_by_precision(float64=6, float32=1), + n_samples=-1, + ) ``` @@ -382,7 +381,8 @@ which checks if: import pytest from pymc.distributions import Blah -from tests.distributions.util import assert_moment_is_expected +from pymc.testing import assert_moment_is_expected + @pytest.mark.parametrize( "param1, param2, size, expected", diff --git a/tests/distributions/util.py b/pymc/testing.py similarity index 90% rename from tests/distributions/util.py rename to pymc/testing.py index 9563d1e4f5d..e56f3fe88d5 100644 --- a/tests/distributions/util.py +++ b/pymc/testing.py @@ -14,29 +14,49 @@ import functools as ft import itertools as it -from contextlib import ExitStack as does_not_raise from typing import Callable, List, Optional import numpy as np -import numpy.random as nr -import numpy.testing as npt import pytensor -import pytensor.tensor as at +import pytensor.tensor as pt import pytest -import scipy.special as sp -import scipy.stats as st +from numpy import random as nr +from numpy import testing as npt from pytensor.compile.mode import Mode +from pytensor.graph.basic import ancestors +from pytensor.graph.rewriting.basic import in2out +from pytensor.tensor.random.op import RandomVariable +from scipy import special as sp +from scipy import stats as st import pymc as pm +from pymc import logcdf, logp from pymc.distributions.shape_utils import change_dist_size from pymc.initial_point import make_initial_point_fn -from pymc.logprob.abstract import logcdf -from pymc.logprob.joint_logprob import joint_logp, logp +from pymc.logprob import joint_logp from pymc.logprob.utils import ParameterValueError -from pymc.pytensorf import compile_pymc, floatX, intX -from tests.helpers import SeededTest, select_by_precision +from pymc.pytensorf import ( + compile_pymc, + floatX, + intX, + local_check_parameter_to_ninf_switch, +) + +# This mode can be used for tests where model compilations takes the bulk of the runtime +# AND where we don't care about posterior numerical or sampling stability (e.g., when +# all that matters are the shape of the draws or deterministic values of observed data). +# DO NOT USE UNLESS YOU HAVE A GOOD REASON TO! +fast_unstable_sampling_mode = ( + pytensor.compile.mode.FAST_COMPILE + # Remove slow rewrite phases + .excluding("canonicalize", "specialize") + # Include necessary rewrites for proper logp handling + .including("remove_TransformedVariables").register( + (in2out(local_check_parameter_to_ninf_switch), -1) + ) +) def product(domains, n_samples=-1): @@ -47,7 +67,8 @@ def product(domains, n_samples=-1): must be "domain-like", as in, have a `.vals` property n_samples: int, maximum samples to return. -1 to return whole product - Returns: + Returns + ------- list of the cartesian product of the domains """ try: @@ -119,22 +140,6 @@ def __neg__(self): return Domain([-v for v in self.vals], self.dtype, (-self.lower, -self.upper), self.shape) -@pytest.mark.parametrize( - "values, edges, expectation", - [ - ([], None, pytest.raises(IndexError)), - ([], (0, 0), pytest.raises(ValueError)), - ([0], None, pytest.raises(ValueError)), - ([0], (0, 0), does_not_raise()), - ([-1, 1], None, pytest.raises(ValueError)), - ([-1, 0, 1], None, does_not_raise()), - ], -) -def test_domain(values, edges, expectation): - with expectation: - Domain(values, edges=edges) - - class ProductDomain: def __init__(self, domains): self.vals = list(it.product(*(d.vals for d in domains))) @@ -203,24 +208,25 @@ def RandomPdMatrix(n): Rplusbig = Domain([0, 0.5, 0.9, 0.99, 1, 1.5, 2, 20, np.inf]) Rminusbig = Domain([-np.inf, -2, -1.5, -1, -0.99, -0.9, -0.5, -0.01, 0]) Unit = Domain([0, 0.001, 0.1, 0.5, 0.75, 0.99, 1]) - Circ = Domain([-np.pi, -2.1, -1, -0.01, 0.0, 0.01, 1, 2.1, np.pi]) - Runif = Domain([-np.inf, -0.4, 0, 0.4, np.inf]) Rdunif = Domain([-np.inf, -1, 0, 1, np.inf], "int64") Rplusunif = Domain([0, 0.5, np.inf]) Rplusdunif = Domain([0, 10, np.inf], "int64") - I = Domain([-np.inf, -3, -2, -1, 0, 1, 2, 3, np.inf], "int64") - NatSmall = Domain([0, 3, 4, 5, np.inf], "int64") Nat = Domain([0, 1, 2, 3, np.inf], "int64") NatBig = Domain([0, 1, 2, 3, 5000, np.inf], "int64") PosNat = Domain([1, 2, 3, np.inf], "int64") - Bool = Domain([0, 0, 1, 1], "int64") +def select_by_precision(float64, float32): + """Helper function to choose reasonable decimal cutoffs for different floatX modes.""" + decimal = float64 if pytensor.config.floatX == "float64" else float32 + return decimal + + def build_model(distfam, valuedomain, vardomains, extra_args=None): if extra_args is None: extra_args = {} @@ -228,9 +234,9 @@ def build_model(distfam, valuedomain, vardomains, extra_args=None): with pm.Model() as m: param_vars = {} for v, dom in vardomains.items(): - v_at = pytensor.shared(np.asarray(dom.vals[0])) - v_at.name = v - param_vars[v] = v_at + v_pt = pytensor.shared(np.asarray(dom.vals[0])) + v_pt.name = v + param_vars[v] = v_pt param_vars.update(extra_args) distfam( "value", @@ -295,10 +301,10 @@ def logp_reference(args): args.update(scipy_args) return scipy_logp(**args) - def _model_input_dict(model, param_vars, pt): + def _model_input_dict(model, param_vars, point): """Create a dict with only the necessary, transformed logp inputs.""" pt_d = {} - for k, v in pt.items(): + for k, v in point.items(): rv_var = model.named_vars.get(k) nv = param_vars.get(k, rv_var) nv = model.rvs_to_values.get(nv, nv) @@ -325,16 +331,16 @@ def _model_input_dict(model, param_vars, pt): # Test supported value and parameters domain matches scipy domains = paramdomains.copy() domains["value"] = domain - for pt in product(domains, n_samples=n_samples): - pt = dict(pt) - pt_d = _model_input_dict(model, param_vars, pt) + for point in product(domains, n_samples=n_samples): + point = dict(point) + pt_d = _model_input_dict(model, param_vars, point) pt_logp = pm.Point(pt_d, model=model) - pt_ref = pm.Point(pt, filter_model_vars=False, model=model) + pt_ref = pm.Point(point, filter_model_vars=False, model=model) npt.assert_almost_equal( logp_pymc(pt_logp), logp_reference(pt_ref), decimal=decimal, - err_msg=str(pt), + err_msg=str(point), ) valid_value = domain.vals[0] @@ -360,7 +366,7 @@ def _model_input_dict(model, param_vars, pt): if invalid_edge is None: continue test_params = valid_params.copy() # Shallow copy should be okay - test_params[invalid_param] = at.as_tensor_variable(invalid_edge) + test_params[invalid_param] = pt.as_tensor_variable(invalid_edge) # We need to remove `Assert`s introduced by checks like # `assert_negative_support` and disable test values; # otherwise, we won't be able to create the `RandomVariable` @@ -441,9 +447,6 @@ def check_logcdf( Whether to run test 2., which checks that pymc distribution logcdf returns -inf for invalid parameter values outside the supported domain edge - Returns - ------- - """ # Test pymc and scipy distributions match for values and parameters # within the supported domain edges (excluding edges) @@ -459,8 +462,8 @@ def check_logcdf( if decimal is None: decimal = select_by_precision(float64=6, float32=3) - for pt in product(domains, n_samples=n_samples): - params = dict(pt) + for point in product(domains, n_samples=n_samples): + params = dict(point) scipy_eval = scipy_logcdf(**params) value = params.pop("value") @@ -496,7 +499,7 @@ def check_logcdf( for invalid_edge in invalid_edges: if invalid_edge is not None: test_params = valid_params.copy() # Shallow copy should be okay - test_params[invalid_param] = at.as_tensor_variable(invalid_edge) + test_params[invalid_param] = pt.as_tensor_variable(invalid_edge) # We need to remove `Assert`s introduced by checks like # `assert_negative_support` and disable test values; # otherwise, we won't be able to create the @@ -559,8 +562,8 @@ def check_selfconsistency_discrete_logcdf( dist_logcdf = model.compile_fn(logcdf(rv, value)) dist_logp = model.compile_fn(logp(rv, value)) - for pt in product(domains, n_samples=n_samples): - params = dict(pt) + for point in product(domains, n_samples=n_samples): + params = dict(point) value = params.pop("value") values = np.arange(domain.lower, value + 1) @@ -573,7 +576,7 @@ def check_selfconsistency_discrete_logcdf( dist_logcdf({"value": value}), sp.logsumexp([dist_logp({"value": value}) for value in values]), decimal=decimal, - err_msg=str(pt), + err_msg=str(point), ) @@ -598,7 +601,7 @@ def assert_moment_is_expected(model, expected, check_finite_logp=True): logp_moment = ( joint_logp( (model["x"],), - rvs_to_values={model["x"]: at.constant(moment)}, + rvs_to_values={model["x"]: pt.constant(moment)}, rvs_to_transforms={}, )[0] .sum() @@ -607,7 +610,7 @@ def assert_moment_is_expected(model, expected, check_finite_logp=True): assert np.isfinite(logp_moment) -def pymc_random( +def continuous_random_tester( dist, paramdomains, ref_rand, @@ -629,12 +632,12 @@ def pymc_random( pymc_rand = compile_pymc([], model_dist) domains = paramdomains.copy() - for pt in product(domains, n_samples=100): - pt = pm.Point(pt, model=model) - pt.update(model_args) + for point in product(domains, n_samples=100): + point = pm.Point(point, model=model) + point.update(model_args) # Update the shared parameter variables in `param_vars` - for k, v in pt.items(): + for k, v in point.items(): nv = param_vars.get(k, model.named_vars.get(k)) if nv.name in param_vars: param_vars[nv.name].set_value(v) @@ -645,13 +648,13 @@ def pymc_random( f = fails while p <= alpha and f > 0: s0 = pymc_rand() - s1 = floatX(ref_rand(size=size, **pt)) + s1 = floatX(ref_rand(size=size, **point)) _, p = st.ks_2samp(np.atleast_1d(s0).flatten(), np.atleast_1d(s1).flatten()) f -= 1 - assert p > alpha, str(pt) + assert p > alpha, str(point) -def pymc_random_discrete( +def discrete_random_tester( dist, paramdomains, valuedomain=None, @@ -668,12 +671,12 @@ def pymc_random_discrete( pymc_rand = compile_pymc([], model_dist) domains = paramdomains.copy() - for pt in product(domains, n_samples=100): - pt = pm.Point(pt, model=model) + for point in product(domains, n_samples=100): + point = pm.Point(point, model=model) p = alpha # Update the shared parameter variables in `param_vars` - for k, v in pt.items(): + for k, v in point.items(): nv = param_vars.get(k, model.named_vars.get(k)) if nv.name in param_vars: param_vars[nv.name].set_value(v) @@ -683,7 +686,7 @@ def pymc_random_discrete( f = fails while p <= alpha and f > 0: o = pymc_rand() - e = intX(ref_rand(size=size, **pt)) + e = intX(ref_rand(size=size, **point)) o = np.atleast_1d(o).flatten() e = np.atleast_1d(e).flatten() bins = min(20, max(len(set(e)), len(set(o)))) @@ -695,12 +698,29 @@ def pymc_random_discrete( else: _, p = st.chisquare(observed + 1, expected + 1) f -= 1 - assert p > alpha, str(pt) + assert p > alpha, str(point) + + +class SeededTest: + random_seed = 20160911 + random_state = None + + @classmethod + def setup_class(cls): + nr.seed(cls.random_seed) + + def setup_method(self): + nr.seed(self.random_seed) + + def get_random_state(self, reset=False): + if self.random_state is None or reset: + self.random_state = nr.RandomState(self.random_seed) + return self.random_state class BaseTestDistributionRandom(SeededTest): """ - This class provides a base for tests that new RandomVariables are correctly + Base class for tests that new RandomVariables are correctly implemented, and that the mapping of parameters between the PyMC Distribution and the respective RandomVariable is correct. @@ -762,7 +782,7 @@ class BaseTestDistributionRandom(SeededTest): reference_dist: Optional[Callable] = None reference_dist_params: Optional[dict] = None expected_rv_op_params: Optional[dict] = None - checks_to_run = [] + checks_to_run: List[str] = [] size = 15 decimal = select_by_precision(float64=6, float32=3) @@ -864,3 +884,8 @@ def seeded_numpy_distribution_builder(dist_name: str) -> Callable: return lambda self: ft.partial( getattr(np.random.RandomState, dist_name), self.get_random_state() ) + + +def assert_no_rvs(var): + assert not any(isinstance(v.owner.op, RandomVariable) for v in ancestors([var]) if v.owner) + return var diff --git a/tests/distributions/test_continuous.py b/tests/distributions/test_continuous.py index dc4c88b567f..3713566bdbd 100644 --- a/tests/distributions/test_continuous.py +++ b/tests/distributions/test_continuous.py @@ -32,7 +32,7 @@ from pymc.logprob.joint_logprob import logp from pymc.logprob.utils import ParameterValueError from pymc.pytensorf import floatX -from tests.distributions.util import ( +from pymc.testing import ( BaseTestDistributionRandom, Circ, Domain, @@ -45,11 +45,11 @@ assert_moment_is_expected, check_logcdf, check_logp, - pymc_random, + continuous_random_tester, seeded_numpy_distribution_builder, seeded_scipy_distribution_builder, + select_by_precision, ) -from tests.helpers import select_by_precision from tests.logprob.utils import create_pytensor_params, scipy_logprob_tester try: @@ -2259,7 +2259,7 @@ def dist(cls, **kwargs): pdf_points = st.norm.pdf(x_points, loc=mu, scale=sigma) return super().dist(x_points=x_points, pdf_points=pdf_points, **kwargs) - pymc_random( + continuous_random_tester( TestedInterpolated, {}, extra_args={"rng": pytensor.shared(rng)}, diff --git a/tests/distributions/test_discrete.py b/tests/distributions/test_discrete.py index 7b15a6b5875..7c929e3f1f0 100644 --- a/tests/distributions/test_discrete.py +++ b/tests/distributions/test_discrete.py @@ -33,8 +33,7 @@ from pymc.logprob.joint_logprob import logp from pymc.logprob.utils import ParameterValueError from pymc.pytensorf import floatX -from pymc.vartypes import discrete_types -from tests.distributions.util import ( +from pymc.testing import ( BaseTestDistributionRandom, Bool, Domain, @@ -58,6 +57,7 @@ seeded_numpy_distribution_builder, seeded_scipy_distribution_builder, ) +from pymc.vartypes import discrete_types from tests.logprob.utils import create_pytensor_params, scipy_logprob_tester diff --git a/tests/distributions/test_distribution.py b/tests/distributions/test_distribution.py index 7251676a2f2..44af5a7d428 100644 --- a/tests/distributions/test_distribution.py +++ b/tests/distributions/test_distribution.py @@ -50,8 +50,8 @@ from pymc.logprob.joint_logprob import logp from pymc.model import Model from pymc.sampling import draw, sample +from pymc.testing import assert_moment_is_expected from pymc.util import _FutureWarningValidatingScratchpad -from tests.distributions.util import assert_moment_is_expected class TestBugfixes: diff --git a/tests/distributions/test_mixture.py b/tests/distributions/test_mixture.py index dc6a4c74466..5afae7cb9ca 100644 --- a/tests/distributions/test_mixture.py +++ b/tests/distributions/test_mixture.py @@ -62,13 +62,13 @@ ) from pymc.sampling.mcmc import sample from pymc.step_methods import Metropolis -from tests.distributions.util import ( +from pymc.testing import ( Domain, + SeededTest, Simplex, assert_moment_is_expected, - pymc_random, + continuous_random_tester, ) -from tests.helpers import SeededTest def generate_normal_mixture_data(w, mu, sigma, size=1000): @@ -850,7 +850,7 @@ def ref_rand(size, w, mu, sigma): component = np.random.choice(w.size, size=size, p=w) return np.random.normal(mu[component], sigma[component], size=size) - pymc_random( + continuous_random_tester( NormalMixture, { "w": Simplex(2), @@ -861,7 +861,7 @@ def ref_rand(size, w, mu, sigma): size=1000, ref_rand=ref_rand, ) - pymc_random( + continuous_random_tester( NormalMixture, { "w": Simplex(3), diff --git a/tests/distributions/test_multivariate.py b/tests/distributions/test_multivariate.py index 5e5385a1457..df79e6b841c 100644 --- a/tests/distributions/test_multivariate.py +++ b/tests/distributions/test_multivariate.py @@ -42,7 +42,7 @@ from pymc.math import kronecker from pymc.pytensorf import compile_pymc, floatX, intX from pymc.sampling.forward import draw -from tests.distributions.util import ( +from pymc.testing import ( BaseTestDistributionRandom, Domain, Nat, @@ -54,10 +54,10 @@ Vector, assert_moment_is_expected, check_logp, - pymc_random, + continuous_random_tester, seeded_numpy_distribution_builder, + select_by_precision, ) -from tests.helpers import select_by_precision def betafn(a): @@ -2010,7 +2010,7 @@ def ref_rand(size, n, eta): beta = eta - 1 + n / 2 return (st.beta.rvs(size=(size, shape), a=beta, b=beta) - 0.5) * 2 - pymc_random( + continuous_random_tester( pm.LKJCorr, { "n": Domain([2, 10, 50], edges=(None, None)), diff --git a/tests/distributions/test_simulator.py b/tests/distributions/test_simulator.py index 7e06d2eb0f3..7172526601b 100644 --- a/tests/distributions/test_simulator.py +++ b/tests/distributions/test_simulator.py @@ -32,7 +32,7 @@ from pymc.initial_point import make_initial_point_fn from pymc.pytensorf import compile_pymc from pymc.smc.kernels import IMH -from tests.helpers import SeededTest +from pymc.testing import SeededTest class TestSimulator(SeededTest): diff --git a/tests/distributions/test_timeseries.py b/tests/distributions/test_timeseries.py index d49268e1a7a..3a045778001 100644 --- a/tests/distributions/test_timeseries.py +++ b/tests/distributions/test_timeseries.py @@ -44,8 +44,7 @@ from pymc.pytensorf import floatX from pymc.sampling.forward import draw, sample_posterior_predictive from pymc.sampling.mcmc import sample -from tests.distributions.util import assert_moment_is_expected -from tests.helpers import select_by_precision +from pymc.testing import assert_moment_is_expected, select_by_precision # Turn all warnings into errors for this module # Ignoring NumPy deprecation warning tracked in https://github.com/pymc-devs/pytensor/issues/146 diff --git a/tests/distributions/test_transform.py b/tests/distributions/test_transform.py index 9d47a59e10b..3f0767f4613 100644 --- a/tests/distributions/test_transform.py +++ b/tests/distributions/test_transform.py @@ -27,20 +27,20 @@ from pymc.logprob.joint_logprob import joint_logp from pymc.pytensorf import floatX, jacobian -from tests.checks import close_to, close_to_logical -from tests.distributions.util import ( +from pymc.testing import ( Circ, MultiSimplex, R, Rminusbig, Rplusbig, + SeededTest, Simplex, SortedVector, Unit, UnitSortedVector, Vector, ) -from tests.helpers import SeededTest +from tests.checks import close_to, close_to_logical # some transforms (stick breaking) require addition of small slack in order to be numerically # stable. The minimal addable slack for float32 is higher thus we need to be less strict diff --git a/tests/distributions/test_truncated.py b/tests/distributions/test_truncated.py index 39ebbb6cc98..ec7e886aeea 100644 --- a/tests/distributions/test_truncated.py +++ b/tests/distributions/test_truncated.py @@ -29,7 +29,7 @@ from pymc.logprob.joint_logprob import logp from pymc.logprob.transforms import IntervalTransform from pymc.logprob.utils import ParameterValueError -from tests.distributions.util import assert_moment_is_expected +from pymc.testing import assert_moment_is_expected class IcdfNormalRV(NormalRV): diff --git a/tests/helpers.py b/tests/helpers.py index 343861787b3..f3aa4f914e3 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -24,34 +24,14 @@ import pytensor from pytensor.gradient import verify_grad as at_verify_grad -from pytensor.graph import ancestors -from pytensor.graph.rewriting.basic import in2out -from pytensor.tensor.random.op import RandomVariable import pymc as pm -from pymc.pytensorf import local_check_parameter_to_ninf_switch +from pymc.testing import fast_unstable_sampling_mode from tests.checks import close_to from tests.models import mv_simple, mv_simple_coarse -class SeededTest: - random_seed = 20160911 - random_state = None - - @classmethod - def setup_class(cls): - nr.seed(cls.random_seed) - - def setup_method(self): - nr.seed(self.random_seed) - - def get_random_state(self, reset=False): - if self.random_state is None or reset: - self.random_state = nr.RandomState(self.random_seed) - return self.random_state - - class LoggingHandler(BufferingHandler): def __init__(self, matcher): # BufferingHandler takes a "capacity" argument @@ -112,12 +92,6 @@ def match_value(self, k, dv, v): return result -def select_by_precision(float64, float32): - """Helper function to choose reasonable decimal cutoffs for different floatX modes.""" - decimal = float64 if pytensor.config.floatX == "float64" else float32 - return decimal - - @contextlib.contextmanager def not_raises(): yield @@ -137,21 +111,6 @@ def assert_random_state_equal(state1, state2): assert field1 == field2 -# This mode can be used for tests where model compilations takes the bulk of the runtime -# AND where we don't care about posterior numerical or sampling stability (e.g., when -# all that matters are the shape of the draws or deterministic values of observed data). -# DO NOT USE UNLESS YOU HAVE A GOOD REASON TO! -fast_unstable_sampling_mode = ( - pytensor.compile.mode.FAST_COMPILE - # Remove slow rewrite phases - .excluding("canonicalize", "specialize") - # Include necessary rewrites for proper logp handling - .including("remove_TransformedVariables").register( - (in2out(local_check_parameter_to_ninf_switch), -1) - ) -) - - class StepMethodTester: def setup_class(self): self.temp_dir = tempfile.mkdtemp() @@ -213,8 +172,3 @@ def continuous_steps(self, step, step_kwargs): assert {m.rvs_to_values[c1], m.rvs_to_values[c2]} == set( step([c1, c2], **step_kwargs).vars ) - - -def assert_no_rvs(var): - assert not any(isinstance(v.owner.op, RandomVariable) for v in ancestors([var]) if v.owner) - return var diff --git a/tests/logprob/test_censoring.py b/tests/logprob/test_censoring.py index a5798e898ce..c18d902b39f 100644 --- a/tests/logprob/test_censoring.py +++ b/tests/logprob/test_censoring.py @@ -43,7 +43,7 @@ from pymc.logprob import factorized_joint_logprob from pymc.logprob.transforms import LogTransform, TransformValuesRewrite -from tests.helpers import assert_no_rvs +from pymc.testing import assert_no_rvs from tests.logprob.utils import joint_logprob diff --git a/tests/logprob/test_composite_logprob.py b/tests/logprob/test_composite_logprob.py index 9f805a62f6b..bfdc2d070e7 100644 --- a/tests/logprob/test_composite_logprob.py +++ b/tests/logprob/test_composite_logprob.py @@ -41,7 +41,7 @@ from pymc.logprob.censoring import MeasurableClip from pymc.logprob.rewriting import construct_ir_fgraph -from tests.helpers import assert_no_rvs +from pymc.testing import assert_no_rvs from tests.logprob.utils import joint_logprob diff --git a/tests/logprob/test_cumsum.py b/tests/logprob/test_cumsum.py index 6a25f6d9839..4d0b6c45650 100644 --- a/tests/logprob/test_cumsum.py +++ b/tests/logprob/test_cumsum.py @@ -40,7 +40,7 @@ import pytest import scipy.stats as st -from tests.helpers import assert_no_rvs +from pymc.testing import assert_no_rvs from tests.logprob.utils import joint_logprob diff --git a/tests/logprob/test_joint_logprob.py b/tests/logprob/test_joint_logprob.py index f0336821b76..74bfdd3f53c 100644 --- a/tests/logprob/test_joint_logprob.py +++ b/tests/logprob/test_joint_logprob.py @@ -58,7 +58,7 @@ from pymc.logprob.abstract import logprob from pymc.logprob.joint_logprob import factorized_joint_logprob, joint_logp from pymc.logprob.utils import rvs_to_value_vars, walk_model -from tests.helpers import assert_no_rvs +from pymc.testing import assert_no_rvs from tests.logprob.utils import joint_logprob diff --git a/tests/logprob/test_mixture.py b/tests/logprob/test_mixture.py index bad48d5d116..546e5611bb6 100644 --- a/tests/logprob/test_mixture.py +++ b/tests/logprob/test_mixture.py @@ -49,7 +49,7 @@ from pymc.logprob.mixture import MixtureRV, expand_indices from pymc.logprob.rewriting import construct_ir_fgraph from pymc.logprob.utils import dirac_delta -from tests.helpers import assert_no_rvs +from pymc.testing import assert_no_rvs from tests.logprob.utils import joint_logprob, scipy_logprob diff --git a/tests/logprob/test_scan.py b/tests/logprob/test_scan.py index 551bb514719..26f809fc026 100644 --- a/tests/logprob/test_scan.py +++ b/tests/logprob/test_scan.py @@ -51,7 +51,7 @@ convert_outer_out_to_in, get_random_outer_outputs, ) -from tests.helpers import assert_no_rvs +from pymc.testing import assert_no_rvs from tests.logprob.utils import joint_logprob diff --git a/tests/logprob/test_tensor.py b/tests/logprob/test_tensor.py index a0b88107684..4262942cfe1 100644 --- a/tests/logprob/test_tensor.py +++ b/tests/logprob/test_tensor.py @@ -48,7 +48,7 @@ from pymc.logprob import factorized_joint_logprob from pymc.logprob.rewriting import logprob_rewrites_db from pymc.logprob.tensor import naive_bcast_rv_lift -from tests.helpers import assert_no_rvs +from pymc.testing import assert_no_rvs from tests.logprob.utils import joint_logprob diff --git a/tests/logprob/test_transforms.py b/tests/logprob/test_transforms.py index bb508c4cdbc..92e015d63d4 100644 --- a/tests/logprob/test_transforms.py +++ b/tests/logprob/test_transforms.py @@ -63,7 +63,7 @@ TransformValuesRewrite, transformed_variable, ) -from tests.helpers import assert_no_rvs +from pymc.testing import assert_no_rvs from tests.logprob.utils import joint_logprob diff --git a/tests/logprob/test_utils.py b/tests/logprob/test_utils.py index 8fe398195c3..7300e2d6823 100644 --- a/tests/logprob/test_utils.py +++ b/tests/logprob/test_utils.py @@ -57,7 +57,7 @@ rvs_to_value_vars, walk_model, ) -from tests.helpers import assert_no_rvs +from pymc.testing import assert_no_rvs from tests.logprob.utils import create_pytensor_params, scipy_logprob_tester diff --git a/tests/ode/test_ode.py b/tests/ode/test_ode.py index 599ed256ded..4b146eac2de 100644 --- a/tests/ode/test_ode.py +++ b/tests/ode/test_ode.py @@ -23,7 +23,7 @@ import pymc as pm from pymc.ode import DifferentialEquation -from tests.helpers import fast_unstable_sampling_mode +from pymc.testing import fast_unstable_sampling_mode def test_simulate(): diff --git a/tests/sampler_fixtures.py b/tests/sampler_fixtures.py index 9a7f1aead37..73af1eba091 100644 --- a/tests/sampler_fixtures.py +++ b/tests/sampler_fixtures.py @@ -21,8 +21,8 @@ import pymc as pm from pymc.backends.arviz import to_inference_data +from pymc.testing import SeededTest from pymc.util import get_var_name -from tests.helpers import SeededTest class KnownMean: diff --git a/tests/sampling/test_forward.py b/tests/sampling/test_forward.py index 751cc9a6994..fba03614eac 100644 --- a/tests/sampling/test_forward.py +++ b/tests/sampling/test_forward.py @@ -40,7 +40,7 @@ get_vars_in_point_list, observed_dependent_deterministics, ) -from tests.helpers import SeededTest, fast_unstable_sampling_mode +from pymc.testing import SeededTest, fast_unstable_sampling_mode class TestDraw(SeededTest): diff --git a/tests/sampling/test_mcmc.py b/tests/sampling/test_mcmc.py index 8f7e06fb5c9..dbdc0d2a407 100644 --- a/tests/sampling/test_mcmc.py +++ b/tests/sampling/test_mcmc.py @@ -45,7 +45,7 @@ Metropolis, Slice, ) -from tests.helpers import SeededTest, fast_unstable_sampling_mode +from pymc.testing import SeededTest, fast_unstable_sampling_mode from tests.models import simple_init diff --git a/tests/smc/test_smc.py b/tests/smc/test_smc.py index c606031c866..a087212ac03 100644 --- a/tests/smc/test_smc.py +++ b/tests/smc/test_smc.py @@ -26,7 +26,8 @@ from pymc.backends.base import MultiTrace from pymc.pytensorf import floatX from pymc.smc.kernels import IMH, systematic_resampling -from tests.helpers import SeededTest, assert_random_state_equal +from pymc.testing import SeededTest +from tests.helpers import assert_random_state_equal class TestSMC(SeededTest): diff --git a/tests/step_methods/test_compound.py b/tests/step_methods/test_compound.py index ba9d90634d0..93cbecbc0d9 100644 --- a/tests/step_methods/test_compound.py +++ b/tests/step_methods/test_compound.py @@ -31,7 +31,8 @@ get_stats_dtypes_shapes_from_steps, infer_warn_stats_info, ) -from tests.helpers import StepMethodTester, fast_unstable_sampling_mode +from pymc.testing import fast_unstable_sampling_mode +from tests.helpers import StepMethodTester from tests.models import simple_2model_continuous diff --git a/tests/step_methods/test_metropolis.py b/tests/step_methods/test_metropolis.py index cee18284081..8da8faac481 100644 --- a/tests/step_methods/test_metropolis.py +++ b/tests/step_methods/test_metropolis.py @@ -31,12 +31,9 @@ MultivariateNormalProposal, NormalProposal, ) +from pymc.testing import fast_unstable_sampling_mode from tests import sampler_fixtures as sf -from tests.helpers import ( - RVsAssignmentStepsTester, - StepMethodTester, - fast_unstable_sampling_mode, -) +from tests.helpers import RVsAssignmentStepsTester, StepMethodTester from tests.models import mv_simple, mv_simple_discrete, simple_categorical diff --git a/tests/test_data.py b/tests/test_data.py index 09d175de4b1..7c4da8db899 100644 --- a/tests/test_data.py +++ b/tests/test_data.py @@ -28,7 +28,7 @@ from pymc.data import is_minibatch from pymc.pytensorf import GeneratorOp, floatX -from tests.helpers import SeededTest +from pymc.testing import SeededTest class TestData(SeededTest): diff --git a/tests/test_math.py b/tests/test_math.py index 3d2120a5f70..624692caa5a 100644 --- a/tests/test_math.py +++ b/tests/test_math.py @@ -39,7 +39,8 @@ softmax, ) from pymc.pytensorf import floatX -from tests.helpers import SeededTest, verify_grad +from pymc.testing import SeededTest +from tests.helpers import verify_grad def test_kronecker(): diff --git a/tests/test_model.py b/tests/test_model.py index 2589e23c4c6..37d061a6160 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -45,9 +45,9 @@ from pymc.logprob.joint_logprob import joint_logp from pymc.logprob.transforms import IntervalTransform from pymc.model import Point, ValueGradFunction, modelcontext +from pymc.testing import SeededTest from pymc.util import _FutureWarningValidatingScratchpad from pymc.variational.minibatch_rv import MinibatchRandomVariable -from tests.helpers import SeededTest from tests.models import simple_model diff --git a/tests/test_model_graph.py b/tests/test_model_graph.py index 08733958b81..40b7293d4c3 100644 --- a/tests/test_model_graph.py +++ b/tests/test_model_graph.py @@ -25,7 +25,7 @@ from pymc.exceptions import ImputationWarning from pymc.model_graph import ModelGraph, model_to_graphviz, model_to_networkx -from tests.helpers import SeededTest +from pymc.testing import SeededTest def school_model(): diff --git a/tests/test_pytensorf.py b/tests/test_pytensorf.py index 3fe9440be3e..ddaf86ed47c 100644 --- a/tests/test_pytensorf.py +++ b/tests/test_pytensorf.py @@ -48,8 +48,8 @@ rvs_to_value_vars, walk_model, ) +from pymc.testing import assert_no_rvs from pymc.vartypes import int_types -from tests.helpers import assert_no_rvs @pytest.mark.parametrize( diff --git a/tests/test_testing.py b/tests/test_testing.py new file mode 100644 index 00000000000..b23e97a1d2b --- /dev/null +++ b/tests/test_testing.py @@ -0,0 +1,34 @@ +# Copyright 2023 The PyMC Developers +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from contextlib import ExitStack as does_not_raise + +import pytest + +from pymc.testing import Domain + + +@pytest.mark.parametrize( + "values, edges, expectation", + [ + ([], None, pytest.raises(IndexError)), + ([], (0, 0), pytest.raises(ValueError)), + ([0], None, pytest.raises(ValueError)), + ([0], (0, 0), does_not_raise()), + ([-1, 1], None, pytest.raises(ValueError)), + ([-1, 0, 1], None, does_not_raise()), + ], +) +def test_domain(values, edges, expectation): + with expectation: + Domain(values, edges=edges) diff --git a/tests/tuning/test_starting.py b/tests/tuning/test_starting.py index cdff83a22e2..4e7a3540ee2 100644 --- a/tests/tuning/test_starting.py +++ b/tests/tuning/test_starting.py @@ -20,10 +20,10 @@ from pymc.exceptions import ImputationWarning from pymc.step_methods.metropolis import tune +from pymc.testing import select_by_precision from pymc.tuning import find_MAP from tests import models from tests.checks import close_to -from tests.helpers import select_by_precision from tests.models import non_normal, simple_arbitrary_det, simple_model diff --git a/tests/variational/test_minibatch_rv.py b/tests/variational/test_minibatch_rv.py index 7f0a1d4dc45..8246c16ca39 100644 --- a/tests/variational/test_minibatch_rv.py +++ b/tests/variational/test_minibatch_rv.py @@ -20,8 +20,8 @@ import pymc as pm from pymc import Normal, draw +from pymc.testing import select_by_precision from pymc.variational.minibatch_rv import create_minibatch_rv -from tests.helpers import select_by_precision from tests.test_data import gen1, gen2