added feature to get equivalence classes

Author: Theomat

grape/cli/prune.py

@@ -6,6 +6,7 @@
 from grape.cli import dsl_loader
 from grape.program import Primitive, Variable
 from grape.evaluator import Evaluator
+from grape.pruning.equivalence_class_manager import EquivalenceClassManager
 from grape.pruning.obs_equiv_pruner import prune
 
 
@@ -78,6 +79,12 @@ def is_python_file(file_path: str) -> bool:
         type=str,
         help="your starting automaton file",
     )
+    parser.add_argument(
+        "--classes",
+        type=str,
+        default=None,
+        help="save equivalence classes ina JSON file",
+    )
 
     return parser.parse_args()
 
@@ -90,9 +97,11 @@ def main():
     inputs = sample_inputs(args.samples, sample_dict, equal_dict)
 
     evaluator = Evaluator(dsl, inputs, equal_dict, skip_exceptions)
+    manager = EquivalenceClassManager()
     grammar, type_req = prune(
         dsl,
         evaluator,
+        manager,
         args.size,
         target_type,
         args.automaton,
@@ -110,6 +119,10 @@ def main():
 
     dump_automaton_to_file(grammar, args.output)
 
+    if args.classes is not None:
+        with open(args.classes, "w") as fd:
+            fd.write(manager.to_json())
+
 
 if __name__ == "__main__":
     main()

grape/pruning/commutativity_pruner.py

@@ -2,6 +2,7 @@
 from grape.dsl import DSL
 from grape.evaluator import Evaluator
 from grape.program import Function, Primitive, Program, Variable
+from grape.pruning.equivalence_class_manager import EquivalenceClassManager
 import grape.types as types
 
 T = TypeVar("T")
@@ -26,15 +27,17 @@ def __produce_all_variants__(
             current[i] += 1
 
 
-def get_rewrites(
-    dsl: DSL, primitive: str, swapped_indices: tuple[int, int]
-) -> list[tuple[Program, Program, str]]:
-    constraints: list[tuple[Program, Program, str]] = []
-    stype = dsl[primitive][0]
+def __add_rewrite__(
+    dsl: DSL,
+    primitive: str,
+    swapped_indices: tuple[int, int],
+    manager: EquivalenceClassManager,
+):
+    stype = dsl.primitives[primitive][0]
     args_type, _ = types.parse(stype)
     swapped_type = args_type[swapped_indices[0]]
 
-    nargs = len(types.arguments(dsl.primitives[primitive]))
+    nargs = len(types.arguments(dsl.primitives[primitive][0]))
 
     for p1, (stype1, _) in dsl.primitives.copy().items():
         args1, rtype1 = types.parse(stype1)
@@ -73,7 +76,7 @@ def get_rewrites(
             )
             equiv_to.arguments[swapped_indices[0]] = second_arg
             equiv_to.arguments[swapped_indices[1]] = first_arg
-            constraints.append((deleted, equiv_to))
+            manager.add_merge(deleted, equiv_to)
         # Add additional constraint for variable type
 
         second_arg = Variable(swapped_indices[0])
@@ -84,14 +87,11 @@ def get_rewrites(
         equiv_to = Function(Primitive(primitive), [Variable(i) for i in range(nargs)])
         equiv_to.arguments[swapped_indices[0]] = second_arg
         equiv_to.arguments[swapped_indices[1]] = first_arg
-        constraints.append((deleted, equiv_to))
-
-    return constraints
+        manager.add_merge(deleted, equiv_to)
 
 
 def prune(
-    dsl: DSL,
-    evaluator: Evaluator,
+    dsl: DSL, evaluator: Evaluator, manager: EquivalenceClassManager
 ) -> list[tuple[str, list[int]]]:
     commutatives = []
     for prim, (stype, _) in dsl.primitives.items():
@@ -117,8 +117,10 @@ def prune(
                 continue
             variant = Function(Primitive(prim), new_args)
             if evaluator.eval(variant, stype) is not None:
+                swapped = [i for i, x in enumerate(new_args) if x.no != i]
                 commutatives.append(
                     (prim, [i for i, x in enumerate(new_args) if x.no != i])
                 )
+                __add_rewrite__(dsl, prim, (swapped[0], swapped[1]), manager)
     evaluator.clean_memoisation()
     return commutatives

grape/pruning/equivalence_class_manager.py

@@ -0,0 +1,42 @@
+import json
+from grape.program import Program
+
+
+class EquivalenceClassManager:
+    def __init__(self):
+        self.classes: dict[Program, set[Program]] = {}
+
+    def new_class(self, representative: Program):
+        """
+        Create a new class of equivalence.
+        Assumes class does not already exist.
+        """
+        assert representative not in self.classes
+        self.classes[representative] = set()
+
+    def add_to_class(self, program: Program, representative: Program):
+        """
+        Add a program to an already existing equivalence class.
+        Assumes class already exists.
+        """
+        self.classes[representative].add(program)
+
+    def add_merge(self, program: Program, representative: Program):
+        """
+        Add a program to an already existing equivalence class.
+        Assumes nothing so it creates a new class if it does not exist.
+        """
+        if representative not in self.classes:
+            self.new_class(representative)
+        self.add_to_class(program, representative)
+
+    def to_json(self) -> str:
+        str_classes = sorted(
+            [
+                {"representative": str(key), "elements": list(map(str, value))}
+                for key, value in self.classes.items()
+            ],
+            key=lambda x: (x["representative"], len(x["elements"])),
+            reverse=True,
+        )
+        return json.dumps(str_classes)

grape/pruning/obs_equiv_pruner.py

@@ -15,6 +15,7 @@
 )
 from grape.automaton.tree_automaton import DFTA
 import grape.pruning.commutativity_pruner as commutativity_pruner
+from grape.pruning.equivalence_class_manager import EquivalenceClassManager
 import grape.types as types
 
 from tqdm import tqdm
@@ -59,13 +60,14 @@ def __get_base_grammar__(
     has_base_grammar: bool,
     dsl: DSL,
     evaluator: Evaluator,
+    manager: EquivalenceClassManager,
     max_size: int,
     base_automaton_file: str,
     type_req: str,
 ):
     base_grammar = grammar_by_saturation(dsl, type_req)
     if not has_base_grammar:
-        commutatives = commutativity_pruner.prune(dsl, evaluator)
+        commutatives = commutativity_pruner.prune(dsl, evaluator, manager)
         grammar = grammar_by_saturation(
             dsl,
             type_req,
@@ -89,6 +91,7 @@ def __get_base_grammar__(
 def prune(
     dsl: DSL,
     evaluator: Evaluator,
+    manager: EquivalenceClassManager,
     max_size: int,
     rtype: str | None,
     base_automaton_file: str,
@@ -102,7 +105,13 @@ def prune(
         base_automaton_file is None or len(base_automaton_file) == 0
     )
     grammar, base_expected_trees, enum_ntrees = __get_base_grammar__(
-        has_base_grammar, dsl, evaluator, max_size, base_automaton_file, type_req
+        has_base_grammar,
+        dsl,
+        evaluator,
+        manager,
+        max_size,
+        base_automaton_file,
+        type_req,
     )
     base_ntrees = sum(base_expected_trees.values())
 
@@ -142,6 +151,8 @@ def estimate_total(size: int) -> tuple[int, float]:
             program = gen.send(should_keep)
             representative = evaluator.eval(program, type_req)
             should_keep = representative is None
+            if not should_keep:
+                manager.add_merge(program, representative)
             n += 1
             if n & 15 == 0:
                 pbar.update(16)