Distribute test_step into step_methods

Author: Armavica

.github/workflows/tests.yml

@@ -148,7 +148,7 @@ jobs:
         python-version: ["3.8"]
         test-subset:
           - pymc/tests/variational/test_approximations.py pymc/tests/variational/test_callbacks.py pymc/tests/variational/test_inference.py pymc/tests/variational/test_opvi.py pymc/tests/test_initial_point.py
-          - pymc/tests/test_model.py pymc/tests/test_step.py
+          - pymc/tests/test_model.py pymc/tests/step_methods/test_compound.py pymc/tests/step_methods/hmc/test_hmc.py
           - pymc/tests/gp/test_cov.py pymc/tests/gp/test_gp.py pymc/tests/gp/test_mean.py pymc/tests/gp/test_util.py pymc/tests/ode/test_ode.py pymc/tests/ode/test_utils.py pymc/tests/test_smc.py pymc/tests/test_parallel_sampling.py
           - pymc/tests/test_sampling.py pymc/tests/step_methods/test_metropolis.py pymc/tests/step_methods/test_slicer.py pymc/tests/step_methods/hmc/test_nuts.py
 

pymc/tests/helpers.py

@@ -13,6 +13,9 @@
 #   limitations under the License.
 
 import contextlib
+import shutil
+import tempfile
+import warnings
 
 from logging.handlers import BufferingHandler
 
@@ -24,7 +27,11 @@
 from aesara.graph.rewriting.basic import in2out
 from aesara.sandbox.rng_mrg import MRG_RandomStream as RandomStream
 
+import pymc as pm
+
 from pymc.aesaraf import at_rng, local_check_parameter_to_ninf_switch, set_at_rng
+from pymc.tests.checks import close_to
+from pymc.tests.models import mv_simple, mv_simple_coarse
 
 
 class SeededTest:
@@ -148,3 +155,66 @@ def assert_random_state_equal(state1, state2):
         (in2out(local_check_parameter_to_ninf_switch), -1)
     )
 )
+
+
+class StepMethodTester:
+    def setup_class(self):
+        self.temp_dir = tempfile.mkdtemp()
+
+    def teardown_class(self):
+        shutil.rmtree(self.temp_dir)
+
+    def check_stat(self, check, idata, name):
+        group = idata.posterior
+        for (var, stat, value, bound) in check:
+            s = stat(group[var].sel(chain=0), axis=0)
+            close_to(s, value, bound, name)
+
+    def check_stat_dtype(self, step, idata):
+        # TODO: This check does not confirm the announced dtypes are correct as the
+        #  sampling machinery will convert them automatically.
+        for stats_dtypes in getattr(step, "stats_dtypes", []):
+            for stat, dtype in stats_dtypes.items():
+                if stat == "tune":
+                    continue
+                assert idata.sample_stats[stat].dtype == np.dtype(dtype)
+
+    def step_continuous(self, step_fn, draws):
+        start, model, (mu, C) = mv_simple()
+        unc = np.diag(C) ** 0.5
+        check = (("x", np.mean, mu, unc / 10), ("x", np.std, unc, unc / 10))
+        _, model_coarse, _ = mv_simple_coarse()
+        with model:
+            step = step_fn(C, model_coarse)
+            with warnings.catch_warnings():
+                warnings.filterwarnings("ignore", "More chains .* than draws .*", UserWarning)
+                idata = pm.sample(
+                    tune=1000,
+                    draws=draws,
+                    chains=1,
+                    step=step,
+                    initvals=start,
+                    model=model,
+                    random_seed=1,
+                )
+            self.check_stat(check, idata, step.__class__.__name__)
+            self.check_stat_dtype(idata, step)
+
+
+class RVsAssignmentStepsTester:
+    """
+    Test that step methods convert input RVs to respective value vars
+    Step methods are tested with one and two variables to cover compound
+    the special branches in `BlockedStep.__new__`
+    """
+
+    def continuous_steps(self, step, step_kwargs):
+        with pm.Model() as m:
+            c1 = pm.HalfNormal("c1")
+            c2 = pm.HalfNormal("c2")
+
+            with aesara.config.change_flags(mode=fast_unstable_sampling_mode):
+                assert [m.rvs_to_values[c1]] == step([c1], **step_kwargs).vars
+            assert {m.rvs_to_values[c1], m.rvs_to_values[c2]} == set(
+                step([c1, c2], **step_kwargs).vars
+            )

pymc/tests/step_methods/hmc/test_hmc.py

