Add metropolis sampling state

Author: lucianopaz

pymc/step_methods/metropolis.py

@@ -12,6 +12,8 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 from collections.abc import Callable
+from dataclasses import field
+from typing import Any
 
 import numpy as np
 import numpy.random as nr
@@ -40,7 +42,8 @@
     StatsType,
     metrop_select,
 )
-from pymc.step_methods.compound import Competence
+from pymc.step_methods.compound import Competence, StepMethodState
+from pymc.step_methods.state import dataclass_state
 
 __all__ = [
     "Metropolis",
@@ -111,18 +114,40 @@ def __call__(self, num_draws=None, rng: np.random.Generator | None = None):
             return np.dot(self.chol, b)
 
 
+@dataclass_state
+class MetropolisState(StepMethodState):
+    scaling: np.ndarray
+    tune: bool
+    steps_until_tune: float
+    tune_interval: float
+    accepted_sum: np.ndarray
+    accept_rate_iter: np.ndarray
+    accepted_iter: np.ndarray
+    enum_dims: np.ndarray
+
+    discrete: np.ndarray = field(metadata={"frozen": True})
+    any_discrete: bool = field(metadata={"frozen": True})
+    all_discrete: bool = field(metadata={"frozen": True})
+    elemwise_update: bool = field(metadata={"frozen": True})
+    _untuned_settings: dict[str, np.ndarray | float] = field(metadata={"frozen": True})
+    mode: Any = field(metadata={"frozen": True})
+
+
 class Metropolis(ArrayStepShared):
     """Metropolis-Hastings sampling step"""
 
     name = "metropolis"
 
+    default_blocked = False
     stats_dtypes_shapes = {
         "accept": (np.float64, []),
         "accepted": (np.float64, []),
         "tune": (bool, []),
         "scaling": (np.float64, []),
     }
 
+    _state_class = MetropolisState
+
     def __init__(
         self,
         vars=None,
@@ -346,6 +371,15 @@ def tune(scale, acc_rate):
     )
 
 
+@dataclass_state
+class BinaryMetropolisState(StepMethodState):
+    tune: bool
+    accepted: int
+    scaling: float
+    tune_interval: int
+    steps_until_tune: int
+
+
 class BinaryMetropolis(ArrayStep):
     """Metropolis-Hastings optimized for binary variables
 
@@ -375,7 +409,9 @@ class BinaryMetropolis(ArrayStep):
         "p_jump": (np.float64, []),
     }
 
-    def __init__(self, vars, scaling=1.0, tune=True, tune_interval=100, model=None):
+    _state_class = BinaryMetropolisState
+
+    def __init__(self, vars, scaling=1.0, tune=True, tune_interval=100, model=None, rng=None):
         model = pm.modelcontext(model)
 
         self.scaling = scaling
@@ -389,7 +425,7 @@ def __init__(self, vars, scaling=1.0, tune=True, tune_interval=100, model=None):
         if not all([v.dtype in pm.discrete_types for v in vars]):
             raise ValueError("All variables must be Bernoulli for BinaryMetropolis")
 
-        super().__init__(vars, [model.compile_logp()])
+        super().__init__(vars, [model.compile_logp()], rng=rng)
 
     def astep(self, apoint: RaveledVars, *args) -> tuple[RaveledVars, StatsType]:
         logp = args[0]
@@ -445,6 +481,14 @@ def competence(var):
         return Competence.INCOMPATIBLE
 
 
+@dataclass_state
+class BinaryGibbsMetropolisState(StepMethodState):
+    tune: bool
+    transit_p: int
+    shuffle_dims: bool
+    order: list
+
+
 class BinaryGibbsMetropolis(ArrayStep):
     """A Metropolis-within-Gibbs step method optimized for binary variables
 
@@ -472,7 +516,9 @@ class BinaryGibbsMetropolis(ArrayStep):
         "tune": (bool, []),
     }
 
-    def __init__(self, vars, order="random", transit_p=0.8, model=None):
+    _state_class = BinaryGibbsMetropolisState
+
+    def __init__(self, vars, order="random", transit_p=0.8, model=None, rng=None):
         model = pm.modelcontext(model)
 
         # Doesn't actually tune, but it's required to emit a sampler stat
@@ -498,7 +544,7 @@ def __init__(self, vars, order="random", transit_p=0.8, model=None):
         if not all([v.dtype in pm.discrete_types for v in vars]):
             raise ValueError("All variables must be binary for BinaryGibbsMetropolis")
 
-        super().__init__(vars, [model.compile_logp()])
+        super().__init__(vars, [model.compile_logp()], rng=rng)
 
     def reset_tuning(self):
         # There are no tuning parameters in this step method.
@@ -557,6 +603,13 @@ def competence(var):
         return Competence.INCOMPATIBLE
 
 
