A Discrete Switching System

from gtsam import DiscreteBayesNet, DiscreteKeys, DiscreteFactorGraph, Ordering
from gtsam.symbol_shorthand import S
from gtsam.symbol_shorthand import M

def P(*args):
    """ Create a DiscreteKeys instances from a variable number of DiscreteKey pairs."""
    # TODO: We can make life easier by providing variable argument functions in C++ itself.
    dks = DiscreteKeys()
    for key in args:
        dks.push_back(key)
    return dks

import graphviz


class show(graphviz.Source):
    """ Display an object with a dot method as a graph."""

    def __init__(self, obj):
        """Construct from object with 'dot' method."""
        # This small class takes an object, calls its dot function, and uses the
        # resulting string to initialize a graphviz.Source instance. This in turn
        # has a _repr_mimebundle_ method, which then renders it in the notebook.
        super().__init__(obj.dot())

nrStates = 3
K = 5

bayesNet = DiscreteBayesNet()
for k in range(1, K):
    key = S(k), nrStates
    key_plus = S(k+1), nrStates
    mode = M(k), 2
    bayesNet.add(key_plus, P(mode, key), "9/1/0 1/8/1 0/1/9  1/9/0 0/1/9 9/0/1")

bayesNet

show(bayesNet)

# Create a factor graph out of the Bayes net.
factorGraph = DiscreteFactorGraph(bayesNet)
show(factorGraph)

# Create a BayesTree out of the factor graph.
ordering = Ordering()
# First eliminate "continuous" states in time order
for k in range(1, K+1):
    ordering.push_back(S(k))
for k in range(1, K):
    ordering.push_back(M(k))
print(ordering)
bayesTree = factorGraph.eliminateMultifrontal(ordering)

show(bayesTree)