Pass TestSemiCategoryRepresentation

1 files changed, 69 insertions, 94 deletions

diff --git a/src/act4e_solutions/semicategory_representation.py b/src/act4e_solutions/semicategory_representation.py
index e99f669..d6977bc 100644
--- a/src/act4e_solutions/semicategory_representation.py
+++ b/src/act4e_solutions/semicategory_representation.py
@@ -5,8 +5,10 @@ from act4e_interfaces import EnumerableSet
 
 from .sets_representation import MyFiniteSet
 
-from itertools import tee
 from dataclasses import dataclass
+from functools import partial
+from itertools import product
+
 
 E = TypeVar("E")
 OD = TypeVar("OD") # Type of the objects
@@ -15,107 +17,91 @@ MD = TypeVar("MD") # Type of the morphisms
 ROD = I.RichObject[OD]
 RMD = I.RichMorphism[MD]
 
-class Graph:
-    def __init__(self):
-        self.nodes = {}
-        self.edges = {}
-
-    def add_node(self, name):
-        self.nodes[name] = []
-
-    def add_edge(self, source, target, name):
-        self.nodes[source].append((target, name))
-        if name not in self.edges:
-            self.edges[name] = []
-        self.edges[name].append((source, target))
 
-    def get_edge_nodes(self, name):
-        return self.edges[name][0]
+@dataclass(frozen=True)
+class Morph(I.RichMorphism[MD]):
+    source: str
+    target: str
 
-    def find_all_paths(self, source, target, length):
-        paths = []
-        def dfs(node, path, depth):
-            if depth == length:
-                if node == target:
-                    paths.append(path)
-                return
 
-            for neighbor, edge_name in self.nodes[node]:
-                new_path = path + [edge_name]
-                dfs(neighbor, new_path, depth + 1)
-
-        dfs(source, [], 0)
-        return paths
+@dataclass(frozen=True)
+class Obj(I.RichObject[OD]):
+    identity: Optional[Morph]
 
 
 class MySemiCategory(Generic[OD, MD], I.SemiCategory[ROD, RMD]):
-    """ Skeleton for a class implementing SemiCategory."""
-    _objects: Dict[str, ROD]
-    _morphisms: Dict[str, RMD]
-    _graph: Graph
+    _objects: Dict[str, Obj]
+    _morphisms: Dict[int, Dict[Tuple[str, str], List[Morph]]]
     _composition: Callable[[OD, OD, OD, MD], MD]
     _equations: Dict[List[str], str]
-    _identity: RMD
 
     def __init__(
         self,
-        objects,
-        morphisms,
-        graph,
         composition: Callable[[OD, OD, OD, MD, MD], MD],
         equations: Dict[str, str],
-        identity: Optional[RMD] = None
     ):
-        self._objects = objects
-        self._morphisms = morphisms
-        self._graph = graph
+        self._objects = {}
+        self._morphisms = {0: {}}
         self._composition = composition
         self._equations = equations
-        self._identity = identity
+
+    def add_object(self, ob: Obj):
+        self._objects[ob.label] = ob
+
+    def add_morphism(self, mor: Morph):
+        if not (mor.source, mor.target) in self._morphisms[0].keys():
+            self._morphisms[0][(mor.source, mor.target)] = []
+        self._morphisms[0][(mor.source, mor.target)].append(mor)
 
     def objects(self, uptolevel: Optional[int] = None) -> EnumerableSet[ROD]: 
-        # FIXME: why uptolevel?
         return MyFiniteSet(self._objects.values())
 
     def hom(self, ob1: ROD, ob2: ROD, uptolevel: Optional[int] = None) -> EnumerableSet[RMD]:
