Arithmetic refactor

mathics/builtin/arithfns/basic.py

@@ -7,7 +7,10 @@
 
 """
 
-from mathics.builtin.arithmetic import create_infix
+import sympy
+
+from mathics.builtin.arithmetic import _MPMathFunction, create_infix
+from mathics.builtin.base import BinaryOperator, Builtin, PrefixOperator, SympyFunction
 from mathics.core.atoms import (
     Complex,
     Integer,
@@ -45,7 +48,6 @@
     Symbol,
     SymbolDivide,
     SymbolHoldForm,
-    SymbolNull,
     SymbolPower,
     SymbolTimes,
 )
@@ -56,10 +58,17 @@
     SymbolInfix,
     SymbolLeft,
     SymbolMinus,
+    SymbolOverflow,
     SymbolPattern,
-    SymbolSequence,
 )
-from mathics.eval.arithmetic import eval_Plus, eval_Times
+from mathics.eval.arithmetic import (
+    associate_powers,
+    eval_Exponential,
+    eval_Plus,
+    eval_Power_inexact,
+    eval_Power_number,
+    eval_Times,
+)
 from mathics.eval.nevaluator import eval_N
 from mathics.eval.numerify import numerify
 
@@ -439,6 +448,8 @@ class Power(BinaryOperator, MPMathFunction):
     rules = {
         "Power[]": "1",
         "Power[x_]": "x",
+        "Power[I,-1]": "-I",
+        "Power[-1, 1/2]": "I",
     }
 
     summary_text = "exponentiate"
@@ -447,15 +458,15 @@ class Power(BinaryOperator, MPMathFunction):
     # Remember to up sympy doc link when this is corrected
     sympy_name = "Pow"
 
+    def eval_exp(self, x, evaluation):
+        "Power[E, x]"
+        return eval_Exponential(x)
+
     def eval_check(self, x, y, evaluation):
         "Power[x_, y_]"
-
-        # Power uses MPMathFunction but does some error checking first
-        if isinstance(x, Number) and x.is_zero:
-            if isinstance(y, Number):
-                y_err = y
-            else:
-                y_err = eval_N(y, evaluation)
+        # if x is zero
+        if x.is_zero:
+            y_err = y if isinstance(y, Number) else eval_N(y, evaluation)
             if isinstance(y_err, Number):
                 py_y = y_err.round_to_float(permit_complex=True).real
                 if py_y > 0:
@@ -469,17 +480,47 @@ def eval_check(self, x, y, evaluation):
                     evaluation.message(
                         "Power", "infy", Expression(SymbolPower, x, y_err)
                     )
-                    return SymbolComplexInfinity
-        if isinstance(x, Complex) and x.real.is_zero:
-            yhalf = Expression(SymbolTimes, y, RationalOneHalf)
-            factor = self.eval(Expression(SymbolSequence, x.imag, y), evaluation)
-            return Expression(
-                SymbolTimes, factor, Expression(SymbolPower, IntegerM1, yhalf)
-            )
-
-        result = self.eval(Expression(SymbolSequence, x, y), evaluation)
-        if result is None or result != SymbolNull:
-            return result
+                return SymbolComplexInfinity
+
+        # If x and y are inexact numbers, use the numerical function
+
+        if x.is_inexact() and y.is_inexact():
+            try:
+                return eval_Power_inexact(x, y)
+            except OverflowError:
+                evaluation.message("General", "ovfl")
+                return Expression(SymbolOverflow)
+
+        # Tries to associate powers a^b^c-> a^(b*c)
+        assoc = associate_powers(x, y)
+        if not assoc.has_form("Power", 2):
+            return assoc
+
+        assoc = numerify(assoc, evaluation)
+        x, y = assoc.elements
+        # If x and y are numbers
+        if isinstance(x, Number) and isinstance(y, Number):
+            try:
+                return eval_Power_number(x, y)
+            except OverflowError:
+                evaluation.message("General", "ovfl")
+                return Expression(SymbolOverflow)
+
+        # if x or y are inexact, leave the expression
+        # as it is:
+        if x.is_inexact() or y.is_inexact():
+            return assoc
+
+        # Finally, try to convert to sympy
+        base_sp, exp_sp = x.to_sympy(), y.to_sympy()
+        if base_sp is None or exp_sp is None:
+            # If base or exp can not be converted to sympy,
+            # returns the result of applying the associative
+            # rule.
+            return assoc
+
+        result = from_sympy(sympy.Pow(base_sp, exp_sp))
+        return result.evaluate_elements(evaluation)
 
 
 class Sqrt(SympyFunction):

mathics/eval/arithmetic.py

@@ -1,7 +1,9 @@
 # -*- coding: utf-8 -*-
 
 """
