From 457d548015904b4f87e042cdf863867c0c35124d Mon Sep 17 00:00:00 2001
From: Varun Agrawal <varagrawal@gmail.com>
Date: Sun, 30 May 2021 11:22:58 -0400
Subject: [PATCH] override the rekey methods so as to update the properties as
 well

---
 gtsam/nonlinear/LinearContainerFactor.cpp | 41 +++++++++++++++++++++++
 gtsam/nonlinear/LinearContainerFactor.h   | 15 ++++++++-
 gtsam/nonlinear/NonlinearFactor.h         |  4 +--
 3 files changed, 57 insertions(+), 3 deletions(-)

diff --git a/gtsam/nonlinear/LinearContainerFactor.cpp b/gtsam/nonlinear/LinearContainerFactor.cpp
index 4f19f36f8..d715eb5c7 100644
--- a/gtsam/nonlinear/LinearContainerFactor.cpp
+++ b/gtsam/nonlinear/LinearContainerFactor.cpp
@@ -164,6 +164,47 @@ NonlinearFactor::shared_ptr LinearContainerFactor::negateToNonlinear() const {
   return NonlinearFactor::shared_ptr(new LinearContainerFactor(antifactor, linearizationPoint_));
 }
 
+/* ************************************************************************* */
+NonlinearFactor::shared_ptr LinearContainerFactor::rekey(
+    const std::map<Key, Key>& rekey_mapping) const {
+  auto rekeyed_base_factor = Base::rekey(rekey_mapping);
+  // Update the keys to the properties as well
+  // Downncast so we have access to members
+  auto new_factor = boost::static_pointer_cast<LinearContainerFactor>(rekeyed_base_factor);
+  // Create a new Values to assign later
+  Values newLinearizationPoint;
+  for (size_t i = 0; i < factor_->size(); ++i) {
+    auto mapping = rekey_mapping.find(factor_->keys()[i]);
+    if (mapping != rekey_mapping.end())
+      new_factor->factor_->keys()[i] = mapping->second;
+      newLinearizationPoint.insert(mapping->second, linearizationPoint_->at(mapping->first));
+  }
+  new_factor->linearizationPoint_ = newLinearizationPoint;
+
+  // upcast back and return
+  return boost::static_pointer_cast<NonlinearFactor>(new_factor);
+}
+
+/* ************************************************************************* */
+NonlinearFactor::shared_ptr LinearContainerFactor::rekey(
+    const KeyVector& new_keys) const {
+  auto rekeyed_base_factor = Base::rekey(new_keys);
+  // Update the keys to the properties as well
+  // Downncast so we have access to members
+  auto new_factor = boost::static_pointer_cast<LinearContainerFactor>(rekeyed_base_factor);
+  new_factor->factor_->keys() = new_factor->keys();
+  // Create a new Values to assign later
+  Values newLinearizationPoint;
+  for(size_t i=0; i<new_keys.size(); ++i) {
+    Key cur_key = linearizationPoint_->keys()[i];
+    newLinearizationPoint.insert(new_keys[i], linearizationPoint_->at(cur_key));
+  }
+  new_factor->linearizationPoint_ = newLinearizationPoint;
+
+  // upcast back and return
+  return boost::static_pointer_cast<NonlinearFactor>(new_factor);
+}
+
 /* ************************************************************************* */
 NonlinearFactorGraph LinearContainerFactor::ConvertLinearGraph(
     const GaussianFactorGraph& linear_graph, const Values& linearizationPoint) {
diff --git a/gtsam/nonlinear/LinearContainerFactor.h b/gtsam/nonlinear/LinearContainerFactor.h
index 8587e6b91..8c5b34f01 100644
--- a/gtsam/nonlinear/LinearContainerFactor.h
+++ b/gtsam/nonlinear/LinearContainerFactor.h
@@ -120,8 +120,21 @@ public:
     return NonlinearFactor::shared_ptr(new LinearContainerFactor(factor_,linearizationPoint_));
   }
 
-  // casting syntactic sugar
+  /**
+   * Creates a shared_ptr clone of the
+   * factor with different keys using
+   * a map from old->new keys
+   */
+  NonlinearFactor::shared_ptr rekey(
+      const std::map<Key, Key>& rekey_mapping) const override;
 
+  /**
+   * Clones a factor and fully replaces its keys
+   * @param new_keys is the full replacement set of keys
+   */
+  NonlinearFactor::shared_ptr rekey(const KeyVector& new_keys) const override;
+
+  /// Casting syntactic sugar
   inline bool hasLinearizationPoint() const { return linearizationPoint_.is_initialized(); }
 
   /**
diff --git a/gtsam/nonlinear/NonlinearFactor.h b/gtsam/nonlinear/NonlinearFactor.h
index 21c05dc2c..6deaaf7fe 100644
--- a/gtsam/nonlinear/NonlinearFactor.h
+++ b/gtsam/nonlinear/NonlinearFactor.h
@@ -126,13 +126,13 @@ public:
    * factor with different keys using
    * a map from old->new keys
    */
-  shared_ptr rekey(const std::map<Key,Key>& rekey_mapping) const;
+  virtual shared_ptr rekey(const std::map<Key,Key>& rekey_mapping) const;
 
   /**
    * Clones a factor and fully replaces its keys
    * @param new_keys is the full replacement set of keys
    */
-  shared_ptr rekey(const KeyVector& new_keys) const;
+  virtual shared_ptr rekey(const KeyVector& new_keys) const;
 
 }; // \class NonlinearFactor