Merge pull request #1237 from skirpichev/series-dir

Misc fixes

diofant/__init__.py

@@ -55,7 +55,7 @@
                       IntegerRing, PythonRational, QQ_gmpy, QQ_python,
                       RationalField, RealAlgebraicField, RealField, ZZ_gmpy,
                       ZZ_python, IntegerModRing)
-from .series import Limit, O, Order, limit, residue, series
+from .series import Limit, O, Order, limit, residue
 from .functions import (E1, Abs, Chi, Ci, DiracDelta, Ei, Eijk,
                         FallingFactorial, Heaviside, Id, KroneckerDelta,
                         LambertW, LeviCivita, Li, Max, Min, Piecewise,
@@ -215,7 +215,7 @@
     'ComplexAlgebraicField', 'ComplexField', 'Domain', 'ExpressionDomain',
     'FF_gmpy', 'FF_python', 'FiniteField', 'IntegerRing', 'PythonRational',
     'QQ_gmpy', 'QQ_python', 'RationalField', 'RealAlgebraicField', 'RealField',
-    'ZZ_gmpy', 'ZZ_python', 'Limit', 'O', 'Order', 'limit', 'residue', 'series',
+    'ZZ_gmpy', 'ZZ_python', 'Limit', 'O', 'Order', 'limit', 'residue',
     'E1', 'Abs', 'Chi', 'Ci', 'DiracDelta', 'Ei', 'Eijk', 'FallingFactorial',
     'Heaviside', 'Id', 'KroneckerDelta', 'LambertW', 'LeviCivita', 'Li', 'Max',
     'Min', 'Piecewise', 'RisingFactorial', 'Shi', 'Si', 'Ynm', 'Ynm_c', 'Znm',

diofant/concrete/expr_with_limits.py

@@ -184,6 +184,23 @@ def free_symbols(self):
                 isyms.update(i.free_symbols)
         return isyms
 
+    def __eq__(self, other):
+        if not isinstance(other, self.func):
+            return False
+
+        return (self.function.xreplace(self.canonical_variables) ==
+                other.function.xreplace(other.canonical_variables) and
+                [[_[0].xreplace(self.canonical_variables) if len(_) > 1 else _[0], *_[1:]]
+                 for _ in self.limits] ==
+                [[_[0].xreplace(other.canonical_variables) if len(_) > 1 else _[0], *_[1:]]
+                 for _ in other.limits])
+    __hash__ = Expr.__hash__
+
+    def _hashable_content(self):
+        return (self.function.xreplace(self.canonical_variables),
+                *((_[0].xreplace(self.canonical_variables) if len(_) > 1 else _[0], *_[1:])
+                  for _ in self.limits))
+
     @property
     def is_number(self):
         """Return True if the Sum has no free symbols, else False."""

diofant/core/basic.py

@@ -30,6 +30,7 @@ class Basic:
     is_Add = False
     is_Mul = False
     is_Pow = False
+    is_Exp = False
     is_Number = False
     is_Float = False
     is_Rational = False
@@ -1074,7 +1075,7 @@ def _eval_rewrite(self, pattern, rule, **hints):
 
         if pattern is None or isinstance(self, pattern):
             if hasattr(self, rule):
-                rewritten = getattr(self, rule)(*args)
+                rewritten = getattr(self, rule)(*args, **hints)
                 if rewritten is not None:
                     return rewritten
         return self.func(*args)

diofant/core/expr.py

@@ -2317,7 +2317,7 @@ def is_comparable(self):
     # ####### SERIES, LEADING TERM, LIMIT, ORDER METHODS ####### #
     ##############################################################
 
-    def series(self, x=None, x0=0, n=6, dir='+', logx=None):
+    def series(self, x=None, x0=0, n=6, dir=None, logx=None):
         """Series expansion of "self" around ``x = x0`` yielding either terms of
         the series one by one (the lazy series given when n=None), else
         all the terms at once when n != None.
@@ -2346,15 +2346,15 @@ def series(self, x=None, x0=0, n=6, dir='+', logx=None):
         >>> [next(term) for i in range(2)]
         [1, -x**2/2]
 
-        For ``dir=+`` (default) the series is calculated from the right and
-        for ``dir=-`` the series from the left.  For infinite ``x0`` (``oo``
+        For ``dir=-1`` (default) the series is calculated from the right and
+        for ``dir=+1`` the series from the left.  For infinite ``x0`` (``oo``
         or ``-oo``), the ``dir`` argument is determined from the direction
-        of the infinity (i.e. ``dir="-"`` for ``oo``).  For smooth functions
+        of the infinity (i.e. ``dir=+1`` for ``oo``).  For smooth functions
         this flag will not alter the results.
 
-        >>> abs(x).series(dir='+')
+        >>> abs(x).series(dir=-1)
         x
-        >>> abs(x).series(dir='-')
+        >>> abs(x).series(dir=+1)
         -x
 
         For rational expressions this method may return original expression.
@@ -2365,6 +2365,7 @@ def series(self, x=None, x0=0, n=6, dir='+', logx=None):
         """
         from .function import expand_mul
         from .symbol import Dummy, Symbol
+        from ..functions import sign
         from ..series import Order
         from ..simplify import collect
 
@@ -2385,31 +2386,31 @@ def series(self, x=None, x0=0, n=6, dir='+', logx=None):
             else:
                 return self
 
-        if len(dir) != 1 or dir not in '+-':
-            raise ValueError("Dir must be '+' or '-'")
-
-        if x0 == oo:
-            return self.aseries(x, n)
-        elif x0 == -oo:
-            return self.subs({x: -x}).aseries(x, n).subs({x: -x})
-
-        # use rep to shift origin to x0 and change sign (if dir is negative)
-        # and undo the process with rep2
-        if x0 or dir == '-':
-            if dir == '-':
-                rep = -x + x0
-                rep2 = -x
-                rep2b = x0
-            else:
-                rep = x + x0
-                rep2 = x
-                rep2b = -x0
-            s = self.subs({x: rep}).series(x, x0=0, n=n, dir='+', logx=logx)
+        x0 = sympify(x0)
+
+        if x0.is_infinite:
+            dir = sign(x0).simplify()
+        elif dir is None:
+            dir = Rational(-1)
+        else:
+            dir = sympify(dir)
+
+        if not (isinstance(dir, Expr) and dir.is_nonzero):
+            raise ValueError(f'Direction must be a nonzero Expr, got {dir}')
+
+        dir = dir/abs(dir)
+
+        if x0.is_infinite:
+            return self.subs({x: dir*x}).aseries(x, n).subs({x: x/dir})
+
+        # use rep to shift origin to x0 and change dir to 1, then undo
+        if x0 or dir != -1:
+            s = self.subs({x: x0 - dir*x}).series(x, x0=0, n=n, dir=-1, logx=logx)
             if n is None:  # lseries...
-                return (si.subs({x: rep2 + rep2b}) for si in s)  # pragma: no branch
-            return s.subs({x: rep2 + rep2b})
+                return (si.subs({x: x0/dir - x/dir}) for si in s)  # pragma: no branch
+            return s.subs({x: x0/dir - x/dir})
 
-        # from here on it's x0=0 and dir='+' handling
+        # from here on it's x0=0 and dir=-1 handling
 
         if x.is_positive is x.is_negative is None or x.is_Symbol is not True:
             # replace x with an x that has a positive assumption

diofant/core/function.py

@@ -1501,7 +1501,16 @@ def _eval_derivative(self, s):
                      for v, p in zip(self.variables, self.point)])
 
     def _eval_nseries(self, x, n, logx=None):
