示例#1
0
 def test_get_broadcastable_dist_samples(self, samples_to_broadcast):
     size, samples, broadcast_shape = samples_to_broadcast
     if broadcast_shape is not None:
         size_ = to_tuple(size)
         outs, out_shape = get_broadcastable_dist_samples(
             samples, size=size, return_out_shape=True)
         assert out_shape == broadcast_shape
         for i, o in zip(samples, outs):
             ishape = i.shape
             if ishape[:min([len(size_), len(ishape)])] == size_:
                 expected_shape = (size_ + (1, ) *
                                   (len(broadcast_shape) - len(ishape)) +
                                   ishape[len(size_):])
             else:
                 expected_shape = ishape
             assert o.shape == expected_shape
         assert shapes_broadcasting(*[o.shape
                                      for o in outs]) == broadcast_shape
     else:
         with pytest.raises(ValueError):
             get_broadcastable_dist_samples(samples, size=size)
示例#2
0
 def test_get_broadcastable_dist_samples(self, samples_to_broadcast):
     size, samples, broadcast_shape = samples_to_broadcast
     if broadcast_shape is not None:
         size_ = to_tuple(size)
         outs, out_shape = get_broadcastable_dist_samples(
             samples, size=size, return_out_shape=True
         )
         assert out_shape == broadcast_shape
         for i, o in zip(samples, outs):
             ishape = i.shape
             if ishape[: min([len(size_), len(ishape)])] == size_:
                 expected_shape = (
                     size_
                     + (1,) * (len(broadcast_shape) - len(ishape))
                     + ishape[len(size_) :]
                 )
             else:
                 expected_shape = ishape
             assert o.shape == expected_shape
         assert shapes_broadcasting(*[o.shape for o in outs]) == broadcast_shape
     else:
         with pytest.raises(ValueError):
             get_broadcastable_dist_samples(samples, size=size)
示例#3
0
def generate_samples(generator, *args, **kwargs):
    """Generate samples from the distribution of a random variable.

    Parameters
    ----------
    generator: function
        Function to generate the random samples. The function is
        expected take parameters for generating samples and
        a keyword argument ``size`` which determines the shape
        of the samples.
        The args and kwargs (stripped of the keywords below) will be
        passed to the generator function.

    keyword arguments
    ~~~~~~~~~~~~~~~~~

    dist_shape: int or tuple of int
        The shape of the random variable (i.e., the shape attribute).
    size: int or tuple of int
        The required shape of the samples.
    broadcast_shape: tuple of int or None
        The shape resulting from the broadcasting of the parameters.
        If not specified it will be inferred from the shape of the
        parameters. This may be required when the parameter shape
        does not determine the shape of a single sample, for example,
        the shape of the probabilities in the Categorical distribution.
    not_broadcast_kwargs: dict or None
        Key word argument dictionary to provide to the random generator, which
        must not be broadcasted with the rest of the args and kwargs.

    Any remaining args and kwargs are passed on to the generator function.
    """
    dist_shape = kwargs.pop("dist_shape", ())
    size = kwargs.pop("size", None)
    broadcast_shape = kwargs.pop("broadcast_shape", None)
    not_broadcast_kwargs = kwargs.pop("not_broadcast_kwargs", None)
    if not_broadcast_kwargs is None:
        not_broadcast_kwargs = dict()

    # Parse out raw input parameters for the generator
    args = tuple(p[0] if isinstance(p, tuple) else p for p in args)
    for key in kwargs:
        p = kwargs[key]
        kwargs[key] = p[0] if isinstance(p, tuple) else p

    # Convert size and dist_shape to tuples
    size_tup = to_tuple(size)
    dist_shape = to_tuple(dist_shape)
    if dist_shape[: len(size_tup)] == size_tup:
        # dist_shape is prepended with size_tup. This is not a consequence
        # of the parameters being drawn size_tup times! By chance, the
        # distribution's shape has its first elements equal to size_tup.
        # This means that we must prepend the size_tup to dist_shape, and
        # check if that broadcasts well with the parameters
        _dist_shape = size_tup + dist_shape
    else:
        _dist_shape = dist_shape

    if broadcast_shape is None:
        # If broadcast_shape is not explicitly provided, it is inferred as the
        # broadcasted shape of the input parameter and dist_shape, taking into
        # account the potential size prefix
        inputs = args + tuple(kwargs.values())
        broadcast_shape = broadcast_dist_samples_shape(
            [np.asarray(i).shape for i in inputs] + [_dist_shape], size=size_tup
        )
        # We do this instead of broadcast_distribution_samples to avoid
        # creating a dummy array with dist_shape in memory
        inputs = get_broadcastable_dist_samples(
            inputs,
            size=size_tup,
            must_bcast_with=broadcast_shape,
        )
        # We modify the arguments with their broadcasted counterparts
        args = tuple(inputs[: len(args)])
        for offset, key in enumerate(kwargs):
            kwargs[key] = inputs[len(args) + offset]
    # Update kwargs with the keyword arguments that were not broadcasted
    kwargs.update(not_broadcast_kwargs)

    # We ensure that broadcast_shape is a tuple
    broadcast_shape = to_tuple(broadcast_shape)

    try:
        dist_bcast_shape = broadcast_dist_samples_shape(
            [_dist_shape, broadcast_shape],
            size=size,
        )
    except (ValueError, TypeError):
        raise TypeError(
            """Attempted to generate values with incompatible shapes:
            size: {size}
            size_tup: {size_tup}
            broadcast_shape[:len(size_tup)] == size_tup: {size_prepended}
            dist_shape: {dist_shape}
            broadcast_shape: {broadcast_shape}
        """.format(
                size=size,
                size_tup=size_tup,
                dist_shape=dist_shape,
                broadcast_shape=broadcast_shape,
                size_prepended=broadcast_shape[: len(size_tup)] == size_tup,
            )
        )
    if dist_bcast_shape[: len(size_tup)] == size_tup:
        samples = generator(size=dist_bcast_shape, *args, **kwargs)
    else:
        samples = generator(size=size_tup + dist_bcast_shape, *args, **kwargs)

    return np.asarray(samples)
