Fraser/dream-coder
Program Synthesis Data Generated program synthesis datasets used to train dreamcoder. Currently just supports text & list data.
6730
1from collections import defaultdict2 3from dreamcoder.frontier import *4from dreamcoder.program import *5from dreamcoder.type import *6from dreamcoder.utilities import *7 8import time9 10class GrammarFailure(Exception):11 pass12 13class SketchEnumerationFailure(Exception):14 pass15 16class NoCandidates(Exception):17 pass18 19 20class Grammar(object):21 def __init__(self, logVariable, productions, continuationType=None):22 self.logVariable = logVariable23 self.productions = productions24 25 self.continuationType = continuationType26 27 self.expression2likelihood = dict((p, l) for l, _, p in productions)28 self.expression2likelihood[Index(0)] = self.logVariable29 30 def randomWeights(self, r):31 """returns a new grammar with random weights drawn from r. calls `r` w/ old weight"""32 return Grammar(logVariable=r(self.logVariable),33 productions=[(r(l),t,p)34 for l,t,p in self.productions ],35 continuationType=self.continuationType)36 37 def strip_primitive_values(self):38 return Grammar(logVariable=self.logVariable,39 productions=[(l,t,strip_primitive_values(p))40 for l,t,p in self.productions ],41 continuationType=self.continuationType)42 43 def unstrip_primitive_values(self):44 return Grammar(logVariable=self.logVariable,45 productions=[(l,t,unstrip_primitive_values(p))46 for l,t,p in self.productions ],47 continuationType=self.continuationType)48 49 def __setstate__(self, state):50 """51 Legacy support for loading grammar objects without the imperative type filled in52 """53 assert 'logVariable' in state54 assert 'productions' in state55 if 'continuationType' in state:56 continuationType = state['continuationType']57 else:58 if any( 'turtle' in str(t) for l,t,p in state['productions'] ):59 continuationType = baseType("turtle")60 elif any( 'tower' in str(t) for l,t,p in state['productions'] ):61 continuationType = baseType("tower")62 else:63 continuationType = None64 65 self.__init__(state['logVariable'], state['productions'], continuationType=continuationType)66 67 @staticmethod68 def fromProductions(productions, logVariable=0.0, continuationType=None):69 """Make a grammar from primitives and their relative logpriors."""70 return Grammar(logVariable, [(l, p.infer(), p)71 for l, p in productions],72 continuationType=continuationType)73 74 @staticmethod75 def uniform(primitives, continuationType=None):76 return Grammar(0.0, [(0.0, p.infer(), p) for p in primitives], continuationType=continuationType)77 78 def __len__(self): return len(self.productions)79 80 def __str__(self):81 def productionKey(xxx_todo_changeme):82 (l, t, p) = xxx_todo_changeme83 return not isinstance(p, Primitive), l is not None and -l84 if self.continuationType is not None:85 lines = ["continuation : %s"%self.continuationType]86 else:87 lines = []88 lines += ["%f\tt0\t$_" % self.logVariable]89 for l, t, p in sorted(self.productions, key=productionKey):90 if l is not None:91 l = "%f\t%s\t%s" % (l, t, p)92 else:93 l = "-Inf\t%s\t%s" % (t, p)94 if not t.isArrow() and isinstance(p, Invented):95 try:96 l += "\teval = %s" % (p.evaluate([]))97 except BaseException:98 pass99 100 lines.append(l)101 return "\n".join(lines)102 103 def json(self):104 j = {"logVariable": self.logVariable,105 "productions": [{"expression": str(p), "logProbability": l}106 for l, _, p in self.productions]}107 if self.continuationType is not None:108 j["continuationType"] = self.continuationType.json()109 return j110 111 def _immutable_code(self): return self.logVariable, tuple(self.productions)112 113 def __eq__(self, o): return self._immutable_code() == o._immutable_code()114 115 def __ne__(self, o): return not (self == o)116 117 def __hash__(self): return hash(self._immutable_code())118 119 @property120 def primitives(self):121 return [p for _, _, p in self.productions]122 123 def removeProductions(self, ps):124 return Grammar(125 self.logVariable, [126 (l, t, p) for (127 l, t, p) in self.productions if p not in ps],128 continuationType=self.continuationType)129 130 def buildCandidates(self, request, context, environment,131 # Should the log probabilities be normalized?132 normalize=True,133 # Should be returned a table mapping primitives to134 # their candidate entry?135 returnTable=False,136 # Should we return probabilities vs log probabilities?137 returnProbabilities=False,138 # Must be a leaf (have no arguments)?139 mustBeLeaf=False):140 """Primitives that are candidates for being used given a requested type141 If returnTable is false (default): returns [((log)likelihood, tp, primitive, context)]142 if returntable is true: returns {primitive: ((log)likelihood, tp, context)}"""143 if returnProbabilities:144 assert normalize145 146 candidates = []147 variableCandidates = []148 for l, t, p in self.productions:149 try:150 newContext, t = t.instantiate(context)151 newContext = newContext.unify(t.returns(), request)152 t = t.apply(newContext)153 if mustBeLeaf and t.isArrow():154 continue155 candidates.append((l, t, p, newContext))156 except UnificationFailure:157 continue158 for j, t in enumerate(environment):159 try:160 newContext = context.unify(t.returns(), request)161 t = t.apply(newContext)162 if mustBeLeaf and t.isArrow():163 continue164 variableCandidates.append((t, Index(j), newContext))165 except UnificationFailure:166 continue167 168 if self.continuationType == request:169 terminalIndices = [v.i for t,v,k in variableCandidates if not t.isArrow()]170 if terminalIndices:171 smallestIndex = Index(min(terminalIndices))172 variableCandidates = [(t,v,k) for t,v,k in variableCandidates173 if t.isArrow() or v == smallestIndex]174 175 candidates += [(self.logVariable - log(len(variableCandidates)), t, p, k)176 for t, p, k in variableCandidates]177 if candidates == []:178 raise NoCandidates()179 #eprint("candidates inside buildCandidates before norm:")180 #eprint(candidates)181 182 if normalize:183 z = lse([l for l, t, p, k in candidates])184 if returnProbabilities:185 candidates = [(exp(l - z), t, p, k)186 for l, t, p, k in candidates]187 else:188 candidates = [(l - z, t, p, k) for l, t, p, k in candidates]189 190 #eprint("candidates inside buildCandidates after norm:")191 #eprint(candidates)192 193 if returnTable:194 return {p: (l, t, k) for l, t, p, k in candidates}195 else:196 return candidates197 198 199 def sample(self, request, maximumDepth=6, maxAttempts=None):200 attempts = 0201 202 while True:203 try:204 _, e = self._sample(205 request, Context.EMPTY, [], maximumDepth=maximumDepth)206 return e207 except NoCandidates:208 if maxAttempts is not None:209 attempts += 1210 if attempts > maxAttempts:211 return None212 continue213 214 def _sample(self, request, context, environment, maximumDepth):215 if request.isArrow():216 context, expression = self._sample(217 request.arguments[1], context, [218 request.arguments[0]] + environment, maximumDepth)219 return context, Abstraction(expression)220 221 candidates = self.buildCandidates(request, context, environment,222 normalize=True,223 returnProbabilities=True,224 # Force it to terminate in a225 # leaf; a primitive with no226 # function arguments227 mustBeLeaf=maximumDepth <= 1)228 #eprint("candidates:")229 #eprint(candidates)230 newType, chosenPrimitive, context = sampleDistribution(candidates)231 232 # Sample the arguments233 xs = newType.functionArguments()234 returnValue = chosenPrimitive235 236 for x in xs:237 x = x.apply(context)238 context, x = self._sample(x, context, environment, maximumDepth - 1)239 returnValue = Application(returnValue, x)240 241 return context, returnValue242 243 def likelihoodSummary(self, context, environment, request, expression, silent=False):244 if request.isArrow():245 if not isinstance(expression, Abstraction):246 if not silent:247 eprint("Request is an arrow but I got", expression)248 return context, None249 return self.likelihoodSummary(context,250 [request.arguments[0]] + environment,251 request.arguments[1],252 expression.body,253 silent=silent)254 # Build the candidates255 candidates = self.buildCandidates(request, context, environment,256 normalize=False,257 returnTable=True)258 259 # A list of everything that would have been possible to use here260 possibles = [p for p in candidates.keys() if not p.isIndex]261 numberOfVariables = sum(p.isIndex for p in candidates.keys())262 if numberOfVariables > 0:263 possibles += [Index(0)]264 265 f, xs = expression.applicationParse()266 267 if f not in candidates:268 if self.continuationType is not None and f.isIndex:269 ls = LikelihoodSummary()270 ls.constant = NEGATIVEINFINITY271 return ls272 273 if not silent:274 eprint(f, "Not in candidates")275 eprint("Candidates is", candidates)276 #eprint("grammar:", grammar.productions)277 eprint("request is", request)278 eprint("xs", xs)279 eprint("environment", environment)280 assert False281 return context, None282 283 thisSummary = LikelihoodSummary()284 thisSummary.record(f, possibles,285 constant= -math.log(numberOfVariables) if f.isIndex else 0)286 287 _, tp, context = candidates[f]288 argumentTypes = tp.functionArguments()289 if len(xs) != len(argumentTypes):290 eprint("PANIC: not enough arguments for the type")291 eprint("request", request)292 eprint("tp", tp)293 eprint("expression", expression)294 eprint("xs", xs)295 eprint("argumentTypes", argumentTypes)296 # This should absolutely never occur297 raise GrammarFailure((context, environment, request, expression))298 299 for argumentType, argument in zip(argumentTypes, xs):300 argumentType = argumentType.apply(context)301 context, newSummary = self.likelihoodSummary(302 context, environment, argumentType, argument, silent=silent)303 if newSummary is None:304 return context, None305 thisSummary.join(newSummary)306 307 return context, thisSummary308 309 def bestFirstEnumeration(self, request):310 from heapq import heappush, heappop311 312 pq = []313 314 def choices(parentCost, xs):315 for c, x in xs:316 heappush(pq, (parentCost + c, x))317 318 def g(parentCost, request, _=None,319 context=None, environment=[],320 k=None):321 """322 k is a continuation.323 k: Expects to be called with MDL, context, expression.324 """325 326 assert k is not None327 if context is None:328 context = Context.EMPTY329 330 if request.isArrow():331 g(parentCost,332 request.arguments[1],333 context=context,334 environment=[request.arguments[0]] + environment,335 k=lambda MDL,336 newContext,337 p: k(MDL,338 newContext,339 Abstraction(p)))340 else:341 candidates = self.buildCandidates(request,342 context,343 environment,344 normalize=True,345 returnProbabilities=False,346 returnTable=True)347 choices(parentCost,348 [(-f_ll_tp_newContext[1][0],349 lambda: ga(parentCost - f_ll_tp_newContext[1][0],350 f_ll_tp_newContext[0],351 f_ll_tp_newContext[1][1].functionArguments(),352 context=f_ll_tp_newContext[1][2],353 environment=environment,354 k=k)) for f_ll_tp_newContext in iter(candidates.items())])355 356 def ga(costSoFar, f, argumentTypes, _=None,357 context=None, environment=None,358 k=None):359 if argumentTypes == []:360 k(costSoFar, context, f)361 else:362 t1 = argumentTypes[0].apply(context)363 g(costSoFar, t1, context=context, environment=environment,364 k=lambda newCost, newContext, argument:365 ga(newCost, Application(f, argument), argumentTypes[1:],366 context=newContext, environment=environment,367 k=k))368 369 def receiveResult(MDL, _, expression):370 heappush(pq, (MDL, expression))371 372 g(0., request, context=Context.EMPTY, environment=[], k=receiveResult)373 frontier = []374 while len(frontier) < 10**3:375 MDL, action = heappop(pq)376 if isinstance(action, Program):377 expression = action378 frontier.append(expression)379 #eprint("Enumerated program",expression,-MDL,self.closedLogLikelihood(request, expression))380 else:381 action()382 383 def closedLikelihoodSummary(self, request, expression, silent=False):384 try:385 context, summary = self.likelihoodSummary(Context.EMPTY, [], request, expression, silent=silent)386 except GrammarFailure as e:387 failureExport = 'failures/grammarFailure%s.pickle' % (388 time.time() + getPID())389 eprint("PANIC: Grammar failure, exporting to ", failureExport)390 with open(failureExport, 'wb') as handle:391 pickle.dump((e, self, request, expression), handle)392 assert False393 394 return summary395 396 def logLikelihood(self, request, expression):397 summary = self.closedLikelihoodSummary(request, expression)398 if summary is None:399 eprint(400 "FATAL: program [ %s ] does not have a likelihood summary." %401 expression, "r = ", request, "\n", self)402 assert False403 return summary.logLikelihood(self)404 405 def rescoreFrontier(self, frontier):406 return Frontier([FrontierEntry(e.program,407 logPrior=self.logLikelihood(frontier.task.request, e.program),408 logLikelihood=e.logLikelihood)409 for e in frontier],410 frontier.task)411 412 def productionUses(self, frontiers):413 """Returns the expected number of times that each production was used. {production: expectedUses}"""414 frontiers = [self.rescoreFrontier(f).normalize()415 for f in frontiers if not f.empty]416 uses = {p: 0. for p in self.primitives}417 for f in frontiers:418 for e in f:419 summary = self.closedLikelihoodSummary(f.task.request,420 e.program)421 for p, u in summary.uses:422 uses[p] += u * math.exp(e.logPosterior)423 return uses424 425 def insideOutside(self, frontiers, pseudoCounts, iterations=1):426 # Replace programs with (likelihood summary, uses)427 frontiers = [ Frontier([ FrontierEntry((summary, summary.toUses()),428 logPrior=summary.logLikelihood(self),429 logLikelihood=e.logLikelihood)430 for e in f431 for summary in [self.closedLikelihoodSummary(f.task.request, e.program)] ],432 task=f.task)433 for f in frontiers ]434 435 g = self436 for i in range(iterations):437 u = Uses()438 for f in frontiers:439 f = f.normalize()440 for e in f:441 _, eu = e.program442 u += math.exp(e.logPosterior) * eu443 444 lv = math.log(u.actualVariables + pseudoCounts) - \445 math.log(u.possibleVariables + pseudoCounts)446 g = Grammar(lv,447 [ (math.log(u.actualUses.get(p,0.) + pseudoCounts) - \448 math.log(u.possibleUses.get(p,0.) + pseudoCounts),449 t,p)450 for _,t,p in g.productions ],451 continuationType=self.continuationType)452 if i < iterations - 1:453 frontiers = [Frontier([ FrontierEntry((summary, uses),454 logPrior=summary.logLikelihood(g),455 logLikelihood=e.logLikelihood)456 for e in f457 for (summary, uses) in [e.program] ],458 task=f.task)459 for f in frontiers ]460 return g461 462 def frontierMDL(self, frontier):463 return max( e.logLikelihood + self.logLikelihood(frontier.task.request, e.program)464 for e in frontier ) 465 466 467 def enumeration(self,context,environment,request,upperBound,468 maximumDepth=20,469 lowerBound=0.):470 '''Enumerates all programs whose MDL satisfies: lowerBound <= MDL < upperBound'''471 if upperBound < 0 or maximumDepth == 1:472 return473 474 if request.isArrow():475 v = request.arguments[0]476 for l, newContext, b in self.enumeration(context, [v] + environment,477 request.arguments[1],478 upperBound=upperBound,479 lowerBound=lowerBound,480 maximumDepth=maximumDepth):481 yield l, newContext, Abstraction(b)482 483 else:484 candidates = self.buildCandidates(request, context, environment,485 normalize=True)486 487 for l, t, p, newContext in candidates:488 mdl = -l489 if not (mdl < upperBound):490 continue491 492 xs = t.functionArguments()493 for aL, aK, application in\494 self.enumerateApplication(newContext, environment, p, xs,495 upperBound=upperBound + l,496 lowerBound=lowerBound + l,497 maximumDepth=maximumDepth - 1):498 yield aL + l, aK, application499 500 def enumerateApplication(self, context, environment,501 function, argumentRequests,502 # Upper bound on the description length of all of503 # the arguments504 upperBound,505 # Lower bound on the description length of all of506 # the arguments507 lowerBound=0.,508 maximumDepth=20,509 originalFunction=None,510 argumentIndex=0):511 if upperBound < 0. or maximumDepth == 1:512 return513 if originalFunction is None:514 originalFunction = function515 516 if argumentRequests == []:517 if lowerBound <= 0. and 0. < upperBound:518 yield 0., context, function519 else:520 return521 else:522 argRequest = argumentRequests[0].apply(context)523 laterRequests = argumentRequests[1:]524 for argL, newContext, arg in self.enumeration(context, environment, argRequest,525 upperBound=upperBound,526 lowerBound=0.,527 maximumDepth=maximumDepth):528 if violatesSymmetry(originalFunction, arg, argumentIndex):529 continue530 531 newFunction = Application(function, arg)532 for resultL, resultK, result in self.enumerateApplication(newContext, environment, newFunction,533 laterRequests,534 upperBound=upperBound + argL,535 lowerBound=lowerBound + argL,536 maximumDepth=maximumDepth,537 originalFunction=originalFunction,538 argumentIndex=argumentIndex + 1):539 yield resultL + argL, resultK, result540 541 def sketchEnumeration(self,context,environment,request,sk,upperBound,542 maximumDepth=20,543 lowerBound=0.):544 '''Enumerates all sketch instantiations whose MDL satisfies: lowerBound <= MDL < upperBound'''545 if upperBound < 0. or maximumDepth == 1:546 return547 548 if sk.isHole:549 yield from self.enumeration(context, environment, request, upperBound,550 maximumDepth=maximumDepth,551 lowerBound=lowerBound)552 elif request.isArrow():553 assert sk.isAbstraction554 v = request.arguments[0]555 for l, newContext, b in self.sketchEnumeration(context, [v] + environment,556 request.arguments[1],557 sk.body,558 upperBound=upperBound,559 lowerBound=lowerBound,560 maximumDepth=maximumDepth):561 yield l, newContext, Abstraction(b)562 563 else:564 f, xs = sk.applicationParse()565 if f.isIndex:566 ft = environment[f.i].apply(context)567 elif f.isInvented or f.isPrimitive:568 context, ft = f.tp.instantiate(context)569 elif f.isAbstraction:570 assert False, "sketch is not in beta longform"571 elif f.isHole:572 assert False, "hole as function not yet supported"573 elif f.isApplication:574 assert False, "should never happen - bug in applicationParse"575 else: assert False576 577 try: context = context.unify(ft.returns(), request) 578 except UnificationFailure:579 print("Exception: sketch is ill-typed")580 return #so that we can continue evaluating581 # raise SketchEnumerationFailure() #"sketch is ill-typed"582 ft = ft.apply(context)583 argumentRequests = ft.functionArguments()584 585 assert len(argumentRequests) == len(xs)586 587 yield from self.sketchApplication(context, environment,588 f, xs, argumentRequests,589 upperBound=upperBound,590 lowerBound=lowerBound,591 maximumDepth=maximumDepth - 1)592 593 594 def sketchApplication(self, context, environment,595 function, arguments, argumentRequests,596 # Upper bound on the description length of all of597 # the arguments598 upperBound,599 # Lower bound on the description length of all of600 # the arguments601 lowerBound=0.,602 maximumDepth=20):603 if upperBound < 0. or maximumDepth == 1:604 return605 606 if argumentRequests == []:607 if lowerBound <= 0. and 0. < upperBound:608 yield 0., context, function609 else:610 return611 else:612 argRequest = argumentRequests[0].apply(context)613 laterRequests = argumentRequests[1:]614 firstSketch = arguments[0]615 laterSketches = arguments[1:]616 for argL, newContext, arg in self.sketchEnumeration(context, environment, argRequest,617 firstSketch,618 upperBound=upperBound,619 lowerBound=0.,620 maximumDepth=maximumDepth):621 622 newFunction = Application(function, arg)623 for resultL, resultK, result in self.sketchApplication(newContext, environment, newFunction,624 laterSketches, laterRequests,625 upperBound=upperBound + argL,626 lowerBound=lowerBound + argL,627 maximumDepth=maximumDepth):628 629 yield resultL + argL, resultK, result630 631 def sketchLogLikelihood(self, request, full, sk, context=Context.EMPTY, environment=[]):632 """633 calculates mdl of full program 'full' from sketch 'sk'634 """635 if sk.isHole:636 _, summary = self.likelihoodSummary(context, environment, request, full)637 if summary is None:638 eprint(639 "FATAL: program [ %s ] does not have a likelihood summary." %640 full, "r = ", request, "\n", self)641 assert False642 return summary.logLikelihood(self), context643 644 elif request.isArrow():645 assert sk.isAbstraction and full.isAbstraction646 #assert sk.f == full.f #is this right? or do i need to recurse?647 v = request.arguments[0]648 return self.sketchLogLikelihood(request.arguments[1], full.body, sk.body, context=context, environment=[v] + environment)649 650 else:651 sk_f, sk_xs = sk.applicationParse()652 full_f, full_xs = full.applicationParse()653 if sk_f.isIndex:654 assert sk_f == full_f, "sketch and full program don't match on an index"655 ft = environment[sk_f.i].apply(context)656 elif sk_f.isInvented or sk_f.isPrimitive:657 assert sk_f == full_f, "sketch and full program don't match on a primitive"658 context, ft = sk_f.tp.instantiate(context)659 elif sk_f.isAbstraction:660 assert False, "sketch is not in beta longform"661 elif sk_f.isHole:662 assert False, "hole as function not yet supported"663 elif sk_f.isApplication:664 assert False, "should never happen - bug in applicationParse"665 else: assert False666 667 try: context = context.unify(ft.returns(), request) 668 except UnificationFailure: assert False, "sketch is ill-typed"669 ft = ft.apply(context)670 argumentRequests = ft.functionArguments()671 672 assert len(argumentRequests) == len(sk_xs) == len(full_xs) #this might not be true if holes??673 674 return self.sketchllApplication(context, environment,675 sk_f, sk_xs, full_f, full_xs, argumentRequests)676 677 def sketchllApplication(self, context, environment,678 sk_function, sk_arguments, full_function, full_arguments, argumentRequests):679 if argumentRequests == []:680 return torch.tensor([0.]).cuda(), context #does this make sense?681 else:682 argRequest = argumentRequests[0].apply(context)683 laterRequests = argumentRequests[1:]684 685 sk_firstSketch = sk_arguments[0]686 full_firstSketch = full_arguments[0]687 sk_laterSketches = sk_arguments[1:]688 full_laterSketches = full_arguments[1:]689 690 argL, newContext = self.sketchLogLikelihood(argRequest, full_firstSketch, sk_firstSketch, context=context, environment=environment)691 692 #finish this...693 sk_newFunction = Application(sk_function, sk_firstSketch) # is this redundant? maybe 694 full_newFunction = Application(full_function, full_firstSketch)695 696 resultL, context = self.sketchllApplication(newContext, environment, sk_newFunction, sk_laterSketches,697 full_newFunction, full_laterSketches, laterRequests)698 699 return resultL + argL, context700 701 702 def enumerateNearby(self, request, expr, distance=3.0):703 """Enumerate programs with local mutations in subtrees with small description length"""704 if distance <= 0:705 yield expr706 else:707 def mutations(tp, loss):708 for l, _, expr in self.enumeration(709 Context.EMPTY, [], tp, distance - loss):710 yield expr, l711 yield from Mutator(self, mutations).execute(expr, request)712 713 714 def enumerateHoles(self, request, expr, k=3, return_obj=Hole):715 """Enumerate programs with a single hole within mdl distance"""716 #TODO: make it possible to enumerate sketches with multiple holes717 def mutations(tp, loss, is_left_application=False):718 """719 to allow applications lhs to become a hole, 720 remove the condition below and ignore all the is_left_application kwds 721 """722 if not is_left_application: 723 yield return_obj(), 0724 top_k = []725 for expr, l in Mutator(self, mutations).execute(expr, request):726 if len(top_k) > 0:727 i, v = min(enumerate(top_k), key=lambda x:x[1][1])728 if l > v[1]:729 if len(top_k) >= k:730 top_k[i] = (expr, l)731 else:732 top_k.append((expr, l))733 elif len(top_k) < k:734 top_k.append((expr, l))735 else:736 top_k.append((expr, l))737 return sorted(top_k, key=lambda x:-x[1])738 739 def untorch(self):740 return Grammar(self.logVariable.data.tolist()[0], 741 [ (l.data.tolist()[0], t, p)742 for l, t, p in self.productions],743 continuationType=self.continuationType)744 745class LikelihoodSummary(object):746 '''Summarizes the terms that will be used in a likelihood calculation'''747 748 def __init__(self):749 self.uses = {}750 self.normalizers = {}751 self.constant = 0.752 753 def __str__(self):754 return """LikelihoodSummary(constant = %f,755uses = {%s},756normalizers = {%s})""" % (self.constant,757 ", ".join(758 "%s: %d" % (k,759 v) for k,760 v in self.uses.items()),761 ", ".join(762 "%s: %d" % (k,763 v) for k,764 v in self.normalizers.items()))765 766 def record(self, actual, possibles, constant=0.):767 # Variables are all normalized to be $0768 if isinstance(actual, Index):769 actual = Index(0)770 771 # Make it something that we can hash772 possibles = frozenset(sorted(possibles, key=hash))773 774 self.constant += constant775 self.uses[actual] = self.uses.get(actual, 0) + 1776 self.normalizers[possibles] = self.normalizers.get(possibles, 0) + 1777 778 def join(self, other):779 self.constant += other.constant780 for k, v in other.uses.items():781 self.uses[k] = self.uses.get(k, 0) + v782 for k, v in other.normalizers.items():783 self.normalizers[k] = self.normalizers.get(k, 0) + v784 785 def logLikelihood(self, grammar):786 return self.constant + \787 sum(count * grammar.expression2likelihood[p] for p, count in self.uses.items()) - \788 sum(count * lse([grammar.expression2likelihood[p] for p in ps])789 for ps, count in self.normalizers.items())790 def logLikelihood_overlyGeneral(self, grammar):791 """Calculates log likelihood of this summary, given that the summary might refer to productions that don't occur in the grammar"""792 return self.constant + \793 sum(count * grammar.expression2likelihood[p] for p, count in self.uses.items()) - \794 sum(count * lse([grammar.expression2likelihood.get(p,NEGATIVEINFINITY) for p in ps])795 for ps, count in self.normalizers.items()) 796 def numerator(self, grammar):797 return self.constant + \798 sum(count * grammar.expression2likelihood[p] for p, count in self.uses.items())799 def denominator(self, grammar):800 return \801 sum(count * lse([grammar.expression2likelihood[p] for p in ps])802 for ps, count in self.normalizers.items())803 def toUses(self):804 from collections import Counter805 806 possibleVariables = sum( count if Index(0) in ps else 0807 for ps, count in self.normalizers.items() )808 actualVariables = self.uses.get(Index(0), 0.)809 actualUses = {k: v810 for k, v in self.uses.items()811 if not k.isIndex }812 possibleUses = dict(Counter(p813 for ps, count in self.normalizers.items()814 for p_ in ps815 if not p_.isIndex816 for p in [p_]*count ))817 return Uses(possibleVariables, actualVariables,818 possibleUses, actualUses)819 820 821class Uses(object):822 '''Tracks uses of different grammar productions'''823 824 def __init__(self, possibleVariables=0., actualVariables=0.,825 possibleUses={}, actualUses={}):826 self.actualVariables = actualVariables827 self.possibleVariables = possibleVariables828 self.possibleUses = possibleUses829 self.actualUses = actualUses830 831 def __str__(self):832 return "Uses(actualVariables = %f, possibleVariables = %f, actualUses = %s, possibleUses = %s)" %\833 (self.actualVariables, self.possibleVariables, self.actualUses, self.possibleUses)834 835 def __repr__(self): return str(self)836 837 def __mul__(self, a):838 return Uses(a * self.possibleVariables,839 a * self.actualVariables,840 {p: a * u for p, u in self.possibleUses.items()},841 {p: a * u for p, u in self.actualUses.items()})842 843 def __imul__(self, a):844 self.possibleVariables *= a845 self.actualVariables *= a846 for p in self.possibleUses:847 self.possibleUses[p] *= a848 for p in self.actualUses:849 self.actualUses[p] *= a850 return self851 852 def __rmul__(self, a):853 return self * a854 855 def __radd__(self, o):856 if o == 0:857 return self858 return self + o859 860 def __add__(self, o):861 if o == 0:862 return self863 864 def merge(x, y):865 z = x.copy()866 for k, v in y.items():867 z[k] = v + x.get(k, 0.)868 return z869 return Uses(self.possibleVariables + o.possibleVariables,870 self.actualVariables + o.actualVariables,871 merge(self.possibleUses, o.possibleUses),872 merge(self.actualUses, o.actualUses))873 874 def __iadd__(self, o):875 self.possibleVariables += o.possibleVariables876 self.actualVariables += o.actualVariables877 for k, v in o.possibleUses.items():878 self.possibleUses[k] = self.possibleUses.get(k, 0.) + v879 for k, v in o.actualUses.items():880 self.actualUses[k] = self.actualUses.get(k, 0.) + v881 return self882 883 @staticmethod884 def join(z, *weightedUses):885 """Consumes weightedUses"""886 if not weightedUses:887 Uses.empty888 if len(weightedUses) == 1:889 return weightedUses[0][1]890 for w, u in weightedUses:891 u *= exp(w - z)892 total = Uses()893 total.possibleVariables = sum(894 u.possibleVariables for _, u in weightedUses)895 total.actualVariables = sum(u.actualVariables for _, u in weightedUses)896 total.possibleUses = defaultdict(float)897 total.actualUses = defaultdict(float)898 for _, u in weightedUses:899 for k, v in u.possibleUses.items():900 total.possibleUses[k] += v901 for k, v in u.actualUses.items():902 total.actualUses[k] += v903 return total904 905 906Uses.empty = Uses()907 908class ContextualGrammar:909 def __init__(self, noParent, variableParent, library):910 self.noParent, self.variableParent, self.library = noParent, variableParent, library911 912 self.productions = [(None,t,p) for _,t,p in self.noParent.productions ]913 self.primitives = [p for _,_2,p in self.productions ]914 915 self.continuationType = noParent.continuationType916 assert variableParent.continuationType == self.continuationType917 918 assert set(noParent.primitives) == set(variableParent.primitives)919 assert set(variableParent.primitives) == set(library.keys())920 for e,gs in library.items():921 assert len(gs) == len(e.infer().functionArguments())922 for g in gs:923 assert set(g.primitives) == set(library.keys())924 assert g.continuationType == self.continuationType925 926 def untorch(self):927 return ContextualGrammar(self.noParent.untorch(), self.variableParent.untorch(),928 {e: [g.untorch() for g in gs ]929 for e,gs in self.library.items() })930 931 def randomWeights(self, r):932 """returns a new grammar with random weights drawn from r. calls `r` w/ old weight"""933 return ContextualGrammar(self.noParent.randomWeights(r),934 self.variableParent.randomWeights(r),935 {e: [g.randomWeights(r) for g in gs]936 for e,gs in self.library.items() })937 def __str__(self):938 lines = ["No parent:",str(self.noParent),"",939 "Variable parent:",str(self.variableParent),"",940 ""]941 for e,gs in self.library.items():942 for j,g in enumerate(gs):943 lines.extend(["Parent %s, argument index %s"%(e,j),944 str(g),945 ""])946 return "\n".join(lines)947 948 def json(self):949 return {"noParent": self.noParent.json(),950 "variableParent": self.variableParent.json(),951 "productions": [{"program": str(e),952 "arguments": [gp.json() for gp in gs ]}953 for e,gs in self.library.items() ]}954 955 @staticmethod956 def fromGrammar(g):957 return ContextualGrammar(g, g,958 {e: [g]*len(e.infer().functionArguments())959 for e in g.primitives })960 961 962 class LS: # likelihood summary963 def __init__(self, owner):964 self.noParent = LikelihoodSummary()965 self.variableParent = LikelihoodSummary()966 self.library = {e: [LikelihoodSummary() for _ in gs] for e,gs in owner.library.items() }967 968 def record(self, parent, parentIndex, actual, possibles, constant):969 if parent is None: ls = self.noParent970 elif parent.isIndex: ls = self.variableParent971 else: ls = self.library[parent][parentIndex]972 ls.record(actual, possibles, constant=constant)973 974 def join(self, other):975 self.noParent.join(other.noParent)976 self.variableParent.join(other.variableParent)977 for e,gs in self.library.items():978 for g1,g2 in zip(gs, other.library[e]):979 g1.join(g2)980 981 def logLikelihood(self, owner):982 return self.noParent.logLikelihood(owner.noParent) + \983 self.variableParent.logLikelihood(owner.variableParent) + \984 sum(r.logLikelihood(g)985 for e, rs in self.library.items()986 for r,g in zip(rs, owner.library[e]) ) 987 def numerator(self, owner):988 return self.noParent.numerator(owner.noParent) + \989 self.variableParent.numerator(owner.variableParent) + \990 sum(r.numerator(g)991 for e, rs in self.library.items()992 for r,g in zip(rs, owner.library[e]) ) 993 def denominator(self, owner):994 return self.noParent.denominator(owner.noParent) + \995 self.variableParent.denominator(owner.variableParent) + \996 sum(r.denominator(g)997 for e, rs in self.library.items()998 for r,g in zip(rs, owner.library[e]) ) 999 1000 def likelihoodSummary(self, parent, parentIndex, context, environment, request, expression):1001 if request.isArrow():1002 assert expression.isAbstraction1003 return self.likelihoodSummary(parent, parentIndex,1004 context,1005 [request.arguments[0]] + environment,1006 request.arguments[1],1007 expression.body)1008 if parent is None: g = self.noParent1009 elif parent.isIndex: g = self.variableParent1010 else: g = self.library[parent][parentIndex] 1011 candidates = g.buildCandidates(request, context, environment,1012 normalize=False, returnTable=True)1013 1014 # A list of everything that would have been possible to use here1015 possibles = [p for p in candidates.keys() if not p.isIndex]1016 numberOfVariables = sum(p.isIndex for p in candidates.keys())1017 if numberOfVariables > 0:1018 possibles += [Index(0)]1019 1020 f, xs = expression.applicationParse()1021 1022 assert f in candidates1023 1024 thisSummary = self.LS(self)1025 thisSummary.record(parent, parentIndex,1026 f, possibles,1027 constant= -math.log(numberOfVariables) if f.isIndex else 0)1028 1029 _, tp, context = candidates[f]1030 argumentTypes = tp.functionArguments()1031 assert len(xs) == len(argumentTypes)1032 1033 for i, (argumentType, argument) in enumerate(zip(argumentTypes, xs)):1034 argumentType = argumentType.apply(context)1035 context, newSummary = self.likelihoodSummary(f, i,1036 context, environment, argumentType, argument)1037 thisSummary.join(newSummary)1038 1039 return context, thisSummary1040 1041 def closedLikelihoodSummary(self, request, expression):1042 return self.likelihoodSummary(None,None,1043 Context.EMPTY,[],1044 request, expression)[1]1045 1046 def logLikelihood(self, request, expression):1047 return self.closedLikelihoodSummary(request, expression).logLikelihood(self)1048 1049 def sample(self, request, maximumDepth=8, maxAttempts=None):1050 attempts = 01051 while True:1052 try:1053 _, e = self._sample(None, None, Context.EMPTY, [], request, maximumDepth)1054 return e1055 except NoCandidates:1056 if maxAttempts is not None:1057 attempts += 11058 if attempts > maxAttempts: return None1059 continue1060 1061 def _sample(self, parent, parentIndex, context, environment, request, maximumDepth):1062 if request.isArrow():1063 context, body = self._sample(parent, parentIndex, context,1064 [request.arguments[0]] + environment,1065 request.arguments[1],1066 maximumDepth)1067 return context, Abstraction(body)1068 if parent is None: g = self.noParent1069 elif parent.isIndex: g = self.variableParent1070 else: g = self.library[parent][parentIndex]1071 candidates = g.buildCandidates(request, context, environment,1072 normalize=True, returnProbabilities=True,1073 mustBeLeaf=(maximumDepth <= 1))1074 newType, chosenPrimitive, context = sampleDistribution(candidates)1075 1076 xs = newType.functionArguments()1077 returnValue = chosenPrimitive1078 1079 for j,x in enumerate(xs):1080 x = x.apply(context)1081 context, x = self._sample(chosenPrimitive, j, context, environment, x, maximumDepth - 1)1082 returnValue = Application(returnValue, x)1083 1084 return context, returnValue1085 1086 def expectedUsesMonteCarlo(self, request, debug=None):1087 import numpy as np1088 n = 01089 u = [0.]*len(self.primitives)1090 primitives = list(sorted(self.primitives, key=str))1091 noInventions = all( not p.isInvented for p in primitives )1092 primitive2index = {primitive: i1093 for i, primitive in enumerate(primitives)1094 if primitive.isInvented or noInventions }1095 eprint(primitive2index)1096 ns = 100001097 with timing(f"calculated expected uses using Monte Carlo simulation w/ {ns} samples"):1098 for _ in range(ns):1099 p = self.sample(request, maxAttempts=0)1100 if p is None: continue1101 n += 11102 if debug and n < 10:1103 eprint(debug, p)1104 for _, child in p.walk():1105 if child not in primitive2index: continue1106 u[primitive2index[child]] += 1.01107 u = np.array(u)/n1108 if debug:1109 eprint(f"Got {n} samples. Feature vector:\n{u}")1110 eprint(f"Likely used primitives: {[p for p,i in primitive2index.items() if u[i] > 0.5]}")1111 eprint(f"Likely used primitive indices: {[i for p,i in primitive2index.items() if u[i] > 0.5]}")1112 return u1113 1114 def featureVector(self, _=None, requests=None, onlyInventions=True, normalize=True):1115 """1116 Returns the probabilities licensed by the type system.1117 This is like the grammar productions, but with irrelevant junk removed.1118 Its intended use case is for clustering; it should be strictly better than the raw transition matrix.1119 """1120 if requests is None:1121 if self.continuationType: requests = {self.continuationType}1122 elif any( 'REAL' == str(p) for p in self.primitives ): requests = set()1123 elif any( 'STRING' == str(p) for p in self.primitives ): requests = {tlist(tcharacter)}1124 else: requests = set()1125 requests = {r.returns() for r in requests}1126 features = []1127 logWeights = []1128 for l,t,p in sorted(self.noParent.productions,1129 key=lambda z: str(z[2])):1130 if onlyInventions and not p.isInvented: continue1131 if any( canUnify(r, t.returns()) for r in requests ) or len(requests) == 0:1132 logWeights.append(l)1133 features.append(logWeights)1134 for parent in sorted(self.primitives, key=str):1135 if onlyInventions and not parent.isInvented: continue1136 if parent not in self.library: continue1137 argumentTypes = parent.infer().functionArguments()1138 for j,g in enumerate(self.library[parent]):1139 argumentType = argumentTypes[j]1140 logWeights = []1141 for l,t,p in sorted(g.productions,1142 key=lambda z: str(z[2])):1143 if onlyInventions and not p.isInvented: continue1144 if canUnify(argumentType.returns(), t.returns()):1145 logWeights.append(l)1146 features.append(logWeights)1147 1148 if normalize:1149 features = [ [math.exp(w - z) for w in lw ]1150 for lw in features1151 if lw1152 for z in [lse(lw)] ]1153 import numpy as np1154 return np.array([f1155 for lw in features1156 for f in lw])1157 1158 def enumeration(self,context,environment,request,upperBound,1159 parent=None, parentIndex=None,1160 maximumDepth=20,1161 lowerBound=0.):1162 '''Enumerates all programs whose MDL satisfies: lowerBound <= MDL < upperBound'''1163 if upperBound < 0 or maximumDepth == 1:1164 return1165 1166 if request.isArrow():1167 v = request.arguments[0]1168 for l, newContext, b in self.enumeration(context, [v] + environment,1169 request.arguments[1],1170 parent=parent, parentIndex=parentIndex,1171 upperBound=upperBound,1172 lowerBound=lowerBound,1173 maximumDepth=maximumDepth):1174 yield l, newContext, Abstraction(b)1175 else:1176 if parent is None: g = self.noParent1177 elif parent.isIndex: g = self.variableParent1178 else: g = self.library[parent][parentIndex]1179 1180 candidates = g.buildCandidates(request, context, environment,1181 normalize=True)1182 1183 for l, t, p, newContext in candidates:1184 mdl = -l1185 if not (mdl < upperBound):1186 continue1187 1188 xs = t.functionArguments()1189 for aL, aK, application in\1190 self.enumerateApplication(newContext, environment, p, xs,1191 parent=p,1192 upperBound=upperBound + l,1193 lowerBound=lowerBound + l,1194 maximumDepth=maximumDepth - 1):1195 yield aL + l, aK, application1196 1197 def enumerateApplication(self, context, environment,1198 function, argumentRequests,1199 # Upper bound on the description length of all of1200 # the arguments