test_metropolis.py

#   Copyright 2024 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 warnings

import arviz as az
import numpy as np
import numpy.testing as npt
import pytensor
import pytest

from pytensor.compile.mode import Mode

import pymc as pm

from pymc.step_methods.metropolis import (
    BinaryGibbsMetropolis,
    BinaryMetropolis,
    CategoricalGibbsMetropolis,
    DEMetropolis,
    DEMetropolisZ,
    Metropolis,
    MultivariateNormalProposal,
    NormalProposal,
)
from pymc.step_methods.state import equal_dataclass_values
from pymc.testing import fast_unstable_sampling_mode
from tests import sampler_fixtures as sf
from tests.helpers import RVsAssignmentStepsTester, StepMethodTester, equal_sampling_states
from tests.models import (
    mv_simple,
    mv_simple_discrete,
    simple_binary,
    simple_categorical,
    simple_model,
)

SEED = sum(ord(c) for c in "test_metropolis")


class TestMetropolisUniform(sf.MetropolisFixture, sf.UniformFixture):
    n_samples = 50000
    tune = 10000
    burn = 0
    chains = 4
    min_n_eff = 10000
    rtol = 0.1
    atol = 0.05
    ks_thin = 10
    step_args = {"rng": np.random.default_rng(SEED)}


class TestMetropolis:
    def test_proposal_choice(self):
        with pytensor.config.change_flags(mode=fast_unstable_sampling_mode):
            _, model, _ = mv_simple()
            with model:
                initial_point = model.initial_point()
                initial_point_size = sum(initial_point[n.name].size for n in model.value_vars)

                s = np.ones(initial_point_size)
                sampler = Metropolis(S=s)
                assert isinstance(sampler.proposal_dist, NormalProposal)
                s = np.diag(s)
                sampler = Metropolis(S=s)
                assert isinstance(sampler.proposal_dist, MultivariateNormalProposal)
                s[0, 0] = -s[0, 0]
                with pytest.raises(np.linalg.LinAlgError):
                    sampler = Metropolis(S=s)

    def test_mv_proposal(self):
        np.random.seed(42)
        cov = np.random.randn(5, 5)
        cov = cov.dot(cov.T)
        prop = MultivariateNormalProposal(cov)
        samples = np.array([prop() for _ in range(10000)])
        npt.assert_allclose(np.cov(samples.T), cov, rtol=0.2)

    def test_tuning_reset(self):
        """Re-use of the step method instance with cores=1 must not leak tuning information between chains."""
        with pm.Model() as pmodel:
            D = 3
            pm.Normal("n", 0, 2, size=(D,))
            idata = pm.sample(
                tune=600,
                draws=500,
                step=Metropolis(tune=True, scaling=0.1, rng=SEED),
                cores=1,
                chains=3,
                discard_tuned_samples=False,
            )
        for c in idata.posterior.chain:
            # check that the tuned settings changed and were reset
            assert idata.warmup_sample_stats["scaling"].sel(chain=c).values[0] == 0.1
            tuned = idata.warmup_sample_stats["scaling"].sel(chain=c).values[-1]
            assert tuned != 0.1
            np.testing.assert_array_equal(idata.sample_stats["scaling"].sel(chain=c).values, tuned)

    @pytest.mark.parametrize(
        "batched_dist",
        (
            pm.Binomial.dist(n=5, p=0.9),  # scalar case
            pm.Binomial.dist(n=np.arange(40) + 1, p=np.linspace(0.1, 0.9, 40), shape=(40,)),
            pm.Binomial.dist(
                n=(np.arange(20) + 1)[::-1],
                p=np.linspace(0.1, 0.9, 20),
                shape=(
                    2,
                    20,
                ),
            ),
            pm.Dirichlet.dist(a=np.ones(3) * (np.arange(40) + 1)[:, None], shape=(40, 3)),
            pm.Dirichlet.dist(a=np.ones(3) * (np.arange(20) + 1)[:, None], shape=(2, 20, 3)),
        ),
    )
    def test_elemwise_update(self, batched_dist):
        with pm.Model() as m:
            m.register_rv(batched_dist, name="batched_dist")
            step = pm.Metropolis([batched_dist], rng=SEED)
            assert step.elemwise_update == (batched_dist.ndim > 0)
            trace = pm.sample(draws=1000, chains=2, step=step, random_seed=428)

        assert az.rhat(trace).max()["batched_dist"].values < 1.1
        assert az.ess(trace).min()["batched_dist"].values > 50

    def test_elemwise_update_different_scales(self):
        mu = [1, 2, 3, 4, 5, 100, 1_000, 10_000]
        with pm.Model() as m:
            x = pm.Poisson("x", mu=mu)
            step = pm.Metropolis([x], rng=SEED)
            trace = pm.sample(draws=1000, chains=2, step=step, random_seed=128).posterior

        np.testing.assert_allclose(trace["x"].mean(("draw", "chain")), mu, rtol=0.1)
        np.testing.assert_allclose(trace["x"].var(("draw", "chain")), mu, rtol=0.2)

    def test_multinomial_no_elemwise_update(self):
        with pm.Model() as m:
            batched_dist = pm.Multinomial("batched_dist", n=5, p=np.ones(4) / 4, shape=(10, 4))
            with pytensor.config.change_flags(mode=fast_unstable_sampling_mode):
                step = pm.Metropolis([batched_dist], rng=SEED)
                assert not step.elemwise_update