示例#4
0
    def random(self, point=None, size=None):
        """
        Draw random values from defined ``MixtureSameFamily`` distribution.

        Parameters
        ----------
        point : dict, optional
            Dict of variable values on which random values are to be
            conditioned (uses default point if not specified).
        size : int, optional
            Desired size of random sample (returns one sample if not
            specified).

        Returns
        -------
        array
        """
        sample_shape = to_tuple(size)
        mixture_axis = self.mixture_axis

        # First we draw values for the mixture component weights
        (w,) = draw_values([self.w], point=point, size=size)

        # We now draw random choices from those weights.
        # However, we have to ensure that the number of choices has the
        # sample_shape present.
        w_shape = w.shape
        batch_shape = self.comp_dists.shape[: mixture_axis + 1]
        param_shape = np.broadcast(np.empty(w_shape), np.empty(batch_shape)).shape
        event_shape = self.comp_dists.shape[mixture_axis + 1 :]

        if np.asarray(self.shape).size != 0:
            comp_dists_ndim = len(self.comp_dists.shape)

            # If event_shape of both comp_dists and supplied shape matches,
            # broadcast only batch_shape
            # else broadcast the entire given shape with batch_shape.
            if list(self.shape[mixture_axis - comp_dists_ndim + 1 :]) == list(event_shape):
                dist_shape = np.broadcast(
                    np.empty(self.shape[:mixture_axis]), np.empty(param_shape[:mixture_axis])
                ).shape
            else:
                dist_shape = np.broadcast(
                    np.empty(self.shape), np.empty(param_shape[:mixture_axis])
                ).shape
        else:
            dist_shape = param_shape[:mixture_axis]

        # Try to determine the size that must be used to get the mixture
        # components (i.e. get random choices using w).
        # 1. There must be size independent choices based on w.
        # 2. There must also be independent draws for each non singleton axis
        # of w.
        # 3. There must also be independent draws for each dimension added by
        # self.shape with respect to the w.ndim. These usually correspond to
        # observed variables with batch shapes
        wsh = (1,) * (len(dist_shape) - len(w_shape) + 1) + w_shape[:mixture_axis]
        psh = (1,) * (len(dist_shape) - len(param_shape) + 1) + param_shape[:mixture_axis]
        w_sample_size = []
        # Loop through the dist_shape to get the conditions 2 and 3 first
        for i in range(len(dist_shape)):
            if dist_shape[i] != psh[i] and wsh[i] == 1:
                # self.shape[i] is a non singleton dimension (usually caused by
                # observed data)
                sh = dist_shape[i]
            else:
                sh = wsh[i]
            w_sample_size.append(sh)

        if sample_shape is not None and w_sample_size[: len(sample_shape)] != sample_shape:
            w_sample_size = sample_shape + tuple(w_sample_size)

        choices = random_choice(p=w, size=w_sample_size)

        # We now draw samples from the mixture components random method
        comp_samples = self.comp_dists.random(point=point, size=size)
        if comp_samples.shape[: len(sample_shape)] != sample_shape:
            comp_samples = np.broadcast_to(
                comp_samples,
                shape=sample_shape + comp_samples.shape,
            )

        # At this point the shapes of the arrays involved are:
        # comp_samples.shape = (sample_shape, batch_shape, mixture_axis, event_shape)
        # choices.shape = (sample_shape, batch_shape)
        #
        # To be able to take the choices along the mixture_axis of the
        # comp_samples, we have to add in dimensions to the right of the
        # choices array.
        # We also need to make sure that the batch_shapes of both the comp_samples
        # and choices broadcast with each other.

        choices = np.reshape(choices, choices.shape + (1,) * (1 + len(event_shape)))

        choices, comp_samples = get_broadcastable_dist_samples([choices, comp_samples], size=size)

        # We now take the choices of the mixture components along the mixture_axis
        # but we use the negative index representation to be able to handle the
        # sample_shape
        samples = np.take_along_axis(
            comp_samples, choices, axis=mixture_axis - len(self.comp_dists.shape)
        )

        # The `samples` array still has the `mixture_axis`, so we must remove it:
        output = samples[(..., 0) + (slice(None),) * len(event_shape)]
        return output