-        raise NotImplementedError
+        if x in self.point:
+            v = self.variables[self.point.index(x)]
+        else:
+            v = x
+        arg = self.expr.nseries(v, n=n, logx=logx)
+        rv = Add(*[self.func(a, *zip(self.variables, self.point))
+                   for a in Add.make_args(arg.removeO())])
+        if o := arg.getO():
+            rv += o.subs({v: x})
+        return rv
 
 
 def diff(f, *args, **kwargs):

diofant/core/power.py

@@ -211,7 +211,10 @@ def __new__(cls, b, e, evaluate=None):
                 obj = b._eval_power(e)
                 if obj is not None:
                     return obj
-        return Expr.__new__(cls, b, e)
+        obj = Expr.__new__(cls, b, e)
+        if b is E:
+            obj.is_Exp = True
+        return obj
 
     def _eval_is_commutative(self):
         return self.base.is_commutative and self.exp.is_commutative
@@ -410,7 +413,7 @@ def _eval_is_extended_real(self):
                 return (2*I*e/pi).is_even
 
         if b.is_extended_real is None:
-            if b.func == self.func and b.base is E and b.exp.is_imaginary:
+            if b.func == self.func and b.is_Exp and b.exp.is_imaginary:
                 return e.is_imaginary
         if e.is_extended_real is None:
             return
