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)
