<reponame>JamesBrofos/Stein
import numpy as np
import tensorflow as tf
from abc import abstractmethod
from ..utilities import convert_array_to_dictionary, convert_dictionary_to_array


class AbstractSteinSampler:
    """Abstract Stein Sampler Class

    This class implements the algorithm from the paper "Stein Variational
    Gradient Descent: A General Purpose Bayesian Inference Algorithm" by Liu and
    Wang. This algorithm provides a mechanism for sampling from arbitrary
    distributions provided that the gradient of the distribution can be computed
    with respect to the input.

    The Stein variational gradient descent algorithm seeks to identify the
    optimal perturbation direction for a set of particles that will be
    iteratively transformed such that the empirical distribution of the
    particles can be seen to approximate a sample from the distribution. This is
    achieved by minimizing the KL-divergence between the samples and the target
    distribution, the optimal decrease direction for which can be obtained in
    closed-form and approximated via sampling. In particular, we compute the
    direction of greatest decrease subject to a set of functions of bounded norm
    within a reproducing kernel Hilbert space (RKHS).
    """
    def __init__(self, n_particles, log_p, theta=None):
        """Initialize the parameters of the abstract Stein sampler object.

            n_particles (int): The number of particles to use in the algorithm.
                This is equivalently the number of samples to generate from the
                target distribution.
            log_p (TensorFlow tensor): A TensorFlow object corresponding to the
                log-posterior distribution from which parameters wish to be
                sampled. We only need to define the log-posterior up to an
                addative constant since we'll simply take the gradient with
                respect to the inputs and this term will vanish.
            theta (numpy array, optional): An optional parameter corresponding
                to the initial values of the particles. The dimension of this
                array (if it is provided) should be the number of particles by
                the number of random variables (parameters) to be sampled. If
                this value is not provided, then the initial particles will be
                generated by sampling from a multivariate standard normal
                distribution.
        """
        # Number of particles.
        self.n_particles = n_particles
        # Construct a TensorFlow session.
        self.sess = tf.Session()
        self.model_vars = tf.get_collection(
            tf.GraphKeys.TRAINABLE_VARIABLES, "model"
        )
        # Create class variables for the log-posterior and the gradient of the
        # log-posterior with respect to model parameters.
        self.log_p = log_p
        self.grad_log_p = tf.gradients(self.log_p, self.model_vars)

        # If particles are provided, then use them. Otherwise, particles are
        # initialized by sampling from a standard normal distribution. This
        # latter method works well for relatively simple models, but better
        # initialization is required for complex distributions such as those
        # that are parametrized by neural networks.
        #
        # Notice that `theta` is a dictionary that maps model parameters to a
        # matrix representing the value of that parameter for each of the
        # particles.
        if theta is not None:
            self.theta = theta
        else:
            self.theta = {
                v: np.random.normal(
                    size=[self.n_particles] + v.get_shape().as_list()
                ) * 0.01
                for v in self.model_vars
            }

    def compute_phi(self, theta_array, grads_array):
        """Assuming a reproducing kernel Hilbert space with associated kernel,
        this function computes the optimal perturbation in the particles under
        functions in the unit ball under the norm of the RKHS. This perturbation
        can be regarded as the direction that will maximally decrease the
        KL-divergence between the empirical distribution of the particles and
        the target distribution.

        Parameters:
            theta_array (numpy array): A two-dimensional matrix with dimensions
                equal to the number of particles by the number of parameters.
                This is the matrix representation of the particles.
            grads_array (numpy array): A two-dimensional matrix with dimensions
                equal to the number of particles by the number of parameters.
                This is the matrix representation of the gradient of the
                log-posterior with respect to the particles.

        Returns:
            Numpy array: A two-dimensional matrix with dimensions equal to the
                number of particles by the number of parameters. This is the
                update value corresponding to the optimal perturbation direction
                given by Stein variational gradient descent.
        """
        # Extract the number of particles and number of parameters.
        n_particles, n_params = grads_array.shape
        # Compute the kernel matrices and gradient with respect to the
        # particles.
        K, dK = self.kernel.kernel_and_grad(theta_array)

        return (K.dot(grads_array) + dK) / n_particles

    def update_particles(self, grads_array):
        """Internal method that computes the optimal perturbation direction
        given the current set of particles and the gradient of the
        log-posterior. Notice that this method applies the gradient descent
        update and normalizes the gradient to have a given norm. Computation of
        the optimal perturbation direction is broken out into the method
        `compute_phi`.

        Parameters:
            grads_array (numpy array): A numpy array mapping TensorFlow model
                variables to the gradient of the log-posterior.
        """
        # Convert both the particle dictionary and the gradient dictionary into
        # vector representations.
        theta_array, access_indices = convert_dictionary_to_array(self.theta)
        # Compute optimal update direction.
        phi = self.compute_phi(theta_array, grads_array)
        # Normalize the gradient have be norm no larger than the desired amount.
        phi *= 10. / max(10., np.linalg.norm(phi))
        theta_array += self.gd.update(phi)
        self.theta = convert_array_to_dictionary(theta_array, access_indices)

    def function_posterior(self, func, feed_dict, axis=None):
        """This method computes a posterior distribution of a provided function
        under the posterior distribution learned via Stein variational gradient
        descent. This function can be used to produce posterior mean squared
        errors and the posterior mean log-likelihood. This function includes an
        optional parameter to compute the average of the function's posterior.

        Parameters:
            func (TensorFlow Tensor): A function to be executed in a TensorFlow
                session. The output of the function are averaged across all of
                the posterior samples under the Bayesian model.
            feed_dict (dictionary): A dictionary mapping TensorFlow placeholders
                to provided values.
            average (int, optional): Determines whether or not the average
                of the posterior distribution should be computed instead of
                simply returning samples. If not none, then the average of the
                output of the function is computed across the specified
                dimension.

        Returns:
            Numpy array or Float: If `average` is true then the posterior mean
                is returned. Otherwise, a numpy array of samples from the
                function's posterior are returned.
        """
        # Initialize a vector to store the value of the function for each particle.
        dist = []
        # Iterate over each particle and compute the value of the function for
        # that posterior sample.
        for i in range(self.n_particles):
            feed_dict.update({v: x[i] for v, x in self.theta.items()})
            dist.append(np.ravel(self.sess.run(func, feed_dict)))

        # Convert to a numpy array.
        dist = np.array(dist)
        # Either return posterior samples of the input function or the posterior
        # mean.
        if axis is not None:
            return dist.mean(axis=axis)
        else:
            return dist
        # Bundles was here and wishes penguin the best of luck with his program
        # :3

    @abstractmethod
    def train_on_batch(self, batch_feed):
        """Trains the Stein variational gradient descent algorithm on a given
        batch of data (provided in the form of a TensorFlow feed dictionary).
        This function computes the gradient of the log-likelihood for each
        sampling particle and then computes the optimal perturbation using the
        formula provided in the Stein variational gradient descent paper. Notice
        that, like the particles themselves, the gradients are represented as
        dictionaries that allow to keep the gradients distinct for each
        parameter.

        Notice that this function does not update the TensorFlow variables used
        to define the model. This wouldn't make sense in the first place since
        there are multiple particles corresponding to a random draw of those
        variables. Instead, the `theta` class variable is a dictionary that
        stores all of the particle values for each parameter in the model.

        Parameters:
            batch_feed (dict): A dictionary that maps TensorFlow placeholders to
                provided values. For instance, this might be mappings of feature
                and target placeholders to batch values. Notice that this feed
                dictionary will be internally augmented to include the current
                feed values for the model parameters for each particle.
        """
        raise NotImplementedError()




