<reponame>feynmanliang/beanmachine
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

"""
The file contains MiniBM, a minimal implementation of Bean Machine PPL with a Metropolis
Hastings implementation and a coin flipping model at the end. It is standalone, in that
MiniBM does not depend on the Bean Machine framework at all. The only two dependencies
for MiniBM are the PyTorch library and tqdm (for progress bar).
"""

from __future__ import annotations

import itertools
import random
from collections import defaultdict
from functools import wraps
from typing import Any, Callable, Dict, List, NamedTuple, Optional, Tuple

import torch
import torch.distributions as dist
from tqdm.auto import tqdm


class RVIdentifier(NamedTuple):
    """
    A struct whose attributes uniquely identifies a random variable in Bean Machine.

    Args:
        wrapper: A reference to the decorated random variable function
        args: Arguments taken by the
    """

    wrapper: Callable
    args: Tuple

    @property
    def function(self):
        """A pointer to the original function that returns the distribution object"""
        return self.wrapper.__wrapped__  # calls the original function


def random_variable(f: Callable[Any, dist.Distribution]):
    """A decorator that convert a Python function that returns a distribution into a
    function that evaluates to a Bean Machine random variable.

    In Bean Machine, a @random_variable function can be used in two ways:

    1. When being invoked outside of an inference scope, it returns an RVIdentifier
       without evaluating the original function.
    2. During inference, or when being invoked from another random variable, the
       function will update the graph (if needed) and return its value at state of
       inference.

    For example::

        @random_variable
        def foo():
            return dist.Normal(0., 1.0)

        print(foo())  # RVIdentifier(wrapper=foo, args=())

        @random_variable
        def bar():
            mean = foo()  # evaluates to value of theta() during inference
            return dist.Normal(mean, 1.0)


    Args:
        f: A function that returns a PyTorch Distribution object
    """

    @wraps(f)
    def wrapper(*args):
        rvid = RVIdentifier(wrapper, args)
        # Bean Machine inference methods use the World class to store and control the
        # state of inference
        world = get_world_context()
        if world is None:
            # We're not in an active inference. Return an ID for the random variable
            return rvid
        else:
            # Update the graph and return current value of the random variable in world
            return world.update_graph(rvid)

    return wrapper


RVDict = Dict[RVIdentifier, torch.Tensor]  # alias for typing
WORLD_STACK: List[World] = []


def get_world_context() -> Optional[World]:
    """Returns the active World (if any) or None"""
    return WORLD_STACK[-1] if WORLD_STACK else None


class World:
    """
    A World is Bean Machine's internal representation of a state of the model. At the
    high level, it stores a value for each of the random variables. It can also be used
    as a context manager to control the behavior of random variables. For example::

        @random_variable
        def foo():
            return dist.Normal(0., 1.0)

        @random_variable
        def bar():
            return dist.Normal(foo(), 1.0)

        # initialize world and add bar() and its ancesters to it
        world = World.initialize_world([bar()])
        world[bar()]  # returns the value of bar() in world
        world[foo()]  # since foo() is bar()'s parent, it is also initialized in world

        # World is also used within inference as a context manager to control the
        # behavior of random variable
        with world:
            foo()  # returns the value of foo() in world, which equals to world[foo()]
    """

    def __init__(self, observations: Optional[RVDict] = None):
        self.observations: RVDict = observations or {}
        self.variables: RVDict = {}

    def __getitem__(self, node: RVIdentifier) -> torch.Tensor:
        return self.variables[node]

    def __enter__(self) -> World:
        WORLD_STACK.append(self)
        return self

    def __exit__(self, *args) -> None:
        WORLD_STACK.pop()

    def update_graph(self, node: RVIdentifier) -> torch.Tensor:
        """Update the graphy by adding node to self (if needed) and retuurn the value
        of node in self."""
        if node not in self.variables:
            # parent nodes will be invoked when calling node.get_distribution
            distribution = self.get_distribution(node)
            if node in self.observations:
                self.variables[node] = self.observations[node]
            else:
                self.variables[node] = distribution.sample()
        return self.variables[node]

    def replace(self, values: RVDict) -> World:
        """Return a new world where the values of the random variables are replaced by
        the provided values"""
        new_world = World(self.observations)
        new_world.variables = {**self.variables, **values}
        return new_world

    def log_prob(self) -> torch.Tensor:
        """Return the joint log prob on all random variables in the world"""
        log_prob = torch.tensor(0.0)
        for node, value in self.variables.items():
            distribution = self.get_distribution(node)
            log_prob += distribution.log_prob(value).sum()
        return log_prob

    def get_distribution(self, node: RVIdentifier) -> dist.Distribution:
        """A utility method that activate the current world and invoke the function
        associated with node. Bean Machine requires random variable functions to return
        a distribution object, so this method will also return a distribution object."""
        with self:
            return node.function(*node.args)

    @staticmethod
    def initialize_world(
        queries: List[RVIdentifier], observations: Optional[RVDict] = None
    ) -> World:
        """Initializes and returns a new world. Starting from the queries and
        observations, the parent nodes will be added recursively to the world."""
        observations = observations or {}
        world = World(observations)
        for node in itertools.chain(queries, observations):
            world.update_graph(node)
        return world


