Metropolis-Hastings in Python

The Metropolis–Hastings algorithm can be implemented in python as an Abstract Base Class using the code below. The code was designed to:

• Take states of a Generic Type ( Generic[T])
• Allow custom approval and likelihood functions to be used

import datetime  
from abc import ABC, abstractmethod  
from typing import Generic, List, TypeVar  
  
import numpy as np  
from matplotlib import pyplot as plt  
from tqdm import trange  
  
T = TypeVar("T")  
  
  
class AbstractMetropolisHastings(ABC, Generic[T]):  
    # Based on https://en.wikipedia.org/wiki/Metropolis%E2%80%93Hastings_algorithm#Description  
  
    def __init__(self, initial_configuration: T):  
        self.configuration_history: List[T] = [initial_configuration]  
        self.accepted_configuration_count: int = 0  
        self.rejected_configuration_count: int = 0  
  
    @property  
    def current_configuration(self) -> T:  
        return self.configuration_history[-1]  
  
    @abstractmethod  
    def generator_function(self) -> T:  
        pass  
  
    @abstractmethod  
    def state_likelihood(self, T) -> float:  
        # This is proportional to the state probability  
        pass  
  
    def approval_function(self, new_configuration: T) -> bool:  
        return (  
            self.state_likelihood(new_configuration)  
            >= self.state_likelihood(self.current_configuration) * np.random.random()  
        )  
  
    def run_single_iteration(self, limit_tries=10**5) -> T:  
        tries = 0  
        while True:  
            new_state = self.generator_function()  
            if self.approval_function(new_state):  
                self.configuration_history.append(new_state)  
                self.accepted_configuration_count += 1  
                return new_state  
  
            self.rejected_configuration_count += 1  
            tries += 1  
            if tries >= limit_tries:  
                # Useful for debugging  
                tries = 0  
                limit_tries *= int(1.1)  
                print(f"{new_state:e}", end=", ")  
  
    def run_iterations(self, n: int) -> None:  
        pbar = trange(n, desc="Bar desc", leave=True)  
        for _ in pbar:  
            self.run_single_iteration()  
  
            # Update the progress bar roughly once a second  
            seconds_passed = datetime.datetime.now().timestamp() - pbar.start_t  
            n = self.rejected_configuration_count + self.accepted_configuration_count  
            iterations_per_second = 1 + int(n / seconds_passed)  
            update_frequency = 2 ** (int(np.log(iterations_per_second) / np.log(2)) - 1)  
  
            if n % update_frequency == 0:  
  
                rejection_rate = np.divide(  
                    self.rejected_configuration_count,  
                    self.accepted_configuration_count  
                    + self.rejected_configuration_count,  
                )  
                pbar.set_description(  
                    f"Rejected {100*rejection_rate:.1f}%",  
                    refresh=True,  
                )  
  
    def plot(self) -> None:  
        plt.plot(self.configuration_history)

Example Usage

The code required to apply the Metropolis–Hastings algorithm to the Variational Monte Carlo problem is as follows

class VMC(AbstractMetropolisHastings[float]):  
    a = [1, 0]  
    p = [.0002, 0, 0, 0]  
    sigma = 5  
  
    def generator_function(self):  
        return self.current_configuration + np.random.normal(0, self.sigma)  
  
    def state_likelihood(self, x: float):  
        # This is proportional to the state probability  
        p = self.p  
        aux_exp = lambda mu, sigma: np.prod([-1, mu, x - sigma, x - sigma])  
        return np.exp(aux_exp(p[0], p[1]) + aux_exp(p[2], p[3]))

The code above defines

• The state space to be of type float. More complex examples might have states in ${\mathbb{R}}^3$ by passing in the type Tuple[float, float, float], etc
• The likelihood function to be $e^{-p_0(x-p_1{)}^2-p_2(x-p_3{)}^2}$.

With the methodology defined, we can run an example

vmc = VMC(initial_configuration=10**2)
vmc.run_iterations(10**2)
vmc.plot()

This starts the simulation at state $x=100$, runs $100$ iterations, and then plots the result. As an example, we get the result