@@ -591,7 +594,7 @@ def _check(ct1, ct2, old):
                     new_l.append(Pow(self.base, Add(*o_al), evaluate=False))
                     return Mul(*new_l)
 
-        if old.is_Pow and old.base is E and self.exp.is_extended_real and self.base.is_positive:
+        if old.is_Exp and self.exp.is_extended_real and self.base.is_positive:
             ct1 = old.exp.as_independent(Symbol, as_Add=False)
             ct2 = (self.exp*log(self.base)).as_independent(
                 Symbol, as_Add=False)
@@ -956,7 +959,7 @@ def as_real_imag(self, deep=True, **hints):
             rp, tp = self.func(r, self.exp), t*self.exp
 
             return rp*cos(tp), rp*sin(tp)
-        elif self.base is E:
+        elif self.is_Exp:
             from ..functions import exp
             re, im = self.exp.as_real_imag()
             if deep:
@@ -1152,7 +1155,7 @@ def _eval_nseries(self, x, n, logx):
         from ..functions import arg, exp, floor, log
         from ..series import Order, limit
         from ..simplify import powsimp
-        if self.base is E:
+        if self.is_Exp:
             e_series = self.exp.nseries(x, n=n, logx=logx)
             if e_series.is_Order:
                 return 1 + e_series
@@ -1213,7 +1216,7 @@ def _eval_as_leading_term(self, x):
         from ..series import Order
         if not self.exp.has(x):
             return self.func(self.base.as_leading_term(x), self.exp)
-        elif self.base is E:
+        elif self.is_Exp:
             if self.exp.is_Mul:
                 k, arg = self.exp.as_independent(x)
             else:
@@ -1229,17 +1232,17 @@ def _eval_as_leading_term(self, x):
 
     def _eval_rewrite_as_sin(self, base, exp):
         from ..functions import sin
-        if self.base is E:
+        if self.is_Exp:
             return sin(I*self.exp + pi/2) - I*sin(I*self.exp)
 
     def _eval_rewrite_as_cos(self, base, exp):
         from ..functions import cos
-        if self.base is E:
+        if self.is_Exp:
             return cos(I*self.exp) + I*cos(I*self.exp + pi/2)
 
     def _eval_rewrite_as_tanh(self, base, exp):
         from ..functions import tanh
-        if self.base is E:
+        if self.is_Exp:
             return (1 + tanh(self.exp/2))/(1 - tanh(self.exp/2))
 
     def as_content_primitive(self, radical=False):

diofant/functions/combinatorial/factorials.py

@@ -92,7 +92,7 @@ def _eval_rewrite_as_gamma(self, n):
         from .. import gamma
         return gamma(n + 1)
 
-    def _eval_rewrite_as_tractable(self, n):
+    def _eval_rewrite_as_tractable(self, n, **kwargs):
         from .. import exp, loggamma
         return exp(loggamma(n + 1))
 
@@ -350,7 +350,7 @@ def _eval_rewrite_as_gamma(self, x, k):
         from .. import gamma
         return gamma(x + k) / gamma(x)
 
-    def _eval_rewrite_as_tractable(self, x, k):
+    def _eval_rewrite_as_tractable(self, x, k, **kwargs):
         return self._eval_rewrite_as_gamma(x, k).rewrite('tractable')
 
     def _eval_is_integer(self):
@@ -554,14 +554,14 @@ def _eval_expand_func(self, **hints):
 
         return self.func(*self.args)
 
-    def _eval_rewrite_as_factorial(self, n, k):
+    def _eval_rewrite_as_factorial(self, n, k, **kwargs):
         return factorial(n)/(factorial(k)*factorial(n - k))
 
     def _eval_rewrite_as_gamma(self, n, k):
         from .. import gamma
         return gamma(n + 1)/(gamma(k + 1)*gamma(n - k + 1))
 
