def comp_dists(self, comp_dists): self._comp_dists = comp_dists if isinstance(comp_dists, Distribution): self._comp_dist_shapes = to_tuple(comp_dists.shape) self._broadcast_shape = self._comp_dist_shapes self.comp_is_distribution = True else: # Now we check the comp_dists distribution shape, see what # the broadcast shape would be. This shape will be the dist_shape # used by generate samples (the shape of a single random sample) # from the mixture self._comp_dist_shapes = [to_tuple(d.shape) for d in comp_dists] # All component distributions must broadcast with each other try: self._broadcast_shape = np.broadcast( *(np.empty(shape) for shape in self._comp_dist_shapes)).shape except Exception: raise TypeError("Supplied comp_dists shapes do not broadcast " "with each other. comp_dists shapes are: " "{}".format(self._comp_dist_shapes)) # We wrap the _comp_dist.random by adding the kwarg raw_size_, # which will be the size attribute passed to _comp_samples. # _comp_samples then calls generate_samples, which may change the # size value to make it compatible with scipy.stats.*.rvs self._generators = [] for comp_dist in comp_dists: generator = Mixture._comp_dist_random_wrapper(comp_dist.random) self._generators.append(generator) self.comp_is_distribution = False
def _set_values(cls, lower, upper, size, shape, initval): if size is None: size = shape lower = np.asarray(lower) lower = floatX(np.where(lower == None, -np.inf, lower)) upper = np.asarray(upper) upper = floatX(np.where(upper == None, np.inf, upper)) if initval is None: _size = np.broadcast_shapes(to_tuple(size), np.shape(lower), np.shape(upper)) _lower = np.broadcast_to(lower, _size) _upper = np.broadcast_to(upper, _size) initval = np.where( (_lower == -np.inf) & (_upper == np.inf), 0, np.where( _lower == -np.inf, _upper - 1, np.where(_upper == np.inf, _lower + 1, (_lower + _upper) / 2), ), ) lower = as_tensor_variable(floatX(lower)) upper = as_tensor_variable(floatX(upper)) return lower, upper, initval
def test_sample_generate_values(fixture_model, fixture_sizes): model, RVs = fixture_model size = to_tuple(fixture_sizes) with model: prior = pm.sample_prior_predictive(samples=fixture_sizes) for rv in RVs: assert prior[rv.name].shape == size + tuple(rv.distribution.shape)
def test_density_dist_default_moment_multivariate(with_random, size): def _random(mu, rng=None, size=None): return rng.normal(mu, scale=1, size=to_tuple(size) + mu.shape) if with_random: random = _random else: random = None mu_val = np.random.normal(loc=2, scale=1, size=5).astype(aesara.config.floatX) with pm.Model(): mu = pm.Normal("mu", size=5) a = pm.DensityDist("a", mu, random=random, ndims_params=[1], ndim_supp=1, size=size) if with_random: evaled_moment = get_moment(a).eval({mu: mu_val}) assert evaled_moment.shape == to_tuple(size) + (5, ) assert np.all(evaled_moment == 0) else: with pytest.raises( TypeError, match= "Cannot safely infer the size of a multivariate random variable's moment.", ): evaled_moment = get_moment(a).eval({mu: mu_val})
def random_choice(p, size): """Return draws from categorical probability functions Args: p: array Probability of each class. If p.ndim > 1, the last axis is interpreted as the probability of each class, and numpy.random.choice is iterated for every other axis element. size: int or tuple Shape of the desired output array. If p is multidimensional, size should broadcast with p.shape[:-1]. Returns: random sample: array """ k = p.shape[-1] if p.ndim > 1: # If p is an nd-array, the last axis is interpreted as the class # probability. We must iterate over the elements of all the other # dimensions. # We first ensure that p is broadcasted to the output's shape size = to_tuple(size) + (1, ) p = np.broadcast_arrays(p, np.empty(size))[0] out_shape = p.shape[:-1] # np.random.choice accepts 1D p arrays, so we semiflatten p to # iterate calls using the last axis as the category probabilities p = np.reshape(p, (-1, p.shape[-1])) samples = np.array([np.random.choice(k, p=p_) for p_ in p]) # We reshape to the desired output shape samples = np.reshape(samples, out_shape) else: samples = np.random.choice(k, p=p, size=size) return samples
def generate_poisson_mixture_data(w, mu, size=1000): component = np.random.choice(w.size, size=size, p=w) mu = np.atleast_1d(mu) out_size = to_tuple(size) + mu.shape[:-1] mu_ = np.array([mu[..., comp] for comp in component.ravel()]) mu_ = np.reshape(mu_, out_size) x = np.random.poisson(mu_, size=out_size) return x
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 get_steps( steps: Optional[Union[int, np.ndarray, TensorVariable]], *, shape: Optional[Shape] = None, dims: Optional[Dims] = None, observed: Optional[Any] = None, step_shape_offset: int = 0, ): """Extract number of steps from shape / dims / observed information Parameters ---------- steps: User specified steps for timeseries distribution shape: User specified shape for timeseries distribution dims: User specified dims for timeseries distribution observed: User specified observed data from timeseries distribution step_shape_offset: Difference between last shape dimension and number of steps in timeseries distribution, defaults to 0 Returns ------- steps Steps, if specified directly by user, or inferred from the last dimension of shape / dims / observed. When two sources of step information are provided, a symbolic Assert is added to ensure they are consistent. """ inferred_steps = None if shape is not None: shape = to_tuple(shape) if shape[-1] is not ...: inferred_steps = shape[-1] - step_shape_offset if inferred_steps is None and dims is not None: dims = convert_dims(dims) if dims[-1] is not ...: model = modelcontext(None) inferred_steps = model.dim_lengths[dims[-1]] - step_shape_offset if inferred_steps is None and observed is not None: observed = convert_observed_data(observed) inferred_steps = observed.shape[-1] - step_shape_offset if inferred_steps is None: inferred_steps = steps # If there are two sources of information for the steps, assert they are consistent elif steps is not None: inferred_steps = Assert(msg="Steps do not match last shape dimension")( inferred_steps, at.eq(inferred_steps, steps) ) return inferred_steps
def test_density_dist_custom_moment_multivariate(size): def moment(rv, size, mu): return (at.ones(size)[..., None] * mu).astype(rv.dtype) mu_val = np.random.normal(loc=2, scale=1, size=5).astype(aesara.config.floatX) with pm.Model(): mu = pm.Normal("mu", size=5) a = pm.DensityDist("a", mu, get_moment=moment, ndims_params=[1], ndim_supp=1, size=size) evaled_moment = get_moment(a).eval({mu: mu_val}) assert evaled_moment.shape == to_tuple(size) + (5,) assert np.all(evaled_moment == mu_val)
def test_density_dist_custom_moment_univariate(size): def moment(rv, size, mu): return (at.ones(size) * mu).astype(rv.dtype) mu_val = np.array(np.random.normal(loc=2, scale=1)).astype(aesara.config.floatX) with pm.Model(): mu = pm.Normal("mu") a = pm.DensityDist("a", mu, get_moment=moment, size=size) evaled_moment = get_moment(a).eval({mu: mu_val}) assert evaled_moment.shape == to_tuple(size) assert np.all(evaled_moment == mu_val)
def generate_normal_mixture_data(w, mu, sigma, size=1000): component = np.random.choice(w.size, size=size, p=w) mu, sigma = np.broadcast_arrays(mu, sigma) out_size = to_tuple(size) + mu.shape[:-1] mu_ = np.array([mu[..., comp] for comp in component.ravel()]) sigma_ = np.array([sigma[..., comp] for comp in component.ravel()]) mu_ = np.reshape(mu_, out_size) sigma_ = np.reshape(sigma_, out_size) x = np.random.normal(mu_, sigma_, size=out_size) return x
def test_normal_mixture_nd(self, nd, ncomp): nd = to_tuple(nd) ncomp = int(ncomp) comp_shape = nd + (ncomp,) test_mus = np.random.randn(*comp_shape) test_taus = np.random.gamma(1, 1, size=comp_shape) observed = generate_normal_mixture_data( w=np.ones(ncomp) / ncomp, mu=test_mus, sigma=1 / np.sqrt(test_taus), size=10 ) with Model() as model0: mus = Normal("mus", shape=comp_shape) taus = Gamma("taus", alpha=1, beta=1, shape=comp_shape) ws = Dirichlet("ws", np.ones(ncomp), shape=(ncomp,)) mixture0 = NormalMixture("m", w=ws, mu=mus, tau=taus, shape=nd, comp_shape=comp_shape) obs0 = NormalMixture( "obs", w=ws, mu=mus, tau=taus, comp_shape=comp_shape, observed=observed ) with Model() as model1: mus = Normal("mus", shape=comp_shape) taus = Gamma("taus", alpha=1, beta=1, shape=comp_shape) ws = Dirichlet("ws", np.ones(ncomp), shape=(ncomp,)) comp_dist = [ Normal.dist(mu=mus[..., i], tau=taus[..., i], shape=nd) for i in range(ncomp) ] mixture1 = Mixture("m", w=ws, comp_dists=comp_dist, shape=nd) obs1 = Mixture("obs", w=ws, comp_dists=comp_dist, observed=observed) with Model() as model2: # Test that results are correct without comp_shape being passed to the Mixture. # This used to fail in V3 mus = Normal("mus", shape=comp_shape) taus = Gamma("taus", alpha=1, beta=1, shape=comp_shape) ws = Dirichlet("ws", np.ones(ncomp), shape=(ncomp,)) mixture2 = NormalMixture("m", w=ws, mu=mus, tau=taus, shape=nd) obs2 = NormalMixture("obs", w=ws, mu=mus, tau=taus, observed=observed) testpoint = model0.compute_initial_point() testpoint["mus"] = test_mus testpoint["taus_log__"] = np.log(test_taus) for logp0, logp1, logp2 in zip( model0.compile_logp(vars=[mixture0, obs0], sum=False)(testpoint), model1.compile_logp(vars=[mixture1, obs1], sum=False)(testpoint), model2.compile_logp(vars=[mixture2, obs2], sum=False)(testpoint), ): assert_allclose(logp0, logp1) assert_allclose(logp0, logp2)
def change_size(cls, rv, new_size, expand=False): weights = rv.tag.weights components = rv.tag.components if expand: component = rv.tag.components[0] # Old size is equal to `shape[:-ndim_supp]`, with care needed for `ndim_supp == 0` size_dims = component.ndim - component.owner.op.ndim_supp if len(rv.tag.components) == 1: # If we have a single component, new size should ignore the mixture axis # dimension, as that is not touched by `_resize_components` size_dims -= 1 old_size = components[0].shape[:size_dims] new_size = to_tuple(new_size) + tuple(old_size) components = cls._resize_components(new_size, *components) return cls.rv_op(weights, *components, size=None)
def __init__(self, w, comp_dists, mixture_axis=-1, *args, **kwargs): self.w = at.as_tensor_variable(w) if not isinstance(comp_dists, Distribution): raise TypeError( "The MixtureSameFamily distribution only accepts Distribution " f"instances as its components. Got {type(comp_dists)} instead." ) self.comp_dists = comp_dists if mixture_axis < 0: mixture_axis = len(comp_dists.shape) + mixture_axis if mixture_axis < 0: raise ValueError( "`mixture_axis` is supposed to be in shape of components' distribution. " f"Got {mixture_axis + len(comp_dists.shape)} axis instead out of the bounds." ) comp_shape = to_tuple(comp_dists.shape) self.shape = comp_shape[:mixture_axis] + comp_shape[mixture_axis + 1:] self.mixture_axis = mixture_axis kwargs.setdefault("dtype", self.comp_dists.dtype) # Compute the mode so we don't always have to pass a initval defaults = kwargs.pop("defaults", []) event_shape = self.comp_dists.shape[mixture_axis + 1:] _w = at.shape_padleft( at.shape_padright(w, len(event_shape)), len(self.comp_dists.shape) - w.ndim - len(event_shape), ) mode = take_along_axis( self.comp_dists.mode, at.argmax(_w, keepdims=True), axis=mixture_axis, ) self.mode = mode[(..., 0) + (slice(None), ) * len(event_shape)] if not all_discrete(comp_dists): mean = at.as_tensor_variable(self.comp_dists.mean) self.mean = (_w * mean).sum(axis=mixture_axis) if "mean" not in defaults: defaults.append("mean") defaults.append("mode") super().__init__(defaults=defaults, *args, **kwargs)
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)
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 _random(mu, rng=None, size=None): return rng.normal(mu, scale=1, size=to_tuple(size) + mu.shape)
def test_normal_mixture_nd(self, nd, ncomp): nd = to_tuple(nd) ncomp = int(ncomp) comp_shape = nd + (ncomp,) test_mus = np.random.randn(*comp_shape) test_taus = np.random.gamma(1, 1, size=comp_shape) observed = generate_normal_mixture_data( w=np.ones(ncomp) / ncomp, mu=test_mus, sd=1 / np.sqrt(test_taus), size=10 ) with Model() as model0: mus = Normal("mus", shape=comp_shape) taus = Gamma("taus", alpha=1, beta=1, shape=comp_shape) ws = Dirichlet("ws", np.ones(ncomp), shape=(ncomp,)) mixture0 = NormalMixture("m", w=ws, mu=mus, tau=taus, shape=nd, comp_shape=comp_shape) obs0 = NormalMixture( "obs", w=ws, mu=mus, tau=taus, shape=nd, comp_shape=comp_shape, observed=observed ) with Model() as model1: mus = Normal("mus", shape=comp_shape) taus = Gamma("taus", alpha=1, beta=1, shape=comp_shape) ws = Dirichlet("ws", np.ones(ncomp), shape=(ncomp,)) comp_dist = [ Normal.dist(mu=mus[..., i], tau=taus[..., i], shape=nd) for i in range(ncomp) ] mixture1 = Mixture("m", w=ws, comp_dists=comp_dist, shape=nd) obs1 = Mixture("obs", w=ws, comp_dists=comp_dist, shape=nd, observed=observed) with Model() as model2: # Expected to fail if comp_shape is not provided, # nd is multidim and it does not broadcast with ncomp. If by chance # it does broadcast, an error is raised if the mixture is given # observed data. # Furthermore, the Mixture will also raise errors when the observed # data is multidimensional but it does not broadcast well with # comp_dists. mus = Normal("mus", shape=comp_shape) taus = Gamma("taus", alpha=1, beta=1, shape=comp_shape) ws = Dirichlet("ws", np.ones(ncomp), shape=(ncomp,)) if len(nd) > 1: if nd[-1] != ncomp: with pytest.raises(ValueError): NormalMixture("m", w=ws, mu=mus, tau=taus, shape=nd) mixture2 = None else: mixture2 = NormalMixture("m", w=ws, mu=mus, tau=taus, shape=nd) else: mixture2 = NormalMixture("m", w=ws, mu=mus, tau=taus, shape=nd) observed_fails = False if len(nd) >= 1 and nd != (1,): try: np.broadcast(np.empty(comp_shape), observed) except Exception: observed_fails = True if observed_fails: with pytest.raises(ValueError): NormalMixture("obs", w=ws, mu=mus, tau=taus, shape=nd, observed=observed) obs2 = None else: obs2 = NormalMixture("obs", w=ws, mu=mus, tau=taus, shape=nd, observed=observed) testpoint = model0.recompute_initial_point() testpoint["mus"] = test_mus testpoint["taus"] = test_taus assert_allclose(model0.logp(testpoint), model1.logp(testpoint)) assert_allclose(mixture0.logp(testpoint), mixture1.logp(testpoint)) assert_allclose(obs0.logp(testpoint), obs1.logp(testpoint)) if mixture2 is not None and obs2 is not None: assert_allclose(model0.logp(testpoint), model2.logp(testpoint)) if mixture2 is not None: assert_allclose(mixture0.logp(testpoint), mixture2.logp(testpoint)) if obs2 is not None: assert_allclose(obs0.logp(testpoint), obs2.logp(testpoint))
def random(self, point=None, size=None): """ Draw random values from MvGaussianRandomWalk. Parameters ---------- point: dict, optional Dict of variable values on which random values are to be conditioned (uses default point if not specified). size: int or tuple of ints, optional Desired size of random sample (returns one sample if not specified). Returns ------- array Examples ------- .. code-block:: python with pm.Model(): mu = np.array([1.0, 0.0]) cov = np.array([[1.0, 0.0], [0.0, 2.0]]) # draw one sample from a 2-dimensional Gaussian random walk with 10 timesteps sample = MvGaussianRandomWalk(mu, cov, shape=(10, 2)).random() # draw three samples from a 2-dimensional Gaussian random walk with 10 timesteps sample = MvGaussianRandomWalk(mu, cov, shape=(10, 2)).random(size=3) # draw four samples from a 2-dimensional Gaussian random walk with 10 timesteps, # indexed with a (2, 2) array sample = MvGaussianRandomWalk(mu, cov, shape=(10, 2)).random(size=(2, 2)) """ # for each draw specified by the size input, we need to draw time_steps many # samples from MvNormal. size = to_tuple(size) multivariate_samples = self.innov.random(point=point, size=size) # this has shape (size, self.shape) if len(self.shape) == 2: # have time dimension in first slot of shape. Therefore the time # component can be accessed with the index equal to the length of size. time_axis = len(size) multivariate_samples = multivariate_samples.cumsum(axis=time_axis) if time_axis != 0: # this for loop covers the case where size is a tuple for idx in np.ndindex(size): multivariate_samples[idx] = ( multivariate_samples[idx] - multivariate_samples[idx][0] ) else: # size was passed as None multivariate_samples = multivariate_samples - multivariate_samples[0] # if the above statement fails, then only a spatial dimension was passed in for self.shape. # Therefore don't subtract off the initial value since otherwise you get all zeros # as your output. return multivariate_samples