class MetropolisHastings:
    """A naive implementation of the `Metropolis-Hastings algorithm
    <https://en.wikipedia.org/wiki/Metropolis%E2%80%93Hastings_algorithm>`_"""

    def infer(
        self,
        queries: List[RVIdentifier],
        observations: Optional[RVDict],
        num_samples: int,
    ) -> RVDict:
        world = World.initialize_world(queries, observations)

        samples = defaultdict(list)
        # the main inference loop
        for _ in tqdm(range(num_samples)):
            latent_nodes = world.variables.keys() - world.observations.keys()
            random.shuffle(latent_nodes)
            # randomly select a node to be updated at a time
            for node in latent_nodes:
                proposer_distribution = world.get_distribution(node)
                new_value = proposer_distribution.sample()
                new_world = world.replace({node: new_value})
                backward_distribution = new_world.get_distribution(node)

                # log P(x, y)
                old_log_prob = world.log_prob()
                # log P(x', y)
                new_log_prob = new_world.log_prob()
                # log g(x'|x)
                forward_log_prob = proposer_distribution.log_prob(new_value).sum()
                # log g(x|x')
                backward_log_prob = backward_distribution.log_prob(world[node]).sum()

                accept_log_prob = (
                    new_log_prob + backward_log_prob - old_log_prob - forward_log_prob
                )
                if torch.bernoulli(accept_log_prob.exp().clamp(max=1)):
                    # accept the new state
                    world = new_world

            # collect the samples before moving to the next iteration
            for node in queries:
                samples[node].append(world[node])
        # stack the list of tensors into a single tensor
        samples = {node: torch.stack(samples[node]) for node in samples}
        return samples


def main():
    # coin fliping model adapted from our tutorial
    # (https://beanmachine.org/docs/overview/tutorials/Coin_flipping/CoinFlipping/)
    @random_variable
    def weight():
        return dist.Beta(2, 2)

    @random_variable
    def y():
        return dist.Bernoulli(weight()).expand((N,))

    # data generation
    true_weight = 0.75
    true_y = dist.Bernoulli(true_weight)
    N = 100
    y_obs = true_y.sample((N,))

    print("Head rate:", y_obs.mean())

    # running inference
    samples = MetropolisHastings().infer([weight()], {y(): y_obs}, num_samples=500)
    print("Estimated weight of the coin:", samples[weight()].mean())


if __name__ == "__main__":
    main()

# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

"""
The file contains MiniBM, a minimal implementation of Bean Machine PPL with a Metropolis
Hastings implementation and a coin flipping model at the end. It is standalone, in that
MiniBM does not depend on the Bean Machine framework at all. The only two dependencies
for MiniBM are the PyTorch library and tqdm (for progress bar).
"""

from __future__ import annotations

import itertools
import random
from collections import defaultdict
from functools import wraps
from typing import Any, Callable, Dict, List, NamedTuple, Optional, Tuple

import torch
import torch.distributions as dist
from tqdm.auto import tqdm


class RVIdentifier(NamedTuple):
    """
    A struct whose attributes uniquely identifies a random variable in Bean Machine.

    Args:
        wrapper: A reference to the decorated random variable function
        args: Arguments taken by the
    """

    wrapper: Callable
    args: Tuple

    @property
    def function(self):
        """A pointer to the original function that returns the distribution object"""
        return self.wrapper.__wrapped__  # calls the original function


