Merge pull request #1890 from NewThinker-Jeffrey/jeffrey/isam2_marginalization

gtsam/inference/VariableIndex.h

@@ -87,9 +87,11 @@ class GTSAM_EXPORT VariableIndex {
   const FactorIndices& operator[](Key variable) const {
     KeyMap::const_iterator item = index_.find(variable);
     if(item == index_.end())
-      throw std::invalid_argument("Requested non-existent variable from VariableIndex");
+      throw std::invalid_argument("Requested non-existent variable '" +
+                                  DefaultKeyFormatter(variable) +
+                                  "' from VariableIndex");
     else
-    return item->second;
+      return item->second;
   }
 
   /// Return true if no factors associated with a variable

gtsam_unstable/nonlinear/BayesTreeMarginalizationHelper.h

@@ -0,0 +1,358 @@
+/* ----------------------------------------------------------------------------
+
+ * GTSAM Copyright 2010, Georgia Tech Research Corporation,
+ * Atlanta, Georgia 30332-0415
+ * All Rights Reserved
+ * Authors: Frank Dellaert, et al. (see THANKS for the full author list)
+
+ * See LICENSE for the license information
+
+ * -------------------------------------------------------------------------- */
+
+/**
+ * @file    BayesTreeMarginalizationHelper.h
+ * @brief   Helper functions for marginalizing variables from a Bayes Tree.
+ *
+ * @author  Jeffrey (Zhiwei Wang)
+ * @date    Oct 28, 2024
+ */
+
+// \callgraph
+#pragma once
+
+#include <unordered_map>
+#include <unordered_set>
+#include <deque>
+#include <gtsam/inference/BayesTree.h>
+#include <gtsam/inference/BayesTreeCliqueBase.h>
+#include <gtsam/base/debug.h>
+#include "gtsam_unstable/dllexport.h"
+
+namespace gtsam {
+
+/**
+ * This class provides helper functions for marginalizing variables from a Bayes Tree.
+ */
+template <typename BayesTree>
+class GTSAM_UNSTABLE_EXPORT BayesTreeMarginalizationHelper {
+
+public:
+  using Clique = typename BayesTree::Clique;
+  using sharedClique = typename BayesTree::sharedClique;
+
+  /**
+   * This function identifies variables that need to be re-eliminated before
+   * performing marginalization.
+   * 
+   * Re-elimination is necessary for a clique containing marginalizable
+   * variables if:
+   * 
+   * 1. Some non-marginalizable variables appear before marginalizable ones
+   *    in that clique;
+   * 2. Or it has a child node depending on a marginalizable variable AND the
+   *    subtree rooted at that child contains non-marginalizables.
+   * 
+   * In addition, for any descendant node depending on a marginalizable
+   * variable, if the subtree rooted at that descendant contains
+   * non-marginalizable variables (i.e., it lies on a path from one of the
+   * aforementioned cliques that require re-elimination to a node containing
+   * non-marginalizable variables at the leaf side), then it also needs to
+   * be re-eliminated.
+   * 
+   * @param[in] bayesTree The Bayes tree
+   * @param[in] marginalizableKeys Keys to be marginalized
+   * @return Set of additional keys that need to be re-eliminated
+   */
+  static std::unordered_set<Key>
+  gatherAdditionalKeysToReEliminate(
+      const BayesTree& bayesTree,
+      const KeyVector& marginalizableKeys) {
+    const bool debug = ISDEBUG("BayesTreeMarginalizationHelper");
+
+    std::unordered_set<const Clique*> additionalCliques =
+        gatherAdditionalCliquesToReEliminate(bayesTree, marginalizableKeys);    
+
+    std::unordered_set<Key> additionalKeys;
+    for (const Clique* clique : additionalCliques) {
+      addCliqueToKeySet(clique, &additionalKeys);
+    }
+
+    if (debug) {
+      std::cout << "BayesTreeMarginalizationHelper: Additional keys to re-eliminate: ";
+      for (const Key& key : additionalKeys) {
+        std::cout << DefaultKeyFormatter(key) << " ";
+      }
+      std::cout << std::endl;
+    }
+
+    return additionalKeys;
+  }
+
+ protected:
+  /**
+   * This function identifies cliques that need to be re-eliminated before
+   * performing marginalization.
+   * See the docstring of @ref gatherAdditionalKeysToReEliminate().
+   */
+  static std::unordered_set<const Clique*>
+  gatherAdditionalCliquesToReEliminate(
+      const BayesTree& bayesTree,
+      const KeyVector& marginalizableKeys) {
+    std::unordered_set<const Clique*> additionalCliques;
+    std::unordered_set<Key> marginalizableKeySet(
+        marginalizableKeys.begin(), marginalizableKeys.end());
+    CachedSearch cachedSearch;
+
+    // Check each clique that contains a marginalizable key
+    for (const Clique* clique :
+         getCliquesContainingKeys(bayesTree, marginalizableKeySet)) {
+      if (additionalCliques.count(clique)) {
+        // The clique has already been visited. This can happen when an
+        // ancestor of the current clique also contain some marginalizable
+        // varaibles and it's processed beore the current.
+        continue;
+      }
+
+      if (needsReelimination(clique, marginalizableKeySet, &cachedSearch)) {
+        // Add the current clique
+        additionalCliques.insert(clique);
+
+        // Then add the dependent cliques
+        gatherDependentCliques(clique, marginalizableKeySet, &additionalCliques,
+                               &cachedSearch);
+      }
+    }
+    return additionalCliques;
+  }
+
+  /**
+   * Gather the cliques containing any of the given keys.
+   * 
+   * @param[in] bayesTree The Bayes tree
+   * @param[in] keysOfInterest Set of keys of interest
+   * @return Set of cliques that contain any of the given keys
+   */
+  static std::unordered_set<const Clique*> getCliquesContainingKeys(
+      const BayesTree& bayesTree,
+      const std::unordered_set<Key>& keysOfInterest) {
+    std::unordered_set<const Clique*> cliques;
+    for (const Key& key : keysOfInterest) {
+      cliques.insert(bayesTree[key].get());
+    }
+    return cliques;
+  }
+
+  /**
+   * A struct to cache the results of the below two functions.
+   */
+  struct CachedSearch {
+    std::unordered_map<const Clique*, bool> wholeMarginalizableCliques;
+    std::unordered_map<const Clique*, bool> wholeMarginalizableSubtrees;
+  };
+
+  /**
+   * Check if all variables in the clique are marginalizable.
+   * 
+   * Note we use a cache map to avoid repeated searches.
+   */
+  static bool isWholeCliqueMarginalizable(
+      const Clique* clique,
+      const std::unordered_set<Key>& marginalizableKeys,
+      CachedSearch* cache) {
+    auto it = cache->wholeMarginalizableCliques.find(clique);
+    if (it != cache->wholeMarginalizableCliques.end()) {
+      return it->second;
+    } else {
+      bool ret = true;
+      for (Key key : clique->conditional()->frontals()) {
+        if (!marginalizableKeys.count(key)) {
+          ret = false;
+          break;
+        }
+      }
+      cache->wholeMarginalizableCliques.insert({clique, ret});
+      return ret;
+    }
+  }
+
+  /**
+   * Check if all variables in the subtree are marginalizable.
+   * 
+   * Note we use a cache map to avoid repeated searches.
+   */
+  static bool isWholeSubtreeMarginalizable(
+      const Clique* subtree,
+      const std::unordered_set<Key>& marginalizableKeys,
+      CachedSearch* cache) {
+    auto it = cache->wholeMarginalizableSubtrees.find(subtree);
+    if (it != cache->wholeMarginalizableSubtrees.end()) {
+      return it->second;
+    } else {
+      bool ret = true;
+      if (isWholeCliqueMarginalizable(subtree, marginalizableKeys, cache)) {
+        for (const sharedClique& child : subtree->children) {
+          if (!isWholeSubtreeMarginalizable(child.get(), marginalizableKeys, cache)) {
+            ret = false;
+            break;
+          }
+        }
+      } else {
+        ret = false;
+      }
+      cache->wholeMarginalizableSubtrees.insert({subtree, ret});
+      return ret;
+    }
+  }
+
+  /**
+   * Check if a clique contains variables that need reelimination due to
+   * elimination ordering conflicts.
+   * 
+   * @param[in] clique The clique to check
+   * @param[in] marginalizableKeys Set of keys to be marginalized
+   * @return true if any variables in the clique need re-elimination
+   */
+  static bool needsReelimination(
+      const Clique* clique,
+      const std::unordered_set<Key>& marginalizableKeys,
+      CachedSearch* cache) {
+    bool hasNonMarginalizableAhead = false;
+
+    // Check each frontal variable in order
+    for (Key key : clique->conditional()->frontals()) {
+      if (marginalizableKeys.count(key)) {
+        // If we've seen non-marginalizable variables before this one,
+        // we need to reeliminate
+        if (hasNonMarginalizableAhead) {
+          return true;
+        }
+
+        // Check if any child depends on this marginalizable key and the
+        // subtree rooted at that child contains non-marginalizables.
+        for (const sharedClique& child : clique->children) {
+          if (hasDependency(child.get(), key) &&
+              !isWholeSubtreeMarginalizable(child.get(), marginalizableKeys, cache)) {
+            return true;
+          }
+        }
+      } else {
+        hasNonMarginalizableAhead = true;
+      }
+    }
+    return false;
+  }
+
+  /**
+   * Gather all dependent nodes that lie on a path from the root clique
+   * to a clique containing a non-marginalizable variable at the leaf side.
+   *
+   * @param[in] rootClique The root clique
+   * @param[in] marginalizableKeys Set of keys to be marginalized
+   */
+  static void gatherDependentCliques(
+      const Clique* rootClique,
+      const std::unordered_set<Key>& marginalizableKeys,
+      std::unordered_set<const Clique*>* additionalCliques,
+      CachedSearch* cache) {
+    std::vector<const Clique*> dependentChildren;
+    dependentChildren.reserve(rootClique->children.size());
+    for (const sharedClique& child : rootClique->children) {
+      if (additionalCliques->count(child.get())) {
+        // This child has already been visited. This can happen if the
+        // child itself contains a marginalizable variable and it's
+        // processed before the current rootClique.
+        continue;
+      }
+      if (hasDependency(child.get(), marginalizableKeys)) {
+        dependentChildren.push_back(child.get());
+      }
+    }
+    gatherDependentCliquesFromChildren(
+        dependentChildren, marginalizableKeys, additionalCliques, cache);
+  }
+
+  /**
+   * A helper function for the above gatherDependentCliques().
+   */
+  static void gatherDependentCliquesFromChildren(
+      const std::vector<const Clique*>& dependentChildren,
+      const std::unordered_set<Key>& marginalizableKeys,
+      std::unordered_set<const Clique*>* additionalCliques,
+      CachedSearch* cache) {
+    std::deque<const Clique*> descendants(
+        dependentChildren.begin(), dependentChildren.end());
+    while (!descendants.empty()) {
+      const Clique* descendant = descendants.front();
+      descendants.pop_front();
+
+      // If the subtree rooted at this descendant contains non-marginalizables,
+      // it must lie on a path from the root clique to a clique containing
+      // non-marginalizables at the leaf side.
+      if (!isWholeSubtreeMarginalizable(descendant, marginalizableKeys, cache)) {
+        additionalCliques->insert(descendant);
+
+        // Add children of the current descendant to the set descendants.
+        for (const sharedClique& child : descendant->children) {
+          if (additionalCliques->count(child.get())) {
+            // This child has already been visited.
+            continue;
+          } else {
+            descendants.push_back(child.get());
+          }
+        }
+      }
+    }
+  }
+
+  /**
+   * Add all frontal variables from a clique to a key set.
+   * 
+   * @param[in] clique Clique to add keys from
+   * @param[out] additionalKeys Pointer to the output key set
+   */
+  static void addCliqueToKeySet(
+      const Clique* clique,
+      std::unordered_set<Key>* additionalKeys) {
+    for (Key key : clique->conditional()->frontals()) {
+      additionalKeys->insert(key);
+    }
+  }
+
+  /**
+   * Check if the clique depends on the given key.
+   * 
+   * @param[in] clique Clique to check
+   * @param[in] key Key to check for dependencies
+   * @return true if clique depends on the key
+   */
+  static bool hasDependency(
+      const Clique* clique, Key key) {
+    auto& conditional = clique->conditional();
+    if (std::find(conditional->beginParents(),
+        conditional->endParents(), key)
+        != conditional->endParents()) {
+      return true;
+    } else {
+      return false;
+    }
+  }
+
+  /**
+   * Check if the clique depends on any of the given keys.
+   */
+  static bool hasDependency(
+      const Clique* clique, const std::unordered_set<Key>& keys) {
+    auto& conditional = clique->conditional();
+    for (auto it = conditional->beginParents();
+        it != conditional->endParents(); ++it) {
+      if (keys.count(*it)) {
+        return true;
+      }
+    }
+
+    return false;
+  }
+};
+// BayesTreeMarginalizationHelper
+
+}/// namespace gtsam

gtsam_unstable/nonlinear/IncrementalFixedLagSmoother.cpp

@@ -20,32 +20,11 @@
  */
 
 #include <gtsam_unstable/nonlinear/IncrementalFixedLagSmoother.h>
+#include <gtsam_unstable/nonlinear/BayesTreeMarginalizationHelper.h>
 #include <gtsam/base/debug.h>
 
 namespace gtsam {
 
-/* ************************************************************************* */
-void recursiveMarkAffectedKeys(const Key& key,
-    const ISAM2Clique::shared_ptr& clique, std::set<Key>& additionalKeys) {
-
-  // Check if the separator keys of the current clique contain the specified key
-  if (std::find(clique->conditional()->beginParents(),
-      clique->conditional()->endParents(), key)
-      != clique->conditional()->endParents()) {
-
-    // Mark the frontal keys of the current clique
-    for(Key i: clique->conditional()->frontals()) {
-      additionalKeys.insert(i);
-    }
-
-    // Recursively mark all of the children
-    for(const ISAM2Clique::shared_ptr& child: clique->children) {
-      recursiveMarkAffectedKeys(key, child, additionalKeys);
-    }
-  }
-  // If the key was not found in the separator/parents, then none of its children can have it either
-}
-
 /* ************************************************************************* */
 void IncrementalFixedLagSmoother::print(const std::string& s,
     const KeyFormatter& keyFormatter) const {
@@ -114,14 +93,9 @@ FixedLagSmoother::Result IncrementalFixedLagSmoother::update(
     std::cout << std::endl;
   }
 
-  // Mark additional keys between the marginalized keys and the leaves
-  std::set<Key> additionalKeys;
-  for(Key key: marginalizableKeys) {
-    ISAM2Clique::shared_ptr clique = isam_[key];
-    for(const ISAM2Clique::shared_ptr& child: clique->children) {
-      recursiveMarkAffectedKeys(key, child, additionalKeys);
-    }
-  }
+  std::unordered_set<Key> additionalKeys =
+      BayesTreeMarginalizationHelper<ISAM2>::gatherAdditionalKeysToReEliminate(
+          isam_, marginalizableKeys);
   KeyList additionalMarkedKeys(additionalKeys.begin(), additionalKeys.end());
 
   // Update iSAM2

gtsam_unstable/nonlinear/tests/testIncrementalFixedLagSmoother.cpp

@@ -49,6 +49,41 @@ bool check_smoother(const NonlinearFactorGraph& fullgraph, const Values& fullini
   return assert_equal(expected, actual);
 }
 
+/* ************************************************************************* */
+void PrintSymbolicTreeHelper(
+    const ISAM2Clique::shared_ptr& clique, const std::string indent = "") {
+
+  // Print the current clique
+  std::cout << indent << "P( ";
+  for(Key key: clique->conditional()->frontals()) {
+    std::cout << DefaultKeyFormatter(key) << " ";
+  }
+  if (clique->conditional()->nrParents() > 0)
+    std::cout << "| ";
+  for(Key key: clique->conditional()->parents()) {
+    std::cout << DefaultKeyFormatter(key) << " ";
+  }
+  std::cout << ")" << std::endl;
+
+  // Recursively print all of the children
+  for(const ISAM2Clique::shared_ptr& child: clique->children) {
+    PrintSymbolicTreeHelper(child, indent + " ");
+  }
+}
+
+/* ************************************************************************* */
+void PrintSymbolicTree(const ISAM2& isam,
+    const std::string& label) {
+  std::cout << label << std::endl;
+  if (!isam.roots().empty()) {
+    for(const ISAM2::sharedClique& root: isam.roots()) {
+      PrintSymbolicTreeHelper(root);
+    }
+  } else
+    std::cout << "{Empty Tree}" << std::endl;
+}
+
+
 /* ************************************************************************* */
 TEST( IncrementalFixedLagSmoother, Example )
 {
@@ -64,7 +99,7 @@ TEST( IncrementalFixedLagSmoother, Example )
 
   // Create a Fixed-Lag Smoother
   typedef IncrementalFixedLagSmoother::KeyTimestampMap Timestamps;
-  IncrementalFixedLagSmoother smoother(7.0, ISAM2Params());
+  IncrementalFixedLagSmoother smoother(12.0, ISAM2Params());
 
   // Create containers to keep the full graph
   Values fullinit;
@@ -158,6 +193,9 @@ TEST( IncrementalFixedLagSmoother, Example )
     Values newValues;
     Timestamps newTimestamps;
 
+    // Add the odometry factor twice to ensure the removeFactor test below works,
+    // where we need to keep the connectivity of the graph.
+    newFactors.push_back(BetweenFactor<Point2>(key1, key2, Point2(1.0, 0.0), odometerNoise));
     newFactors.push_back(BetweenFactor<Point2>(key1, key2, Point2(1.0, 0.0), odometerNoise));
     newValues.insert(key2, Point2(double(i)+0.1, -0.1));
     newTimestamps[key2] = double(i);
@@ -188,6 +226,7 @@ TEST( IncrementalFixedLagSmoother, Example )
 	  newFactors.push_back(BetweenFactor<Point2>(key1, key2, Point2(1.0, 0.0), odometerNoise));
 	  newValues.insert(key2, Point2(double(i)+0.1, -0.1));
 	  newTimestamps[key2] = double(i);
+	  ++i;
 
 	  fullgraph.push_back(newFactors);
 	  fullinit.insert(newValues);
@@ -210,6 +249,10 @@ TEST( IncrementalFixedLagSmoother, Example )
 
 	  const NonlinearFactorGraph smootherFactorsBeforeRemove = smoother.getFactors();
 
+    std::cout << "fullgraph.size() = " << fullgraph.size() << std::endl;
+    std::cout << "smootherFactorsBeforeRemove.size() = "
+              << smootherFactorsBeforeRemove.size() << std::endl;
+
 	  // remove factor
 	  smoother.update(emptyNewFactors, emptyNewValues, emptyNewTimestamps,factorToRemove);
 
@@ -231,6 +274,67 @@ TEST( IncrementalFixedLagSmoother, Example )
 	    }
 	  }
   }
+
+  {
+    SETDEBUG("BayesTreeMarginalizationHelper", true);
+    PrintSymbolicTree(smoother.getISAM2(), "Bayes Tree Before marginalization test:");
+
+    // Do pressure test on marginalization. Enlarge max_i to enhance the test.
+    const int max_i = 500;
+    while(i <= max_i) {
+      Key key_0 = MakeKey(i);
+      Key key_1 = MakeKey(i-1);
+      Key key_2 = MakeKey(i-2);
+      Key key_3 = MakeKey(i-3);
+      Key key_4 = MakeKey(i-4);
+      Key key_5 = MakeKey(i-5);
+      Key key_6 = MakeKey(i-6);
+      Key key_7 = MakeKey(i-7);
+      Key key_8 = MakeKey(i-8);
+      Key key_9 = MakeKey(i-9);
+      Key key_10 = MakeKey(i-10);
+
+      NonlinearFactorGraph newFactors;
+      Values newValues;
+      Timestamps newTimestamps;
+
+      // To make a complex graph
+      newFactors.push_back(BetweenFactor<Point2>(key_1, key_0, Point2(1.0, 0.0), odometerNoise));
+      if (i % 2 == 0)
+        newFactors.push_back(BetweenFactor<Point2>(key_2, key_1, Point2(1.0, 0.0), odometerNoise));
+      if (i % 3 == 0)
+        newFactors.push_back(BetweenFactor<Point2>(key_3, key_2, Point2(1.0, 0.0), odometerNoise));
+      if (i % 4 == 0)
+        newFactors.push_back(BetweenFactor<Point2>(key_4, key_3, Point2(1.0, 0.0), odometerNoise));
+      if (i % 5 == 0)
+        newFactors.push_back(BetweenFactor<Point2>(key_5, key_4, Point2(1.0, 0.0), odometerNoise));
+      if (i % 6 == 0)
+        newFactors.push_back(BetweenFactor<Point2>(key_6, key_5, Point2(1.0, 0.0), odometerNoise));
+      if (i % 7 == 0)
+        newFactors.push_back(BetweenFactor<Point2>(key_7, key_6, Point2(1.0, 0.0), odometerNoise));
+      if (i % 8 == 0)
+        newFactors.push_back(BetweenFactor<Point2>(key_8, key_7, Point2(1.0, 0.0), odometerNoise));
+      if (i % 9 == 0)
+        newFactors.push_back(BetweenFactor<Point2>(key_9, key_8, Point2(1.0, 0.0), odometerNoise));
+      if (i % 10 == 0)
+        newFactors.push_back(BetweenFactor<Point2>(key_10, key_9, Point2(1.0, 0.0), odometerNoise));
+
+      newValues.insert(key_0, Point2(double(i)+0.1, -0.1));
+      newTimestamps[key_0] = double(i);
+
+      fullgraph.push_back(newFactors);
+      fullinit.insert(newValues);
+
+      // Update the smoother
+      smoother.update(newFactors, newValues, newTimestamps);
+
+      // Check
+      CHECK(check_smoother(fullgraph, fullinit, smoother, key_0));
+      PrintSymbolicTree(smoother.getISAM2(), "Bayes Tree marginalization test: i = " + std::to_string(i));
+
+      ++i;
+    }
+  }
 }
 
 /* ************************************************************************* */