-arithmetic-related evaluation functions.
+helper functions for arithmetic evaluation, which do not
+depends on the evaluation context. Conversions to Sympy are
+used just as a last resource.
 
 Many of these do do depend on the evaluation context. Conversions to Sympy are
 used just as a last resource.
@@ -320,6 +322,28 @@ def eval_complex_sign(n: BaseElement) -> Optional[BaseElement]:
     return sign or eval_complex_sign(expr)
 
 
+def eval_Sign_number(n: Number) -> Number:
+    """
+    Evals the absolute value of a number.
+    """
+    if n.is_zero:
+        return Integer0
+    if isinstance(n, (Integer, Rational, Real)):
+        return Integer1 if n.value > 0 else IntegerM1
+    if isinstance(n, Complex):
+        abs_sq = eval_add_numbers(
+            *(eval_multiply_numbers(x, x) for x in (n.real, n.imag))
+        )
+        criteria = eval_add_numbers(abs_sq, IntegerM1)
+        if test_zero_arithmetic_expr(criteria):
+            return n
+        if n.is_inexact():
+            return eval_multiply_numbers(n, eval_Power_number(abs_sq, RealM0p5))
+        if test_zero_arithmetic_expr(criteria, numeric=True):
+            return n
+        return eval_multiply_numbers(n, eval_Power_number(abs_sq, RationalMOneHalf))
+
+
 def eval_mpmath_function(
     mpmath_function: Callable, *args: Number, prec: Optional[int] = None
 ) -> Optional[Number]:
@@ -345,6 +369,31 @@ def eval_mpmath_function(
         return call_mpmath(mpmath_function, tuple(mpmath_args), prec)
 
 
+def eval_Exponential(exp: BaseElement) -> BaseElement:
+    """
+    Eval E^exp
+    """
+    # If both base and exponent are exact quantities,
+    # use sympy.
+
+    if not exp.is_inexact():
+        exp_sp = exp.to_sympy()
+        if exp_sp is None:
+            return None
+        return from_sympy(sympy.Exp(exp_sp))
+
+    prec = exp.get_precision()
+    if prec is not None:
+        if exp.is_machine_precision():
+            number = mpmath.exp(exp.to_mpmath())
+            result = from_mpmath(number)
+            return result
+        else:
+            with mpmath.workprec(prec):
+                number = mpmath.exp(exp.to_mpmath())
+                return from_mpmath(number, prec)
+
+
 def eval_Plus(*items: BaseElement) -> BaseElement:
     "evaluate Plus for general elements"
     numbers, items_tuple = segregate_numbers_from_sorted_list(*items)
@@ -643,9 +692,112 @@ def eval_Times(*items: BaseElement) -> BaseElement:
     )
 
 
