def test_marginalization():
c = distributions.conditional.Discrete([[1, 2, 3]], [[4, 5, 6]],
np.array([
np.array([1, 0, 0]),
np.array([0, 1, 0]),
np.array([0, 0, 1])
]))
d = distributions.Discrete([[4, 5, 6]], np.array([0, 0, 1]))
e = c.marginalize(d)
assert np.array_equal(e.sample(), np.array([3]))
def marginalize(
self, distribution: "distributions.Discrete"
) -> "distributions.Discrete":
if not np.array_equal(distribution.values, self._conditioned_values):
raise ValueError(
"Values of the given distribution do not align with the conditioned "
"values.")
conditioned_shape = tuple(len(x) for x in self._conditioned_values)
new_p = np.zeros(tuple(len(x) for x in self._values))
for p, di in zip(distribution.probabilities.ravel(),
np.ndindex(conditioned_shape)):
new_p += np.asarray(p * self._probabilities[di], dtype=np.float32)
return distributions.Discrete(self._values, new_p)
def test_probability():
d = distributions.Discrete(
[[1, 2, 3], [4, 5, 6]],
[[0.2, 0.2, 0.6], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]])
assert d.probability(np.array([1, 4])) == 0.2
assert d.probability(np.array([1, 5])) == 0.2
assert d.probability(np.array([1, 6])) == 0.6
assert d.probability(np.array([2, 6])) == 0.0
assert np.array_equal(
d.probability(np.array([[1, 4], [1, 5], [1, 6], [2, 6]])),
np.array([0.2, 0.2, 0.6, 0.0]))
assert d.probability(np.array([0.5, 3.5])) == 0.2
assert d.probability(np.array([0.5, 5.5])) == 0.6
assert d.probability(np.array([0.5, 6.5])) == 0.6
def test_discrete():
d = distributions.Discrete([[1, 2, 3]], [1.0, 0.0, 0.0])
assert np.array_equal(d.sample(), np.array([1]))
assert d.sample().shape == (1, )
assert d.sample(5).shape == (5, 1)
def test_marginalize():
d = distributions.Discrete(
[[1, 2, 3], [4, 5, 6]],
[[0.2, 0.2, 0.6], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]])
m = d.marginalize(1)
assert np.array_equal(m.probabilities, np.array([1.0, 0.0, 0.0]))
def test_multivariate():
d = distributions.Discrete([[1, 2, 3], [4, 5, 6]],
[[0, 0, 0], [0, 0, 1], [0, 0, 0]])
assert np.array_equal(d.sample(), np.array([2, 6]))
def by_evidence(self, evidence: np.ndarray) -> "distributions.Discrete":
i = tuple(
np.where(self._conditioned_values[i] == x)[0][0]
for i, x in enumerate(evidence))
return distributions.Discrete(self._values, self._probabilities[i])