def test_multiple_samplers(self, caplog):
     with Model():
         prob = Beta("prob", alpha=5.0, beta=3.0)
         Binomial("outcome", n=1, p=prob)
         caplog.clear()
         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_multiple_samplers(self):
     with Model():
         prob = Beta('prob', alpha=5., beta=3.)
         Binomial('outcome', n=1, p=prob)
         with pytest.warns(None) as warns:
             sample(3, tune=2, discard_tuned_samples=False,
                    n_init=None)
         messages = [warn.message.args[0] for warn in warns]
         assert any("contains only 3" in msg for msg in messages)
         assert all('boolean index did not' not in msg for msg in messages)

 def test_checks_population_size(self):
     """Test that population samplers check the population size."""
     with Model() as model:
         n = Normal("n", mu=0, sigma=1)
         for stepper in TestPopulationSamplers.steppers:
             step = stepper()
             with pytest.raises(ValueError):
                 trace = sample(draws=100, chains=1, step=step)
             trace = sample(draws=100, chains=4, step=step)
     pass

    def test_float32(self):
        with Model() as model:
            x = Normal('x', testval=np.array(1., dtype='float32'))
            obs = Normal('obs', mu=x, sigma=1., observed=np.random.randn(5).astype('float32'))

        assert x.dtype == 'float32'
        assert obs.dtype == 'float32'

        for sampler in self.samplers:
            with model:
                sample(10, sampler())

    def test_float32(self):
        theano.config.floatX = 'float32'
        theano.config.warn_float64 = 'warn'

        with Model() as model:
            x = Normal('x', testval=np.array(1., dtype='float32'))
            obs = Normal('obs', mu=x, sd=1., observed=np.random.randn(5).astype('float32'))

        assert x.dtype == 'float32'
        assert obs.dtype == 'float32'

        for sampler in self.samplers:
            with model:
                sample(10, sampler())

    def check_trace(self, step_method):
        """Tests whether the trace for step methods is exactly the same as on master.

        Code changes that effect how random numbers are drawn may change this, and require
        `master_samples` to be updated, but such changes should be noted and justified in the
        commit.

        This method may also be used to benchmark step methods across commits, by running, for
        example

        ```
        BENCHMARK=100000 ./scripts/test.sh -s pymc3/tests/test_step.py:TestStepMethods
        ```

        on multiple commits.
        """
        n_steps = 100
        with Model() as model:
            x = Normal("x", mu=0, sigma=1)
            y = Normal("y", mu=x, sigma=1, observed=1)
            if step_method.__name__ == "SMC":
                trace = sample(
                    draws=200, random_seed=1, progressbar=False, step=step_method(parallel=False)
                )
            elif step_method.__name__ == "NUTS":
                step = step_method(scaling=model.test_point)
                trace = sample(
                    0,
                    tune=n_steps,
                    discard_tuned_samples=False,
                    step=step,
                    random_seed=1,
                    chains=1,
                )
            else:
                trace = sample(
                    0,
                    tune=n_steps,
                    discard_tuned_samples=False,
                    step=step_method(),
                    random_seed=1,
                    chains=1,
                )

        assert_array_almost_equal(
            trace["x"],
            self.master_samples[step_method],
            decimal=select_by_precision(float64=6, float32=4),
        )

    def test_linalg(self):
        with Model():
            a = Normal('a', shape=2)
            a = tt.switch(a > 0, np.inf, a)
            b = tt.slinalg.solve(floatX(np.eye(2)), a)
            Normal('c', mu=b, shape=2)
            with pytest.warns(None) as warns:
                trace = sample(20, init=None, tune=5)
            assert np.any(trace['diverging'])
            assert any('diverging samples after tuning' in str(warn.message)
                       for warn in warns)
            assert any('contains only' in str(warn.message) for warn in warns)

            with pytest.raises(SamplingError):
                sample(20, init=None, nuts_kwargs={'on_error': 'raise'})

def test_constant_step():

    with Model() as model:
        x = Normal('x', 0, 1)
        start = {'x':-1}
        tr = sample(10, step=Constant([x]), start=start)
        assert_almost_equal(tr['x'], start['x'], decimal=10)

