def samples_to_broadcast(fixture_sizes, fixture_shapes): samples = [np.empty(s) for s in fixture_shapes] try: broadcast_shape = broadcast_dist_samples_shape(fixture_shapes, size=fixture_sizes) except ValueError: broadcast_shape = None return fixture_sizes, samples, broadcast_shape
def random(self, point=None, size=None, **kwargs): if self.rand is not None: not_broadcast_kwargs = dict(point=point) not_broadcast_kwargs.update(**kwargs) if self.wrap_random_with_dist_shape: size = to_tuple(size) with _DrawValuesContextBlocker(): test_draw = generate_samples( self.rand, size=None, not_broadcast_kwargs=not_broadcast_kwargs, ) test_shape = test_draw.shape if self.shape[:len(size)] == size: dist_shape = size + self.shape else: dist_shape = self.shape broadcast_shape = broadcast_dist_samples_shape( [dist_shape, test_shape], size=size) broadcast_shape = broadcast_shape[:len(broadcast_shape) - len(test_shape)] samples = generate_samples( self.rand, broadcast_shape=broadcast_shape, size=size, not_broadcast_kwargs=not_broadcast_kwargs, ) else: samples = self.rand(point=point, size=size, **kwargs) if self.check_shape_in_random: expected_shape = self.shape if size is None else to_tuple( size) + self.shape if not expected_shape == samples.shape: raise RuntimeError( "DensityDist encountered a shape inconsistency " "while drawing samples using the supplied random " "function. Was expecting to get samples of shape " "{expected} but got {got} instead.\n" "Whenever possible wrap_random_with_dist_shape = True " "is recommended.\n" "Be aware that the random callable provided as the " "DensityDist random method cannot " "adapt to shape changes in the distribution's " "shape, which sometimes are necessary for sampling " "when the model uses pymc3.Data or aesara shared " "tensors, or when the DensityDist has observed " "values.\n" "This check can be disabled by passing " "check_shape_in_random=False when the DensityDist " "is initialized.".format( expected=expected_shape, got=samples.shape, )) return samples else: raise ValueError( "Distribution was not passed any random method. " "Define a custom random method and pass it as kwarg random")
def samples_to_broadcast(fixture_sizes, fixture_shapes): samples = [np.empty(s) for s in fixture_shapes] try: broadcast_shape = broadcast_dist_samples_shape( fixture_shapes, size=fixture_sizes ) except ValueError: broadcast_shape = None return fixture_sizes, samples, broadcast_shape
def samples_to_broadcast_to(request, samples_to_broadcast): to_shape = request.param size, samples, broadcast_shape = samples_to_broadcast if broadcast_shape is not None: try: broadcast_shape = broadcast_dist_samples_shape( [broadcast_shape, to_tuple(to_shape)], size=size) except ValueError: broadcast_shape = None return to_shape, size, samples, broadcast_shape
def samples_to_broadcast_to(request, samples_to_broadcast): to_shape = request.param size, samples, broadcast_shape = samples_to_broadcast if broadcast_shape is not None: try: broadcast_shape = broadcast_dist_samples_shape( [broadcast_shape, to_tuple(to_shape)], size=size ) except ValueError: broadcast_shape = None return to_shape, size, samples, broadcast_shape
def test_broadcast_dist_samples_shape(self, fixture_sizes, fixture_shapes): size = fixture_sizes shapes = fixture_shapes size_ = to_tuple(size) shapes_ = [ s if s[:min([len(size_), len(s)])] != size_ else s[len(size_):] for s in shapes ] try: expected_out = np.broadcast(*[np.empty(s) for s in shapes_]).shape except ValueError: expected_out = None if expected_out is not None and any( s[:min([len(size_), len(s)])] == size_ for s in shapes): expected_out = size_ + expected_out if expected_out is None: with pytest.raises(ValueError): broadcast_dist_samples_shape(shapes, size=size) else: out = broadcast_dist_samples_shape(shapes, size=size) assert out == expected_out
def test_broadcast_dist_samples_shape(self, fixture_sizes, fixture_shapes): size = fixture_sizes shapes = fixture_shapes size_ = to_tuple(size) shapes_ = [ s if s[: min([len(size_), len(s)])] != size_ else s[len(size_) :] for s in shapes ] try: expected_out = np.broadcast(*[np.empty(s) for s in shapes_]).shape except ValueError: expected_out = None if expected_out is not None and any( (s[: min([len(size_), len(s)])] == size_ for s in shapes) ): expected_out = size_ + expected_out if expected_out is None: with pytest.raises(ValueError): broadcast_dist_samples_shape(shapes, size=size) else: out = broadcast_dist_samples_shape(shapes, size=size) assert out == expected_out
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)