Fraser/dream-coder

1from dreamcoder.fragmentUtilities import *2from dreamcoder.grammar import *3from dreamcoder.program import *4 5from itertools import chain6import time7 8 9class FragmentGrammar(object):10    def __init__(self, logVariable, productions):11        self.logVariable = logVariable12        self.productions = productions13        self.likelihoodCache = {}14 15    def clearCache(self):16        self.likelihoodCache = {}17 18    def __repr__(self):19        return "FragmentGrammar(logVariable={self.logVariable}, productions={self.productions}".format(20            self=self)21 22    def __str__(self):23        def productionKey(xxx_todo_changeme):24            (l, t, p) = xxx_todo_changeme25            return not isinstance(p, Primitive), -l26        return "\n".join(["%f\tt0\t$_" % self.logVariable] + ["%f\t%s\t%s" % (l, t, p)27                                                              for l, t, p in sorted(self.productions, key=productionKey)])28 29    def buildCandidates(self, context, environment, request):30        candidates = []31        variableCandidates = []32        for l, t, p in self.productions:33            try:34                newContext, t = t.instantiate(context)35                newContext = newContext.unify(t.returns(), request)36                candidates.append((l, newContext,37                                   t.apply(newContext),38                                   p))39            except UnificationFailure:40                continue41        for j, t in enumerate(environment):42            try:43                newContext = context.unify(t.returns(), request)44                variableCandidates.append((newContext,45                                           t.apply(newContext),46                                           Index(j)))47            except UnificationFailure:48                continue49        if variableCandidates:50            z = math.log(len(variableCandidates))51            for newContext, newType, index in variableCandidates:52                candidates.append(53                    (self.logVariable - z, newContext, newType, index))54 55        z = lse([candidate[0] for candidate in candidates])56        return [(l - z, c, t, p) for l, c, t, p in candidates]57 58    def logLikelihood(self, request, expression):59        _, l, _ = self._logLikelihood(Context.EMPTY, [], request, expression)60        if invalid(l):61            f = 'failures/likelihoodFailure%s.pickle' % (time() + getPID())62            eprint("PANIC: Invalid log likelihood. expression:",63                   expression, "tp:", request, "Exported to:", f)64            with open(f, 'wb') as handle:65                pickle.dump((self, request, expression), handle)66            assert False67        return l68 69    def closedUses(self, request, expression):70        _, l, u = self._logLikelihood(Context.EMPTY, [], request, expression)71        return l, u72 73    def _logLikelihood(self, context, environment, request, expression):74        '''returns (context, log likelihood, uses)'''75 76        # We can cash likelihood calculations faster whenever they don't involve type inference77        # This is because they are guaranteed to not modify the context,78        polymorphic = request.isPolymorphic or any(79            v.isPolymorphic for v in environment)80        # For some reason polymorphic caching slows it down81        shouldDoCaching = not polymorphic82 83        # Caching84        if shouldDoCaching:85            if polymorphic:86                inTypes = canonicalTypes(87                    [request.apply(context)] + [v.apply(context) for v in environment])88            else:89                inTypes = canonicalTypes([request] + environment)90            cacheKey = (tuple(inTypes), expression)91            if cacheKey in self.likelihoodCache:92                outTypes, l, u = self.likelihoodCache[cacheKey]93                context, instantiatedTypes = instantiateTypes(94                    context, outTypes)95                outRequest = instantiatedTypes[0]96                outEnvironment = instantiatedTypes[1:]97                # eprint("request:", request.apply(context), "environment:",98                #        [ v.apply(context) for v in environment ])99                # eprint("will be unified with: out request:",outRequest,"out environment",outEnvironment)100                if polymorphic:101                    context = context.unify(request, outRequest)102                    for v, vp in zip(environment, outEnvironment):103                        context = context.unify(v, vp)104                return context, l, u105 106        if request.isArrow():107            if not isinstance(expression, Abstraction):108                return (context, NEGATIVEINFINITY, Uses.empty)109            return self._logLikelihood(context,110                                       [request.arguments[0]] + environment,111                                       request.arguments[1],112                                       expression.body)113 114        # Not a function type115 116        # Construct and normalize the candidate productions117        candidates = self.buildCandidates(context, environment, request)118 119        # Consider each way of breaking the expression up into a120        # function and arguments121        totalLikelihood = NEGATIVEINFINITY122        weightedUses = []123 124        possibleVariables = float(int(any(isinstance(candidate, Index)125                                          for _, _, _, candidate in candidates)))126        possibleUses = {candidate: 1. for _, _, _, candidate in candidates127                        if not isinstance(candidate, Index)}128 129        for f, xs in expression.applicationParses():130            for candidateLikelihood, newContext, tp, production in candidates:131                variableBindings = {}132                # This is a variable in the environment133                if production.isIndex:134                    if production != f:135                        continue136                else:137                    try:138                        newContext, fragmentType, variableBindings = \139                            Matcher.match(newContext, production, f, len(xs))140                        # This is necessary because the types of the variable141                        # bindings and holes need to match up w/ request142                        fragmentTypeTemplate = request143                        for _ in xs:144                            newContext, newVariable = newContext.makeVariable()145                            fragmentTypeTemplate = arrow(146                                newVariable, fragmentTypeTemplate)147                        newContext = newContext.unify(148                            fragmentType, fragmentTypeTemplate)149                        # update the unified type150                        tp = fragmentType.apply(newContext)151                    except MatchFailure:152                        continue153 154                argumentTypes = tp.functionArguments()155                if len(xs) != len(argumentTypes):156                    # I think that this is some kind of bug. But I can't figure it out right now.157                    # As a hack, count this as though it were a failure158                    continue159                    #raise GrammarFailure('len(xs) != len(argumentTypes): tp={}, xs={}'.format(tp, xs))160 161                thisLikelihood = candidateLikelihood162                if isinstance(production, Index):163                    theseUses = Uses(possibleVariables=possibleVariables,164                                     actualVariables=1.,165                                     possibleUses=possibleUses.copy(),166                                     actualUses={})167                else:168                    theseUses = Uses(possibleVariables=possibleVariables,169                                     actualVariables=0.,170                                     possibleUses=possibleUses.copy(),171                                     actualUses={production: 1.})172 173                # Accumulate likelihood from free variables and holes and174                # arguments175                for freeType, freeExpression in chain(176                        variableBindings.values(), zip(argumentTypes, xs)):177                    freeType = freeType.apply(newContext)178                    newContext, expressionLikelihood, newUses = self._logLikelihood(179                        newContext, environment, freeType, freeExpression)180                    if expressionLikelihood is NEGATIVEINFINITY:181                        thisLikelihood = NEGATIVEINFINITY182                        break183 184                    thisLikelihood += expressionLikelihood185                    theseUses += newUses186 187                if thisLikelihood is NEGATIVEINFINITY:188                    continue189 190                weightedUses.append((thisLikelihood, theseUses))191                totalLikelihood = lse(totalLikelihood, thisLikelihood)192 193                # Any of these new context objects should be equally good194                context = newContext195 196        if totalLikelihood is NEGATIVEINFINITY:197            return context, totalLikelihood, Uses.empty198        assert weightedUses != []199 200        allUses = Uses.join(totalLikelihood, *weightedUses)201 202        # memoize result203        if shouldDoCaching:204            outTypes = [request.apply(context)] + \205                [v.apply(context) for v in environment]206            outTypes = canonicalTypes(outTypes)207            self.likelihoodCache[cacheKey] = (208                outTypes, totalLikelihood, allUses)209 210        return context, totalLikelihood, allUses211 212    def expectedUses(self, frontiers):213        if len(list(frontiers)) == 0:214            return Uses()215        likelihoods = [[(l + entry.logLikelihood, u)216                        for entry in frontier217                        for l, u in [self.closedUses(frontier.task.request, entry.program)]]218                       for frontier in frontiers]219        zs = (lse([l for l, _ in ls]) for ls in likelihoods)220        return sum(math.exp(l - z) * u221                   for z, frontier in zip(zs, likelihoods)222                   for l, u in frontier)223 224    def insideOutside(self, frontiers, pseudoCounts):225        uses = self.expectedUses(frontiers)226        return FragmentGrammar(log(uses.actualVariables +227                                   pseudoCounts) -228                               log(max(uses.possibleVariables, 1.)), [(log(uses.actualUses.get(p, 0.) +229                                                                           pseudoCounts) -230                                                                       log(uses.possibleUses.get(p, 0.) +231                                                                           pseudoCounts), t, p) for _, t, p in self.productions])232 233    def jointFrontiersLikelihood(self, frontiers):234        return sum(lse([entry.logLikelihood + self.logLikelihood(frontier.task.request, entry.program)235                        for entry in frontier])236                   for frontier in frontiers)237 238    def jointFrontiersMDL(self, frontiers, CPUs=1):239        return sum(240            parallelMap(241                CPUs,242                lambda frontier: max(243                    entry.logLikelihood +244                    self.logLikelihood(245                        frontier.task.request,246                        entry.program) for entry in frontier),247                frontiers))248 249    def __len__(self): return len(self.productions)250 251    @staticmethod252    def fromGrammar(g):253        return FragmentGrammar(g.logVariable, g.productions)254 255    def toGrammar(self):256        return Grammar(self.logVariable, [(l, q.infer(), q)257                                          for l, t, p in self.productions258                                          for q in [defragment(p)]])259 260    @property261    def primitives(self): return [p for _, _, p in self.productions]262 263    @staticmethod264    def uniform(productions):265        return FragmentGrammar(0., [(0., p.infer(), p) for p in productions])266 267    def normalize(self):268        z = lse([l for l, t, p in self.productions] + [self.logVariable])269        return FragmentGrammar(self.logVariable - z,270                               [(l - z, t, p) for l, t, p in self.productions])271 272    def makeUniform(self):273        return FragmentGrammar(0., [(0., p.infer(), p)274                                    for _, _, p in self.productions])275 276    def rescoreFrontier(self, frontier):277        return Frontier([FrontierEntry(e.program,278                                       logPrior=self.logLikelihood(frontier.task.request, e.program),279                                       logLikelihood=e.logLikelihood)280                         for e in frontier],281                        frontier.task)282 283    @staticmethod284    def induceFromFrontiers(285            g0,286            frontiers,287            _=None,288            topK=1,289            topk_use_only_likelihood=False,290            pseudoCounts=1.0,291            aic=1.0,292            structurePenalty=0.001,293            a=0,294            CPUs=1):295        _ = topk_use_only_likelihood # not used in python compressor296        originalFrontiers = frontiers297        frontiers = [frontier for frontier in frontiers if not frontier.empty]298        eprint("Inducing a grammar from", len(frontiers), "frontiers")299 300        bestGrammar = FragmentGrammar.fromGrammar(g0)301        oldJoint = bestGrammar.jointFrontiersMDL(frontiers, CPUs=1)302 303        # "restricted frontiers" only contain the top K according to the best grammar304        def restrictFrontiers():305            return parallelMap(306                CPUs,307                lambda f: bestGrammar.rescoreFrontier(f).topK(topK),308                frontiers)309        restrictedFrontiers = []310 311        def grammarScore(g):312            g = g.makeUniform().insideOutside(restrictedFrontiers, pseudoCounts)313            likelihood = g.jointFrontiersMDL(restrictedFrontiers)314            structure = sum(primitiveSize(p) for p in g.primitives)315            score = likelihood - aic * len(g) - structurePenalty * structure316            g.clearCache()317            if invalid(score):318                # FIXME: This should never occur but it does anyway319                score = float('-inf')320            return score, g321 322        if aic is not POSITIVEINFINITY:323            restrictedFrontiers = restrictFrontiers()324            bestScore, _ = grammarScore(bestGrammar)325            eprint("Starting score", bestScore)326            while True:327                restrictedFrontiers = restrictFrontiers()328                fragments = [f329                             for f in proposeFragmentsFromFrontiers(restrictedFrontiers, a, CPUs=CPUs)330                             if f not in bestGrammar.primitives331                             and defragment(f) not in bestGrammar.primitives]332                eprint("Proposed %d fragments." % len(fragments))333 334                candidateGrammars = [335                    FragmentGrammar.uniform(336                        bestGrammar.primitives +337                        [fragment]) for fragment in fragments]338                if not candidateGrammars:339                    break340 341                scoredFragments = parallelMap(CPUs, grammarScore, candidateGrammars,342                                              # Each process handles up to 100343                                              # grammars at a time, a "job"344                                              chunksize=max(345                                                  1, min(len(candidateGrammars) // CPUs, 100)),346                                              # maxTasks: Maximum number of jobs allocated to a process347                                              # This means that after evaluating this*chunk many grammars,348                                              # we killed the process, freeing up its memory.349                                              # In exchange we pay the cost of spawning a new process.350                                              # We should play with this number,351                                              # figuring out how big we can make it without352                                              # running out of memory.353                                              maxtasksperchild=5)354                newScore, newGrammar = max(scoredFragments, key=lambda sg: sg[0])355 356                if newScore <= bestScore:357                    break358                dS = newScore - bestScore359                bestScore, bestGrammar = newScore, newGrammar360                newPrimitiveLikelihood, newType, newPrimitive = bestGrammar.productions[-1]361                expectedUses = bestGrammar.expectedUses(362                    restrictedFrontiers).actualUses.get(newPrimitive, 0)363                eprint(364                    "New primitive of type %s\t%s\t\n(score = %f; dScore = %f; <uses> = %f)" %365                    (newType, newPrimitive, newScore, dS, expectedUses))366 367                # Rewrite the frontiers in terms of the new fragment368                concretePrimitive = defragment(newPrimitive)369                bestGrammar.productions[-1] = (newPrimitiveLikelihood,370                                               concretePrimitive.tp,371                                               concretePrimitive)372                frontiers = parallelMap(373                    CPUs, lambda frontier: bestGrammar.rescoreFrontier(374                        RewriteFragments.rewriteFrontier(375                            frontier, newPrimitive)), frontiers)376                eprint(377                    "\t(<uses> in rewritten frontiers: %f)" %378                    (bestGrammar.expectedUses(frontiers).actualUses[concretePrimitive]))379        else:380            eprint("Skipping fragment proposals")381 382        if False:383            # Reestimate the parameters using the entire frontiers384            bestGrammar = bestGrammar.makeUniform().insideOutside(frontiers, pseudoCounts)385        elif True:386            # Reestimate the parameters using the best programs387            restrictedFrontiers = restrictFrontiers()388            bestGrammar = bestGrammar.makeUniform().insideOutside(389                restrictedFrontiers, pseudoCounts)390        else:391            # Use parameters that were found during search392            pass393 394        eprint("Old joint = %f\tNew joint = %f\n" %395               (oldJoint, bestGrammar.jointFrontiersMDL(frontiers, CPUs=CPUs)))396        # Return all of the frontiers, which have now been rewritten to use the397        # new fragments398        frontiers = {f.task: f for f in frontiers}399        frontiers = [frontiers.get(f.task, f)400                     for f in originalFrontiers]401 402        productionUses = bestGrammar.expectedUses(403            [f for f in frontiers if not f.empty]).actualUses404        productionUses = {405            p: productionUses.get(406                p, 0.) for p in bestGrammar.primitives}407        possibleUses = bestGrammar.expectedUses(408            [f for f in frontiers if not f.empty]).possibleUses409        possibleUses = {410            p: possibleUses.get(411                p, 0.) for p in bestGrammar.primitives}412 413        for p in bestGrammar.primitives:414            eprint("%f / %f\t%s" % (productionUses[p],415                                    possibleUses[p],416                                    p))417 418        bestGrammar.clearCache()419 420        grammar = bestGrammar.toGrammar()421 422        if False and \423           any(productionUses.get(p, 0) < 0.5 for p in grammar.primitives if p.isInvented):424            uselessProductions = [ p for p in grammar.primitives                 425                                   if p.isInvented and productionUses.get(p, 0) < 0.5]426            eprint("The following invented primitives are no longer needed, removing them...")427            eprint("\t" + "\t\n".join(map(str, uselessProductions)))428            grammar = grammar.removeProductions(uselessProductions)429 430        return grammar, frontiers431