-
Notifications
You must be signed in to change notification settings - Fork 0
/
test_fglib.py
41 lines (31 loc) · 976 Bytes
/
test_fglib.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
from fglib import graphs, nodes, inference, rv
from random import choice
import networkx as nx
# Create factor graph
fg = graphs.FactorGraph()
# Create variable nodes
x1 = nodes.VNode("x1", rv.Discrete) # with 2 states (Bernoulli)
x2 = nodes.VNode("x2", rv.Discrete) # with 3 states
x3 = nodes.VNode("x3", rv.Discrete)
x4 = nodes.VNode("x4", rv.Discrete)
# Create factor nodes (with joint distributions)
dist_fa = [[3, 4],
[3, 9]]
fa = nodes.FNode("fa", rv.Discrete(dist_fa, x1, x2))
dist_fb = [[3, 4],
[5, 1]]
fb = nodes.FNode("fb", rv.Discrete(dist_fb, x2, x3))
dist_fc = [[7, 8],
[3, 9]]
fc = nodes.FNode("fc", rv.Discrete(dist_fc, x2, x4))
# Add nodes to factor graph
fg.set_nodes([x1, x2, x3, x4])
fg.set_nodes([fa, fb, fc])
# Add edges to factor graph
fg.set_edge(x1, fa)
fg.set_edge(fa, x2)
fg.set_edge(x2, fb)
fg.set_edge(fb, x3)
fg.set_edge(x2, fc)
fg.set_edge(fc, x4)
fg.save_graph_fig(num=1, pos=pos, fig_name='tree-1')