-        morphisms, paths = [], []
-
-        for i in range(1, uptolevel + 2):
-            for p in self._graph.find_all_paths(ob1.label, ob2.label, i):
-                # Remove empty paths
-                if not p:
-                    continue
-
-                # Remove paths with identities (except the path [id])
-                is_identity = lambda e: e.startswith("id_")
-                if any(map(is_identity, p)) and len(p) > 1:
-                    continue
-
-                paths.append(p)
-
-        for path in paths:
-            # Create compositon by following path
-            f = self._morphisms[path[0]]
-            _, prev_target = self._graph.get_edge_nodes(f.label)
-            for g_label in path[1:]:
-                _, target = self._graph.get_edge_nodes(g_label)
-                g = self._morphisms[g_label]
-                f = self.compose(ob1, self._objects[prev_target], self._objects[target], f, g)
-                prev_target = target
-
-            morphisms.append(f)
-        
+        def all_morphisms(source, target, lvl):
+            """ return all morphisms up to level lvl """
+            # Base case, for level zero return given morphisms
+            if lvl <= 0:
+                return self._morphisms[0].get((source, target), [])
+
+            # We need to compute the morphisms for this level
+            if lvl not in self._morphisms.keys():
+                self._morphisms[lvl] = {}
+                # For all intermediate object
+                for interm, interm_ob in self._objects.items():
+                    # Compute all paths through that interm object
+                    fs = self._morphisms[lvl - 1].get((source, interm), [])
+                    gs = all_morphisms(interm, target, lvl - 1)
+
+                    # Build all possible compositions
+                    not_id = lambda m: not m.label.startswith("id_")
+                    for f, g in product(filter(not_id, fs), filter(not_id, gs)):
+                        source_ob = self._objects[source]
+                        target_ob = self._objects[target]
+                        h = self.compose(source_ob, interm_ob, target_ob, f, g)
+
+                        if (source, target) not in self._morphisms[lvl]:
+                            self._morphisms[lvl][source, target] = []
+                        self._morphisms[lvl][source, target].append(h)
+
+            # Compute the union
+            morphisms = []
+            for l in range(lvl + 1):
+                morphisms += self._morphisms[l].get((source, target), [])
+
+            return morphisms
+
+        morphisms = all_morphisms(ob1.label, ob2.label, uptolevel)
         return MyFiniteSet(morphisms)
 
     def compose(self, ob1: ROD, ob2: ROD, ob3: ROD, m1: RMD, m2: RMD) -> RMD:
         m3 = self._composition(ob1.obdata, ob2.obdata, ob3.obdata, m1.mordata, m2.mordata)
-        rm3 = I.RichMorphism(";".join((m1.label, m2.label)), m3)
-        return rm3
+        return Morph(";".join((m1.label, m2.label)), m3, ob1.label, ob2.label)
 
     def identity(self, ob: ROD) -> RMD:
-        if not self._identity:
+        if not self._objects[ob.label]:
             raise I.InvalidValue()
-        return self._identity
+
+        return self._objects[ob.label].identity
 
 
 class SolSemiCategoryRepresentation(I.SemiCategoryRepresentation):
@@ -130,37 +116,26 @@ class SolSemiCategoryRepresentation(I.SemiCategoryRepresentation):
         if not all(k in data.keys() for k in ["objects", "morphisms", "equations"]):
             raise I.InvalidFormat()
 
-        # Create an empty diagram
-        graph = Graph()
-
-        # Prepare dictionaries for O(1) access and identity
-        objects, morphisms, identity = {}, {}, None
+        sc = MySemiCategory(compose, data["equations"])
 
         # Get the objects
         for obname, ob in data["objects"].items():
             if not "obdata" in ob.keys():
                 raise I.InvalidFormat()
 
-            objects[obname] = I.RichObject(obname, ob["obdata"])
-            graph.add_node(obname)
+            identity = None
+            if "identity" in ob.keys() and ob["identity"]:
+                identity = Morph(f"id_{obname}", ob["identity"]["mordata"], obname, obname)
+                sc.add_morphism(identity)
 
-            if "identity" in ob.keys():
-                # is not None
-                if ob["identity"]:
-                    idname = f"id_{obname}"
-                    morphisms[idname] = I.RichMorphism(idname, ob["identity"]["mordata"])
-                    graph.add_edge(obname, obname, idname)
+            sc.add_object(Obj(obname, ob["obdata"], identity))
             
         # Get the morphisms
         for morname, mor in data["morphisms"].items():
             if not all(k in mor.keys() for k in ["mordata", "source", "target"]):
                 raise I.InvalidFormat()
 
-            morphisms[morname] = I.RichMorphism(morname, mor["mordata"])
-            graph.add_edge(mor["source"], mor["target"], morname)
-
-        equations = data["equations"]
-
-        return MySemiCategory(objects, morphisms, graph, compose, equations, identity)
-
+            src, tar = mor["source"], mor["target"]
+            sc.add_morphism(Morph(morname, mor["mordata"], mor["source"], mor["target"]))
 
+        return sc