DiscreteBayesNet_FG.cpp 4.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120
  1. /* ----------------------------------------------------------------------------
  2. * GTSAM Copyright 2010, Georgia Tech Research Corporation,
  3. * Atlanta, Georgia 30332-0415
  4. * All Rights Reserved
  5. * Authors: Frank Dellaert, et al. (see THANKS for the full author list)
  6. * See LICENSE for the license information
  7. * -------------------------------------------------------------------------- */
  8. /**
  9. * @file DiscreteBayesNet_FG.cpp
  10. * @brief Discrete Bayes Net example using Factor Graphs
  11. * @author Abhijit
  12. * @date Jun 4, 2012
  13. *
  14. * We use the famous Rain/Cloudy/Sprinkler Example of [Russell & Norvig, 2009,
  15. * p529] You may be familiar with other graphical model packages like BNT
  16. * (available at http://bnt.googlecode.com/svn/trunk/docs/usage.html) where this
  17. * is used as an example. The following demo is same as that in the above link,
  18. * except that everything is using GTSAM.
  19. */
  20. #include <gtsam/discrete/DiscreteFactorGraph.h>
  21. #include <gtsam/discrete/DiscreteMarginals.h>
  22. #include <iomanip>
  23. using namespace std;
  24. using namespace gtsam;
  25. int main(int argc, char **argv) {
  26. // Define keys and a print function
  27. Key C(1), S(2), R(3), W(4);
  28. auto print = [=](DiscreteFactor::sharedValues values) {
  29. cout << boolalpha << "Cloudy = " << static_cast<bool>((*values)[C])
  30. << " Sprinkler = " << static_cast<bool>((*values)[S])
  31. << " Rain = " << boolalpha << static_cast<bool>((*values)[R])
  32. << " WetGrass = " << static_cast<bool>((*values)[W]) << endl;
  33. };
  34. // We assume binary state variables
  35. // we have 0 == "False" and 1 == "True"
  36. const size_t nrStates = 2;
  37. // define variables
  38. DiscreteKey Cloudy(C, nrStates), Sprinkler(S, nrStates), Rain(R, nrStates),
  39. WetGrass(W, nrStates);
  40. // create Factor Graph of the bayes net
  41. DiscreteFactorGraph graph;
  42. // add factors
  43. graph.add(Cloudy, "0.5 0.5"); // P(Cloudy)
  44. graph.add(Cloudy & Sprinkler, "0.5 0.5 0.9 0.1"); // P(Sprinkler | Cloudy)
  45. graph.add(Cloudy & Rain, "0.8 0.2 0.2 0.8"); // P(Rain | Cloudy)
  46. graph.add(Sprinkler & Rain & WetGrass,
  47. "1 0 0.1 0.9 0.1 0.9 0.001 0.99"); // P(WetGrass | Sprinkler, Rain)
  48. // Alternatively we can also create a DiscreteBayesNet, add
  49. // DiscreteConditional factors and create a FactorGraph from it. (See
  50. // testDiscreteBayesNet.cpp)
  51. // Since this is a relatively small distribution, we can as well print
  52. // the whole distribution..
  53. cout << "Distribution of Example: " << endl;
  54. cout << setw(11) << "Cloudy(C)" << setw(14) << "Sprinkler(S)" << setw(10)
  55. << "Rain(R)" << setw(14) << "WetGrass(W)" << setw(15) << "P(C,S,R,W)"
  56. << endl;
  57. for (size_t a = 0; a < nrStates; a++)
  58. for (size_t m = 0; m < nrStates; m++)
  59. for (size_t h = 0; h < nrStates; h++)
  60. for (size_t c = 0; c < nrStates; c++) {
  61. DiscreteFactor::Values values;
  62. values[C] = c;
  63. values[S] = h;
  64. values[R] = m;
  65. values[W] = a;
  66. double prodPot = graph(values);
  67. cout << setw(8) << static_cast<bool>(c) << setw(14)
  68. << static_cast<bool>(h) << setw(12) << static_cast<bool>(m)
  69. << setw(13) << static_cast<bool>(a) << setw(16) << prodPot
  70. << endl;
  71. }
  72. // "Most Probable Explanation", i.e., configuration with largest value
  73. DiscreteFactor::sharedValues mpe = graph.eliminateSequential()->optimize();
  74. cout << "\nMost Probable Explanation (MPE):" << endl;
  75. print(mpe);
  76. // "Inference" We show an inference query like: probability that the Sprinkler
  77. // was on; given that the grass is wet i.e. P( S | C=0) = ?
  78. // add evidence that it is not Cloudy
  79. graph.add(Cloudy, "1 0");
  80. // solve again, now with evidence
  81. DiscreteBayesNet::shared_ptr chordal = graph.eliminateSequential();
  82. DiscreteFactor::sharedValues mpe_with_evidence = chordal->optimize();
  83. cout << "\nMPE given C=0:" << endl;
  84. print(mpe_with_evidence);
  85. // we can also calculate arbitrary marginals:
  86. DiscreteMarginals marginals(graph);
  87. cout << "\nP(S=1|C=0):" << marginals.marginalProbabilities(Sprinkler)[1]
  88. << endl;
  89. cout << "\nP(R=0|C=0):" << marginals.marginalProbabilities(Rain)[0] << endl;
  90. cout << "\nP(W=1|C=0):" << marginals.marginalProbabilities(WetGrass)[1]
  91. << endl;
  92. // We can also sample from it
  93. cout << "\n10 samples:" << endl;
  94. for (size_t i = 0; i < 10; i++) {
  95. DiscreteFactor::sharedValues sample = chordal->sample();
  96. print(sample);
  97. }
  98. return 0;
  99. }