fixed mypy issues

Author: Theomat

examples/compare_enumeration.py

@@ -15,20 +15,21 @@
     ProgramEnumerator,
     auto_type,
 )
+from synth.syntax.grammars.det_grammar import DetGrammar
 from synth.syntax.grammars.enumeration.constant_delay import (
     enumerate_prob_grammar as cd,
 )
 import tqdm
 import timeout_decorator
 
 SEARCH_ALGOS = {
-    "a_star": as_enumerate_prob_grammar,
-    "bee_search": bs_enumerate_prob_grammar,
+    # "a_star": as_enumerate_prob_grammar,
+    # "bee_search": bs_enumerate_prob_grammar,
     "beap_search": bps_enumerate_prob_grammar,
-    "heap_search": hs_enumerate_prob_grammar,
-    "cd4": lambda x: cd(x, k=4),
-    "cd16": lambda x: cd(x, k=16),
-    "cd64": lambda x: cd(x, k=64),
+    # "heap_search": hs_enumerate_prob_grammar,
+    # "cd4": lambda x: cd(x, k=4),
+    # "cd16": lambda x: cd(x, k=16),
+    # "cd64": lambda x: cd(x, k=64),
 }
 
 parser = argparse.ArgumentParser(
@@ -163,6 +164,8 @@ def enumerative_search(
     Tuple[str, int, int, float, int, int, int, int],
     List[Tuple[str, int, int, float, int, int, int, int]],
 ]:
+    import numpy as np
+
     n = 0
     non_terminals = len(pcfg.rules)
     derivation_rules = sum(len(pcfg.rules[S]) for S in pcfg.rules)
@@ -171,9 +174,11 @@ def enumerative_search(
     pbar = tqdm.tqdm(total=programs, desc=title or name, smoothing=0)
     enumerator = custom_enumerate(pcfg)
     gen = enumerator.generator()
+    det_g = pcfg.grammar
+    assert isinstance(det_g, DetGrammar)
     program = 1
-    datum_each = 100000
-    target_generation_speed = 1000000
+    datum_each = 10000
+    target_generation_speed = 10000
     start = 0
     detailed = []
     try:
@@ -184,9 +189,15 @@ def fun():
         get_next = timeout_decorator.timeout(timeout, timeout_exception=StopIteration)(
             fun
         )
+        last_multiset = None
+        max_dist = -1
         start = time.perf_counter_ns()
         while program is not None:
             program = get_next()
+            new_m = det_g.to_multiset(program)
+            if last_multiset is not None:
+                max_dist = max(np.sum(np.abs(new_m - last_multiset)), max_dist)
+            last_multiset = new_m
             n += 1
             if n % datum_each == 0 or n >= programs:
                 used_time = time.perf_counter_ns() - start
@@ -200,7 +211,7 @@ def fun():
                     (
                         name,
                         non_terminals,
-                        derivation_rules,
+                        max_dist,
                         used_time / 1e9,
                         n,
                         enumerator.programs_in_queues(),

examples/plot_helper.py

@@ -246,16 +246,16 @@ def plot_dist(
     methods: Dict[str, Dict[int, List]],
     y_data: Tuple[int, str],
     x_axis_name: str,
+    nbins: int = 5,
 ) -> None:
     width = 1.0
     data_length = 0
     a_index, a_name = y_data
     max_a = max(
-        max(max([y[a_index] for y in x]) for x in seed_dico.values())
+        max(max(y[a_index] for y in x) for x in seed_dico.values())
         for seed_dico in methods.values()
     )
     bottom = None
-    nbins = 5
     bins = [max_a]
     while len(bins) <= nbins:
         bins.insert(0, np.sqrt(bins[0] + 1))

synth/pbe/solvers/pbe_solver.py

@@ -72,7 +72,7 @@ def _close_task_solving_(
 
     def solve(
         self, task: Task[PBE], enumerator: ProgramEnumerator[None], timeout: float = 60
-    ) -> Generator[Program, bool, None]:
+    ) -> Generator[Program, None, bool]:
         """
         Solve the given task by enumerating programs with the given enumerator.
         When the timeout is reached, this function returns.
@@ -101,6 +101,7 @@ def solve(
             except StopIteration as e:
                 self._close_task_solving_(task, enumerator, time, False, program)
                 raise e
+        return False
 
     def _test_(self, task: Task[PBE], program: Program) -> bool:
         """

synth/pbe/solvers/restart_pbe_solver.py

@@ -61,7 +61,7 @@ def _close_task_solving_(
 
     def solve(
         self, task: Task[PBE], enumerator: ProgramEnumerator[None], timeout: float = 60
-    ) -> Generator[Program, bool, None]:
+    ) -> Generator[Program, None, bool]:
         with chrono.clock(f"solve.{self.name()}.{self.subsolver.name()}") as c:  # type: ignore
             self._enumerator = enumerator
             self._init_task_solving_(task, self._enumerator, timeout)
@@ -73,15 +73,15 @@ def solve(
                     self._close_task_solving_(
                         task, self._enumerator, time, False, program
                     )
-                    return
+                    return False
                 self._programs += 1
                 if self._test_(task, program):
                     should_stop = yield program
                     if should_stop:
                         self._close_task_solving_(
                             task, self._enumerator, time, True, program
                         )
-                        return
+                        return True
                 self._score = self.subsolver._score
                 # Saves data
                 if self._score > 0:
@@ -92,6 +92,7 @@ def solve(
                     self._enumerator = self._restart_(self._enumerator)
                     gen = self._enumerator.generator()
                 program = next(gen)
+        return False
 
     def _should_restart_(self) -> bool:
         return self.restart_criterion(self)

synth/syntax/grammars/det_grammar.py

@@ -11,7 +11,7 @@
     Generic,
 )
 from functools import lru_cache
-import copy
+import numpy as np
 
 from synth.syntax.grammars.grammar import DerivableProgram, Grammar
 from synth.syntax.program import Constant, Function, Primitive, Program, Variable
@@ -23,6 +23,13 @@
 T = TypeVar("T")
 
 
+def __tuplify__(element: Any) -> Any:
+    if isinstance(element, (List, Tuple)):
+        return tuple(__tuplify__(x) for x in element)
+    else:
+        return element
+
+
 class DetGrammar(Grammar, ABC, Generic[U, V, W]):
     """
     Represents a deterministic grammar.
@@ -54,6 +61,14 @@ def __init__(
         self.type_request = self._guess_type_request_()
         if clean:
             self.clean()
+        self._derivation2index = {}
+        self._index2derivation = []
+
+        for S in self.rules:
+            for P, args in self.rules[S].items():
+                elem = __tuplify__((S, P, args))
+                self._derivation2index[elem] = len(self._index2derivation)
+                self._index2derivation.append(elem)
 
     @lru_cache()
     def primitives_used(self) -> Set[Primitive]:
@@ -83,6 +98,16 @@ def __str__(self) -> str:
                 s += "   {}\n".format(self.__rule_to_str__(P, out))
         return s
 
+    def to_multiset(self, program: Program) -> np.ndarray:
+        out = np.zeros((len(self._derivation2index)))
+
+        def reduce(acc, S, P, args):
+            elem = __tuplify__((S, P, args))
+            out[self._derivation2index[elem]] += 1
+
+        self.reduce_derivations(reduce, None, program, None)
+        return out
+
     def __repr__(self) -> str:
         return self.__str__()