From c8bafab4307d7b2da98bf50b4f85e04cac53ba9e Mon Sep 17 00:00:00 2001
From: Varun Agrawal <varagrawal@gmail.com>
Date: Tue, 28 Jan 2025 16:21:26 -0500
Subject: [PATCH] add scale() method in DiscreteFactor

---
 gtsam/discrete/DiscreteFactor.cpp | 8 ++++++++
 gtsam/discrete/DiscreteFactor.h   | 8 ++++++++
 2 files changed, 16 insertions(+)

diff --git a/gtsam/discrete/DiscreteFactor.cpp b/gtsam/discrete/DiscreteFactor.cpp
index c8e2e0fd2..faae02af2 100644
--- a/gtsam/discrete/DiscreteFactor.cpp
+++ b/gtsam/discrete/DiscreteFactor.cpp
@@ -71,4 +71,12 @@ AlgebraicDecisionTree<Key> DiscreteFactor::errorTree() const {
   return AlgebraicDecisionTree<Key>(dkeys, errors);
 }
 
+/* ************************************************************************ */
+DiscreteFactor::shared_ptr DiscreteFactor::scale() const {
+  // Max over all the potentials by pretending all keys are frontal:
+  shared_ptr denominator = this->max(this->size());
+  // Normalize the product factor to prevent underflow.
+  return this->operator/(denominator);
+}
+
 }  // namespace gtsam
diff --git a/gtsam/discrete/DiscreteFactor.h b/gtsam/discrete/DiscreteFactor.h
index 92bbd18dd..fafb4dbf5 100644
--- a/gtsam/discrete/DiscreteFactor.h
+++ b/gtsam/discrete/DiscreteFactor.h
@@ -158,6 +158,14 @@ class GTSAM_EXPORT DiscreteFactor : public Factor {
   /// Create new factor by maximizing over all values with the same separator.
   virtual DiscreteFactor::shared_ptr max(const Ordering& keys) const = 0;
 
+  /**
+   * @brief Scale the factor values by the maximum
+   * to prevent underflow/overflow.
+   * 
+   * @return DiscreteFactor::shared_ptr 
+   */
+  DiscreteFactor::shared_ptr scale() const;
+
   /**
    * Get the number of non-zero values contained in this factor.
    * It could be much smaller than `prod_{key}(cardinality(key))`.