From 00be610e18ced3a1226a5271a6f1ad275d5c3539 Mon Sep 17 00:00:00 2001
From: Varun Agrawal <varagrawal@gmail.com>
Date: Tue, 7 Jun 2022 18:39:33 -0400
Subject: [PATCH] Add Switching class and make it linear

---
 gtsam/hybrid/tests/Switching.h | 139 +++++++++++++++++++++++++++++++--
 1 file changed, 132 insertions(+), 7 deletions(-)

diff --git a/gtsam/hybrid/tests/Switching.h b/gtsam/hybrid/tests/Switching.h
index c081b8e87..85c86f64e 100644
--- a/gtsam/hybrid/tests/Switching.h
+++ b/gtsam/hybrid/tests/Switching.h
@@ -18,18 +18,136 @@
 
 #include <gtsam/base/Matrix.h>
 #include <gtsam/discrete/DecisionTreeFactor.h>
+#include <gtsam/discrete/DiscreteDistribution.h>
 #include <gtsam/hybrid/GaussianMixtureFactor.h>
 #include <gtsam/hybrid/HybridGaussianFactorGraph.h>
 #include <gtsam/inference/Symbol.h>
 #include <gtsam/linear/JacobianFactor.h>
+#include <gtsam/slam/BetweenFactor.h>
 
 #pragma once
 
