Implemented A*

Author: Theomat

examples/pbe/solve.py

@@ -24,6 +24,7 @@
     hs_enumerate_bucket_prob_grammar,
     hs_enumerate_bucket_prob_u_grammar,
     cd_enumerate_prob_grammar,
+    as_enumerate_prob_grammar,
     ProgramEnumerator,
     Type,
     CFG,
@@ -62,6 +63,7 @@
         lambda x: hs_enumerate_bucket_prob_u_grammar(x, 3),
     ),
     "bee_search": (bs_enumerate_prob_grammar, None),
+    "a_star": (as_enumerate_prob_grammar, None),
 }
 
 PRUNING = {"dfta", "obs-eq"}

synth/syntax/__init__.py

@@ -50,5 +50,6 @@
     hs_enumerate_bucket_prob_grammar,
     hs_enumerate_bucket_prob_u_grammar,
     cd_enumerate_prob_grammar,
+    as_enumerate_prob_grammar,
     split,
 )

synth/syntax/grammars/__init__.py

@@ -12,6 +12,7 @@
     hs_enumerate_bucket_prob_grammar,
     hs_enumerate_bucket_prob_u_grammar,
     cd_enumerate_prob_grammar,
+    as_enumerate_prob_grammar,
     split,
 )
 from synth.syntax.grammars.u_grammar import UGrammar

synth/syntax/grammars/enumeration/__init__.py

@@ -18,5 +18,8 @@
 from synth.syntax.grammars.enumeration.constant_delay import (
     enumerate_prob_grammar as cd_enumerate_prob_grammar,
 )
+from synth.syntax.grammars.enumeration.a_star import (
+    enumerate_prob_grammar as as_enumerate_prob_grammar,
+)
 
 enumerate_prob_grammar = cd_enumerate_prob_grammar

synth/syntax/grammars/enumeration/a_star.py

@@ -0,0 +1,133 @@
+from heapq import heappush, heappop
+from typing import (
+    Generator,
+    Generic,
+    List,
+    Optional,
+    Tuple,
+    TypeVar,
+    Union,
+)
+from dataclasses import dataclass, field
+
+import numpy as np
+
+from synth.filter.filter import Filter
+from synth.syntax.grammars.enumeration.program_enumerator import ProgramEnumerator
+from synth.syntax.grammars.tagged_u_grammar import ProbUGrammar
+from synth.syntax.program import Function, Program
+from synth.syntax.grammars.tagged_det_grammar import ProbDetGrammar, DerivableProgram
+from synth.syntax.type_system import Type
+
+U = TypeVar("U")
+V = TypeVar("V")
+W = TypeVar("W")
+
+
+def _build_(
+    elems: List[Tuple[DerivableProgram, Tuple[Type, U]]], G: ProbDetGrammar[U, V, W]
+) -> Program:
+    P, S = elems.pop(0)
+    nargs = G.arguments_length_for(S, P)
+    if nargs == 0:
+        return P
+    else:
+        args = []
+        while nargs > 0:
+            args.append(_build_(elems, G))
+            nargs -= 1
+        return Function(P, args)
+
+
+@dataclass(order=True, frozen=True)
+class HeapElement(Generic[U]):
+    priority: float
+    to_expand: List[Tuple[Type, U]] = field(compare=False)
+    parts: List[Tuple[DerivableProgram, Tuple[Type, U]]] = field(compare=False)
+
+    def __repr__(self) -> str:
+        return f"({self.priority}, {self.parts})"
+
+    def make_program(self, g: ProbDetGrammar[U, V, W]) -> Program:
+        return _build_(self.parts, g)
+
+
+class AStar(
+    ProgramEnumerator[None],
+    Generic[U, V, W],
+):
+    def __init__(
+        self,
+        G: ProbDetGrammar[U, V, W],
+        filter: Optional[Filter[Program]] = None,
+    ) -> None:
+        super().__init__(filter)
+        self.current: Optional[Program] = None
+
+        self.G = G
+        self.start = G.start
+        self.rules = G.rules
+
+        self.frontier: List[HeapElement[U]] = []
+
+    def probability(self, program: Program) -> float:
+        return self.G.probability(program)
+
+    @classmethod
+    def name(cls) -> str:
+        return "a-star"
+
+    def generator(self) -> Generator[Program, None, None]:
+        """
+        A generator which outputs the next most probable program
+        """
+        first = (self.G.start[0], self.G.start[1][0])  # type: ignore
+        heappush(self.frontier, HeapElement(0, [first], []))  # type: ignore
+
+        while self.frontier:
+            elem = heappop(self.frontier)
+            if len(elem.to_expand) == 0:
+                p = elem.make_program(self.G)
+                if self._should_keep_subprogram(p):
+                    yield p
+            else:
+                partS = elem.to_expand.pop()
+                S = (partS[0], (partS[1], None))
+                for P in self.G.rules[S]:  # type: ignore
+                    args = self.G.rules[S][P][0]  # type: ignore
+                    p = self.G.probabilities[S][P]  # type: ignore
+                    new_el = HeapElement(
+                        elem.priority + p,
+                        elem.to_expand + list(args),
+                        elem.parts + [(P, S)],  # type: ignore
+                    )
+                    heappush(self.frontier, new_el)
+
+    def merge_program(self, representative: Program, other: Program) -> None:
+        """
+        Merge other into representative.
+        In other words, other will no longer be generated through heap search
+        """
+        pass
+
+    def programs_in_banks(self) -> int:
+        return 0
+
+    def programs_in_queues(self) -> int:
+        return len(self.frontier)
+
+    def clone(self, G: Union[ProbDetGrammar, ProbUGrammar]) -> "AStar[U, V, W]":
+        assert isinstance(G, ProbDetGrammar)
+        enum = self.__class__(G)
+        return enum
+
+
+def enumerate_prob_grammar(G: ProbDetGrammar[U, V, W]) -> AStar[U, V, W]:
+    Gp: ProbDetGrammar = ProbDetGrammar(
+        G.grammar,
+        {
+            S: {P: -np.log(p) for P, p in val.items() if p > 0}
+            for S, val in G.probabilities.items()
+        },
+    )
+    return AStar(Gp)

