diff --git a/gtsam/discrete/TableFactor.cpp b/gtsam/discrete/TableFactor.cpp index bf9662e346..a59095d406 100644 --- a/gtsam/discrete/TableFactor.cpp +++ b/gtsam/discrete/TableFactor.cpp @@ -87,7 +87,15 @@ static Eigen::SparseVector ComputeSparseTable( }); sparseTable.reserve(nrValues); - std::set allKeys(dt.keys().begin(), dt.keys().end()); + KeySet allKeys(dt.keys().begin(), dt.keys().end()); + + // Compute denominators to be used in computing sparse table indices + std::map denominators; + double denom = sparseTable.size(); + for (const DiscreteKey& dkey : dkeys) { + denom /= dkey.second; + denominators.insert(std::pair(dkey.first, denom)); + } /** * @brief Functor which is called by the DecisionTree for each leaf. @@ -102,13 +110,13 @@ static Eigen::SparseVector ComputeSparseTable( auto op = [&](const Assignment& assignment, double p) { if (p > 0) { // Get all the keys involved in this assignment - std::set assignmentKeys; + KeySet assignmentKeys; for (auto&& [k, _] : assignment) { assignmentKeys.insert(k); } // Find the keys missing in the assignment - std::vector diff; + KeyVector diff; std::set_difference(allKeys.begin(), allKeys.end(), assignmentKeys.begin(), assignmentKeys.end(), std::back_inserter(diff)); @@ -127,12 +135,10 @@ static Eigen::SparseVector ComputeSparseTable( // Generate index and add to the sparse vector. Eigen::Index idx = 0; - size_t previousCardinality = 1; // We go in reverse since a DecisionTree has the highest label first for (auto&& it = updatedAssignment.rbegin(); it != updatedAssignment.rend(); it++) { - idx += previousCardinality * it->second; - previousCardinality *= dt.cardinality(it->first); + idx += it->second * denominators.at(it->first); } sparseTable.coeffRef(idx) = p; } @@ -252,9 +258,19 @@ DecisionTreeFactor TableFactor::operator*(const DecisionTreeFactor& f) const { DecisionTreeFactor TableFactor::toDecisionTreeFactor() const { DiscreteKeys dkeys = discreteKeys(); - std::vector table; - for (auto i = 0; i < sparse_table_.size(); i++) { - table.push_back(sparse_table_.coeff(i)); + // If no keys, then return empty DecisionTreeFactor + if (dkeys.size() == 0) { + AlgebraicDecisionTree tree; + // We can have an empty sparse_table_ or one with a single value. + if (sparse_table_.size() != 0) { + tree = AlgebraicDecisionTree(sparse_table_.coeff(0)); + } + return DecisionTreeFactor(dkeys, tree); + } + + std::vector table(sparse_table_.size(), 0.0); + for (SparseIt it(sparse_table_); it; ++it) { + table[it.index()] = it.value(); } AlgebraicDecisionTree tree(dkeys, table); diff --git a/gtsam/discrete/tests/testTableFactor.cpp b/gtsam/discrete/tests/testTableFactor.cpp index a455faaaa4..e6c71e15ce 100644 --- a/gtsam/discrete/tests/testTableFactor.cpp +++ b/gtsam/discrete/tests/testTableFactor.cpp @@ -173,6 +173,36 @@ TEST(TableFactor, Conversion) { TableFactor tf(dtf.discreteKeys(), dtf); EXPECT(assert_equal(dtf, tf.toDecisionTreeFactor())); + + // Test for correct construction when keys are not in reverse order. + // This is possible in conditionals e.g. P(x1 | x0) + DiscreteKey X(1, 2), Y(0, 2); + DiscreteConditional dtf2( + X, {Y}, std::vector{0.33333333, 0.6, 0.66666667, 0.4}); + + TableFactor tf2(dtf2); + // GTSAM_PRINT(dtf2); + // GTSAM_PRINT(tf2); + // GTSAM_PRINT(tf2.toDecisionTreeFactor()); + + // Check for ADT equality since the order of keys is irrelevant + EXPECT(assert_equal>(dtf2, + tf2.toDecisionTreeFactor())); +} + +/* ************************************************************************* */ +TEST(TableFactor, Empty) { + DiscreteKey X(1, 2); + + TableFactor single = *TableFactor({X}, "1 1").sum(1); + // Should not throw a segfault + EXPECT(assert_equal(*DecisionTreeFactor(X, "1 1").sum(1), + single.toDecisionTreeFactor())); + + TableFactor empty = *TableFactor({X}, "0 0").sum(1); + // Should not throw a segfault + EXPECT(assert_equal(*DecisionTreeFactor(X, "0 0").sum(1), + empty.toDecisionTreeFactor())); } /* ************************************************************************* */