def test_discrete(self):
pmf1 = np.array([0.36, 0.48, 0.16])
pmf2 = np.array([1, 1, 1]) / 3
dist1 = dists.FiniteDist(pmf1)
dist2 = dists.FiniteDist(pmf2)
self.assertAlmostEqual(dists.kl(dist1, dist2), 0.0852996)
self.assertAlmostEqual(dists.kl(dist2, dist1), 0.097455)
def process_update(self, belief: dists.FiniteDist,
t: int) -> dists.FiniteDist:
(n, _, m) = self._dynamics.shape
belief_pmf = belief.pmf()
next_belief_given_input = np.zeros(n, m)
for i in range(m):
next_belief_given_input[:,
i] = self._dynamics.shape[:, :,
i] @ belief_pmf
input_dist = self._policy.input_channel(t).marginal(
dists.FiniteDist(belief_pmf))
return dists.FiniteDist(next_belief_given_input @ input_dist.pmf())
def test_marginal(self):
channel = channels.DiscreteChannel(np.array([[0.2, 0.5], [0.8, 0.5]]))
input = dists.FiniteDist(np.array([0.5, 0.5]))
marginal = channel.marginal(input)
self.assertTrue(np.allclose(marginal.pmf(), np.array([0.35, 0.65])))
def conditional(self, chan_input: int) -> dists.FiniteDist:
"""
Computes the distribution of Y assuming X = x.
:param chan_input: The integer x.
:return: A finite distribution over a set of size m representing the resulting distribution of Y.
"""
return dists.FiniteDist(self._y_given_x[:, chan_input])
def test_pmf(self):
pmf = np.array([0.2, 0.4, 0.4])
dist = dists.FiniteDist(pmf)
for i in range(3):
self.assertEqual(dist.pmf(i), pmf[i])
self.assertTrue(np.allclose(dist.pmf(), pmf))
def test_numel(self):
pmf = np.array([0.2, 0.4, 0.4])
dist = dists.FiniteDist(pmf)
for i in range(3):
self.assertEqual(dist.pmf(i), pmf[i])
self.assertEqual(dist.numel(), 3)
def joint(self, chan_input: dists.FiniteDist) -> dists.FiniteDist:
"""
Computes the joint distribution for (X, Y).
:param chan_input: A finite distribution over a set of size n.
:return: A finite distribution over a set of size n * m. The probability of the event {Y = i, X = j} can be
accessed using a pmf method call with val=(i, j), shape=(m, n).
"""
return dists.FiniteDist((self._y_given_x * chan_input.pmf()).flatten())
def marginal(self, prior: dists.FiniteDist) -> dists.FiniteDist:
"""
Computes the distribution of Y resulting from a prior over X.
:param prior: The assumed prior on X.
:return: The marginal distribution of Y.
"""
return dists.FiniteDist(self._y_given_x @ prior.pmf())
def test_mutual_info(self):
input = dists.FiniteDist(np.array([1 / 2, 1 / 4, 1 / 8, 1 / 8]))
channel = channels.DiscreteChannel(
np.array([[1 / 4, 1 / 4, 1 / 4,
1 / 4], [1 / 8, 1 / 2, 1 / 4, 1 / 4],
[1 / 8, 1 / 4, 1 / 2, 1 / 2], [1 / 2, 0, 0, 0]]))
mi = channel.mutual_info(input, 2)
self.assertAlmostEqual(mi, 3 / 8)
def test_joint(self):
channel = channels.DiscreteChannel(np.array([[0.2, 0.5], [0.8, 0.5]]))
input = dists.FiniteDist(np.array([0.5, 0.5]))
joint = channel.joint(input)
self.assertTrue(
np.allclose(joint.pmf(), np.array([0.1, 0.25, 0.4, 0.25])))
self.assertAlmostEqual(joint.pmf((0, 1), (2, 2)), 0.25)
self.assertAlmostEqual(joint.pmf((1, 0), (2, 2)), 0.4)
def posterior(self, prior: dists.FiniteDist,
output: int) -> dists.FiniteDist:
"""
Computes the posterior distribution over X given Y = y.
:param prior: A finite distribution over n elements representing assumed prior distribution over X.
:param output: The index of the observed value y.
:return: A finite distribution over n elements representing the posterior distribution over X.
"""
return dists.FiniteDist((self._y_given_x[output, :] * prior.pmf()) /
self.marginal(prior).pmf(output))
def test_sample(self):
np.random.seed(0)
pmf = np.array([0.2, 0.4, 0.4])
dist = dists.FiniteDist(pmf)
n = 10000
samples = dist.sample(n)
sample_dist = np.array([(samples == 0).sum() / n,
(samples == 1).sum() / n,
(samples == 2).sum() / n])
self.assertTrue(np.allclose(sample_dist, pmf, atol=1e-2))
def test_value_iter(self):
lava = Lava(5, 2, dists.FiniteDist(np.array([1, 0, 0, 0, 0])), 10)
policy = DiscretePolicy(lava)
val = policy.solve()
for t in range(9):
self.assertTrue(np.allclose(policy._input_given_state[:, :, t], np.array([[0, 0, 0, 1, 0],
[0, 0, 1, 0, 1],
[1, 1, 0, 0, 0]])))
self.assertTrue(np.allclose(policy._input_given_state[:, :, -1], np.array([[0, 0, 0, 1, 0],
[1, 0, 1, 0, 1],
[0, 1, 0, 0, 0]])))
self.assertTrue(np.allclose(val, -54))
def sensor(self, state: int, t: int = 0) -> dists.FiniteDist:
return dists.FiniteDist(self._sensor_tensor[:, state])
def dynamics(self, state: int, input: int, t: int = 0) -> dists.FiniteDist:
return dists.FiniteDist(self._dynamics_tensor[:, state, input])
def test_mean(self):
pmf = np.array([0.2, 0.4, 0.4])
dist = dists.FiniteDist(pmf)
self.assertAlmostEqual(dist.mean(), 1.2)
def test_posterior(self):
prior = dists.FiniteDist(np.array([0.4, 0.6]))
channel = channels.DiscreteChannel(np.array([[0.5, 1], [0.5, 0]]))
posterior = channel.posterior(prior, 0)
self.assertTrue(np.allclose(posterior.pmf(), np.array([0.25, 0.75])))
def test_cov(self):
pmf = np.array([0.2, 0.4, 0.4])
dist = dists.FiniteDist(pmf)
self.assertAlmostEqual(dist.cov(), 0.56)