def random_variable(f: Callable[Any, dist.Distribution]):
    """A decorator that convert a Python function that returns a distribution into a
    function that evaluates to a Bean Machine random variable.

    In Bean Machine, a @random_variable function can be used in two ways:

    1. When being invoked outside of an inference scope, it returns an RVIdentifier
       without evaluating the original function.
    2. During inference, or when being invoked from another random variable, the
       function will update the graph (if needed) and return its value at state of
       inference.

    For example::

        @random_variable
        def foo():
            return dist.Normal(0., 1.0)

        print(foo())  # RVIdentifier(wrapper=foo, args=())

        @random_variable
        def bar():
            mean = foo()  # evaluates to value of theta() during inference
            return dist.Normal(mean, 1.0)


    Args:
        f: A function that returns a PyTorch Distribution object
    """

    @wraps(f)
    def wrapper(*args):
        rvid = RVIdentifier(wrapper, args)
        # Bean Machine inference methods use the World class to store and control the
        # state of inference
        world = get_world_context()
        if world is None:
            # We're not in an active inference. Return an ID for the random variable
            return rvid
        else:
            # Update the graph and return current value of the random variable in world
            return world.update_graph(rvid)

    return wrapper


RVDict = Dict[RVIdentifier, torch.Tensor]  # alias for typing
WORLD_STACK: List[World] = []


def get_world_context() -> Optional[World]:
    """Returns the active World (if any) or None"""
    return WORLD_STACK[-1] if WORLD_STACK else None


class World:
    """
    A World is Bean Machine's internal representation of a state of the model. At the
    high level, it stores a value for each of the random variables. It can also be used
    as a context manager to control the behavior of random variables. For example::

        @random_variable
        def foo():
            return dist.Normal(0., 1.0)

        @random_variable
        def bar():
            return dist.Normal(foo(), 1.0)

        # initialize world and add bar() and its ancesters to it
        world = World.initialize_world([bar()])
        world[bar()]  # returns the value of bar() in world
        world[foo()]  # since foo() is bar()'s parent, it is also initialized in world

        # World is also used within inference as a context manager to control the
        # behavior of random variable
        with world:
            foo()  # returns the value of foo() in world, which equals to world[foo()]
    """

    def __init__(self, observations: Optional[RVDict] = None):
        self.observations: RVDict = observations or {}
        self.variables: RVDict = {}

    def __getitem__(self, node: RVIdentifier) -> torch.Tensor:
        return self.variables[node]

    def __enter__(self) -> World:
        WORLD_STACK.append(self)
        return self

    def __exit__(self, *args) -> None:
        WORLD_STACK.pop()

    def update_graph(self, node: RVIdentifier) -> torch.Tensor:
        """Update the graphy by adding node to self (if needed) and retuurn the value
        of node in self."""
        if node not in self.variables:
            # parent nodes will be invoked when calling node.get_distribution
            distribution = self.get_distribution(node)
            if node in self.observations:
                self.variables[node] = self.observations[node]
            else:
                self.variables[node] = distribution.sample()
        return self.variables[node]

    def replace(self, values: RVDict) -> World:
        """Return a new world where the values of the random variables are replaced by
        the provided values"""
        new_world = World(self.observations)
        new_world.variables = {**self.variables, **values}
        return new_world

    def log_prob(self) -> torch.Tensor:
        """Return the joint log prob on all random variables in the world"""
        log_prob = torch.tensor(0.0)
        for node, value in self.variables.items():
            distribution = self.get_distribution(node)
            log_prob += distribution.log_prob(value).sum()
        return log_prob

    def get_distribution(self, node: RVIdentifier) -> dist.Distribution:
        """A utility method that activate the current world and invoke the function
        associated with node. Bean Machine requires random variable functions to return
        a distribution object, so this method will also return a distribution object."""
        with self:
            return node.function(*node.args)

    @staticmethod
    def initialize_world(
        queries: List[RVIdentifier], observations: Optional[RVDict] = None
    ) -> World:
        """Initializes and returns a new world. Starting from the queries and
        observations, the parent nodes will be added recursively to the world."""
        observations = observations or {}
        world = World(observations)
        for node in itertools.chain(queries, observations):
            world.update_graph(node)
        return world


class MetropolisHastings:
    """A naive implementation of the `Metropolis-Hastings algorithm
    <https://en.wikipedia.org/wiki/Metropolis%E2%80%93Hastings_algorithm>`_"""

    def infer(
        self,
        queries: List[RVIdentifier],
        observations: Optional[RVDict],
        num_samples: int,
    ) -> RVDict:
        world = World.initialize_world(queries, observations)

        samples = defaultdict(list)
        # the main inference loop
        for _ in tqdm(range(num_samples)):
            latent_nodes = world.variables.keys() - world.observations.keys()
            random.shuffle(latent_nodes)
            # randomly select a node to be updated at a time
            for node in latent_nodes:
                proposer_distribution = world.get_distribution(node)
                new_value = proposer_distribution.sample()
                new_world = world.replace({node: new_value})
                backward_distribution = new_world.get_distribution(node)

                # log P(x, y)
                old_log_prob = world.log_prob()
                # log P(x', y)
                new_log_prob = new_world.log_prob()
                # log g(x'|x)
                forward_log_prob = proposer_distribution.log_prob(new_value).sum()
                # log g(x|x')
                backward_log_prob = backward_distribution.log_prob(world[node]).sum()

                accept_log_prob = (
                    new_log_prob + backward_log_prob - old_log_prob - forward_log_prob
                )
                if torch.bernoulli(accept_log_prob.exp().clamp(max=1)):
                    # accept the new state
                    world = new_world

            # collect the samples before moving to the next iteration
            for node in queries:
                samples[node].append(world[node])
        # stack the list of tensors into a single tensor
        samples = {node: torch.stack(samples[node]) for node in samples}
        return samples


