def test_logp_helper_exceptions():
with pytest.raises(TypeError, match="When RV is not a pure distribution"):
logp(at.exp(Normal.dist()), [1, 2])
with pytest.raises(NotImplementedError,
match="PyMC could not infer logp of input variable"):
logp(at.cos(Normal.dist()), 1)
def marginal_mixture_logprob(op, values, rng, weights, *components, **kwargs):
(value, ) = values
# single component
if len(components) == 1:
# Need to broadcast value across mixture axis
mix_axis = -components[0].owner.op.ndim_supp - 1
components_logp = logp(components[0], at.expand_dims(value, mix_axis))
else:
components_logp = at.stack(
[logp(component, value) for component in components],
axis=-1,
)
mix_logp = at.logsumexp(at.log(weights) + components_logp, axis=-1)
mix_logp = check_parameters(
mix_logp,
0 <= weights,
weights <= 1,
at.isclose(at.sum(weights, axis=-1), 1),
msg="0 <= weights <= 1, sum(weights) == 1",
)
return mix_logp
def test_single_multivariate_component_deterministic_weights(self, weights, component, size):
# This test needs seeding to avoid repetitions
rngs = [
aesara.shared(np.random.default_rng(seed))
for seed in self.get_random_state().randint(2**30, size=2)
]
mix = Mixture.dist(weights, component, size=size, rngs=rngs)
mix_eval = mix.eval()
# Test shape
# component shape is either (4, 2, 3), (2, 3)
# weights shape is either (4, 2) or (2,)
if size is not None:
expected_shape = size + (3,)
elif component.ndim == 3 or weights.ndim == 2:
expected_shape = (4, 3)
else:
expected_shape = (3,)
assert mix_eval.shape == expected_shape
# Test draws
totals = mix_eval.sum(-1)
expected_large_count = (weights == 1)[..., 1]
assert np.all((totals == 10_000) == expected_large_count)
repetitions = np.unique(mix_eval[..., 0]).size < totals.size
assert not repetitions
# Test logp
mix_logp_eval = logp(mix, mix_eval).eval()
expected_logp_shape = expected_shape[:-1]
assert mix_logp_eval.shape == expected_logp_shape
bcast_weights = np.broadcast_to(weights, (*expected_logp_shape, 2))
expected_logp = logp(component, mix_eval[..., None, :]).eval()[bcast_weights == 1]
expected_logp = expected_logp.reshape(expected_logp_shape)
assert np.allclose(mix_logp_eval, expected_logp)
def ar_logp(op, values, rhos, sigma, init_dist, steps, noise_rng, **kwargs):
(value,) = values
ar_order = op.ar_order
constant_term = op.constant_term
# Convolve rhos with values
if constant_term:
expectation = at.add(
rhos[..., 0, None],
*(
rhos[..., i + 1, None] * value[..., ar_order - (i + 1) : -(i + 1)]
for i in range(ar_order)
),
)
else:
expectation = at.add(
*(
rhos[..., i, None] * value[..., ar_order - (i + 1) : -(i + 1)]
for i in range(ar_order)
)
)
# Compute and collapse logp across time dimension
innov_logp = at.sum(
logp(Normal.dist(0, sigma[..., None]), value[..., ar_order:] - expectation), axis=-1
)
init_logp = logp(init_dist, value[..., :ar_order])
if init_dist.owner.op.ndim_supp == 0:
init_logp = at.sum(init_logp, axis=-1)
return init_logp + innov_logp
def marginal_mixture_logprob(op, values, rng, weights, *components, **kwargs):
(value, ) = values
# single component
if len(components) == 1:
# Need to broadcast value across mixture axis
mix_axis = -components[0].owner.op.ndim_supp - 1
components_logp = logp(components[0], at.expand_dims(value, mix_axis))
else:
components_logp = at.stack(
[logp(component, value) for component in components],
axis=-1,
)
mix_logp = at.logsumexp(at.log(weights) + components_logp, axis=-1)
# Squeeze stack dimension
# There is a Aesara bug in squeeze with negative axis
# https://github.com/aesara-devs/aesara/issues/830
# mix_logp = at.squeeze(mix_logp, axis=-1)
mix_logp = at.squeeze(mix_logp, axis=mix_logp.ndim - 1)
mix_logp = check_parameters(
mix_logp,
0 <= weights,
weights <= 1,
at.isclose(at.sum(weights, axis=-1), 1),
msg="0 <= weights <= 1, sum(weights) == 1",
)
return mix_logp
def test_logp_helper():
value = at.vector("value")
x = Normal.dist(0, 1)
x_logp = logp(x, value)
np.testing.assert_almost_equal(x_logp.eval({value: [0, 1]}), sp.norm(0, 1).logpdf([0, 1]))
x_logp = logp(x, [0, 1])
np.testing.assert_almost_equal(x_logp.eval(), sp.norm(0, 1).logpdf([0, 1]))
def test_multivariate_init_dist(self):
init_dist = Dirichlet.dist(a=np.full((5, 2), [1, 10]))
x = AR.dist(rho=[0, 0], init_dist=init_dist, steps=0)
x_eval = x.eval()
assert x_eval.shape == (5, 2)
init_dist_eval = init_dist.eval()
init_dist_logp_eval = logp(init_dist, init_dist_eval).eval()
x_logp_eval = logp(x, init_dist_eval).eval()
assert x_logp_eval.shape == (5, )
assert np.allclose(x_logp_eval, init_dist_logp_eval)
def test_with_mvnormal(self):
# 10 batch, 3-variate Gaussian
mu = np.random.randn(self.mixture_comps, 3)
mat = np.random.randn(3, 3)
cov = mat @ mat.T
chol = np.linalg.cholesky(cov)
w = np.ones(self.mixture_comps) / self.mixture_comps
with Model() as model:
comp_dists = MvNormal.dist(mu=mu, chol=chol, shape=(self.mixture_comps, 3))
mixture = Mixture("mixture", w=w, comp_dists=comp_dists, shape=(3,))
prior = sample_prior_predictive(samples=self.n_samples, return_inferencedata=False)
assert prior["mixture"].shape == (self.n_samples, 3)
assert draw(mixture, draws=self.size).shape == (self.size, 3)
if aesara.config.floatX == "float32":
rtol = 1e-4
else:
rtol = 1e-7
initial_point = model.compute_initial_point()
comp_logp = logp(comp_dists, initial_point["mixture"].reshape(1, 3))
log_sum_exp = logsumexp(comp_logp.eval() + np.log(w), axis=0, keepdims=True).sum()
assert_allclose(
model.compile_logp()(initial_point),
log_sum_exp,
rtol,
)
def test_with_multinomial(self, batch_shape):
p = np.random.uniform(size=(*batch_shape, self.mixture_comps, 3))
n = 100 * np.ones((*batch_shape, 1))
w = np.ones(self.mixture_comps) / self.mixture_comps
mixture_axis = len(batch_shape)
with Model() as model:
comp_dists = Multinomial.dist(p=p, n=n, shape=(*batch_shape, self.mixture_comps, 3))
mixture = Mixture(
"mixture",
w=w,
comp_dists=comp_dists,
shape=(*batch_shape, 3),
)
prior = sample_prior_predictive(samples=self.n_samples, return_inferencedata=False)
assert prior["mixture"].shape == (self.n_samples, *batch_shape, 3)
assert draw(mixture, draws=self.size).shape == (self.size, *batch_shape, 3)
if aesara.config.floatX == "float32":
rtol = 1e-4
else:
rtol = 1e-7
initial_point = model.compute_initial_point()
comp_logp = logp(comp_dists, initial_point["mixture"].reshape(*batch_shape, 1, 3))
log_sum_exp = logsumexp(
comp_logp.eval() + np.log(w), axis=mixture_axis, keepdims=True
).sum()
assert_allclose(
model.compile_logp()(initial_point),
log_sum_exp,
rtol,
)
def logp(value, distribution, lower, upper):
"""
Calculate log-probability of Bounded distribution at specified value.
Parameters
----------
value: numeric
Value for which log-probability is calculated.
distribution: TensorVariable
Distribution which is being bounded
lower: numeric
Lower bound for the distribution being bounded.
upper: numeric
Upper bound for the distribution being bounded.
Returns
-------
TensorVariable
"""
res = at.switch(
at.or_(at.lt(value, lower), at.gt(value, upper)),
-np.inf,
logp(distribution, value),
)
return check_parameters(
res,
lower <= upper,
msg="lower <= upper",
)
def test_list_univariate_components_deterministic_weights(self, weights, components, size):
# Size can't be smaller than what is implied by replication dimensions
if size is not None and len(size) < max(components[0].ndim, weights.ndim - 1):
return
mix = Mixture.dist(weights, components, size=size)
mix_eval = mix.eval()
# Test shape
# components[0] shape is either (4, 3), (3,) or ()
# weights shape is either (3, 2) or (2,)
if size is not None:
expected_shape = size
elif components[0].ndim == 2:
expected_shape = (4, 3)
elif components[0].ndim == 1 or weights.ndim == 2:
expected_shape = (3,)
else:
expected_shape = ()
assert mix_eval.shape == expected_shape
# Test draws
expected_positive = np.zeros_like(mix_eval)
if expected_positive.ndim > 0:
expected_positive[..., :] = (weights == 1)[..., 1]
assert np.all((mix_eval > 0) == expected_positive)
repetitions = np.unique(mix_eval).size < mix_eval.size
assert not repetitions
# Test logp
mix_logp_eval = logp(mix, mix_eval).eval()
assert mix_logp_eval.shape == expected_shape
bcast_weights = np.broadcast_to(weights, (*expected_shape, 2))
expected_logp = np.stack(
(
logp(components[0], mix_eval).eval(),
logp(components[1], mix_eval).eval(),
),
axis=-1,
)[bcast_weights == 1]
expected_logp = expected_logp.reshape(expected_shape)
assert np.allclose(mix_logp_eval, expected_logp)
def test_logcdf_transformed_argument():
with Model() as m:
sigma = HalfFlat("sigma")
x = Normal("x", 0, sigma)
Potential("norm_term", -logcdf(x, 1.0))
sigma_value_log = -1.0
sigma_value = np.exp(sigma_value_log)
x_value = 0.5
observed = m.logp_nojac({"sigma_log__": sigma_value_log, "x": x_value})
expected = logp(TruncatedNormal.dist(0, sigma_value, lower=None, upper=1.0), x_value).eval()
assert np.isclose(observed, expected)
def test_list_multivariate_components_deterministic_weights(self, weights, components, size):
mix = Mixture.dist(weights, components, size=size)
mix_eval = mix.eval()
# Test shape
# components[0] shape is either (4, 3) or (3,)
# weights shape is either (4, 2) or (2,)
if size is not None:
expected_shape = size + (3,)
elif components[0].ndim == 2 or weights.ndim == 2:
expected_shape = (4, 3)
else:
expected_shape = (3,)
assert mix_eval.shape == expected_shape
# Test draws
expected_positive = np.zeros_like(mix_eval)
expected_positive[..., :] = (weights == 1)[..., 1, None]
assert np.all((mix_eval > 0) == expected_positive)
repetitions = np.unique(mix_eval).size < mix_eval.size
assert not repetitions
# Test logp
# MvNormal logp is currently limited to 2d values
expectation = pytest.raises(ValueError) if mix_eval.ndim > 2 else does_not_raise()
with expectation:
mix_logp_eval = logp(mix, mix_eval).eval()
assert mix_logp_eval.shape == expected_shape[:-1]
bcast_weights = np.broadcast_to(weights, (*expected_shape[:-1], 2))
expected_logp = np.stack(
(
logp(components[0], mix_eval).eval(),
logp(components[1], mix_eval).eval(),
),
axis=-1,
)[bcast_weights == 1]
expected_logp = expected_logp.reshape(expected_shape[:-1])
assert np.allclose(mix_logp_eval, expected_logp)
def logp(
value: at.Variable,
mu: at.Variable,
sigma: at.Variable,
init: at.Variable,
steps: at.Variable,
) -> at.TensorVariable:
"""Calculate log-probability of Gaussian Random Walk distribution at specified value."""
# Calculate initialization logp
init_logp = logp(init, value[..., 0])
# Make time series stationary around the mean value
stationary_series = value[..., 1:] - value[..., :-1]
# Add one dimension to the right, so that mu and sigma broadcast safely along
# the steps dimension
series_logp = logp(Normal.dist(mu[..., None], sigma[..., None]), stationary_series)
return check_parameters(
init_logp + series_logp.sum(axis=-1),
steps > 0,
msg="steps > 0",
)
def test_logp_helper_derived_rv():
assert np.isclose(
logp(at.exp(Normal.dist()), 5).eval(),
logp(LogNormal.dist(), 5).eval(),
)
