From 143022c1393ad4d3d770d7ab25a3d5473842e82e Mon Sep 17 00:00:00 2001
From: Frank Dellaert <dellaert@gmail.com>
Date: Sat, 31 Dec 2022 18:07:17 -0500
Subject: [PATCH] Tiny Bayes net example

---
 gtsam/hybrid/tests/TinyHybridExample.h    | 63 +++++++++++++++++++++++
 gtsam/hybrid/tests/testHybridBayesNet.cpp | 10 +++-
 2 files changed, 72 insertions(+), 1 deletion(-)
 create mode 100644 gtsam/hybrid/tests/TinyHybridExample.h

diff --git a/gtsam/hybrid/tests/TinyHybridExample.h b/gtsam/hybrid/tests/TinyHybridExample.h
new file mode 100644
index 000000000..c3e7b2f5b
--- /dev/null
+++ b/gtsam/hybrid/tests/TinyHybridExample.h
@@ -0,0 +1,63 @@
+/* ----------------------------------------------------------------------------
+
+ * GTSAM Copyright 2010-2022, 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 TinyHybrdiExample.h
+ *  @date Mar 11, 2022
+ *  @author Varun Agrawal
+ *  @author Fan Jiang
+ */
+
+#include <gtsam/hybrid/HybridBayesNet.h>
+#include <gtsam/inference/Symbol.h>
+#pragma once
+
+namespace gtsam {
+namespace tiny {
+
+using symbol_shorthand::M;
+using symbol_shorthand::X;
+using symbol_shorthand::Z;
+
+/**
+ * Create a tiny two variable hybrid model which represents
+ * the generative probability P(z, x, n) = P(z | x, n)P(x)P(n).
+ */
+static HybridBayesNet createHybridBayesNet(int num_measurements = 1) {
+  // Create hybrid Bayes net.
+  HybridBayesNet bayesNet;
+
+  // Create mode key: 0 is low-noise, 1 is high-noise.
+  const DiscreteKey mode{M(0), 2};
+
+  // Create Gaussian mixture Z(0) = X(0) + noise for each measurement.
+  for (int i = 0; i < num_measurements; i++) {
+    const auto conditional0 = boost::make_shared<GaussianConditional>(
+        GaussianConditional::FromMeanAndStddev(Z(i), I_1x1, X(0), Z_1x1, 0.5));
+    const auto conditional1 = boost::make_shared<GaussianConditional>(
+        GaussianConditional::FromMeanAndStddev(Z(i), I_1x1, X(0), Z_1x1, 3));
+    GaussianMixture gm({Z(i)}, {X(0)}, {mode}, {conditional0, conditional1});
+    bayesNet.emplaceMixture(gm);  // copy :-(
+  }
+
+  // Create prior on X(0).
+  const auto prior_on_x0 =
+      GaussianConditional::FromMeanAndStddev(X(0), Vector1(5.0), 5.0);
+  bayesNet.emplaceGaussian(prior_on_x0);  // copy :-(
+
+  // Add prior on mode.
+  bayesNet.emplaceDiscrete(mode, "4/6");
+
+  return bayesNet;
+}
+
+}  // namespace tiny
+}  // namespace gtsam
diff --git a/gtsam/hybrid/tests/testHybridBayesNet.cpp b/gtsam/hybrid/tests/testHybridBayesNet.cpp
index ce483b593..5cc6566fb 100644
--- a/gtsam/hybrid/tests/testHybridBayesNet.cpp
+++ b/gtsam/hybrid/tests/testHybridBayesNet.cpp
@@ -23,6 +23,7 @@
 #include <gtsam/hybrid/HybridBayesTree.h>
 #include <gtsam/nonlinear/NonlinearFactorGraph.h>
 
+#include "TinyHybridExample.h"
 #include "Switching.h"
 
 // Include for test suite
@@ -63,7 +64,7 @@ TEST(HybridBayesNet, Add) {
 
 /* ****************************************************************************/
 // Test evaluate for a pure discrete Bayes net P(Asia).
-TEST(HybridBayesNet, evaluatePureDiscrete) {
+TEST(HybridBayesNet, EvaluatePureDiscrete) {
   HybridBayesNet bayesNet;
   bayesNet.emplaceDiscrete(Asia, "99/1");
   HybridValues values;
@@ -71,6 +72,13 @@ TEST(HybridBayesNet, evaluatePureDiscrete) {
   EXPECT_DOUBLES_EQUAL(0.99, bayesNet.evaluate(values), 1e-9);
 }
 
+/* ****************************************************************************/
+// Test creation of a tiny hybrid Bayes net.
+TEST(HybridBayesNet, Tiny) {
+  auto bayesNet = tiny::createHybridBayesNet();
+  EXPECT_LONGS_EQUAL(3, bayesNet.size());
+}
+
 /* ****************************************************************************/
 // Test evaluate for a hybrid Bayes net P(X0|X1) P(X1|Asia) P(Asia).
 TEST(HybridBayesNet, evaluateHybrid) {