+# Here I used the convention of calling eval_* to functions that can produce a new expression, or None
+# if the result can not be evaluated, or is trivial. For example, if we call eval_Power_number(Integer2, RationalOneHalf)
+# it returns ``None`` instead of ``Expression(SymbolPower, Integer2, RationalOneHalf)``.
+# The reason is that these functions are written to be part of replacement rules, to be applied during the evaluation process.
+# In that process, a rule is considered applied if produces an expression that is different from the original one, or
+# if the replacement function returns (Python's) ``None``.
+#
+# For example, when the expression ``Power[4, 1/2]`` is evaluated, a (Builtin) rule  ``Power[base_, exp_]->eval_repl_rule(base, expr)``
+# is applied. If the rule matches, `repl_rule` is called with arguments ``(4, 1/2)`` and produces `2`. As `Integer2.sameQ(Power[4, 1/2])`
+# is False, then no new rules for `Power` are checked, and a new round of evaluation is atempted.
+#
+# On the other hand, if ``Power[3, 1/2]``, ``repl_rule`` can do two possible things: one is return ``Power[3, 1/2]``. If it does,
+# the rule is considered applied. Then, the evaluation method checks if `Power[3, 1/2].sameQ(Power[3, 1/2])`. In this case it is true,
+# and then the expression is kept as it is.
+# The other possibility is to return  (Python's) `None`. In that case, the evaluator considers that the rule failed to be applied,
+# and look for another rule associated to ``Power``. To return ``None`` produces then a faster evaluation, since no ``sameQ`` call is needed,
+# and do not prevent that other rules are attempted.
+#
+# The bad part of using ``None`` as a return is that I would expect that ``eval`` produces always a valid Expression, so if at some point of
+# the code I call ``eval_Power_number(Integer3, RationalOneHalf)`` I get ``Expression(SymbolPower, Integer3, RationalOneHalf)``.
+#
+# From my point of view, it would make more sense to use the following convention:
+#  * if the method has signature ``eval_method(...)->BaseElement:`` then use the prefix ``eval_``
+#  * if the method has the siguature ``apply_method(...)->Optional[BaseElement]`` use the prefix ``apply_`` or maybe ``repl_``.
+#
+# In any case, let's keep the current convention.
+#
+#
+
+
+def associate_powers(expr: BaseElement, power: BaseElement = Integer1) -> BaseElement:
+    """
+    base^a^b^c^...^power -> base^(a*b*c*...power)
+    provided one of the following cases
+    * `a`, `b`, ... `power` are all integer numbers
+    * `a`, `b`,... are Rational/Real number with absolute value <=1,
+      and the other powers are not integer numbers.
+    * `a` is not a Rational/Real number, and b, c, ... power are all
+      integer numbers.
+    """
+    powers = []
+    base = expr
+    if power is not Integer1:
+        powers.append(power)
+
+    while base.has_form("Power", 2):
+        previous_base, outer_power = base, power
+        base, power = base.elements
+        if len(powers) == 0:
+            if power is not Integer1:
+                powers.append(power)
+            continue
+        if power is IntegerM1:
+            powers.append(power)
+            continue
+        if isinstance(power, (Rational, Real)):
+            if abs(power.value) < 1:
+                powers.append(power)
+                continue
+        # power is not rational/real and outer_power is integer,
+        elif isinstance(outer_power, Integer):
+            if power is not Integer1:
+                powers.append(power)
+            if isinstance(power, Integer):
+                continue
+            else:
+                break
+        # in any other case, use the previous base and
+        # exit the loop
+        base = previous_base
+        break
+
+    if len(powers) == 0:
+        return base
+    elif len(powers) == 1:
+        return Expression(SymbolPower, base, powers[0])
+    result = Expression(SymbolPower, base, Expression(SymbolTimes, *powers))
+    return result
+
+
+def distribute_factor(expr: BaseElement, factor: BaseElement) -> BaseElement:
+    """
+    q * (a + b  + c) -> (q a  + q b + q c)
+    """
+    if not expr.has_form("Plus", None):
+        return expr
+    terms = (Expression(SymbolTimes, factor, term) for term in expr.elements)
+    return Expression(SymbolPlus, *terms)
+
+
+def distribute_powers(expr: BaseElement) -> BaseElement:
+    """
+    (a b c)^p -> (a^p b^p c^p)
+    """
+    if not expr.has_form("Power", 2):
+        return expr
+    base, exp = expr.elements
+    if not base.has_form("Times", None):
+        return expr
+    factors = (Expression(SymbolPower, factor, exp) for factor in base.elements)
+    return Expression(SymbolTimes, *factors)
+
+
 def eval_add_numbers(
-    *numbers: Number,
-) -> BaseElement:
+    *numbers: List[Number],
+) -> Number:
     """
     Add the elements in ``numbers``.
     """
@@ -671,6 +823,16 @@ def eval_add_numbers(
         return from_sympy(sum(item.to_sympy() for item in numbers))
 
 
+def eval_complex_conjugate(z: Number) -> Number:
+    """
+    Evaluates the complex conjugate of z.
+    """
+    if isinstance(z, Complex):
+        re, im = z.real, z.imag
+        return Complex(re, eval_negate_number(im))
+    return z
+
+
 def eval_inverse_number(n: Number) -> Number:
     """
     Eval 1/n
@@ -729,6 +891,24 @@ def eval_negate_number(n: Number) -> Number:
     return eval_multiply_numbers(IntegerM1, n)
 
 
+def flat_arithmetic_operators(expr: Expression) -> Expression:
+    """
+    operation[a_number, b, operation[c_number, d], e]-> operation[a, c, b, c, d, e]
+    """
+    # items is a dict with two keys: True and False.
+    # In True we store numeric items, and in False the symbolic ones.
+    items = {True: [], False: []}
+    head = expr.get_head()
+    for element in expr.elements:
+        # If the element is also head[elements],
+        # take its elements, and append to the main expression.
+        if element.get_head() is head:
+            for item in flat_arithmetic_operators(element).elements:
+                item[isinstance(item, Number)].append(item)
+        item[isinstance(item, Number)].append(item)
+    return Expression(head, *items[True], *items[False])
+
+
 def segregate_numbers(
     *elements: BaseElement,
 ) -> Tuple[List[Number], List[BaseElement]]: