sathandlers.py 9.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323
  1. from collections import defaultdict
  2. from sympy.assumptions.ask import Q
  3. from sympy.core import (Add, Mul, Pow, Number, NumberSymbol, Symbol)
  4. from sympy.core.numbers import ImaginaryUnit
  5. from sympy.functions.elementary.complexes import Abs
  6. from sympy.logic.boolalg import (Equivalent, And, Or, Implies)
  7. from sympy.matrices.expressions import MatMul
  8. # APIs here may be subject to change
  9. ### Helper functions ###
  10. def allargs(symbol, fact, expr):
  11. """
  12. Apply all arguments of the expression to the fact structure.
  13. Parameters
  14. ==========
  15. symbol : Symbol
  16. A placeholder symbol.
  17. fact : Boolean
  18. Resulting ``Boolean`` expression.
  19. expr : Expr
  20. Examples
  21. ========
  22. >>> from sympy import Q
  23. >>> from sympy.assumptions.sathandlers import allargs
  24. >>> from sympy.abc import x, y
  25. >>> allargs(x, Q.negative(x) | Q.positive(x), x*y)
  26. (Q.negative(x) | Q.positive(x)) & (Q.negative(y) | Q.positive(y))
  27. """
  28. return And(*[fact.subs(symbol, arg) for arg in expr.args])
  29. def anyarg(symbol, fact, expr):
  30. """
  31. Apply any argument of the expression to the fact structure.
  32. Parameters
  33. ==========
  34. symbol : Symbol
  35. A placeholder symbol.
  36. fact : Boolean
  37. Resulting ``Boolean`` expression.
  38. expr : Expr
  39. Examples
  40. ========
  41. >>> from sympy import Q
  42. >>> from sympy.assumptions.sathandlers import anyarg
  43. >>> from sympy.abc import x, y
  44. >>> anyarg(x, Q.negative(x) & Q.positive(x), x*y)
  45. (Q.negative(x) & Q.positive(x)) | (Q.negative(y) & Q.positive(y))
  46. """
  47. return Or(*[fact.subs(symbol, arg) for arg in expr.args])
  48. def exactlyonearg(symbol, fact, expr):
  49. """
  50. Apply exactly one argument of the expression to the fact structure.
  51. Parameters
  52. ==========
  53. symbol : Symbol
  54. A placeholder symbol.
  55. fact : Boolean
  56. Resulting ``Boolean`` expression.
  57. expr : Expr
  58. Examples
  59. ========
  60. >>> from sympy import Q
  61. >>> from sympy.assumptions.sathandlers import exactlyonearg
  62. >>> from sympy.abc import x, y
  63. >>> exactlyonearg(x, Q.positive(x), x*y)
  64. (Q.positive(x) & ~Q.positive(y)) | (Q.positive(y) & ~Q.positive(x))
  65. """
  66. pred_args = [fact.subs(symbol, arg) for arg in expr.args]
  67. res = Or(*[And(pred_args[i], *[~lit for lit in pred_args[:i] +
  68. pred_args[i+1:]]) for i in range(len(pred_args))])
  69. return res
  70. ### Fact registry ###
  71. class ClassFactRegistry:
  72. """
  73. Register handlers against classes.
  74. Explanation
  75. ===========
  76. ``register`` method registers the handler function for a class. Here,
  77. handler function should return a single fact. ``multiregister`` method
  78. registers the handler function for multiple classes. Here, handler function
  79. should return a container of multiple facts.
  80. ``registry(expr)`` returns a set of facts for *expr*.
  81. Examples
  82. ========
  83. Here, we register the facts for ``Abs``.
  84. >>> from sympy import Abs, Equivalent, Q
  85. >>> from sympy.assumptions.sathandlers import ClassFactRegistry
  86. >>> reg = ClassFactRegistry()
  87. >>> @reg.register(Abs)
  88. ... def f1(expr):
  89. ... return Q.nonnegative(expr)
  90. >>> @reg.register(Abs)
  91. ... def f2(expr):
  92. ... arg = expr.args[0]
  93. ... return Equivalent(~Q.zero(arg), ~Q.zero(expr))
  94. Calling the registry with expression returns the defined facts for the
  95. expression.
  96. >>> from sympy.abc import x
  97. >>> reg(Abs(x))
  98. {Q.nonnegative(Abs(x)), Equivalent(~Q.zero(x), ~Q.zero(Abs(x)))}
  99. Multiple facts can be registered at once by ``multiregister`` method.
  100. >>> reg2 = ClassFactRegistry()
  101. >>> @reg2.multiregister(Abs)
  102. ... def _(expr):
  103. ... arg = expr.args[0]
  104. ... return [Q.even(arg) >> Q.even(expr), Q.odd(arg) >> Q.odd(expr)]
  105. >>> reg2(Abs(x))
  106. {Implies(Q.even(x), Q.even(Abs(x))), Implies(Q.odd(x), Q.odd(Abs(x)))}
  107. """
  108. def __init__(self):
  109. self.singlefacts = defaultdict(frozenset)
  110. self.multifacts = defaultdict(frozenset)
  111. def register(self, cls):
  112. def _(func):
  113. self.singlefacts[cls] |= {func}
  114. return func
  115. return _
  116. def multiregister(self, *classes):
  117. def _(func):
  118. for cls in classes:
  119. self.multifacts[cls] |= {func}
  120. return func
  121. return _
  122. def __getitem__(self, key):
  123. ret1 = self.singlefacts[key]
  124. for k in self.singlefacts:
  125. if issubclass(key, k):
  126. ret1 |= self.singlefacts[k]
  127. ret2 = self.multifacts[key]
  128. for k in self.multifacts:
  129. if issubclass(key, k):
  130. ret2 |= self.multifacts[k]
  131. return ret1, ret2
  132. def __call__(self, expr):
  133. ret = set()
  134. handlers1, handlers2 = self[type(expr)]
  135. for h in handlers1:
  136. ret.add(h(expr))
  137. for h in handlers2:
  138. ret.update(h(expr))
  139. return ret
  140. class_fact_registry = ClassFactRegistry()
  141. ### Class fact registration ###
  142. x = Symbol('x')
  143. ## Abs ##
  144. @class_fact_registry.multiregister(Abs)
  145. def _(expr):
  146. arg = expr.args[0]
  147. return [Q.nonnegative(expr),
  148. Equivalent(~Q.zero(arg), ~Q.zero(expr)),
  149. Q.even(arg) >> Q.even(expr),
  150. Q.odd(arg) >> Q.odd(expr),
  151. Q.integer(arg) >> Q.integer(expr),
  152. ]
  153. ### Add ##
  154. @class_fact_registry.multiregister(Add)
  155. def _(expr):
  156. return [allargs(x, Q.positive(x), expr) >> Q.positive(expr),
  157. allargs(x, Q.negative(x), expr) >> Q.negative(expr),
  158. allargs(x, Q.real(x), expr) >> Q.real(expr),
  159. allargs(x, Q.rational(x), expr) >> Q.rational(expr),
  160. allargs(x, Q.integer(x), expr) >> Q.integer(expr),
  161. exactlyonearg(x, ~Q.integer(x), expr) >> ~Q.integer(expr),
  162. ]
  163. @class_fact_registry.register(Add)
  164. def _(expr):
  165. allargs_real = allargs(x, Q.real(x), expr)
  166. onearg_irrational = exactlyonearg(x, Q.irrational(x), expr)
  167. return Implies(allargs_real, Implies(onearg_irrational, Q.irrational(expr)))
  168. ### Mul ###
  169. @class_fact_registry.multiregister(Mul)
  170. def _(expr):
  171. return [Equivalent(Q.zero(expr), anyarg(x, Q.zero(x), expr)),
  172. allargs(x, Q.positive(x), expr) >> Q.positive(expr),
  173. allargs(x, Q.real(x), expr) >> Q.real(expr),
  174. allargs(x, Q.rational(x), expr) >> Q.rational(expr),
  175. allargs(x, Q.integer(x), expr) >> Q.integer(expr),
  176. exactlyonearg(x, ~Q.rational(x), expr) >> ~Q.integer(expr),
  177. allargs(x, Q.commutative(x), expr) >> Q.commutative(expr),
  178. ]
  179. @class_fact_registry.register(Mul)
  180. def _(expr):
  181. # Implicitly assumes Mul has more than one arg
  182. # Would be allargs(x, Q.prime(x) | Q.composite(x)) except 1 is composite
  183. # More advanced prime assumptions will require inequalities, as 1 provides
  184. # a corner case.
  185. allargs_prime = allargs(x, Q.prime(x), expr)
  186. return Implies(allargs_prime, ~Q.prime(expr))
  187. @class_fact_registry.register(Mul)
  188. def _(expr):
  189. # General Case: Odd number of imaginary args implies mul is imaginary(To be implemented)
  190. allargs_imag_or_real = allargs(x, Q.imaginary(x) | Q.real(x), expr)
  191. onearg_imaginary = exactlyonearg(x, Q.imaginary(x), expr)
  192. return Implies(allargs_imag_or_real, Implies(onearg_imaginary, Q.imaginary(expr)))
  193. @class_fact_registry.register(Mul)
  194. def _(expr):
  195. allargs_real = allargs(x, Q.real(x), expr)
  196. onearg_irrational = exactlyonearg(x, Q.irrational(x), expr)
  197. return Implies(allargs_real, Implies(onearg_irrational, Q.irrational(expr)))
  198. @class_fact_registry.register(Mul)
  199. def _(expr):
  200. # Including the integer qualification means we don't need to add any facts
  201. # for odd, since the assumptions already know that every integer is
  202. # exactly one of even or odd.
  203. allargs_integer = allargs(x, Q.integer(x), expr)
  204. anyarg_even = anyarg(x, Q.even(x), expr)
  205. return Implies(allargs_integer, Equivalent(anyarg_even, Q.even(expr)))
  206. ### MatMul ###
  207. @class_fact_registry.register(MatMul)
  208. def _(expr):
  209. allargs_square = allargs(x, Q.square(x), expr)
  210. allargs_invertible = allargs(x, Q.invertible(x), expr)
  211. return Implies(allargs_square, Equivalent(Q.invertible(expr), allargs_invertible))
  212. ### Pow ###
  213. @class_fact_registry.multiregister(Pow)
  214. def _(expr):
  215. base, exp = expr.base, expr.exp
  216. return [
  217. (Q.real(base) & Q.even(exp) & Q.nonnegative(exp)) >> Q.nonnegative(expr),
  218. (Q.nonnegative(base) & Q.odd(exp) & Q.nonnegative(exp)) >> Q.nonnegative(expr),
  219. (Q.nonpositive(base) & Q.odd(exp) & Q.nonnegative(exp)) >> Q.nonpositive(expr),
  220. Equivalent(Q.zero(expr), Q.zero(base) & Q.positive(exp))
  221. ]
  222. ### Numbers ###
  223. _old_assump_getters = {
  224. Q.positive: lambda o: o.is_positive,
  225. Q.zero: lambda o: o.is_zero,
  226. Q.negative: lambda o: o.is_negative,
  227. Q.rational: lambda o: o.is_rational,
  228. Q.irrational: lambda o: o.is_irrational,
  229. Q.even: lambda o: o.is_even,
  230. Q.odd: lambda o: o.is_odd,
  231. Q.imaginary: lambda o: o.is_imaginary,
  232. Q.prime: lambda o: o.is_prime,
  233. Q.composite: lambda o: o.is_composite,
  234. }
  235. @class_fact_registry.multiregister(Number, NumberSymbol, ImaginaryUnit)
  236. def _(expr):
  237. ret = []
  238. for p, getter in _old_assump_getters.items():
  239. pred = p(expr)
  240. prop = getter(expr)
  241. if prop is not None:
  242. ret.append(Equivalent(pred, prop))
  243. return ret