From eb22de0b50da1d1dcbcd58e815f9fc8979a4315c Mon Sep 17 00:00:00 2001
From: Tom Schierenbeck <tom_sch@uni-bremen.de>
Date: Tue, 2 Jan 2024 08:23:06 +0100
Subject: [PATCH] Mixtures models as factors soon

---
 requirements.txt                     |  2 +-
 src/fglib2/probabilistic_circuits.py | 34 +++++++++++++++
 test/test_interface.py               |  2 +-
 test/test_pc.py                      | 62 ++++++++++++++++++++++++++++
 4 files changed, 98 insertions(+), 2 deletions(-)
 create mode 100644 src/fglib2/probabilistic_circuits.py
 create mode 100644 test/test_pc.py

diff --git a/requirements.txt b/requirements.txt
index 7acd511..7425810 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,5 +2,5 @@ networkx>=3.0
 numpy>=1.24.4
 random_events>=1.2.5
 tabulate>=0.9.0
-probabilistic-model>=1.4.13
+probabilistic-model>=1.5.18
 typing-extensions
\ No newline at end of file
diff --git a/src/fglib2/probabilistic_circuits.py b/src/fglib2/probabilistic_circuits.py
new file mode 100644
index 0000000..17d3fbe
--- /dev/null
+++ b/src/fglib2/probabilistic_circuits.py
@@ -0,0 +1,34 @@
+from typing import Iterable, Optional, Union, List
+
+import numpy as np
+from typing_extensions import Self
+
+from probabilistic_model.probabilistic_model import ProbabilisticModel
+from probabilistic_model.probabilistic_circuit import SmoothSumUnit
+from random_events.variables import Variable, Symbolic
+from .distributions import Multinomial
+
+
+class SumUnitFactor(ProbabilisticModel):
+    """
+    A sum unite (mixture model) that can be used as factor for variables in a factor graph.
+
+    Example use-case:
+        Imagine you have a set of variables that expand over some template, e.g. time.
+        You learn a mixture for each time step and then use the latent variable interpretation of the
+        mixture model to create a factor graph. The factors for the transition model are multinomial distributions
+        over the latent variables. The factors for the emission model are the joint probability trees.
+    """
+
+    def __init__(self, model: SmoothSumUnit):
+        self.model = model
+        latent_variable = Symbolic(f"latent_{str(id(model))}", range(len(self.model.weights)))
+        super().__init__([latent_variable])
+
+    @property
+    def latent_variable(self) -> Symbolic:
+        return self.variables[0]
+
+    def marginal(self, variables: List[Variable]) -> Union[Multinomial, Self]:
+        if variables[0] == self.latent_variable:
+            return Multinomial([self.latent_variable], np.array(self.model.weights))
\ No newline at end of file
diff --git a/test/test_interface.py b/test/test_interface.py
index 6ff6b4e..6a556c4 100644
--- a/test/test_interface.py
+++ b/test/test_interface.py
@@ -52,7 +52,7 @@ def generate_random_likelihood_event_over_variables(variables: List[Symbolic]) -
     """
     return [random.choice(variable.domain) for variable in variables]
 
-
+@unittest.skip("Not implemented")
 class InterfaceTestCase(unittest.TestCase):
 
     variables = [Symbolic(f"x_{i}", range(random.randrange(2, 4))) for i in range(9)]
diff --git a/test/test_pc.py b/test/test_pc.py
new file mode 100644
index 0000000..18a398c
--- /dev/null
+++ b/test/test_pc.py
@@ -0,0 +1,62 @@
+import unittest
+
+import networkx as nx
+import numpy as np
+from random_events.variables import Continuous
+from probabilistic_model.probabilistic_circuit.units import DeterministicSumUnit
+from probabilistic_model.probabilistic_circuit.distributions import UniformDistribution
+from fglib2.distributions import Multinomial
+from fglib2.probabilistic_circuits import SumUnitFactor
+import portion
+from fglib2.graphs import FactorGraph, FactorNode
+import matplotlib.pyplot as plt
+
+
+class JPTTestCase(unittest.TestCase):
+
+    variables = [
+        Continuous("x1"),
+        Continuous("x2"),
+        # Continuous("x3"),
+        # Continuous("x4"),
+        # Continuous("x5")
+        ]
+
+    model: FactorGraph
+
+    def setUp(self):
+        interval_1 = portion.closed(-1.25, -0.75)
+        interval_2 = portion.closed(0.75, 1.25)
+
+        model = FactorGraph()
+
+        for variable in self.variables:
+            distribution = SumUnitFactor(UniformDistribution(variable, interval_1) +
+                                         UniformDistribution(variable, interval_2))
+            factor = FactorNode(distribution)
+            model *= factor
+
+        for f1, f2 in zip(model.factor_nodes[:-1], model.factor_nodes[1:]):
+            model *= FactorNode(Multinomial([f1.variables[0], f2.variables[0]], np.array([[0, 0.5], [0.5, 0]])))
+
+        self.model = model
+
+    def test_creation(self):
+        nx.draw(self.model, with_labels=True)
+        plt.show()
+
+    def test_marginal(self):
+        self.model.sum_product()
+        latent_variables = self.model.variables
+        print(latent_variables)
+        for variable in latent_variables:
+            print(self.model.belief(variable))
+        # self.model.marginal(self.variables[:1])
+
+
+    def test_sum_product(self):
+        self.model.sum_product()
+
+
+if __name__ == '__main__':
+    unittest.main()