tree.py 3.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136
  1. from functools import partial
  2. from sympy.strategies import chain, minimize
  3. import sympy.strategies.branch as branch
  4. from sympy.strategies.branch import yieldify
  5. identity = lambda x: x
  6. def treeapply(tree, join, leaf=identity):
  7. """ Apply functions onto recursive containers (tree).
  8. Explanation
  9. ===========
  10. join - a dictionary mapping container types to functions
  11. e.g. ``{list: minimize, tuple: chain}``
  12. Keys are containers/iterables. Values are functions [a] -> a.
  13. Examples
  14. ========
  15. >>> from sympy.strategies.tree import treeapply
  16. >>> tree = [(3, 2), (4, 1)]
  17. >>> treeapply(tree, {list: max, tuple: min})
  18. 2
  19. >>> add = lambda *args: sum(args)
  20. >>> def mul(*args):
  21. ... total = 1
  22. ... for arg in args:
  23. ... total *= arg
  24. ... return total
  25. >>> treeapply(tree, {list: mul, tuple: add})
  26. 25
  27. """
  28. for typ in join:
  29. if isinstance(tree, typ):
  30. return join[typ](*map(partial(treeapply, join=join, leaf=leaf),
  31. tree))
  32. return leaf(tree)
  33. def greedy(tree, objective=identity, **kwargs):
  34. """ Execute a strategic tree. Select alternatives greedily
  35. Trees
  36. -----
  37. Nodes in a tree can be either
  38. function - a leaf
  39. list - a selection among operations
  40. tuple - a sequence of chained operations
  41. Textual examples
  42. ----------------
  43. Text: Run f, then run g, e.g. ``lambda x: g(f(x))``
  44. Code: ``(f, g)``
  45. Text: Run either f or g, whichever minimizes the objective
  46. Code: ``[f, g]``
  47. Textx: Run either f or g, whichever is better, then run h
  48. Code: ``([f, g], h)``
  49. Text: Either expand then simplify or try factor then foosimp. Finally print
  50. Code: ``([(expand, simplify), (factor, foosimp)], print)``
  51. Objective
  52. ---------
  53. "Better" is determined by the objective keyword. This function makes
  54. choices to minimize the objective. It defaults to the identity.
  55. Examples
  56. ========
  57. >>> from sympy.strategies.tree import greedy
  58. >>> inc = lambda x: x + 1
  59. >>> dec = lambda x: x - 1
  60. >>> double = lambda x: 2*x
  61. >>> tree = [inc, (dec, double)] # either inc or dec-then-double
  62. >>> fn = greedy(tree)
  63. >>> fn(4) # lowest value comes from the inc
  64. 5
  65. >>> fn(1) # lowest value comes from dec then double
  66. 0
  67. This function selects between options in a tuple. The result is chosen that
  68. minimizes the objective function.
  69. >>> fn = greedy(tree, objective=lambda x: -x) # maximize
  70. >>> fn(4) # highest value comes from the dec then double
  71. 6
  72. >>> fn(1) # highest value comes from the inc
  73. 2
  74. Greediness
  75. ----------
  76. This is a greedy algorithm. In the example:
  77. ([a, b], c) # do either a or b, then do c
  78. the choice between running ``a`` or ``b`` is made without foresight to c
  79. """
  80. optimize = partial(minimize, objective=objective)
  81. return treeapply(tree, {list: optimize, tuple: chain}, **kwargs)
  82. def allresults(tree, leaf=yieldify):
  83. """ Execute a strategic tree. Return all possibilities.
  84. Returns a lazy iterator of all possible results
  85. Exhaustiveness
  86. --------------
  87. This is an exhaustive algorithm. In the example
  88. ([a, b], [c, d])
  89. All of the results from
  90. (a, c), (b, c), (a, d), (b, d)
  91. are returned. This can lead to combinatorial blowup.
  92. See sympy.strategies.greedy for details on input
  93. """
  94. return treeapply(tree, {list: branch.multiplex, tuple: branch.chain},
  95. leaf=leaf)
  96. def brute(tree, objective=identity, **kwargs):
  97. return lambda expr: min(tuple(allresults(tree, **kwargs)(expr)),
  98. key=objective)