-using gtsam::symbol_shorthand::C;
-using gtsam::symbol_shorthand::X;
-
 namespace gtsam {
-inline HybridGaussianFactorGraph::shared_ptr makeSwitchingChain(
+
+using symbol_shorthand::C;
+using symbol_shorthand::M;
+using symbol_shorthand::X;
+
+/* ****************************************************************************/
+// Test fixture with switching network.
+// TODO(Varun) Currently this is only linear. We need to add nonlinear support
+// and then update to
+// https://github.com/borglab/gtsam/pull/973/files#diff-58c02b3b197ebf731694946e87762d252e9eaa2f5c6c4ba22d618085b321ca23
+struct Switching {
+  size_t K;
+  DiscreteKeys modes;
+  HybridGaussianFactorGraph linearizedFactorGraph;
+  Values linearizationPoint;
+
+  using MotionModel = BetweenFactor<double>;
+  // using MotionMixture = MixtureFactor<MotionModel>;
+
+  /// Create with given number of time steps.
+  Switching(size_t K, double between_sigma = 1.0, double prior_sigma = 0.1)
+      : K(K) {
+    // Create DiscreteKeys for binary K modes, modes[0] will not be used.
+    modes = addDiscreteModes(K);
+
+    // Create hybrid factor graph.
+    // Add a prior on X(1).
+    auto prior = boost::make_shared<JacobianFactor>(
+        X(1), Matrix11::Ones() / prior_sigma, Vector1::Zero());
+    linearizedFactorGraph.push_back(prior);
+
+    // Add "motion models".
+    linearizedFactorGraph = addMotionModels(K);
+
+    // Add measurement factors
+    for (size_t k = 1; k <= K; k++) {
+      linearizedFactorGraph.emplace_gaussian<JacobianFactor>(
+          X(k), Matrix11::Ones() / 0.1, Vector1::Zero());
+    }
+
+    // Add "mode chain"
+    linearizedFactorGraph = addModeChain(linearizedFactorGraph);
+
+    // Create the linearization point.
+    for (size_t k = 1; k <= K; k++) {
+      linearizationPoint.insert<double>(X(k), static_cast<double>(k));
+    }
+  }
+
+  /// Create DiscreteKeys for K binary modes.
+  DiscreteKeys addDiscreteModes(size_t K) {
+    DiscreteKeys m;
+    for (size_t k = 0; k <= K; k++) {
+      m.emplace_back(M(k), 2);
+    }
+    return m;
+  }
+
+  /// Helper function to add motion models for each [k, k+1] interval.
+  HybridGaussianFactorGraph addMotionModels(size_t K) {
+    HybridGaussianFactorGraph hgfg;
+    for (size_t k = 1; k < K; k++) {
+      auto keys = {X(k), X(k + 1)};
+      auto components = motionModels(k);
+      hgfg.emplace_hybrid<GaussianMixtureFactor>(keys, DiscreteKeys{modes[k]},
+                                                 components);
+    }
+    return hgfg;
+  }
+
+  // Create motion models for a given time step
+  static std::vector<GaussianFactor::shared_ptr> motionModels(
+      size_t k, double sigma = 1.0) {
+    auto noise_model = noiseModel::Isotropic::Sigma(1, sigma);
+    auto still = boost::make_shared<JacobianFactor>(
+             X(k), -Matrix11::Ones() / sigma, X(k + 1),
+             Matrix11::Ones() / sigma, Vector1::Zero()),
+         moving = boost::make_shared<JacobianFactor>(
+             X(k), -Matrix11::Ones() / sigma, X(k + 1),
+             Matrix11::Ones() / sigma, -Vector1::Ones() / sigma);
+    return {boost::dynamic_pointer_cast<GaussianFactor>(still),
+            boost::dynamic_pointer_cast<GaussianFactor>(moving)};
+  }
+
+  // // Add "mode chain" to NonlinearHybridFactorGraph
+  // void addModeChain(HybridNonlinearFactorGraph& fg) {
+  //   auto prior = boost::make_shared<DiscreteDistribution>(modes[1], "1/1");
+  //   fg.push_discrete(prior);
+  //   for (size_t k = 1; k < K - 1; k++) {
+  //     auto parents = {modes[k]};
+  //     auto conditional = boost::make_shared<DiscreteConditional>(
+  //         modes[k + 1], parents, "1/2 3/2");
+  //     fg.push_discrete(conditional);
+  //   }
+  // }
+
+  // Add "mode chain" to GaussianHybridFactorGraph
+  HybridGaussianFactorGraph addModeChain(HybridGaussianFactorGraph& fg) {
+    auto prior = boost::make_shared<DiscreteDistribution>(modes[1], "1/1");
+    fg.push_discrete(prior);
+    for (size_t k = 1; k < K - 1; k++) {
+      auto parents = {modes[k]};
+      auto conditional = boost::make_shared<DiscreteConditional>(
+          modes[k + 1], parents, "1/2 3/2");
+      fg.push_discrete(conditional);
+    }
+    return fg;
+  }
+};
+
+/**
+ * @brief Create a switching system chain. A switching system is a continuous
+ * system which depends on a discrete mode at each time step of the chain.
+ *
+ * @param n The number of chain elements.
+ * @param keyFunc The functional to help specify the continuous key.
+ * @param dKeyFunc The functional to help specify the discrete key.
+ * @return HybridGaussianFactorGraph::shared_ptr
+ */
+HybridGaussianFactorGraph::shared_ptr makeSwitchingChain(
     size_t n, std::function<Key(int)> keyFunc = X,
     std::function<Key(int)> dKeyFunc = C) {
   HybridGaussianFactorGraph hfg;
@@ -54,9 +172,16 @@ inline HybridGaussianFactorGraph::shared_ptr makeSwitchingChain(
   return boost::make_shared<HybridGaussianFactorGraph>(std::move(hfg));
 }
 
-inline std::pair<KeyVector, std::vector<int>> makeBinaryOrdering(
-    std::vector<Key> &input) {
+/**
+ * @brief
+ *
+ * @param input The original ordering.
+ * @return std::pair<KeyVector, std::vector<int>>
+ */
+std::pair<KeyVector, std::vector<int>> makeBinaryOrdering(
+    std::vector<Key>& input) {
   KeyVector new_order;
+
   std::vector<int> levels(input.size());
   std::function<void(std::vector<Key>::iterator, std::vector<Key>::iterator,
                      int)>
@@ -79,7 +204,7 @@ inline std::pair<KeyVector, std::vector<int>> makeBinaryOrdering(
 
   bsg(input.begin(), input.end(), 0);
   std::reverse(new_order.begin(), new_order.end());
-  // std::reverse(levels.begin(), levels.end());
+
   return {new_order, levels};
 }