From edadd352af69cae1292ac09a305d5b1695e94ce6 Mon Sep 17 00:00:00 2001
From: Frank Dellaert <dellaert@gmail.com>
Date: Fri, 24 Dec 2021 12:33:23 -0500
Subject: [PATCH] markdown for Bayes nets

---
 gtsam/discrete/DiscreteBayesNet.cpp           | 12 ++++++-
 gtsam/discrete/DiscreteBayesNet.h             |  9 +++++
 gtsam/discrete/discrete.i                     |  2 ++
 gtsam/discrete/tests/testDiscreteBayesNet.cpp | 33 ++++++++++++++++---
 4 files changed, 50 insertions(+), 6 deletions(-)

diff --git a/gtsam/discrete/DiscreteBayesNet.cpp b/gtsam/discrete/DiscreteBayesNet.cpp
index 219f2d93e..e50f4586f 100644
--- a/gtsam/discrete/DiscreteBayesNet.cpp
+++ b/gtsam/discrete/DiscreteBayesNet.cpp
@@ -38,7 +38,7 @@ namespace gtsam {
   double DiscreteBayesNet::evaluate(const DiscreteValues & values) const {
     // evaluate all conditionals and multiply
     double result = 1.0;
-    for(DiscreteConditional::shared_ptr conditional: *this)
+    for(const DiscreteConditional::shared_ptr& conditional: *this)
       result *= (*conditional)(values);
     return result;
   }
@@ -61,5 +61,15 @@ namespace gtsam {
     return result;
   }
 
+  /* ************************************************************************* */
+  std::string DiscreteBayesNet::_repr_markdown_(
+      const KeyFormatter& keyFormatter) const {
+    using std::endl;
+    std::stringstream ss;
+    ss << "`DiscreteBayesNet` of size " << size() << endl << endl;
+    for(const DiscreteConditional::shared_ptr& conditional: *this)
+      ss << conditional->_repr_markdown_(keyFormatter) << endl;
+    return ss.str();
+  }
 /* ************************************************************************* */
 } // namespace
diff --git a/gtsam/discrete/DiscreteBayesNet.h b/gtsam/discrete/DiscreteBayesNet.h
index 2d92b72e8..5eb656b3b 100644
--- a/gtsam/discrete/DiscreteBayesNet.h
+++ b/gtsam/discrete/DiscreteBayesNet.h
@@ -13,6 +13,7 @@
  * @file DiscreteBayesNet.h
  * @date Feb 15, 2011
  * @author Duy-Nguyen Ta
+ * @author Frank dellaert
  */
 
 #pragma once
@@ -97,6 +98,14 @@ namespace gtsam {
     DiscreteValues sample() const;
 
     ///@}
+    /// @name Wrapper support
+    /// @{
+
+    /// Render as markdown table.
+    std::string _repr_markdown_(
+        const KeyFormatter& keyFormatter = DefaultKeyFormatter) const;
+
+    /// @}
 
   private:
     /** Serialization function */
diff --git a/gtsam/discrete/discrete.i b/gtsam/discrete/discrete.i
index 4618073fa..19d1b8cd9 100644
--- a/gtsam/discrete/discrete.i
+++ b/gtsam/discrete/discrete.i
@@ -93,6 +93,8 @@ class DiscreteBayesNet {
   double operator()(const gtsam::DiscreteValues& values) const;
   gtsam::DiscreteValues optimize() const;
   gtsam::DiscreteValues sample() const;
+  string _repr_markdown_(const gtsam::KeyFormatter& keyFormatter =
+                 gtsam::DefaultKeyFormatter) const;
 };
 
 #include <gtsam/discrete/DiscreteBayesTree.h>
diff --git a/gtsam/discrete/tests/testDiscreteBayesNet.cpp b/gtsam/discrete/tests/testDiscreteBayesNet.cpp
index f0c7d3728..0a3c5d6e1 100644
--- a/gtsam/discrete/tests/testDiscreteBayesNet.cpp
+++ b/gtsam/discrete/tests/testDiscreteBayesNet.cpp
@@ -38,6 +38,9 @@ using namespace boost::assign;
 using namespace std;
 using namespace gtsam;
 
+static const DiscreteKey Asia(0, 2), Smoking(4, 2), Tuberculosis(3, 2),
+    LungCancer(6, 2), Bronchitis(7, 2), Either(5, 2), XRay(2, 2), Dyspnea(1, 2);
+
 /* ************************************************************************* */
 TEST(DiscreteBayesNet, bayesNet) {
   DiscreteBayesNet bayesNet;
@@ -71,8 +74,6 @@ TEST(DiscreteBayesNet, bayesNet) {
 /* ************************************************************************* */
 TEST(DiscreteBayesNet, Asia) {
   DiscreteBayesNet asia;
-  DiscreteKey Asia(0, 2), Smoking(4, 2), Tuberculosis(3, 2), LungCancer(6, 2),
-      Bronchitis(7, 2), Either(5, 2), XRay(2, 2), Dyspnea(1, 2);
 
   asia.add(Asia % "99/1");
   asia.add(Smoking % "50/50");
@@ -151,9 +152,6 @@ TEST(DiscreteBayesNet, Sugar) {
 
 /* ************************************************************************* */
 TEST(DiscreteBayesNet, Dot) {
-  DiscreteKey Asia(0, 2), Smoking(4, 2), Tuberculosis(3, 2), LungCancer(6, 2),
-      Either(5, 2);
-
   DiscreteBayesNet fragment;
   fragment.add(Asia % "99/1");
   fragment.add(Smoking % "50/50");
@@ -172,6 +170,31 @@ TEST(DiscreteBayesNet, Dot) {
          "}");
 }
 
+/* ************************************************************************* */
+// Check markdown representation looks as expected.
+TEST(DiscreteBayesNet, markdown) {
+  DiscreteBayesNet fragment;
+  fragment.add(Asia % "99/1");
+  fragment.add(Smoking | Asia = "8/2 7/3");
+
+  string expected =
+      "`DiscreteBayesNet` of size 2\n"
+      "\n"
+      " $P(Asia)$:\n"
+      "|0|1|\n"
+      "|:-:|:-:|\n"
+      "|0.99|0.01|\n"
+      "\n"
+      " $P(Smoking|Asia)$:\n"
+      "|Asia|0|1|\n"
+      "|:-:|:-:|:-:|\n"
+      "|0|0.8|0.2|\n"
+      "|1|0.7|0.3|\n\n";
+  auto formatter = [](Key key) { return key == 0 ? "Asia" : "Smoking"; };
+  string actual = fragment._repr_markdown_(formatter);
+  EXPECT(actual == expected);
+}
+
 /* ************************************************************************* */
 int main() {
   TestResult tr;