def main():
    # coin fliping model adapted from our tutorial
    # (https://beanmachine.org/docs/overview/tutorials/Coin_flipping/CoinFlipping/)
    @random_variable
    def weight():
        return dist.Beta(2, 2)

    @random_variable
    def y():
        return dist.Bernoulli(weight()).expand((N,))

    # data generation
    true_weight = 0.75
    true_y = dist.Bernoulli(true_weight)
    N = 100
    y_obs = true_y.sample((N,))

    print("Head rate:", y_obs.mean())

    # running inference
    samples = MetropolisHastings().infer([weight()], {y(): y_obs}, num_samples=500)
    print("Estimated weight of the coin:", samples[weight()].mean())


if __name__ == "__main__":
    main()
# Copyright (c) Meta Platforms, Inc. and affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. The file contains MiniBM, a minimal implementation of Bean Machine PPL with a Metropolis Hastings implementation and a coin flipping model at the end. It is standalone, in that MiniBM does not depend on the Bean Machine framework at all. The only two dependencies for MiniBM are the PyTorch library and tqdm (for progress bar). A struct whose attributes uniquely identifies a random variable in Bean Machine.      Args:         wrapper: A reference to the decorated random variable function         args: Arguments taken by the A pointer to the original function that returns the distribution object # calls the original function A decorator that convert a Python function that returns a distribution into a     function that evaluates to a Bean Machine random variable.      In Bean Machine, a @random_variable function can be used in two ways:      1. When being invoked outside of an inference scope, it returns an RVIdentifier        without evaluating the original function.     2. During inference, or when being invoked from another random variable, the        function will update the graph (if needed) and return its value at state of        inference.      For example::          @random_variable         def foo():             return dist.Normal(0., 1.0)          print(foo())  # RVIdentifier(wrapper=foo, args=())          @random_variable         def bar():             mean = foo()  # evaluates to value of theta() during inference             return dist.Normal(mean, 1.0)       Args:         f: A function that returns a PyTorch Distribution object # Bean Machine inference methods use the World class to store and control the # state of inference # We're not in an active inference. Return an ID for the random variable # Update the graph and return current value of the random variable in world # alias for typing Returns the active World (if any) or None A World is Bean Machine's internal representation of a state of the model. At the     high level, it stores a value for each of the random variables. It can also be used     as a context manager to control the behavior of random variables. For example::          @random_variable         def foo():             return dist.Normal(0., 1.0)          @random_variable         def bar():             return dist.Normal(foo(), 1.0)          # initialize world and add bar() and its ancesters to it         world = World.initialize_world([bar()])         world[bar()]  # returns the value of bar() in world         world[foo()]  # since foo() is bar()'s parent, it is also initialized in world          # World is also used within inference as a context manager to control the         # behavior of random variable         with world:             foo()  # returns the value of foo() in world, which equals to world[foo()] Update the graphy by adding node to self (if needed) and retuurn the value         of node in self. # parent nodes will be invoked when calling node.get_distribution Return a new world where the values of the random variables are replaced by         the provided values Return the joint log prob on all random variables in the world A utility method that activate the current world and invoke the function         associated with node. Bean Machine requires random variable functions to return         a distribution object, so this method will also return a distribution object. Initializes and returns a new world. Starting from the queries and         observations, the parent nodes will be added recursively to the world. A naive implementation of the `Metropolis-Hastings algorithm     <https://en.wikipedia.org/wiki/Metropolis%E2%80%93Hastings_algorithm>`_ # the main inference loop # randomly select a node to be updated at a time # log P(x, y) # log P(x', y) # log g(x'|x) # log g(x|x') # accept the new state # collect the samples before moving to the next iteration # stack the list of tensors into a single tensor # coin fliping model adapted from our tutorial # (https://beanmachine.org/docs/overview/tutorials/Coin_flipping/CoinFlipping/) # data generation # running inference