bayesian_search.py

import numpy as np
import matplotlib
from matplotlib.axes import Axes
from matplotlib.animation import FuncAnimation
import matplotlib.pyplot as plt
import numpy as np
matplotlib.rcParams.update({'font.size': 48})

sqrt2pi = np.sqrt(2*np.pi)

def gaussian_2d(x, y, x0, y0, p, sx, sy):
    """2d Gaussian Probability Density Function"""
    left = 1 / 2*np.pi*np.sqrt(1-np.square(p))
    right = np.exp(
        -1*(np.square((x-x0)/sx)
            - 2*p*((x-x0)/sx * (y-y0)/sy)
            + np.square((y-y0)/sy)) /
        2*np.sqrt(1-np.square(p))
    )
    return left*right


class BayesianSearchProblem:
    def __init__(self, fig, ax: Axes, mu_x, mu_y, sigma_x, sigma_y, pearson, detection_prob, grid_size=20, min_grid=-30, max_grid=30):

        self.grid_size: int = grid_size
        self.mu_x: int = mu_x
        self.mu_y: int = mu_y
        self.ax: Axes = ax
        self.detect: float = 0.0
        self.detection_prob: float = detection_prob
        self.iterations: int = 0

        x = np.linspace(min_grid, max_grid, grid_size)
        y = np.linspace(min_grid, max_grid, grid_size)

        cov = pearson * (sigma_x) * (sigma_y)
        sigma_matrix = np.array(
            [[np.square(sigma_x), cov],
             [cov, np.square(sigma_y)]])

        # Create Grid and Sample Point
        location = np.random.multivariate_normal(
            np.array([mu_x, mu_y]), sigma_matrix)
        location_matrix = np.zeros((grid_size, grid_size))

        # Determine which grid it's in
        grid_x, grid_y = (x > location[0]).argmax(), (y > location[1]).argmax()
        location_matrix[grid_x][grid_y] = 1.0
        self.location_matrix: np.array = location_matrix

        # Create Distribution
        # Assign some non-zero dist to very low prob regions
        xx, yy = np.meshgrid(x, y)
        z = gaussian_2d(xx, yy, mu_x, mu_y,
                        pearson, sigma_x, sigma_y) + 1

        # quick and easy but incorrect normalization
        self.z: np.array = z/z.sum()

        # Plot prior distribution
        self.img = self.ax.imshow(
            self.z, origin="lower",
            interpolation="nearest", cmap="gnuplot2")
        fig.colorbar(self.img, ax=ax)

    def __call__(self, i):
        if self.detect:
            # Already detected
            return ()

        # Not yet detected
        p_max_pos: int = self.z.argmax()
        search_x = int(np.floor(p_max_pos / self.grid_size))
        search_y = p_max_pos - search_x*self.grid_size
        detect: int = np.random.binomial(
            1, self.detection_prob) * self.location_matrix[search_x][search_y]

        if detect:
            # We detected!
            ax.set_title(f"Bayesian Search Found In {self.iterations}")
            z = np.zeros((self.grid_size, self.grid_size))
            z[search_x][search_y] = 1.0
            self.z = z
            self.img.set_data(z)
            self.detect = detect
            return self.img,

        else:
            # Update posterior based on evidence of not detecting in cell search_x, search_y
            self.iterations += 1
            ax.set_title(f"Bayesian Search {i}")
            # Posterior of position_x_y
            new_z_xy = (1-self.detection_prob)*self.z[search_x][search_y] / \
                (1-self.detection_prob*self.z[search_x][search_y])
            self.z = self.z / (1-self.detection_prob *
                               self.z[search_x][search_y])
            self.z[search_x][search_y] = new_z_xy
            self.img.set_data(self.z)
            self.img.autoscale()

        return self.img,


if __name__ == "__main__":
    fig, ax = plt.subplots(figsize=(20, 20), dpi=150)
    ax.set_title("Bayesian Search")
    ud = BayesianSearchProblem(fig, ax, 0, 0, 6, 6, 0.7, 0.44, 40, -30, 30)
    anim = FuncAnimation(fig, ud, frames=150, interval=20,
                         blit=True, cache_frame_data=False)

    anim.save('animation.gif', writer='ffmepg', fps=60)