tests/syntax/grammars/enumeration/test_a_star.py

@@ -0,0 +1,71 @@
+from synth.syntax.grammars.enumeration.a_star import (
+    enumerate_prob_grammar,
+)
+from synth.syntax.grammars.cfg import CFG
+from synth.syntax.grammars.ttcfg import TTCFG
+from synth.syntax.grammars.tagged_det_grammar import ProbDetGrammar
+from synth.syntax.dsl import DSL
+from synth.syntax.type_system import (
+    INT,
+    STRING,
+    List,
+    PolymorphicType,
+    PrimitiveType,
+)
+from synth.syntax.type_helper import FunctionType, auto_type
+
+import pytest
+
+
+syntax = {
+    "+": FunctionType(INT, INT, INT),
+    "head": FunctionType(List(PolymorphicType("a")), PolymorphicType("a")),
+    "non_reachable": PrimitiveType("non_reachable"),
+    "1": INT,
+    "2": INT,
+    "non_productive": FunctionType(INT, STRING),
+}
+dsl = DSL(syntax)
+dsl.instantiate_polymorphic_types()
+testdata = [
+    CFG.depth_constraint(dsl, FunctionType(INT, INT), 3),
+    CFG.depth_constraint(dsl, FunctionType(INT, INT), 4),
+]
+
+
+@pytest.mark.parametrize("cfg", testdata)
+def test_unicity_a_star(cfg: TTCFG) -> None:
+    pcfg = ProbDetGrammar.uniform(cfg)
+    seen = set()
+    print(cfg)
+    for program in enumerate_prob_grammar(pcfg):
+        assert program not in seen
+        seen.add(program)
+    # print(pcfg.grammar)
+    assert len(seen) == cfg.programs()
+
+
+@pytest.mark.parametrize("cfg", testdata)
+def test_order_a_star(cfg: TTCFG) -> None:
+    pcfg = ProbDetGrammar.uniform(cfg)
+    last = 1.0
+    for program in enumerate_prob_grammar(pcfg):
+        p = pcfg.probability(program)
+        assert p <= last
+        last = p
+
+
+def test_infinite() -> None:
+    pcfg = ProbDetGrammar.random(
+        CFG.infinite(dsl, testdata[0].type_request, n_gram=1), 1
+    )
+    count = 10000
+    last = 1.0
+    for program in enumerate_prob_grammar(pcfg):
+        count -= 1
+        p = pcfg.probability(program)
+        assert -1e-12 <= last - p, f"failed at program n°{count}:{program}"
+        last = p
+        if count < 0:
+            break
+    assert count == -1