From 98febf2f0cfbaf1e8138a938f620cab1f997bf61 Mon Sep 17 00:00:00 2001
From: Varun Agrawal <varagrawal@gmail.com>
Date: Mon, 7 Nov 2022 18:24:21 -0500
Subject: [PATCH] allow optional discrete transition probability for Switching
 test fixture

---
 gtsam/hybrid/tests/Switching.h | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/gtsam/hybrid/tests/Switching.h b/gtsam/hybrid/tests/Switching.h
index 76721edc4..4bf1678be 100644
--- a/gtsam/hybrid/tests/Switching.h
+++ b/gtsam/hybrid/tests/Switching.h
@@ -133,7 +133,8 @@ struct Switching {
    * @param measurements Vector of measurements for each timestep.
    */
   Switching(size_t K, double between_sigma = 1.0, double prior_sigma = 0.1,
-            std::vector<double> measurements = {})
+            std::vector<double> measurements = {},
+            std::string discrete_transition_prob = "1/2 3/2")
       : K(K) {
     // Create DiscreteKeys for binary K modes.
     for (size_t k = 0; k < K; k++) {
@@ -173,7 +174,7 @@ struct Switching {
     }
 
     // Add "mode chain"
-    addModeChain(&nonlinearFactorGraph);
+    addModeChain(&nonlinearFactorGraph, discrete_transition_prob);
 
     // Create the linearization point.
     for (size_t k = 0; k < K; k++) {
@@ -202,13 +203,14 @@ struct Switching {
    *
    * @param fg The nonlinear factor graph to which the mode chain is added.
    */
-  void addModeChain(HybridNonlinearFactorGraph *fg) {
+  void addModeChain(HybridNonlinearFactorGraph *fg,
+                    std::string discrete_transition_prob = "1/2 3/2") {
     auto prior = boost::make_shared<DiscreteDistribution>(modes[0], "1/1");
     fg->push_discrete(prior);
     for (size_t k = 0; k < K - 2; k++) {
       auto parents = {modes[k]};
       auto conditional = boost::make_shared<DiscreteConditional>(
-          modes[k + 1], parents, "1/2 3/2");
+          modes[k + 1], parents, discrete_transition_prob);
       fg->push_discrete(conditional);
     }
   }
@@ -219,13 +221,14 @@ struct Switching {
    *
    * @param fg The gaussian factor graph to which the mode chain is added.
    */
-  void addModeChain(HybridGaussianFactorGraph *fg) {
+  void addModeChain(HybridGaussianFactorGraph *fg,
+                    std::string discrete_transition_prob = "1/2 3/2") {
     auto prior = boost::make_shared<DiscreteDistribution>(modes[0], "1/1");
     fg->push_discrete(prior);
     for (size_t k = 0; k < K - 2; k++) {
       auto parents = {modes[k]};
       auto conditional = boost::make_shared<DiscreteConditional>(
-          modes[k + 1], parents, "1/2 3/2");
+          modes[k + 1], parents, discrete_transition_prob);
       fg->push_discrete(conditional);
     }
   }