-    def _eval_rewrite_as_tractable(self, n, k):
+    def _eval_rewrite_as_tractable(self, n, k, **kwargs):
         return self._eval_rewrite_as_gamma(n, k).rewrite('tractable')
 
     def _eval_is_integer(self):

diofant/functions/combinatorial/numbers.py

@@ -112,11 +112,10 @@ def eval(cls, n, sym=None):
                                      'only for positive integer indices.')
                 return cls._fibpoly(n).subs({_sym: sym})
 
-    def _eval_rewrite_as_sqrt(self, n, sym=None):
+    def _eval_rewrite_as_sqrt(self, n, sym=None, **kwargs):
         from .. import sqrt
         if sym is None:
             return (GoldenRatio**n - cos(pi*n)/GoldenRatio**n)/sqrt(5)
-
     _eval_rewrite_as_tractable = _eval_rewrite_as_sqrt
 
 
@@ -634,7 +633,7 @@ def _eval_rewrite_as_trigamma(self, n, m=1):
 
     def _eval_rewrite_as_Sum(self, n, m=None):
         from ...concrete import Sum
-        k = Dummy('k', integer=True)
+        k = Dummy('k')
         if m is None:
             m = Integer(1)
         return Sum(k**(-m), (k, 1, n))
@@ -672,7 +671,7 @@ def _eval_expand_func(self, **hints):
 
         return self
 
-    def _eval_rewrite_as_tractable(self, n, m=1):
+    def _eval_rewrite_as_tractable(self, n, m=1, **kwargs):
         from .. import polygamma
         return self.rewrite(polygamma).rewrite('tractable', deep=True)
 

diofant/functions/elementary/complexes.py

@@ -1,4 +1,4 @@
-from ...core import (Add, Derivative, Dummy, E, Eq, Expr, Function, I, Integer,
+from ...core import (Add, Derivative, Dummy, Eq, Expr, Function, I, Integer,
                      Mul, Rational, Symbol, Tuple, factor_terms, nan, oo, pi,
                      zoo)
 from ...core.function import AppliedUndef, ArgumentIndexError
@@ -526,6 +526,10 @@ def _eval_rewrite_as_Piecewise(self, arg):
     def _eval_rewrite_as_sign(self, arg):
         return arg/sign(arg)
 
+    def _eval_rewrite_as_tractable(self, arg, wrt=None, **kwargs):
+        if wrt is not None and (s := sign(arg.limit(wrt, oo))) in (1, -1):
+            return s*arg
+
 
 class arg(Function):
     """Returns the argument (in radians) of a complex number.
@@ -945,7 +949,7 @@ def _polarify(eq, lift, pause=False):
         return r
     elif eq.is_Function:
         return eq.func(*[_polarify(arg, lift, pause=False) for arg in eq.args])
-    elif eq.is_Pow and eq.base is E:
+    elif eq.is_Exp:
         return eq.func(eq.base, _polarify(eq.exp, lift, pause=False))
     elif isinstance(eq, Integral):
         # Don't lift the integration variable
@@ -1029,12 +1033,12 @@ def _unpolarify(eq, exponents_only, pause=False):
         if isinstance(eq, polar_lift):
             return _unpolarify(eq.args[0], exponents_only)
 
-    if eq.is_Pow and eq.base is not E:
+    if eq.is_Pow and not eq.is_Exp:
         expo = _unpolarify(eq.exp, exponents_only)
         base = _unpolarify(eq.base, exponents_only,
                            not (expo.is_integer and not pause))
         return base**expo
-    elif eq.is_Pow and eq.base is E:
+    elif eq.is_Exp:
         return exp(_unpolarify(eq.exp, exponents_only, exponents_only))
     elif isinstance(eq, ExprCondPair):
         return eq.func(_unpolarify(eq.expr, exponents_only, exponents_only),

diofant/functions/elementary/exponential.py

@@ -216,7 +216,7 @@ def eval(cls, arg, base=None):
                 if arg.denominator != 1:
                     return cls(arg.numerator) - cls(arg.denominator)
 
-        if arg.is_Pow and arg.base is E and arg.exp.is_extended_real:
+        if arg.is_Exp and arg.exp.is_extended_real:
             return arg.exp
         elif isinstance(arg, exp_polar):
             return unpolarify(arg.exp)