You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

324 lines
9.2 KiB

5 months ago
from collections import defaultdict
from sympy.assumptions.ask import Q
from sympy.core import (Add, Mul, Pow, Number, NumberSymbol, Symbol)
from sympy.core.numbers import ImaginaryUnit
from sympy.functions.elementary.complexes import Abs
from sympy.logic.boolalg import (Equivalent, And, Or, Implies)
from sympy.matrices.expressions import MatMul
# APIs here may be subject to change
### Helper functions ###
def allargs(symbol, fact, expr):
"""
Apply all arguments of the expression to the fact structure.
Parameters
==========
symbol : Symbol
A placeholder symbol.
fact : Boolean
Resulting ``Boolean`` expression.
expr : Expr
Examples
========
>>> from sympy import Q
>>> from sympy.assumptions.sathandlers import allargs
>>> from sympy.abc import x, y
>>> allargs(x, Q.negative(x) | Q.positive(x), x*y)
(Q.negative(x) | Q.positive(x)) & (Q.negative(y) | Q.positive(y))
"""
return And(*[fact.subs(symbol, arg) for arg in expr.args])
def anyarg(symbol, fact, expr):
"""
Apply any argument of the expression to the fact structure.
Parameters
==========
symbol : Symbol
A placeholder symbol.
fact : Boolean
Resulting ``Boolean`` expression.
expr : Expr
Examples
========
>>> from sympy import Q
>>> from sympy.assumptions.sathandlers import anyarg
>>> from sympy.abc import x, y
>>> anyarg(x, Q.negative(x) & Q.positive(x), x*y)
(Q.negative(x) & Q.positive(x)) | (Q.negative(y) & Q.positive(y))
"""
return Or(*[fact.subs(symbol, arg) for arg in expr.args])
def exactlyonearg(symbol, fact, expr):
"""
Apply exactly one argument of the expression to the fact structure.
Parameters
==========
symbol : Symbol
A placeholder symbol.
fact : Boolean
Resulting ``Boolean`` expression.
expr : Expr
Examples
========
>>> from sympy import Q
>>> from sympy.assumptions.sathandlers import exactlyonearg
>>> from sympy.abc import x, y
>>> exactlyonearg(x, Q.positive(x), x*y)
(Q.positive(x) & ~Q.positive(y)) | (Q.positive(y) & ~Q.positive(x))
"""
pred_args = [fact.subs(symbol, arg) for arg in expr.args]
res = Or(*[And(pred_args[i], *[~lit for lit in pred_args[:i] +
pred_args[i+1:]]) for i in range(len(pred_args))])
return res
### Fact registry ###
class ClassFactRegistry:
"""
Register handlers against classes.
Explanation
===========
``register`` method registers the handler function for a class. Here,
handler function should return a single fact. ``multiregister`` method
registers the handler function for multiple classes. Here, handler function
should return a container of multiple facts.
``registry(expr)`` returns a set of facts for *expr*.
Examples
========
Here, we register the facts for ``Abs``.
>>> from sympy import Abs, Equivalent, Q
>>> from sympy.assumptions.sathandlers import ClassFactRegistry
>>> reg = ClassFactRegistry()
>>> @reg.register(Abs)
... def f1(expr):
... return Q.nonnegative(expr)
>>> @reg.register(Abs)
... def f2(expr):
... arg = expr.args[0]
... return Equivalent(~Q.zero(arg), ~Q.zero(expr))
Calling the registry with expression returns the defined facts for the
expression.
>>> from sympy.abc import x
>>> reg(Abs(x))
{Q.nonnegative(Abs(x)), Equivalent(~Q.zero(x), ~Q.zero(Abs(x)))}
Multiple facts can be registered at once by ``multiregister`` method.
>>> reg2 = ClassFactRegistry()
>>> @reg2.multiregister(Abs)
... def _(expr):
... arg = expr.args[0]
... return [Q.even(arg) >> Q.even(expr), Q.odd(arg) >> Q.odd(expr)]
>>> reg2(Abs(x))
{Implies(Q.even(x), Q.even(Abs(x))), Implies(Q.odd(x), Q.odd(Abs(x)))}
"""
def __init__(self):
self.singlefacts = defaultdict(frozenset)
self.multifacts = defaultdict(frozenset)
def register(self, cls):
def _(func):
self.singlefacts[cls] |= {func}
return func
return _
def multiregister(self, *classes):
def _(func):
for cls in classes:
self.multifacts[cls] |= {func}
return func
return _
def __getitem__(self, key):
ret1 = self.singlefacts[key]
for k in self.singlefacts:
if issubclass(key, k):
ret1 |= self.singlefacts[k]
ret2 = self.multifacts[key]
for k in self.multifacts:
if issubclass(key, k):
ret2 |= self.multifacts[k]
return ret1, ret2
def __call__(self, expr):
ret = set()
handlers1, handlers2 = self[type(expr)]
for h in handlers1:
ret.add(h(expr))
for h in handlers2:
ret.update(h(expr))
return ret
class_fact_registry = ClassFactRegistry()
### Class fact registration ###
x = Symbol('x')
## Abs ##
@class_fact_registry.multiregister(Abs)
def _(expr):
arg = expr.args[0]
return [Q.nonnegative(expr),
Equivalent(~Q.zero(arg), ~Q.zero(expr)),
Q.even(arg) >> Q.even(expr),
Q.odd(arg) >> Q.odd(expr),
Q.integer(arg) >> Q.integer(expr),
]
### Add ##
@class_fact_registry.multiregister(Add)
def _(expr):
return [allargs(x, Q.positive(x), expr) >> Q.positive(expr),
allargs(x, Q.negative(x), expr) >> Q.negative(expr),
allargs(x, Q.real(x), expr) >> Q.real(expr),
allargs(x, Q.rational(x), expr) >> Q.rational(expr),
allargs(x, Q.integer(x), expr) >> Q.integer(expr),
exactlyonearg(x, ~Q.integer(x), expr) >> ~Q.integer(expr),
]
@class_fact_registry.register(Add)
def _(expr):
allargs_real = allargs(x, Q.real(x), expr)
onearg_irrational = exactlyonearg(x, Q.irrational(x), expr)
return Implies(allargs_real, Implies(onearg_irrational, Q.irrational(expr)))
### Mul ###
@class_fact_registry.multiregister(Mul)
def _(expr):
return [Equivalent(Q.zero(expr), anyarg(x, Q.zero(x), expr)),
allargs(x, Q.positive(x), expr) >> Q.positive(expr),
allargs(x, Q.real(x), expr) >> Q.real(expr),
allargs(x, Q.rational(x), expr) >> Q.rational(expr),
allargs(x, Q.integer(x), expr) >> Q.integer(expr),
exactlyonearg(x, ~Q.rational(x), expr) >> ~Q.integer(expr),
allargs(x, Q.commutative(x), expr) >> Q.commutative(expr),
]
@class_fact_registry.register(Mul)
def _(expr):
# Implicitly assumes Mul has more than one arg
# Would be allargs(x, Q.prime(x) | Q.composite(x)) except 1 is composite
# More advanced prime assumptions will require inequalities, as 1 provides
# a corner case.
allargs_prime = allargs(x, Q.prime(x), expr)
return Implies(allargs_prime, ~Q.prime(expr))
@class_fact_registry.register(Mul)
def _(expr):
# General Case: Odd number of imaginary args implies mul is imaginary(To be implemented)
allargs_imag_or_real = allargs(x, Q.imaginary(x) | Q.real(x), expr)
onearg_imaginary = exactlyonearg(x, Q.imaginary(x), expr)
return Implies(allargs_imag_or_real, Implies(onearg_imaginary, Q.imaginary(expr)))
@class_fact_registry.register(Mul)
def _(expr):
allargs_real = allargs(x, Q.real(x), expr)
onearg_irrational = exactlyonearg(x, Q.irrational(x), expr)
return Implies(allargs_real, Implies(onearg_irrational, Q.irrational(expr)))
@class_fact_registry.register(Mul)
def _(expr):
# Including the integer qualification means we don't need to add any facts
# for odd, since the assumptions already know that every integer is
# exactly one of even or odd.
allargs_integer = allargs(x, Q.integer(x), expr)
anyarg_even = anyarg(x, Q.even(x), expr)
return Implies(allargs_integer, Equivalent(anyarg_even, Q.even(expr)))
### MatMul ###
@class_fact_registry.register(MatMul)
def _(expr):
allargs_square = allargs(x, Q.square(x), expr)
allargs_invertible = allargs(x, Q.invertible(x), expr)
return Implies(allargs_square, Equivalent(Q.invertible(expr), allargs_invertible))
### Pow ###
@class_fact_registry.multiregister(Pow)
def _(expr):
base, exp = expr.base, expr.exp
return [
(Q.real(base) & Q.even(exp) & Q.nonnegative(exp)) >> Q.nonnegative(expr),
(Q.nonnegative(base) & Q.odd(exp) & Q.nonnegative(exp)) >> Q.nonnegative(expr),
(Q.nonpositive(base) & Q.odd(exp) & Q.nonnegative(exp)) >> Q.nonpositive(expr),
Equivalent(Q.zero(expr), Q.zero(base) & Q.positive(exp))
]
### Numbers ###
_old_assump_getters = {
Q.positive: lambda o: o.is_positive,
Q.zero: lambda o: o.is_zero,
Q.negative: lambda o: o.is_negative,
Q.rational: lambda o: o.is_rational,
Q.irrational: lambda o: o.is_irrational,
Q.even: lambda o: o.is_even,
Q.odd: lambda o: o.is_odd,
Q.imaginary: lambda o: o.is_imaginary,
Q.prime: lambda o: o.is_prime,
Q.composite: lambda o: o.is_composite,
}
@class_fact_registry.multiregister(Number, NumberSymbol, ImaginaryUnit)
def _(expr):
ret = []
for p, getter in _old_assump_getters.items():
pred = p(expr)
prop = getter(expr)
if prop is not None:
ret.append(Equivalent(pred, prop))
return ret