def test_step_continuous():
    start, model, (mu, C) = mv_simple()

    with model:
        mh = Metropolis()
        slicer = Slice()
        hmc = HamiltonianMC(scaling=C, is_cov=True, blocked=False)
        nuts = NUTS(scaling=C, is_cov=True, blocked=False)

        mh_blocked = Metropolis(S=C,
                                   proposal_dist=MultivariateNormalProposal,
                                   blocked=True)
        slicer_blocked = Slice(blocked=True)
        hmc_blocked = HamiltonianMC(scaling=C, is_cov=True)
        nuts_blocked = NUTS(scaling=C, is_cov=True)

        compound = CompoundStep([hmc_blocked, mh_blocked])


    steps = [slicer, hmc, nuts, mh_blocked, hmc_blocked,
             slicer_blocked, nuts_blocked, compound]

    unc = np.diag(C) ** .5
    check = [('x', np.mean, mu, unc / 10.),
             ('x', np.std, unc, unc / 10.)]

    for st in steps:
        h = sample(8000, st, start, model=model, random_seed=1)
        for (var, stat, val, bound) in check:
            yield check_stat, repr(st), h, var, stat, val, bound

 def test_step_continuous(self):
     start, model, (mu, C) = mv_simple()
     unc = np.diag(C) ** 0.5
     check = (("x", np.mean, mu, unc / 10.0), ("x", np.std, unc, unc / 10.0))
     with model:
         steps = (
             Slice(),
             HamiltonianMC(scaling=C, is_cov=True, blocked=False),
             NUTS(scaling=C, is_cov=True, blocked=False),
             Metropolis(S=C, proposal_dist=MultivariateNormalProposal, blocked=True),
             Slice(blocked=True),
             HamiltonianMC(scaling=C, is_cov=True),
             NUTS(scaling=C, is_cov=True),
             CompoundStep(
                 [
                     HamiltonianMC(scaling=C, is_cov=True),
                     HamiltonianMC(scaling=C, is_cov=True, blocked=False),
                 ]
             ),
         )
     for step in steps:
         trace = sample(
             0,
             tune=8000,
             chains=1,
             discard_tuned_samples=False,
             step=step,
             start=start,
             model=model,
             random_seed=1,
         )
         self.check_stat(check, trace, step.__class__.__name__)

    def check_trace(self, step_method):
        """Tests whether the trace for step methods is exactly the same as on master.

        Code changes that effect how random numbers are drawn may change this, and require
        `master_samples` to be updated, but such changes should be noted and justified in the
        commit.

        This method may also be used to benchmark step methods across commits, by running, for
        example

        ```
        BENCHMARK=100000 ./scripts/test.sh -s pymc3/tests/test_step.py:TestStepMethods
        ```

        on multiple commits.
        """
        test_steps = 100
        n_steps = int(os.getenv('BENCHMARK', 100))
        benchmarking = (n_steps != test_steps)
        if benchmarking:
            tqdm.write('Benchmarking {} with {:,d} samples'.format(step_method.__name__, n_steps))
        else:
            tqdm.write('Checking {} has same trace as on master'.format(step_method.__name__))
        with Model():
            Normal('x', mu=0, sd=1)
            trace = sample(n_steps, step=step_method(), random_seed=1)

        if not benchmarking:
            assert_array_almost_equal(trace.get_values('x'), self.master_samples[step_method])

    def test_sampler_stats(self):
        with Model() as model:
            x = Normal("x", mu=0, sigma=1)
            trace = sample(draws=10, tune=1, chains=1)

        # 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",
        }
        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_ = np.array(
            [
                model.logp(trace.point(i, chain=c))
                for c in trace.chains
                for i in range(len(trace))
            ]
        )
        assert (trace.model_logp == model_logp_).all()

    def check_trace(self, step_method):
        """Tests whether the trace for step methods is exactly the same as on master.

        Code changes that effect how random numbers are drawn may change this, and require
        `master_samples` to be updated, but such changes should be noted and justified in the
        commit.

        This method may also be used to benchmark step methods across commits, by running, for
        example

        ```
        BENCHMARK=100000 ./scripts/test.sh -s pymc3/tests/test_step.py:TestStepMethods
        ```

        on multiple commits.
        """
        n_steps = 100
        with Model():
            x = Normal('x', mu=0, sd=1)
            if step_method.__name__ == 'SMC':
                Deterministic('like', - 0.5 * tt.log(2 * np.pi) - 0.5 * x.T.dot(x))
                trace = smc.ATMIP_sample(n_steps=n_steps, step=step_method(random_seed=1),
                                         n_jobs=1, progressbar=False, stage='0',
                                         homepath=self.temp_dir)
            else:
                trace = sample(n_steps, step=step_method(), random_seed=1)

        print(repr(trace.get_values('x')))
        assert_array_almost_equal(
            trace.get_values('x'),
            self.master_samples[step_method],
            decimal=select_by_precision(float64=6, float32=4))

    def test_linalg(self):
        with Model():
            a = Normal('a', shape=2)
            a = tt.switch(a > 0, np.inf, a)
            b = tt.slinalg.solve(floatX(np.eye(2)), a)
            Normal('c', mu=b, shape=2)
            with pytest.warns(None) as warns:
                trace = sample(20, init=None, tune=5)
            warns = [str(warn.message) for warn in warns]
            print(warns)
            assert np.any(trace['diverging'])
            assert any('diverging samples after tuning' in warn
                       for warn in warns)
            # FIXME This test fails sporadically on py27.
            # It seems that capturing warnings doesn't work as expected.
            # assert any('contains only' in warn for warn in warns)

            with pytest.raises(SamplingError):
                sample(20, init=None, nuts_kwargs={'on_error': 'raise'})

 def test_step_discrete(self):
     start, model, (mu, C) = mv_simple_discrete()
     unc = np.diag(C) ** .5
     check = (('x', np.mean, mu, unc / 10.),
              ('x', np.std, unc, unc / 10.))
     with model:
         steps = (
             Metropolis(S=C, proposal_dist=MultivariateNormalProposal),
         )
     for step in steps:
         trace = sample(20000, step=step, start=start, model=model, random_seed=1)
         self.check_stat(check, trace)

    def check_trace(self, step_method):
        """Tests whether the trace for step methods is exactly the same as on master.

        Code changes that effect how random numbers are drawn may change this, and require
        `master_samples` to be updated, but such changes should be noted and justified in the
        commit.

        This method may also be used to benchmark step methods across commits, by running, for
        example

        ```
        BENCHMARK=100000 ./scripts/test.sh -s pymc3/tests/test_step.py:TestStepMethods
        ```

        on multiple commits.
        """
        n_steps = 100
        with Model() as model:
            x = Normal('x', mu=0, sd=1)
            if step_method.__name__ == 'SMC':
                trace = smc.sample_smc(n_steps=n_steps,
                                       n_chains=2,
                                       start=[{'x':1.}, {'x':-1.}],
                                       random_seed=1,
                                       n_jobs=1, progressbar=False,
                                       homepath=self.temp_dir)

            elif step_method.__name__ == 'NUTS':
                step = step_method(scaling=model.test_point)
                trace = sample(0, tune=n_steps,
                               discard_tuned_samples=False,
                               step=step, random_seed=1, chains=1)
            else:
                trace = sample(0, tune=n_steps,
                               discard_tuned_samples=False,
                               step=step_method(), random_seed=1, chains=1)
        assert_array_almost_equal(
            trace.get_values('x'),
            self.master_samples[step_method],
            decimal=select_by_precision(float64=6, float32=4))

 def test_step_categorical(self):
     start, model, (mu, C) = simple_categorical()
     unc = C ** .5
     check = (('x', np.mean, mu, unc / 10.),
              ('x', np.std, unc, unc / 10.))
     with model:
         steps = (
             CategoricalGibbsMetropolis(model.x, proposal='uniform'),
             CategoricalGibbsMetropolis(model.x, proposal='proportional'),
         )
     for step in steps:
         trace = sample(8000, step=step, start=start, model=model, random_seed=1)
         yield self.check_stat, check, trace, step.__class__.__name__

    def test_parallelized_chains_are_random(self):
        with Model() as model:
            x = Normal('x', 0, 1)
            for stepper in TestPopulationSamplers.steppers:
                step = stepper()

                trace = sample(chains=4, draws=20, tune=0, step=DEMetropolis(),
                               parallelize=True)
                samples = np.array(trace.get_values('x', combine=False))[:,5]

                assert len(set(samples)) == 4, 'Parallelized {} ' \
                    'chains are identical.'.format(stepper)
        pass

 def test_step_elliptical_slice(self):
     start, model, (K, L, mu, std, noise) = mv_prior_simple()
     unc = noise ** 0.5
     check = (('x', np.mean, mu, unc / 10.),
              ('x', np.std, std, unc / 10.))
     with model:
         steps = (
             EllipticalSlice(prior_cov=K),
             EllipticalSlice(prior_chol=L),
         )
     for step in steps:
         trace = sample(5000, step=step, start=start, model=model, random_seed=1)
         yield self.check_stat, check, trace, step.__class__.__name__

 def test_step_discrete(self):
     if theano.config.floatX == "float32":
         return  # Cannot use @skip because it only skips one iteration of the yield
     start, model, (mu, C) = mv_simple_discrete()
     unc = np.diag(C) ** .5
     check = (('x', np.mean, mu, unc / 10.),
              ('x', np.std, unc, unc / 10.))
     with model:
         steps = (
             Metropolis(S=C, proposal_dist=MultivariateNormalProposal),
         )
     for step in steps:
         trace = sample(20000, step=step, start=start, model=model, random_seed=1)
         yield self.check_stat, check, trace, step.__class__.__name__