@@ -0,0 +1,36 @@
+#   Copyright 2020 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.
+
+import pytest
+
+from pymc.step_methods.hmc import HamiltonianMC
+from pymc.tests.helpers import RVsAssignmentStepsTester, StepMethodTester
+
+
+class TestStepHamiltonianMC(StepMethodTester):
+    @pytest.mark.parametrize(
+        "step_fn, draws",
+        [
+            (lambda C, _: HamiltonianMC(scaling=C, is_cov=True, blocked=False), 1000),
+            (lambda C, _: HamiltonianMC(scaling=C, is_cov=True), 1000),
+        ],
+    )
+    def test_step_continuous(self, step_fn, draws):
+        self.step_continuous(step_fn, draws)
+
+
+class TestRVsAssignmentHamiltonianMC(RVsAssignmentStepsTester):
+    @pytest.mark.parametrize("step, step_kwargs", [(HamiltonianMC, {})])
+    def test_continuous_steps(self, step, step_kwargs):
+        self.continuous_steps(step, step_kwargs)

pymc/tests/step_methods/hmc/test_nuts.py

@@ -12,9 +12,20 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 
+import sys
+import warnings
+
+import aesara.tensor as at
+import numpy as np
 import pytest
 
+import pymc as pm
+
+from pymc.aesaraf import floatX
+from pymc.exceptions import SamplingError
+from pymc.step_methods.hmc import NUTS
 from pymc.tests import sampler_fixtures as sf
+from pymc.tests.helpers import RVsAssignmentStepsTester, StepMethodTester
 
 
 class TestNUTSUniform(sf.NutsFixture, sf.UniformFixture):
@@ -81,3 +92,116 @@ class TestNUTSLKJCholeskyCov(sf.NutsFixture, sf.LKJCholeskyCovFixture):
     burn = 0
     chains = 2
     min_n_eff = 200
+
+
+class TestNutsCheckTrace:
+    def test_multiple_samplers(self, caplog):
+        with pm.Model():
+            prob = pm.Beta("prob", alpha=5.0, beta=3.0)
+            pm.Binomial("outcome", n=1, p=prob)
+            caplog.clear()
+            with warnings.catch_warnings():
+                warnings.filterwarnings("ignore", ".*number of samples.*", UserWarning)
+                pm.sample(3, tune=2, discard_tuned_samples=False, n_init=None, chains=1)
+            messages = [msg.msg for msg in caplog.records]
+            assert all("boolean index did not" not in msg for msg in messages)
+
+    def test_bad_init_nonparallel(self):
+        with pm.Model():
+            pm.HalfNormal("a", sigma=1, initval=-1, transform=None)
+            with pytest.raises(SamplingError) as error:
+                pm.sample(chains=1, random_seed=1)
+            error.match("Initial evaluation")
+
+    @pytest.mark.skipif(sys.version_info < (3, 6), reason="requires python3.6 or higher")
+    def test_bad_init_parallel(self):
+        with pm.Model():
+            pm.HalfNormal("a", sigma=1, initval=-1, transform=None)
+            with pytest.raises(SamplingError) as error:
+                pm.sample(cores=2, random_seed=1)
+            error.match("Initial evaluation")
+
+    def test_linalg(self, caplog):
+        with pm.Model():
+            a = pm.Normal("a", size=2, initval=floatX(np.zeros(2)))
+            a = at.switch(a > 0, np.inf, a)
+            b = at.slinalg.solve(floatX(np.eye(2)), a, check_finite=False)
+            pm.Normal("c", mu=b, size=2, initval=floatX(np.r_[0.0, 0.0]))
+            caplog.clear()
+            with warnings.catch_warnings():
+                warnings.filterwarnings("ignore", ".*number of samples.*", UserWarning)
+                trace = pm.sample(20, tune=5, chains=2, return_inferencedata=False, random_seed=526)
+            warns = [msg.msg for msg in caplog.records]
+            assert np.any(trace["diverging"])
+            assert (
+                any("divergence after tuning" in warn for warn in warns)
+                or any("divergences after tuning" in warn for warn in warns)
+                or any("only diverging samples" in warn for warn in warns)
+            )
+
+            with pytest.raises(ValueError) as error:
+                trace.report.raise_ok()
+            error.match("issues during sampling")
+
+            assert not trace.report.ok
+
+    def test_sampler_stats(self):
+        with pm.Model() as model:
+            pm.Normal("x", mu=0, sigma=1)
+            with warnings.catch_warnings():
+                warnings.filterwarnings("ignore", ".*number of samples.*", UserWarning)
+                trace = pm.sample(draws=10, tune=1, chains=1, return_inferencedata=False)
+
+        # Assert stats exist and have the correct shape.
+        expected_stat_names = {
+            "depth",
+            "diverging",
+            "energy",
+            "energy_error",
+            "model_logp",
+            "max_energy_error",
+            "mean_tree_accept",
+            "step_size",
+            "step_size_bar",
+            "tree_size",
+            "tune",
+            "perf_counter_diff",
+            "perf_counter_start",
+            "process_time_diff",
+            "index_in_trajectory",
+            "largest_eigval",
+            "smallest_eigval",
+        }
+        assert trace.stat_names == expected_stat_names
+        for varname in trace.stat_names:
+            assert trace.get_sampler_stats(varname).shape == (10,)
+
+        # Assert model logp is computed correctly: computing post-sampling
+        # and tracking while sampling should give same results.
+        model_logp_fn = model.compile_logp()
+        model_logp_ = np.array(
+            [
+                model_logp_fn(trace.point(i, chain=c))
+                for c in trace.chains
+                for i in range(len(trace))
+            ]
+        )
+        assert (trace.model_logp == model_logp_).all()
+
+
+class TestStepNUTS(StepMethodTester):
+    @pytest.mark.parametrize(
+        "step_fn, draws",
+        [
+            (lambda C, _: NUTS(scaling=C, is_cov=True, blocked=False), 1000),
+            (lambda C, _: NUTS(scaling=C, is_cov=True), 1000),
+        ],
+    )
+    def test_step_continuous(self, step_fn, draws):
+        self.step_continuous(step_fn, draws)
+
+
+class TestRVsAssignmentNUTS(RVsAssignmentStepsTester):
+    @pytest.mark.parametrize("step, step_kwargs", [(NUTS, {})])
+    def test_continuous_steps(self, step, step_kwargs):
+        self.continuous_steps(step, step_kwargs)