class TestDEMetropolis:
    def test_demcmc_tune_parameter(self):
        """Tests that validity of the tune setting is checked"""
        with pm.Model() as model:
            pm.Normal("n", mu=0, sigma=1, size=(2, 3))

            step = DEMetropolis()
            assert step.tune == "scaling"

            step = DEMetropolis(tune=None)
            assert step.tune is None

            step = DEMetropolis(tune="scaling")
            assert step.tune == "scaling"

            step = DEMetropolis(tune="lambda")
            assert step.tune == "lambda"

            with pytest.raises(ValueError):
                DEMetropolis(tune="foo")


class TestDEMetropolisZ:
    def test_tuning_lambda_sequential(self):
        with pm.Model() as pmodel:
            pm.Normal("n", 0, 2, size=(3,))
            idata = pm.sample(
                tune=1000,
                draws=500,
                step=DEMetropolisZ(tune="lambda", lamb=0.92, rng=SEED),
                cores=1,
                chains=3,
                discard_tuned_samples=False,
            )
        for c in idata.posterior.chain:
            # check that the tuned settings changed and were reset
            assert idata.warmup_sample_stats["lambda"].sel(chain=c).values[0] == 0.92
            tuned = idata.warmup_sample_stats["lambda"].sel(chain=c).values[-1]
            assert tuned != 0.92
            np.testing.assert_array_equal(idata.sample_stats["lambda"].sel(chain=c).values, tuned)

    def test_tuning_epsilon_parallel(self):
        with pm.Model() as pmodel:
            pm.Normal("n", 0, 2, size=(3,))
            idata = pm.sample(
                tune=1000,
                draws=500,
                step=DEMetropolisZ(tune="scaling", scaling=0.002, rng=SEED),
                cores=2,
                chains=2,
                discard_tuned_samples=False,
            )
        for c in idata.posterior.chain:
            # check that the tuned settings changed and were reset
            assert idata.warmup_sample_stats["scaling"].sel(chain=c).values[0] == 0.002
            tuned = idata.warmup_sample_stats["scaling"].sel(chain=c).values[-1]
            assert tuned != 0.002
            np.testing.assert_array_equal(idata.sample_stats["scaling"].sel(chain=c).values, tuned)

    def test_tuning_none(self):
        with pm.Model() as pmodel:
            pm.Normal("n", 0, 2, size=(3,))
            idata = pm.sample(
                tune=1000,
                draws=500,
                step=DEMetropolisZ(tune=None, rng=SEED),
                cores=1,
                chains=2,
                discard_tuned_samples=False,
            )
        for c in idata.posterior.chain:
            # check that all tunable parameters remained constant
            assert len(set(idata.warmup_sample_stats["lambda"].sel(chain=c).values)) == 1
            assert len(set(idata.warmup_sample_stats["scaling"].sel(chain=c).values)) == 1

    def test_tuning_reset(self):
        """Re-use of the step method instance with cores=1 must not leak tuning information between chains."""
        with pm.Model() as pmodel:
            D = 3
            pm.Normal("n", 0, 2, size=(D,))
            idata = pm.sample(
                tune=1000,
                draws=500,
                step=DEMetropolisZ(tune="scaling", scaling=0.002, rng=SEED),
                cores=1,
                chains=3,
                discard_tuned_samples=False,
                random_seed=1,
            )
        for c in idata.posterior.chain:
            # check that the tuned settings changed and were reset
            warmup = idata.warmup_sample_stats["scaling"].sel(chain=c).values
            assert warmup[0] == 0.002
            assert warmup[-1] != 0.002
            # check that the variance of the first 50 iterations is much lower than the last 100
            samples = idata.warmup_posterior["n"].sel(chain=c).values
            for d in range(D):
                var_start = np.var(samples[:50, d])
                var_end = np.var(samples[-100:, d])
                assert var_start < 0.1 * var_end

    def test_tune_drop_fraction(self):
        tune = 300
        tune_drop_fraction = 0.85
        draws = 200
        with pm.Model() as pmodel:
            pm.Normal("n", 0, 2, size=(3,))
            step = DEMetropolisZ(tune_drop_fraction=tune_drop_fraction, rng=SEED)
            idata = pm.sample(
                tune=tune, draws=draws, step=step, cores=1, chains=1, discard_tuned_samples=False
            )
            assert len(idata.warmup_posterior.draw) == tune
            assert len(idata.posterior.draw) == draws
            assert len(step._history) == (tune - tune * tune_drop_fraction) + draws

    @pytest.mark.parametrize(
        "variable,has_grad,outcome",
        [("n", True, 1), ("n", False, 1), ("b", True, 0), ("b", False, 0)],
    )
    def test_competence(self, variable, has_grad, outcome):
        with pm.Model() as pmodel:
            pm.Normal("n", 0, 2, size=(3,))
            pm.Binomial("b", n=2, p=0.3)
        assert DEMetropolisZ.competence(pmodel[variable], has_grad=has_grad) == outcome

    @pytest.mark.parametrize("tune_setting", ["foo", True, False])
    def test_invalid_tune(self, tune_setting):
        with pm.Model() as pmodel:
            pm.Normal("n", 0, 2, size=(3,))
            with pytest.raises(ValueError):
                DEMetropolisZ(tune=tune_setting)

    def test_custom_proposal_dist(self):
        with pm.Model() as pmodel:
            D = 3
            pm.Normal("n", 0, 2, size=(D,))
            with warnings.catch_warnings():
                warnings.filterwarnings("ignore", ".*number of samples.*", UserWarning)
                trace = pm.sample(
                    tune=100,
                    draws=50,
                    step=DEMetropolisZ(proposal_dist=NormalProposal),
                    cores=1,
                    chains=3,
                    discard_tuned_samples=False,
                )


class TestStepMetropolis(StepMethodTester):
    def test_step_discrete(self):
        start, model, (mu, C) = mv_simple_discrete()
        unc = np.diag(C) ** 0.5
        check = (("x", np.mean, mu, unc / 10.0), ("x", np.std, unc, unc / 10.0))
        with model:
            step = Metropolis(S=C, proposal_dist=MultivariateNormalProposal, rng=123456)
            idata = pm.sample(
                tune=1000,
                draws=2000,
                chains=1,
                step=step,
                initvals=start,
                model=model,
                random_seed=1,
            )
            self.check_stat(check, idata)
            self.check_stat_dtype(idata, step)

    @pytest.mark.parametrize("proposal", ["uniform", "proportional"])
    def test_step_categorical(self, proposal):
        start, model, (mu, C) = simple_categorical()
        unc = C**0.5
        check = (("x", np.mean, mu, unc / 10.0), ("x", np.std, unc, unc / 10.0))
        with model:
            step = CategoricalGibbsMetropolis([model.x], proposal=proposal, rng=SEED)
            idata = pm.sample(
                tune=1000,
                draws=2000,
                chains=1,
                step=step,
                initvals=start,
                model=model,
                random_seed=1,
            )
            self.check_stat(check, idata)
            self.check_stat_dtype(idata, step)

    @pytest.mark.parametrize(
        "step_fn, draws",
        [
            (
                lambda C, _: Metropolis(
                    S=C, proposal_dist=MultivariateNormalProposal, blocked=True, rng=SEED
                ),
                4000,
            ),
        ],
        ids=str,
    )
    def test_step_continuous(self, step_fn, draws):
        self.step_continuous(step_fn, draws)