import numpy as np
import tensorflow as tf
from abc import abstractmethod
from ..utilities import convert_array_to_dictionary, convert_dictionary_to_array


class AbstractSteinSampler:
    """Abstract Stein Sampler Class

    This class implements the algorithm from the paper "Stein Variational
    Gradient Descent: A General Purpose Bayesian Inference Algorithm" by Liu and
    Wang. This algorithm provides a mechanism for sampling from arbitrary
    distributions provided that the gradient of the distribution can be computed
    with respect to the input.

    The Stein variational gradient descent algorithm seeks to identify the
    optimal perturbation direction for a set of particles that will be
    iteratively transformed such that the empirical distribution of the
    particles can be seen to approximate a sample from the distribution. This is
    achieved by minimizing the KL-divergence between the samples and the target
    distribution, the optimal decrease direction for which can be obtained in
    closed-form and approximated via sampling. In particular, we compute the
    direction of greatest decrease subject to a set of functions of bounded norm
    within a reproducing kernel Hilbert space (RKHS).
    """
    def __init__(self, n_particles, log_p, theta=None):
        """Initialize the parameters of the abstract Stein sampler object.

            n_particles (int): The number of particles to use in the algorithm.
                This is equivalently the number of samples to generate from the
                target distribution.
            log_p (TensorFlow tensor): A TensorFlow object corresponding to the
                log-posterior distribution from which parameters wish to be
                sampled. We only need to define the log-posterior up to an
                addative constant since we'll simply take the gradient with
                respect to the inputs and this term will vanish.
            theta (numpy array, optional): An optional parameter corresponding
                to the initial values of the particles. The dimension of this
                array (if it is provided) should be the number of particles by
                the number of random variables (parameters) to be sampled. If
                this value is not provided, then the initial particles will be
                generated by sampling from a multivariate standard normal
                distribution.
        """
        # Number of particles.
        self.n_particles = n_particles
        # Construct a TensorFlow session.
        self.sess = tf.Session()
        self.model_vars = tf.get_collection(
            tf.GraphKeys.TRAINABLE_VARIABLES, "model"
        )
        # Create class variables for the log-posterior and the gradient of the
        # log-posterior with respect to model parameters.
        self.log_p = log_p
        self.grad_log_p = tf.gradients(self.log_p, self.model_vars)

        # If particles are provided, then use them. Otherwise, particles are
        # initialized by sampling from a standard normal distribution. This
        # latter method works well for relatively simple models, but better
        # initialization is required for complex distributions such as those
        # that are parametrized by neural networks.
        #
        # Notice that `theta` is a dictionary that maps model parameters to a
        # matrix representing the value of that parameter for each of the
        # particles.
        if theta is not None:
            self.theta = theta
        else:
            self.theta = {
                v: np.random.normal(
                    size=[self.n_particles] + v.get_shape().as_list()
                ) * 0.01
                for v in self.model_vars
            }

    def compute_phi(self, theta_array, grads_array):
        """Assuming a reproducing kernel Hilbert space with associated kernel,
        this function computes the optimal perturbation in the particles under
        functions in the unit ball under the norm of the RKHS. This perturbation
        can be regarded as the direction that will maximally decrease the
        KL-divergence between the empirical distribution of the particles and
        the target distribution.

        Parameters:
            theta_array (numpy array): A two-dimensional matrix with dimensions
                equal to the number of particles by the number of parameters.
                This is the matrix representation of the particles.
            grads_array (numpy array): A two-dimensional matrix with dimensions
                equal to the number of particles by the number of parameters.
                This is the matrix representation of the gradient of the
                log-posterior with respect to the particles.

        Returns:
            Numpy array: A two-dimensional matrix with dimensions equal to the
                number of particles by the number of parameters. This is the
                update value corresponding to the optimal perturbation direction
                given by Stein variational gradient descent.
        """
        # Extract the number of particles and number of parameters.
        n_particles, n_params = grads_array.shape
        # Compute the kernel matrices and gradient with respect to the
        # particles.
        K, dK = self.kernel.kernel_and_grad(theta_array)

        return (K.dot(grads_array) + dK) / n_particles

    def update_particles(self, grads_array):
        """Internal method that computes the optimal perturbation direction
        given the current set of particles and the gradient of the
        log-posterior. Notice that this method applies the gradient descent
        update and normalizes the gradient to have a given norm. Computation of
        the optimal perturbation direction is broken out into the method
        `compute_phi`.

        Parameters:
            grads_array (numpy array): A numpy array mapping TensorFlow model
                variables to the gradient of the log-posterior.
        """
        # Convert both the particle dictionary and the gradient dictionary into
        # vector representations.
        theta_array, access_indices = convert_dictionary_to_array(self.theta)
        # Compute optimal update direction.
        phi = self.compute_phi(theta_array, grads_array)
        # Normalize the gradient have be norm no larger than the desired amount.
        phi *= 10. / max(10., np.linalg.norm(phi))
        theta_array += self.gd.update(phi)
        self.theta = convert_array_to_dictionary(theta_array, access_indices)

    def function_posterior(self, func, feed_dict, axis=None):
        """This method computes a posterior distribution of a provided function
        under the posterior distribution learned via Stein variational gradient
        descent. This function can be used to produce posterior mean squared
        errors and the posterior mean log-likelihood. This function includes an
        optional parameter to compute the average of the function's posterior.

        Parameters:
            func (TensorFlow Tensor): A function to be executed in a TensorFlow
                session. The output of the function are averaged across all of
                the posterior samples under the Bayesian model.
            feed_dict (dictionary): A dictionary mapping TensorFlow placeholders
                to provided values.
            average (int, optional): Determines whether or not the average
                of the posterior distribution should be computed instead of
                simply returning samples. If not none, then the average of the
                output of the function is computed across the specified
                dimension.

        Returns:
            Numpy array or Float: If `average` is true then the posterior mean
                is returned. Otherwise, a numpy array of samples from the
                function's posterior are returned.
        """
        # Initialize a vector to store the value of the function for each particle.
        dist = []
        # Iterate over each particle and compute the value of the function for
        # that posterior sample.
        for i in range(self.n_particles):
            feed_dict.update({v: x[i] for v, x in self.theta.items()})
            dist.append(np.ravel(self.sess.run(func, feed_dict)))

        # Convert to a numpy array.
        dist = np.array(dist)
        # Either return posterior samples of the input function or the posterior
        # mean.
        if axis is not None:
            return dist.mean(axis=axis)
        else:
            return dist
        # Bundles was here and wishes penguin the best of luck with his program
        # :3

    @abstractmethod
    def train_on_batch(self, batch_feed):
        """Trains the Stein variational gradient descent algorithm on a given
        batch of data (provided in the form of a TensorFlow feed dictionary).
        This function computes the gradient of the log-likelihood for each
        sampling particle and then computes the optimal perturbation using the
        formula provided in the Stein variational gradient descent paper. Notice
        that, like the particles themselves, the gradients are represented as
        dictionaries that allow to keep the gradients distinct for each
        parameter.

        Notice that this function does not update the TensorFlow variables used
        to define the model. This wouldn't make sense in the first place since
        there are multiple particles corresponding to a random draw of those
        variables. Instead, the `theta` class variable is a dictionary that
        stores all of the particle values for each parameter in the model.

        Parameters:
            batch_feed (dict): A dictionary that maps TensorFlow placeholders to
                provided values. For instance, this might be mappings of feature
                and target placeholders to batch values. Notice that this feed
                dictionary will be internally augmented to include the current
                feed values for the model parameters for each particle.
        """
        raise NotImplementedError()