pymc/tests/step_methods/test_compound.py

@@ -0,0 +1,93 @@
+#   Copyright 2020 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.
+
+import aesara
+import pytest
+
+import pymc as pm
+
+from pymc.step_methods import (
+    NUTS,
+    CompoundStep,
+    DEMetropolis,
+    HamiltonianMC,
+    Metropolis,
+    Slice,
+)
+from pymc.tests.helpers import StepMethodTester, fast_unstable_sampling_mode
+from pymc.tests.models import simple_2model_continuous
+
+
+class TestCompoundStep:
+    samplers = (Metropolis, Slice, HamiltonianMC, NUTS, DEMetropolis)
+
+    def test_non_blocked(self):
+        """Test that samplers correctly create non-blocked compound steps."""
+        with aesara.config.change_flags(mode=fast_unstable_sampling_mode):
+            _, model = simple_2model_continuous()
+            with model:
+                for sampler in self.samplers:
+                    assert isinstance(sampler(blocked=False), CompoundStep)
+
+    def test_blocked(self):
+        with aesara.config.change_flags(mode=fast_unstable_sampling_mode):
+            _, model = simple_2model_continuous()
+            with model:
+                for sampler in self.samplers:
+                    sampler_instance = sampler(blocked=True)
+                    assert not isinstance(sampler_instance, CompoundStep)
+                    assert isinstance(sampler_instance, sampler)
+
+    def test_name(self):
+        with pm.Model() as m:
+            c1 = pm.HalfNormal("c1")
+            c2 = pm.HalfNormal("c2")
+
+            step1 = NUTS([c1])
+            step2 = Slice([c2])
+            step = CompoundStep([step1, step2])
+        assert step.name == "Compound[nuts, slice]"
+
+
+class TestStepCompound(StepMethodTester):
+    @pytest.mark.parametrize(
+        "step_fn, draws",
+        [
+            (
+                lambda C, _: CompoundStep(
+                    [
+                        HamiltonianMC(scaling=C, is_cov=True),
+                        HamiltonianMC(scaling=C, is_cov=True, blocked=False),
+                    ]
+                ),
+                1000,
+            ),
+        ],
+        ids=str,
+    )
+    def test_step_continuous(self, step_fn, draws):
+        self.step_continuous(step_fn, draws)
+
+
+class TestRVsAssignmentCompound:
+    def test_compound_step(self):
+        with pm.Model() as m:
+            c1 = pm.HalfNormal("c1")
+            c2 = pm.HalfNormal("c2")
+
+            with aesara.config.change_flags(mode=fast_unstable_sampling_mode):
+                step1 = NUTS([c1])
+                step2 = NUTS([c2])
+                step = CompoundStep([step1, step2])
+            assert {m.rvs_to_values[c1], m.rvs_to_values[c2]} == set(step.vars)