+@dataclass_state
+class CategoricalGibbsMetropolisState(StepMethodState):
+    shuffle_dims: bool
+    dimcats: list[tuple]
+    tune: bool
+
+
 class CategoricalGibbsMetropolis(ArrayStep):
     """A Metropolis-within-Gibbs step method optimized for categorical variables.
 
@@ -573,6 +626,8 @@ class CategoricalGibbsMetropolis(ArrayStep):
         "tune": (bool, []),
     }
 
+    _state_class = CategoricalGibbsMetropolisState
+
     def __init__(self, vars, proposal="uniform", order="random", model=None, rng=None):
         model = pm.modelcontext(model)
 
@@ -728,6 +783,18 @@ def competence(var):
         return Competence.INCOMPATIBLE
 
 
+@dataclass_state
+class DEMetropolisState(StepMethodState):
+    scaling: np.ndarray
+    lamb: float
+    tune: str | None
+    tune_interval: int
+    steps_until_tune: int
+    accepted: int
+
+    mode: Any = field(metadata={"frozen": True})
+
+
 class DEMetropolis(PopulationArrayStepShared):
     """
     Differential Evolution Metropolis sampling step.
@@ -778,6 +845,8 @@ class DEMetropolis(PopulationArrayStepShared):
         "lambda": (np.float64, []),
     }
 
+    _state_class = DEMetropolisState
+
     def __init__(
         self,
         vars=None,
@@ -789,6 +858,7 @@ def __init__(
         tune_interval=100,
         model=None,
         mode=None,
+        rng=None,
         **kwargs,
     ):
         model = pm.modelcontext(model)
@@ -824,7 +894,7 @@ def __init__(
 
         shared = pm.make_shared_replacements(initial_values, vars, model)
         self.delta_logp = delta_logp(initial_values, model.logp(), vars, shared)
-        super().__init__(vars, shared)
+        super().__init__(vars, shared, rng=rng)
 
     def astep(self, q0: RaveledVars) -> tuple[RaveledVars, StatsType]:
         point_map_info = q0.point_map_info
@@ -843,9 +913,11 @@ def astep(self, q0: RaveledVars) -> tuple[RaveledVars, StatsType]:
 
         # differential evolution proposal
         # select two other chains
-        ir1, ir2 = np.random.choice(self.other_chains, 2, replace=False)
-        r1 = DictToArrayBijection.map(self.population[ir1])
-        r2 = DictToArrayBijection.map(self.population[ir2])
+        if self.other_chains is None:  # pragma: no cover
+            raise RuntimeError("Population sampler has not been linked to the other chains")
+        ir1, ir2 = self.rng.choice(self.other_chains, 2, replace=False)
+        r1 = DictToArrayBijection.map(self.population[ir1])  # type: ignore
+        r2 = DictToArrayBijection.map(self.population[ir2])  # type: ignore
         # propose a jump
         q = floatX(q0d + self.lamb * (r1.data - r2.data) + epsilon)
 
@@ -872,6 +944,21 @@ def competence(var, has_grad):
         return Competence.COMPATIBLE
 
 
+@dataclass_state
+class DEMetropolisZState(StepMethodState):
+    scaling: np.ndarray
+    lamb: float
+    tune: bool
+    tune_target: str | None
+    tune_interval: int
+    steps_until_tune: int
+    accepted: int
+    _history: list
+
+    _untuned_settings: dict[str, np.ndarray | float] = field(metadata={"frozen": True})
+    mode: Any = field(metadata={"frozen": True})
+
+
 class DEMetropolisZ(ArrayStepShared):
     """
     Adaptive Differential Evolution Metropolis sampling step that uses the past to inform jumps.
@@ -925,6 +1012,8 @@ class DEMetropolisZ(ArrayStepShared):
         "lambda": (np.float64, []),
     }
 
+    _state_class = DEMetropolisZState
+
     def __init__(
         self,
         vars=None,
@@ -937,6 +1026,7 @@ def __init__(
         tune_drop_fraction: float = 0.9,
         model=None,
         mode=None,
+        rng=None,
         **kwargs,
     ):
         model = pm.modelcontext(model)
@@ -984,7 +1074,7 @@ def __init__(
 
         shared = pm.make_shared_replacements(initial_values, vars, model)
         self.delta_logp = delta_logp(initial_values, model.logp(), vars, shared)
-        super().__init__(vars, shared)
+        super().__init__(vars, shared, rng=rng)
 
     def reset_tuning(self):
         """Resets the tuned sampler parameters and history to their initial values."""

tests/models.py

@@ -186,3 +186,14 @@ def simple_normal(bounded_prior=False):
         pm.Normal("X_obs", mu=mu_i, sigma=sigma, observed=x0)
 
     return model.initial_point(), model, None
+
+
+def simple_binary():
+    p1 = 0.5
+    p2 = 0.5
+
+    with pm.Model() as model:
+        pm.Bernoulli("d1", p=p1)
+        pm.Bernoulli("d2", p=p2)
+
+    return model.initial_point(), model, (p1, p2)

tests/step_methods/test_metropolis.py

@@ -14,6 +14,8 @@
 
 import warnings
 
+from copy import deepcopy
+
 import arviz as az
 import numpy as np
 import numpy.testing as npt
@@ -24,17 +26,25 @@
 
 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
-from tests.models import mv_simple, mv_simple_discrete, simple_categorical
+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")
 
@@ -47,6 +57,7 @@ class TestMetropolisUniform(sf.MetropolisFixture, sf.UniformFixture):
     min_n_eff = 10000
     rtol = 0.1
     atol = 0.05
+    ks_thin = 10
     step_args = {"rng": np.random.default_rng(SEED)}
 
 
@@ -367,3 +378,45 @@ def test_discrete_steps(self, step, step_kwargs):
     )
     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)
+            sampler_orig = deepcopy(sampler)
+            state_orig = sampler_orig.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)