unit_tests_algorithms.py

import logging
import unittest
import random
from math import sqrt
from scipy.stats import chisquare

from type_system import Type, PolymorphicType, PrimitiveType, Arrow, List, UnknownType, INT, BOOL, STRING
from program import Program, Function, Variable, BasicPrimitive, New
from program_as_list import evaluation_from_compressed, reconstruct_from_compressed

from dsl import DSL
from DSL.deepcoder import semantics,primitive_types

from Algorithms.heap_search import heap_search
from Algorithms.a_star import a_star
from Algorithms.sqrt_sampling import sqrt_sampling, sqrt_sampling_with_sbsur
from Algorithms.threshold_search import bounded_threshold

class TestSum(unittest.TestCase):
    def test_heap_search(self):
        """
        Check if heap_search does not miss any program and if it outputs programs in decreasing order.
        """

        N = 10_000  # number of programs to be generated by heap search
        K = 1000  # number of programs to be sampled from the PCFG

        deepcoder = DSL(semantics, primitive_types)
        type_request = Arrow(List(INT), List(INT))
        deepcoder_CFG = deepcoder.DSL_to_CFG(type_request)
        deepcoder_PCFG = deepcoder_CFG.CFG_to_Random_PCFG()

        gen_heap_search = heap_search(deepcoder_PCFG)
        gen_sampling = deepcoder_PCFG.sampling()

        seen_sampling = set()
        seen_heaps = set()

        current_probability = 1
        for i in range(N):
            program = next(gen_heap_search)
            new_probability = program.probability[
                (deepcoder_PCFG.__hash__(), deepcoder_PCFG.start)
            ]
            self.assertTrue(
                program.probability[(deepcoder_PCFG.__hash__(), deepcoder_PCFG.start)]
                == deepcoder_PCFG.probability_program(deepcoder_PCFG.start, program)
            )
            self.assertLessEqual(new_probability, current_probability)
            current_probability = new_probability
            seen_heaps.add(str(program))

        min_proba = current_probability

        while len(seen_sampling) < K:
            program = next(gen_sampling)
            if (
                deepcoder_PCFG.probability_program(deepcoder_PCFG.start, program)
                >= min_proba
            ):
                seen_sampling.add(str(program))

        diff = seen_sampling - seen_heaps
        self.assertEqual(0, len(diff))

    def test_a_star(self):
        """
        Check if a_star does not miss any program and if it outputs programs in decreasing order.
        """

        N = 10_000  # number of programs to be generated by heap search
        K = 1000  # number of programs to be sampled from the PCFG

        deepcoder = DSL(semantics, primitive_types)
        type_request = Arrow(List(INT), List(INT))
        deepcoder_CFG = deepcoder.DSL_to_CFG(type_request)
        deepcoder_PCFG = deepcoder_CFG.CFG_to_Random_PCFG()

        gen_a_star = a_star(deepcoder_PCFG)
        gen_sampling = deepcoder_PCFG.sampling()

        seen_sampling = set()
        seen_astar = set()

        current_probability = 1
        for i in range(N):
            program = next(gen_a_star)
            program = reconstruct_from_compressed(program, type_request.returns())
            new_probability = deepcoder_PCFG.probability_program(
                deepcoder_PCFG.start, program
            )
            self.assertLessEqual(new_probability, current_probability + 10e-15)
            current_probability = new_probability
            seen_astar.add(str(program))

        min_proba = current_probability

        while len(seen_sampling) < K:
            program = next(gen_sampling)
            if (
                deepcoder_PCFG.probability_program(deepcoder_PCFG.start, program)
                >= min_proba
            ):
                seen_sampling.add(str(program))

        diff = seen_sampling - seen_astar
        self.assertEqual(0, len(diff))

    def test_threshold_search(self):
        """
        Check if threshold search does not miss any program and if it outputs programs above the given threshold
        """
        deepcoder = DSL(semantics, primitive_types)
        type_request = Arrow(List(INT), List(INT))
        deepcoder_CFG = deepcoder.DSL_to_CFG(type_request)
        deepcoder_PCFG = deepcoder_CFG.CFG_to_Random_PCFG()

        threshold = 0.00001
        gen_threshold = bounded_threshold(deepcoder_PCFG, threshold)

        seen_threshold = set()

        while True:
            try:
                program = next(gen_threshold)
                program = reconstruct_from_compressed(program, type_request.returns())
                proba_program = deepcoder_PCFG.probability_program(
                    deepcoder_PCFG.start, program
                )
                self.assertLessEqual(threshold, proba_program)
                seen_threshold.add(str(program))
            except StopIteration:
                break
        K = len(seen_threshold) // 5

        gen_sampling = deepcoder_PCFG.sampling()

        seen_sampling = set()
        while len(seen_sampling) < K:
            program = next(gen_sampling)
            proba_program = deepcoder_PCFG.probability_program(deepcoder_PCFG.start, program)
            if proba_program >= threshold:
                seen_sampling.add(str(program))

        diff = seen_sampling - seen_threshold

        self.assertEqual(0, len(diff))

    def test_sampling(self):
        """
        Check if the sampling algorithm samples according to the correct probabilities using a chi_square test
        """
        K = 20_000  # number of samples from the L-th first programs
        L = 50  # we test the probabilities of the first L programs are ok
        alpha = 0.05  # threshold to reject the "H0 hypothesis"

        deepcoder = DSL(semantics, primitive_types)
        type_request = Arrow(List(INT), List(INT))
        deepcoder_CFG = deepcoder.DSL_to_CFG(type_request)
        deepcoder_PCFG = deepcoder_CFG.CFG_to_Random_PCFG()

        gen_heap_search = heap_search(deepcoder_PCFG)
        gen_sampling = deepcoder_PCFG.sampling()  # generator for sampling

        count = {}
        for _ in range(L):
            program = next(gen_heap_search)
            count[str(program)] = [
                deepcoder_PCFG.probability_program(deepcoder_PCFG.start, program),
                0,
            ]  # expected frequencies versus observed frequencies

        normalisation_factor = sum(count[program][0] for program in count)
        for program in count:
            count[program][0] *= K / normalisation_factor

        i = 0
        while i < K:
            # if (100 * i // K) != (100 * (i + 1) // K):
            #     print(100 * (i + 1) // K, " %")
            program = next(gen_sampling)
            program_hashed = str(program)
            if program_hashed in count:
                count[program_hashed][1] += 1
                i += 1
        f_exp = []
        f_obs = []
        for p in count:
            f_exp.append(count[p][0])
            f_obs.append(count[p][1])

        chisq, p_value = chisquare(f_obs, f_exp=f_exp)
        self.assertLessEqual(alpha, p_value)

    def test_sqrt_sampling(self):
        """
        Check if sqrt_sampling algorithm samples according to the correct probabilities
        """
        L = 10  # we test that the probabilities of the first L programs are ok
        K = 50_000  # number of samples

        t0 = PolymorphicType("t0")
        t1 = PolymorphicType("t1")
        semantics = {
            "RANGE": (),
            "HEAD": (),
            "TAIL": (),
            "SUCC": (),
            "PRED": (),
            "MAP": (),
        }
        primitive_types = {
            "HEAD": Arrow(List(INT), INT),
            "TAIL": Arrow(List(INT), INT),
            "RANGE": Arrow(INT, List(INT)),
            "SUCC": Arrow(INT, INT),
            "PRED": Arrow(INT, INT),
            "MAP": Arrow(Arrow(t0, t1), Arrow(List(t0), List(t1))),
        }
        toy_DSL = DSL(semantics, primitive_types)
        type_request = Arrow(List(INT), List(INT))
        toy_CFG = toy_DSL.DSL_to_CFG(type_request)
        toy_PCFG = toy_CFG.CFG_to_Random_PCFG(alpha = 0.8)

        gen_heap_search = heap_search(toy_PCFG)  # to generate the L first programs
        gen_sqrt_sampling = sqrt_sampling(toy_PCFG)  # generator for sqrt sampling

        count = {}
        for _ in range(L):
            program = next(gen_heap_search)
            count[str(program)] = [
                K * sqrt(toy_PCFG.probability_program(toy_PCFG.start, program)),
                0,
            ]
        i = 0
        while i < K:
            program = next(gen_sqrt_sampling)
            program_hashed = str(program)
            if program_hashed in count:
                count[program_hashed][1] += 1
                i += 1
        ratios = []
        for p in count:
            ratios.append(count[p][1] / count[p][0])

        random_ratios = random.sample(ratios, 5)
        for r in random_ratios:
            self.assertAlmostEqual(ratios[0], r, 1)

    def test_sqrt_sampling_with_sbsur(self):
        """
        Check if sqrt_sampling_with_sbsur algorithm does not miss any program.
        """
        try:
            import sbsur
        except:
            print("did not import sbsur")
            return

        deepcoder = DSL(semantics, primitive_types)
        type_request = Arrow(List(INT), List(INT))
        deepcoder_CFG = deepcoder.DSL_to_CFG(type_request)
        deepcoder_PCFG = deepcoder_CFG.CFG_to_Random_PCFG()

        n = 10_0000

        gen_sqrt_sampling = sqrt_sampling_with_sbsur(deepcoder_PCFG)
        r = type_request.returns()

        seen_programs = set()
        for program in gen_sqrt_sampling:
            prog = reconstruct_from_compressed(program, r)
            self.assertNotIn(str(prog), seen_programs)
            seen_programs.add(str(prog))
            if len(seen_programs) > n:
                break

if __name__ == "__main__":
    unittest.main(verbosity=2)