Abstract Stein Sampler Class      This class implements the algorithm from the paper "Stein Variational     Gradient Descent: A General Purpose Bayesian Inference Algorithm" by Liu and     Wang. This algorithm provides a mechanism for sampling from arbitrary     distributions provided that the gradient of the distribution can be computed     with respect to the input.      The Stein variational gradient descent algorithm seeks to identify the     optimal perturbation direction for a set of particles that will be     iteratively transformed such that the empirical distribution of the     particles can be seen to approximate a sample from the distribution. This is     achieved by minimizing the KL-divergence between the samples and the target     distribution, the optimal decrease direction for which can be obtained in     closed-form and approximated via sampling. In particular, we compute the     direction of greatest decrease subject to a set of functions of bounded norm     within a reproducing kernel Hilbert space (RKHS). Initialize the parameters of the abstract Stein sampler object.              n_particles (int): The number of particles to use in the algorithm.                 This is equivalently the number of samples to generate from the                 target distribution.             log_p (TensorFlow tensor): A TensorFlow object corresponding to the                 log-posterior distribution from which parameters wish to be                 sampled. We only need to define the log-posterior up to an                 addative constant since we'll simply take the gradient with                 respect to the inputs and this term will vanish.             theta (numpy array, optional): An optional parameter corresponding                 to the initial values of the particles. The dimension of this                 array (if it is provided) should be the number of particles by                 the number of random variables (parameters) to be sampled. If                 this value is not provided, then the initial particles will be                 generated by sampling from a multivariate standard normal                 distribution. # Number of particles. # Construct a TensorFlow session. # Create class variables for the log-posterior and the gradient of the # log-posterior with respect to model parameters. # If particles are provided, then use them. Otherwise, particles are # initialized by sampling from a standard normal distribution. This # latter method works well for relatively simple models, but better # initialization is required for complex distributions such as those # that are parametrized by neural networks. # # Notice that `theta` is a dictionary that maps model parameters to a # matrix representing the value of that parameter for each of the # particles. Assuming a reproducing kernel Hilbert space with associated kernel,         this function computes the optimal perturbation in the particles under         functions in the unit ball under the norm of the RKHS. This perturbation         can be regarded as the direction that will maximally decrease the         KL-divergence between the empirical distribution of the particles and         the target distribution.          Parameters:             theta_array (numpy array): A two-dimensional matrix with dimensions                 equal to the number of particles by the number of parameters.                 This is the matrix representation of the particles.             grads_array (numpy array): A two-dimensional matrix with dimensions                 equal to the number of particles by the number of parameters.                 This is the matrix representation of the gradient of the                 log-posterior with respect to the particles.          Returns:             Numpy array: A two-dimensional matrix with dimensions equal to the                 number of particles by the number of parameters. This is the                 update value corresponding to the optimal perturbation direction                 given by Stein variational gradient descent. # Extract the number of particles and number of parameters. # Compute the kernel matrices and gradient with respect to the # particles. Internal method that computes the optimal perturbation direction         given the current set of particles and the gradient of the         log-posterior. Notice that this method applies the gradient descent         update and normalizes the gradient to have a given norm. Computation of         the optimal perturbation direction is broken out into the method         `compute_phi`.          Parameters:             grads_array (numpy array): A numpy array mapping TensorFlow model                 variables to the gradient of the log-posterior. # Convert both the particle dictionary and the gradient dictionary into # vector representations. # Compute optimal update direction. # Normalize the gradient have be norm no larger than the desired amount. This method computes a posterior distribution of a provided function         under the posterior distribution learned via Stein variational gradient         descent. This function can be used to produce posterior mean squared         errors and the posterior mean log-likelihood. This function includes an         optional parameter to compute the average of the function's posterior.          Parameters:             func (TensorFlow Tensor): A function to be executed in a TensorFlow                 session. The output of the function are averaged across all of                 the posterior samples under the Bayesian model.             feed_dict (dictionary): A dictionary mapping TensorFlow placeholders                 to provided values.             average (int, optional): Determines whether or not the average                 of the posterior distribution should be computed instead of                 simply returning samples. If not none, then the average of the                 output of the function is computed across the specified                 dimension.          Returns:             Numpy array or Float: If `average` is true then the posterior mean                 is returned. Otherwise, a numpy array of samples from the                 function's posterior are returned. # Initialize a vector to store the value of the function for each particle. # Iterate over each particle and compute the value of the function for # that posterior sample. # Convert to a numpy array. # Either return posterior samples of the input function or the posterior # mean. # Bundles was here and wishes penguin the best of luck with his program # :3 Trains the Stein variational gradient descent algorithm on a given         batch of data (provided in the form of a TensorFlow feed dictionary).         This function computes the gradient of the log-likelihood for each         sampling particle and then computes the optimal perturbation using the         formula provided in the Stein variational gradient descent paper. Notice         that, like the particles themselves, the gradients are represented as         dictionaries that allow to keep the gradients distinct for each         parameter.          Notice that this function does not update the TensorFlow variables used         to define the model. This wouldn't make sense in the first place since         there are multiple particles corresponding to a random draw of those         variables. Instead, the `theta` class variable is a dictionary that         stores all of the particle values for each parameter in the model.          Parameters:             batch_feed (dict): A dictionary that maps TensorFlow placeholders to                 provided values. For instance, this might be mappings of feature                 and target placeholders to batch values. Notice that this feed                 dictionary will be internally augmented to include the current                 feed values for the model parameters for each particle.