""" 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 Bayes trees. Author: Frank Dellaert """ # pylint: disable=no-name-in-module, invalid-name import unittest import numpy as np from gtsam.symbol_shorthand import A, X from gtsam.utils.test_case import GtsamTestCase import gtsam from gtsam import (DiscreteBayesNet, DiscreteBayesTreeClique, DiscreteConditional, DiscreteFactorGraph, DiscreteValues, Ordering) class TestDiscreteBayesNet(GtsamTestCase): """Tests for Discrete Bayes Nets.""" def test_elimination(self): """Test Multifrontal elimination.""" # Define DiscreteKey pairs. keys = [(j, 2) for j in range(15)] # Create thin-tree Bayes net. bayesNet = DiscreteBayesNet() bayesNet.add(keys[0], [keys[8], keys[12]], "2/3 1/4 3/2 4/1") bayesNet.add(keys[1], [keys[8], keys[12]], "4/1 2/3 3/2 1/4") bayesNet.add(keys[2], [keys[9], keys[12]], "1/4 8/2 2/3 4/1") bayesNet.add(keys[3], [keys[9], keys[12]], "1/4 2/3 3/2 4/1") bayesNet.add(keys[4], [keys[10], keys[13]], "2/3 1/4 3/2 4/1") bayesNet.add(keys[5], [keys[10], keys[13]], "4/1 2/3 3/2 1/4") bayesNet.add(keys[6], [keys[11], keys[13]], "1/4 3/2 2/3 4/1") bayesNet.add(keys[7], [keys[11], keys[13]], "1/4 2/3 3/2 4/1") bayesNet.add(keys[8], [keys[12], keys[14]], "T 1/4 3/2 4/1") bayesNet.add(keys[9], [keys[12], keys[14]], "4/1 2/3 F 1/4") bayesNet.add(keys[10], [keys[13], keys[14]], "1/4 3/2 2/3 4/1") bayesNet.add(keys[11], [keys[13], keys[14]], "1/4 2/3 3/2 4/1") bayesNet.add(keys[12], [keys[14]], "3/1 3/1") bayesNet.add(keys[13], [keys[14]], "1/3 3/1") bayesNet.add(keys[14], "1/3") # Create a factor graph out of the Bayes net. factorGraph = DiscreteFactorGraph(bayesNet) # Create a BayesTree out of the factor graph. ordering = Ordering() for j in range(15): ordering.push_back(j) bayesTree = factorGraph.eliminateMultifrontal(ordering) # Uncomment these for visualization: # print(bayesTree) # for key in range(15): # bayesTree[key].printSignature() # bayesTree.saveGraph("test_DiscreteBayesTree.dot") # The root is P( 8 12 14), we can retrieve it by key: root = bayesTree[8] self.assertIsInstance(root, DiscreteBayesTreeClique) self.assertTrue(root.isRoot()) self.assertIsInstance(root.conditional(), DiscreteConditional) # Test all methods in DiscreteBayesTree self.gtsamAssertEquals(bayesTree, bayesTree) # Check value at 0 zero_values = DiscreteValues() for j in range(15): zero_values[j] = 0 value_at_zeros = bayesTree.evaluate(zero_values) self.assertAlmostEqual(value_at_zeros, 0.0) # Check value at max values_star = factorGraph.optimize() max_value = bayesTree.evaluate(values_star) self.assertAlmostEqual(max_value, 0.002548) # Check operator sugar max_value = bayesTree(values_star) self.assertAlmostEqual(max_value, 0.002548) self.assertFalse(bayesTree.empty()) self.assertEqual(12, bayesTree.size()) @unittest.skip("Too Slow") def test_discrete_bayes_tree_lookup(self): """Check that we can have a multi-frontal lookup table.""" # Make a small planning-like graph: 3 states, 2 actions graph = DiscreteFactorGraph() x1, x2, x3 = (X(1), 3), (X(2), 3), (X(3), 3) a1, a2 = (A(1), 2), (A(2), 2) # Constraint on start and goal graph.add([x1], np.array([1, 0, 0])) graph.add([x3], np.array([0, 0, 1])) # Should I stay or should I go? # "Reward" (exp(-cost)) for an action is 10, and rewards multiply: r = 10 table = np.array([ r, 0, 0, 0, r, 0, # x1 = 0 0, r, 0, 0, 0, r, # x1 = 1 0, 0, r, 0, 0, r # x1 = 2 ]) graph.add([x1, a1, x2], table) graph.add([x2, a2, x3], table) # Eliminate for MPE (maximum probable explanation). ordering = Ordering(keys=[A(2), X(3), X(1), A(1), X(2)]) lookup = graph.eliminateMultifrontal(ordering, gtsam.EliminateForMPE) # Check that the lookup table is correct assert lookup.size() == 2 lookup_x1_a1_x2 = lookup[X(1)].conditional() assert lookup_x1_a1_x2.nrFrontals() == 3 # Check that sum is 1.0 (not 100, as we now normalize to prevent underflow) empty = gtsam.DiscreteValues() self.assertAlmostEqual(lookup_x1_a1_x2.sum(3)(empty), 1.0) # And that only non-zero reward is for x1 a1 x2 == 0 1 1 values = DiscreteValues() values[X(1)] = 0 values[A(1)] = 1 values[X(2)] = 1 self.assertAlmostEqual(lookup_x1_a1_x2(values), 1.0) lookup_a2_x3 = lookup[X(3)].conditional() # Check that the sum depends on x2 and is non-zero only for x2 in {1, 2} sum_x2 = lookup_a2_x3.sum(2) values = DiscreteValues() values[X(2)] = 0 self.assertAlmostEqual(sum_x2(values), 0) values[X(2)] = 1 self.assertAlmostEqual(sum_x2(values), 1.0) # not 10, as we normalize values[X(2)] = 2 self.assertAlmostEqual(sum_x2(values), 2.0) # not 20, as we normalize assert lookup_a2_x3.nrFrontals() == 2 # And that the non-zero rewards are for x2 a2 x3 == 1 1 2 values = DiscreteValues() values[X(2)] = 1 values[A(2)] = 1 values[X(3)] = 2 self.assertAlmostEqual(lookup_a2_x3(values), 1.0) # not 10... def test_direct_from_cliques(self): """Test creating a Bayes tree directly from cliques.""" # Create a BayesNet bayesNet = DiscreteBayesNet() A, B, C = (0, 2), (1, 2), (2, 2) bayesNet.add(A, "1/3") bayesNet.add(B, [A], "1/3 3/1") bayesNet.add(C, [B], "3/1 3/1") # Create cliques directly clique2 = DiscreteBayesTreeClique(DiscreteConditional(C, [B], "3/1 3/1")) clique1 = DiscreteBayesTreeClique(DiscreteConditional(B, [A], "1/3 3/1")) clique0 = DiscreteBayesTreeClique(DiscreteConditional(A, "1/3")) # Create a BayesTree bayesTree = gtsam.DiscreteBayesTree() bayesTree.insertRoot(clique2) bayesTree.addClique(clique1, clique2) bayesTree.addClique(clique0, clique1) # Check that the BayesTree is correct values = DiscreteValues() values[0] = 1 values[1] = 1 values[2] = 1 # regression expected = .046875 self.assertAlmostEqual(expected, bayesNet.evaluate(values)) if __name__ == "__main__": unittest.main()