Use ABC abstract methods in step method interface

Author: brandonwillard

pymc3/step_methods/arraystep.py

@@ -12,21 +12,24 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 
+from abc import ABC, abstractmethod
 from enum import IntEnum, unique
-from typing import Dict, List
+from typing import Dict, List, Tuple, TypeVar, Union
 
 import numpy as np
 
 from aesara.graph.basic import Variable
 from numpy.random import uniform
 
-from pymc3.blocking import DictToArrayBijection, RaveledVars
+from pymc3.blocking import DictToArrayBijection, PointType, RaveledVars
 from pymc3.model import modelcontext
 from pymc3.step_methods.compound import CompoundStep
 from pymc3.util import get_var_name
 
 __all__ = ["ArrayStep", "ArrayStepShared", "metrop_select", "Competence"]
 
+StatsType = TypeVar("StatsType")
+
 
 @unique
 class Competence(IntEnum):
@@ -44,7 +47,7 @@ class Competence(IntEnum):
     IDEAL = 3
 
 
-class BlockedStep:
+class BlockedStep(ABC):
 
     generates_stats = False
     stats_dtypes: List[Dict[str, np.dtype]] = []
@@ -99,6 +102,10 @@ def __new__(cls, *args, **kwargs):
     def __getnewargs_ex__(self):
         return self.__newargs
 
+    @abstractmethod
+    def step(point: PointType, *args, **kwargs) -> Union[PointType, Tuple[PointType, StatsType]]:
+        """Perform a single step of the sampler."""
+
     @staticmethod
     def competence(var, has_grad):
         return Competence.INCOMPATIBLE
@@ -139,7 +146,7 @@ def __init__(self, vars, fs, allvars=False, blocked=True):
         self.allvars = allvars
         self.blocked = blocked
 
-    def step(self, point: Dict[str, np.ndarray]):
+    def step(self, point: PointType):
 
         partial_funcs_and_point = [DictToArrayBijection.mapf(x, start_point=point) for x in self.fs]
         if self.allvars:
@@ -164,8 +171,11 @@ def step(self, point: Dict[str, np.ndarray]):
 
         return point_new
 
-    def astep(self, apoint: RaveledVars, point: Dict[str, np.ndarray]):
-        raise NotImplementedError()
+    @abstractmethod
+    def astep(
+        self, apoint: RaveledVars, point: PointType, *args
+    ) -> Union[RaveledVars, Tuple[RaveledVars, StatsType]]:
+        """Perform a single sample step in a raveled and concatenated parameter space."""
 
 
 class ArrayStepShared(BlockedStep):
@@ -213,9 +223,6 @@ def step(self, point):
 
         return new_point
 
-    def astep(self, apoint: RaveledVars):
-        raise NotImplementedError()
-
 
 class PopulationArrayStepShared(ArrayStepShared):
     """Version of ArrayStepShared that allows samplers to access the states
@@ -278,9 +285,6 @@ def step(self, point):
         self._logp_dlogp_func._extra_are_set = True
         return super().step(point)
 
-    def astep(self, apoint):
-        raise NotImplementedError()
-
 
 def metrop_select(mr, q, q0):
     """Perform rejection/acceptance step for Metropolis class samplers.
@@ -300,6 +304,8 @@ def metrop_select(mr, q, q0):
     q or q0
     """
     # Compare acceptance ratio to uniform random number
+    # TODO XXX: This `uniform` is not given a model-specific RNG state, which
+    # means that sampler runs that use it will not be reproducible.
     if np.isfinite(mr) and np.log(uniform()) < mr:
         return q, True
     else:

pymc3/step_methods/hmc/base_hmc.py

@@ -15,6 +15,7 @@
 import logging
 import time
 
+from abc import abstractmethod
 from collections import namedtuple
 
 import numpy as np
@@ -24,7 +25,8 @@
 from pymc3.blocking import DictToArrayBijection, RaveledVars
 from pymc3.exceptions import SamplingError
 from pymc3.model import Point, modelcontext
-from pymc3.step_methods import arraystep, step_sizes
+from pymc3.step_methods import step_sizes
+from pymc3.step_methods.arraystep import GradientSharedStep
 from pymc3.step_methods.hmc import integration
 from pymc3.step_methods.hmc.quadpotential import QuadPotentialDiagAdapt, quad_potential
 from pymc3.tuning import guess_scaling
@@ -36,7 +38,7 @@
 DivergenceInfo = namedtuple("DivergenceInfo", "message, exec_info, state, state_div")
 
 
-class BaseHMC(arraystep.GradientSharedStep):
+class BaseHMC(GradientSharedStep):
     """Superclass to implement Hamiltonian/hybrid monte carlo."""
 
     default_blocked = True
@@ -85,8 +87,6 @@ def __init__(
         if vars is None:
             vars = self._model.cont_vars
 
-        # vars = inputvars(vars)
-
         super().__init__(vars, blocked=blocked, model=self._model, dtype=dtype, **aesara_kwargs)
 
         self.adapt_step_size = adapt_step_size
@@ -132,12 +132,12 @@ def __init__(
         self._samples_after_tune = 0
         self._num_divs_sample = 0
 
+    @abstractmethod
     def _hamiltonian_step(self, start, p0, step_size):
         """Compute one hamiltonian trajectory and return the next state.
 
         Subclasses must overwrite this method and return a `HMCStepData`.
         """
-        raise NotImplementedError("Abstract method")
 
     def astep(self, q0):
         """Perform a single HMC iteration."""

pymc3/step_methods/metropolis.py

@@ -11,7 +11,7 @@
 #   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 typing import Any, Dict, List, Tuple
+from typing import Any, Callable, Dict, List, Tuple
 
 import aesara
 import numpy as np
@@ -425,7 +425,7 @@ def __init__(self, vars, order="random", transit_p=0.8, model=None):
 
         super().__init__(vars, [model.fastlogp])
 
-    def astep(self, q0: RaveledVars, logp) -> RaveledVars:
+    def astep(self, q0: RaveledVars, logp: Callable[[RaveledVars], np.ndarray]) -> RaveledVars:
 
         order = self.order
         if self.shuffle_dims:
@@ -475,6 +475,7 @@ def competence(var):
 
 class CategoricalGibbsMetropolis(ArrayStep):
     """A Metropolis-within-Gibbs step method optimized for categorical variables.
+
     This step method works for Bernoulli variables as well, but it is not
     optimized for them, like BinaryGibbsMetropolis is. Step method supports
     two types of proposals: A uniform proposal and a proportional proposal,
@@ -574,6 +575,9 @@ def astep_prop(self, q0: RaveledVars, logp) -> RaveledVars:
 
         return q
 
+    def astep(self, q0, logp):
+        raise NotImplementedError()
+
     def metropolis_proportional(self, q, logp, logp_curr, dim, k):
         given_cat = int(q.data[dim])
         log_probs = np.zeros(k)

pymc3/tests/test_hmc.py

@@ -32,7 +32,13 @@ def test_leapfrog_reversible():
     start, model, _ = models.non_normal(n)
     size = sum(start[n.name].size for n in model.value_vars)
     scaling = floatX(np.random.rand(size))
-    step = BaseHMC(vars=model.value_vars, model=model, scaling=scaling)
+
+    class HMC(BaseHMC):
+        def _hamiltonian_step(self, *args, **kwargs):
+            pass
+
+    step = HMC(vars=model.value_vars, model=model, scaling=scaling)
+
     step.integrator._logp_dlogp_func.set_extra_values({})
     astart = DictToArrayBijection.map(start)
     p = RaveledVars(floatX(step.potential.random()), astart.point_map_info)