class TestRVsAssignmentMetropolis(RVsAssignmentStepsTester):
    @pytest.mark.parametrize(
        "step",
        [
            BinaryMetropolis,
            BinaryGibbsMetropolis,
            CategoricalGibbsMetropolis,
        ],
    )
    def test_discrete_steps(self, step):
        with pm.Model() as m:
            d1 = pm.Bernoulli("d1", p=0.5)
            d2 = pm.Bernoulli("d2", p=0.5)

            # Test it can take initial_point, and compile_kwargs as a kwarg
            step_kwargs = {
                "initial_point": {
                    "d1": np.array(0, dtype="int64"),
                    "d2": np.array(1, dtype="int64"),
                },
                "compile_kwargs": {"mode": Mode(linker="py", optimizer=None)},
            }
            assert [m.rvs_to_values[d1]] == step([d1]).vars
            assert {m.rvs_to_values[d1], m.rvs_to_values[d2]} == set(step([d1, d2]).vars)

    @pytest.mark.parametrize(
        "step, step_kwargs", [(Metropolis, {}), (DEMetropolis, {}), (DEMetropolisZ, {})]
    )
    def test_continuous_steps(self, step, step_kwargs):
        self.continuous_steps(step, step_kwargs)


@pytest.mark.parametrize(
    ["step_method", "model_fn"],
    [
        [Metropolis, simple_model],
        [BinaryMetropolis, simple_binary],
        [BinaryGibbsMetropolis, simple_binary],
        [CategoricalGibbsMetropolis, simple_categorical],
        [DEMetropolis, simple_model],
        [DEMetropolisZ, simple_model],
    ],
)
def test_sampling_state(step_method, model_fn):
    with pytensor.config.change_flags(mode=fast_unstable_sampling_mode):
        initial_point, model, _ = model_fn()
        with model:
            sampler = step_method(model.value_vars)
            if hasattr(sampler, "link_population"):
                sampler.link_population([initial_point] * 100, 0)
            state_orig = sampler.sampling_state

            sample1, stat1 = sampler.step(initial_point)
            sampler.tune = False

            final_state1 = sampler.sampling_state

            assert not equal_sampling_states(final_state1, state_orig)

            sampler.sampling_state = state_orig

            assert equal_sampling_states(sampler.sampling_state, state_orig)

            sample2, stat2 = sampler.step(initial_point)
            sampler.tune = False

            final_state2 = sampler.sampling_state

            assert equal_sampling_states(final_state1, final_state2)
            assert equal_dataclass_values(sample1, sample2)
            assert equal_dataclass_values(stat1, stat2)