test_DiscreteDistribution.py

Go to the documentation of this file.

"""
GTSAM Copyright 2010-2021, Georgia Tech Research Corporation,
Atlanta, Georgia 30332-0415
All Rights Reserved
 
See LICENSE for the license information
 
Unit tests for Discrete Priors.
Author: Frank Dellaert
"""
 
# pylint: disable=no-name-in-module, invalid-name
 
import unittest
 
import numpy as np
from gtsam import DecisionTreeFactor, DiscreteKeys, DiscreteDistribution
from gtsam.utils.test_case import GtsamTestCase
 
X = 0, 2
 
 
class TestDiscreteDistribution(GtsamTestCase):
    """Tests for Discrete Priors."""
 
    def test_constructor(self):
        """Test various constructors."""
        keys = DiscreteKeys()
        keys.push_back(X)
        f = DecisionTreeFactor(keys, "0.4 0.6")
        expected = DiscreteDistribution(f)
 
        actual = DiscreteDistribution(X, "2/3")
        self.gtsamAssertEquals(actual, expected)
 
        actual2 = DiscreteDistribution(X, [0.4, 0.6])
        self.gtsamAssertEquals(actual2, expected)
 
    def test_operator(self):
        prior = DiscreteDistribution(X, "2/3")
        self.assertAlmostEqual(prior(0), 0.4)
        self.assertAlmostEqual(prior(1), 0.6)
 
    def test_pmf(self):
        prior = DiscreteDistribution(X, "2/3")
        expected = np.array([0.4, 0.6])
        np.testing.assert_allclose(expected, prior.pmf())
 
    def test_sample(self):
        prior = DiscreteDistribution(X, "2/3")
        actual = prior.sample()
        self.assertIsInstance(actual, int)
 
    def test_markdown(self):
        """Test the _repr_markdown_ method."""
 
        prior = DiscreteDistribution(X, "2/3")
        expected = " *P(0):*\n\n" \
            "|0|value|\n" \
            "|:-:|:-:|\n" \
            "|0|0.4|\n" \
            "|1|0.6|\n" \
 
        actual = prior._repr_markdown_()
        self.assertEqual(actual, expected)
 
 
if __name__ == "__main__":
    unittest.main()