def test__generate_mutual_information_graph(self):
        samples = [
            [0, 0, 0],
            [0, 0, 1],
            [0, 1, 1],
            [1, 1, 1],
        ]

        expected_results = [
            (0, 1, {
                'weight': -0.21576155433883565
            }),
            (0, 2, {
                'weight': -0.084949518397698542
            }),
            (1, 2, {
                'weight': -0.21576155433883565
            }),
        ]

        distribution = mimic.Distribution(samples)

        self.assertEqual(
            expected_results,
            distribution.complete_graph.edges(data=True),
        )
    def test__generate_spanning_graph(self):
        samples = [
            [0, 0, 0],
            [0, 0, 1],
            [0, 1, 1],
            [1, 1, 1],
        ]

        expected_results = [
            (0, 1),
            (1, 2),
        ]

        distribution = mimic.Distribution(samples)

        self.assertEqual(
            expected_results,
            distribution.spanning_graph.edges(),
        )
    def test_generate_samples(self):
        expected_results = np.array([[1, 0, 1, 1]])
        graph = nx.DiGraph()

        graph.add_node(0, probabilities={0: 0, 1: 1})
        graph.add_node(1, probabilities={0: {0: 0, 1: 1}, 1: {0: 1, 1: 0}})
        graph.add_node(2, probabilities={0: {0: 0, 1: 1}, 1: {0: 1, 1: 0}})
        graph.add_node(3, probabilities={0: {0: 0, 1: 1}, 1: {0: 1, 1: 0}})

        graph.add_edges_from([
            (0, 1),
            (1, 2),
            (1, 3),
        ])

        distribution = mimic.Distribution([[0, 0], [0, 0]])
        distribution.bayes_net = graph

        self.assertTrue(
            np.equal(
                expected_results,
                distribution.generate_samples(1),
            ).all())