Add scaling_exp argument to series

Author: pearu

functional_algorithms/context.py

@@ -254,7 +254,6 @@ def constant(self, value, like_expr=UNSPECIFIED):
                 like_expr = self.default_like
         if like_expr is int:
             like_expr = self.symbol("_integer_value", int)
-        print(f"LIKE: {like_expr}")
         return make_constant(self, value, like_expr)
 
     def call(self, func, args):
@@ -524,5 +523,10 @@ def downcast(self, x):
     def is_finite(self, x):
         return Expr(self, "is_finite", (x,))
 
-    def series(self, unit_index, *terms):
-        return make_series(self, unit_index, terms)
+    def series(self, *terms, **params):
+        unit_index = params.get("unit_index", 0)
+        scaling_exp = params.get("scaling_exp", 0)
+        return make_series(self, unit_index, scaling_exp, terms)
+
+    def _series(self, terms, params):
+        return self.series(*terms, **params)

functional_algorithms/expr.py

@@ -121,8 +121,8 @@ def make_apply(context, name, args, result):
     return Expr(context, "apply", (name, *args, result))
 
 
-def make_series(context, order, terms):
-    return Expr(context, "series", (order, *terms))
+def make_series(context, unit_index, scaling_exp, terms):
+    return Expr(context, "series", ((unit_index, scaling_exp), *terms))
 
 
 def normalize(context, operands):
@@ -219,8 +219,8 @@ def make_ref(expr):
             ]
             if all_operands_have_ref_name:
                 # for readability
-                i = expr.operands[0]
-                lst = [expr.kind] + [f"minus{i}" if i < 0 else str(i)] + list(map(make_ref, expr.operands[1:]))
+                params = [f"minus{p}" if p < 0 else str(p) for p in expr.operands[0]]
+                lst = [expr.kind] + params + list(map(make_ref, expr.operands[1:]))
                 ref = "_".join(lst)
             else:
                 ref = f"{expr.kind}_{expr.intkey}"
@@ -360,7 +360,9 @@ def __new__(cls, context, kind, operands):
                     kind = "constant"
 
             if kind == "series":
-                assert isinstance(operands[0], int), type(operands[0])
+                assert isinstance(operands[0], tuple) and len(operands[0]) == 2, type(operands[0])
+                assert isinstance(operands[0][0], int)
+                assert isinstance(operands[0][1], int)
                 assert False not in [isinstance(operand, Expr) for operand in operands[1:]], operands
             else:
                 assert False not in [isinstance(operand, Expr) for operand in operands], operands

functional_algorithms/rewrite.py

@@ -480,7 +480,7 @@ def _binary_op(self, expr, op):
                 return expr.context.constant(r, xlike)
 
     @staticmethod
-    def _add_terms(index_terms1, index_terms2, op=None):
+    def _add_terms(params_terms1, params_terms2, op=None):
         """\
 
         Let `^` indicate the location of unity, then
@@ -517,11 +517,12 @@ def op(x, y):
                 return x + y
 
         swapped = False
-        if index_terms1[0] < index_terms2[0]:
-            index_terms1, index_terms2 = index_terms2, index_terms1
+        if params_terms1[0][0] < params_terms2[0][0]:
+            params_terms1, params_terms2 = params_terms2, params_terms1
             swapped = True
-        index1, terms1 = index_terms1[0], index_terms1[1:]
-        index2, terms2 = index_terms2[0], index_terms2[1:]
+        (index1, sexp1), terms1 = params_terms1[0], params_terms1[1:]
+        (index2, sexp2), terms2 = params_terms2[0], params_terms2[1:]
+        assert sexp1 == sexp2, (sexp1, sexp2)
 
         terms = []
 
@@ -546,10 +547,10 @@ def op(x, y):
             else:
                 terms.append(0)
 
-        return (index1, *terms)
+        return (terms, dict(unit_index=index1, scaling_exp=sexp1))
 
     @staticmethod
-    def _subtract_terms(index_terms1, index_terms2):
+    def _subtract_terms(params_terms1, params_terms2):
 
         def op(x, y):
             if x is None:
@@ -558,10 +559,10 @@ def op(x, y):
                 return x
             return x - y
 
-        return Rewriter._add_terms(index_terms1, index_terms2, op=op)
+        return Rewriter._add_terms(params_terms1, params_terms2, op=op)
 
     @staticmethod
-    def _multiply_terms(index_terms1, index_terms2, op=None):
+    def _multiply_terms(params_terms1, params_terms2, op=None):
         """\
 
         Let `^` indicate the location of unity, then
@@ -585,8 +586,9 @@ def _multiply_terms(index_terms1, index_terms2, op=None):
             def op(x, y):
                 return x * y
 
-        index1, terms1 = index_terms1[0], index_terms1[1:]
-        index2, terms2 = index_terms2[0], index_terms2[1:]
+        (index1, sexp1), terms1 = params_terms1[0], params_terms1[1:]
+        (index2, sexp2), terms2 = params_terms2[0], params_terms2[1:]
+        assert sexp1 == sexp2, (sexp1, sexp2)
 
         terms = []
         for n in range(len(terms1) + len(terms2) - 1):
@@ -600,7 +602,7 @@ def op(x, y):
                             xy += op(x, y)
             assert xy is not None
             terms.append(xy)
-        return (index1 + index2, *terms)
+        return terms, dict(unit_index=index1 + index2, scaling_exp=sexp1)
 
     def add(self, expr):
         result = self._binary_op(expr, lambda x, y: x + y)
@@ -616,10 +618,10 @@ def add(self, expr):
 
         if x.kind == "series":
             if y.kind == "series":
-                return expr.context.series(*self._add_terms(x.operands, y.operands))
-            return expr.context.series(*self._add_terms(x.operands, (0, y)))
+                return expr.context._series(*self._add_terms(x.operands, y.operands))
+            return expr.context._series(*self._add_terms(x.operands, ((0, x.operands[0][1]), y)))
         elif y.kind == "series":
-            return expr.context.series(*self._add_terms((0, x), y.operands))
+            return expr.context._series(*self._add_terms(((0, y.operands[0][1]), x), y.operands))
 
     def subtract(self, expr):
         result = self._binary_op(expr, lambda x, y: x - y)
@@ -636,10 +638,10 @@ def subtract(self, expr):
 
         if x.kind == "series":
             if y.kind == "series":
-                return expr.context.series(*self._subtract_terms(x.operands, y.operands))
-            return expr.context.series(*self._subtract_terms(x.operands, (0, y)))
+                return expr.context._series(*self._subtract_terms(x.operands, y.operands))
+            return expr.context._series(*self._subtract_terms(x.operands, ((0, x.operands[0][1]), y)))
         elif y.kind == "series":
-            return expr.context.series(*self._subtract_terms((0, x), y.operands))
+            return expr.context._series(*self._subtract_terms(((0, y.operands[0][1]), x), y.operands))
 
     def multiply(self, expr):
         result = self._binary_op(expr, lambda x, y: x * y)
@@ -656,10 +658,10 @@ def multiply(self, expr):
 
         if x.kind == "series":
             if y.kind == "series":
-                return expr.context.series(*self._multiply_terms(x.operands, y.operands))
-            return expr.context.series(*self._multiply_terms(x.operands, (0, y)))
+                return expr.context._series(*self._multiply_terms(x.operands, y.operands))
+            return expr.context._series(*self._multiply_terms(x.operands, ((0, x.operands[0][1]), y)))
         elif y.kind == "series":
-            return expr.context.series(*self._multiply_terms((0, x), y.operands))
+            return expr.context._series(*self._multiply_terms(((0, y.operands[0][1]), x), y.operands))
 
     def minimum(self, expr):
         return self._binary_op(expr, lambda x, y: min(x, y))
@@ -679,7 +681,7 @@ def divide(self, expr):
             def op(x, y):
                 return x / y
 
-            return expr.context.series(*self._multiply_terms(x.operands, (0, y), op=op))
+            return expr.context._series(*self._multiply_terms(x.operands, ((0, 0), y), op=op))
 
     def complex(self, expr):
         pass
@@ -805,7 +807,7 @@ def negative(self, expr):
             return x.operands[0]
 
         if x.kind == "series":
-            return expr.context.series(x.operands[0], *(-x_ for x_ in x.operands[1:]))
+            return _expr.make_series(expr.context, *x.operands[0], tuple(-x_ for x_ in x.operands[1:]))
 
     def conjugate(self, expr):
 
@@ -826,7 +828,7 @@ def conjugate(self, expr):
             return x
 
         if x.kind == "series":
-            return expr.context.series(x.operands[0], *(expr.context.conjugate(x_) for x_ in x.operands[1:]))
+            return _expr.make_series(expr.context, *x.operands[0], tuple(expr.context.conjugate(x_) for x_ in x.operands[1:]))
 
     def real(self, expr):
 
@@ -839,7 +841,7 @@ def real(self, expr):
             return x.operands[0]
 
         if x.kind == "series":
-            return expr.context.series(x.operands[0], *(expr.context.real(x_) for x_ in x.operands[1:]))
+            return _expr.make_series(expr.context, *x.operands[0], tuple(expr.context.real(x_) for x_ in x.operands[1:]))
 
     def imag(self, expr):
 
@@ -852,7 +854,7 @@ def imag(self, expr):
             return x.operands[1]
 
         if x.kind == "series":
-            return expr.context.series(x.operands[0], *(expr.context.imag(x_) for x_ in x.operands[1:]))
+            return _expr.make_series(expr.context, *x.operands[0], tuple(expr.context.imag(x_) for x_ in x.operands[1:]))
 
     def _compare(self, expr, relop, relop_index, swap_relop_index):
         x, y = expr.operands
@@ -974,7 +976,7 @@ def square(self, expr):
                 return self._eval(like, "square", value)
 
         if x.kind == "series":
-            return expr.context.series(*self._multiply_terms(x.operands, x.operands))
+            return expr.context._series(*self._multiply_terms(x.operands, x.operands))
 
     def pow(self, expr):
         base, exp = expr.operands
@@ -1091,10 +1093,18 @@ class ReplaceSeries:
     def __rewrite_modifier__(self, expr):
         if expr.kind == "series":
             s = None
-            for t in reversed(expr.operands[1:]):
-                if s is None:
-                    s = t
+            index, sexp = expr.operands[0]
+            for i, t in enumerate(reversed(expr.operands[1:])):
+                i = len(expr.operands[1:]) - 1 - i + index
+                if sexp == 0 or i == 0:
+                    if s is None:
+                        s = t
+                    else:
+                        s += t
                 else:
-                    s += t
+                    if s is None:
+                        s = t * expr.context.constant(2 ** (-i * sexp), t)
+                    else:
+                        s += t * expr.context.constant(2 ** (-i * sexp), t)
             return s
         return expr

functional_algorithms/tests/test_expr.py

@@ -760,61 +760,61 @@ def test_series():
     z = ctx.symbol("z")
     w = ctx.symbol("w")
 
-    s1 = ctx.series(0, x, y)
-    s2 = ctx.series(0, z, w)
-    s3 = ctx.series(1, z, w)
-    s4 = ctx.series(-1, z, w)
-    s5 = ctx.series(-3, z, w)
+    s1 = ctx.series(x, y, unit_index=0)
+    s2 = ctx.series(z, w, unit_index=0)
+    s3 = ctx.series(z, w, unit_index=1)
+    s4 = ctx.series(z, w, unit_index=-1)
+    s5 = ctx.series(z, w, unit_index=-3)
 
     def rewrite(expr):
         return expr.rewrite(fa.rewrite)
 
     assert_equal(s1 + s2, ctx.add(s1, s2))
-    assert_equal(rewrite(s1 + s2), ctx.series(0, x + z, y + w))
-    assert_equal(rewrite(s1 + s3), ctx.series(1, z, x + w, y))
-    assert_equal(rewrite(s3 + s1), ctx.series(1, z, w + x, y))
-    assert_equal(rewrite(s1 + s4), ctx.series(0, x, y + z, w))
-    assert_equal(rewrite(s4 + s1), ctx.series(0, x, z + y, w))
-    assert_equal(rewrite(s1 + s5), ctx.series(0, x, y, 0, z, w))
-    assert_equal(rewrite(s5 + s1), ctx.series(0, x, y, 0, z, w))
-    assert_equal(rewrite(s4 + s5), ctx.series(-1, z, w, z, w))
+    assert_equal(rewrite(s1 + s2), ctx.series(x + z, y + w, unit_index=0))
+    assert_equal(rewrite(s1 + s3), ctx.series(z, x + w, y, unit_index=1))
+    assert_equal(rewrite(s3 + s1), ctx.series(z, w + x, y, unit_index=1))
+    assert_equal(rewrite(s1 + s4), ctx.series(x, y + z, w, unit_index=0))
+    assert_equal(rewrite(s4 + s1), ctx.series(x, z + y, w, unit_index=0))
+    assert_equal(rewrite(s1 + s5), ctx.series(x, y, 0, z, w, unit_index=0))
+    assert_equal(rewrite(s5 + s1), ctx.series(x, y, 0, z, w, unit_index=0))
+    assert_equal(rewrite(s4 + s5), ctx.series(z, w, z, w, unit_index=-1))
 
     assert_equal(s1 - s2, ctx.subtract(s1, s2))
-    assert_equal(rewrite(s1 - s2), ctx.series(0, x - z, y - w))
-    assert_equal(rewrite(s1 - s3), ctx.series(1, -z, x - w, y))
-    assert_equal(rewrite(s3 - s1), ctx.series(1, z, w - x, -y))
-    assert_equal(rewrite(s1 - s4), ctx.series(0, x, y - z, -w))
-    assert_equal(rewrite(s4 - s1), ctx.series(0, -x, z - y, w))
-    assert_equal(rewrite(s1 - s5), ctx.series(0, x, y, 0, -z, -w))
-    assert_equal(rewrite(s5 - s1), ctx.series(0, -x, -y, 0, z, w))
-    assert_equal(rewrite(s4 - s5), ctx.series(-1, z, w, -z, -w))
-    assert_equal(rewrite(s5 - s4), ctx.series(-1, -z, -w, z, w))
-
-    assert_equal(rewrite(s1 + z), ctx.series(0, x + z, y))
-    assert_equal(rewrite(s1 - z), ctx.series(0, x - z, y))
-    assert_equal(rewrite(z + s1), ctx.series(0, z + x, y))
-    assert_equal(rewrite(z - s1), ctx.series(0, z - x, -y))
-    assert_equal(rewrite(s2 + x), ctx.series(0, z + x, w))
-    assert_equal(rewrite(s3 + x), ctx.series(1, z, w + x))
-    assert_equal(rewrite(s4 + x), ctx.series(0, x, z, w))
-    assert_equal(rewrite(s5 + x), ctx.series(0, x, 0, 0, z, w))
+    assert_equal(rewrite(s1 - s2), ctx.series(x - z, y - w, unit_index=0))
+    assert_equal(rewrite(s1 - s3), ctx.series(-z, x - w, y, unit_index=1))
+    assert_equal(rewrite(s3 - s1), ctx.series(z, w - x, -y, unit_index=1))
+    assert_equal(rewrite(s1 - s4), ctx.series(x, y - z, -w, unit_index=0))
+    assert_equal(rewrite(s4 - s1), ctx.series(-x, z - y, w, unit_index=0))
+    assert_equal(rewrite(s1 - s5), ctx.series(x, y, 0, -z, -w, unit_index=0))
+    assert_equal(rewrite(s5 - s1), ctx.series(-x, -y, 0, z, w, unit_index=0))
+    assert_equal(rewrite(s4 - s5), ctx.series(z, w, -z, -w, unit_index=-1))
+    assert_equal(rewrite(s5 - s4), ctx.series(-z, -w, z, w, unit_index=-1))
+
+    assert_equal(rewrite(s1 + z), ctx.series(x + z, y, unit_index=0))
+    assert_equal(rewrite(s1 - z), ctx.series(x - z, y, unit_index=0))
+    assert_equal(rewrite(z + s1), ctx.series(z + x, y, unit_index=0))
+    assert_equal(rewrite(z - s1), ctx.series(z - x, -y, unit_index=0))
+    assert_equal(rewrite(s2 + x), ctx.series(z + x, w, unit_index=0))
+    assert_equal(rewrite(s3 + x), ctx.series(z, w + x, unit_index=1))
+    assert_equal(rewrite(s4 + x), ctx.series(x, z, w, unit_index=0))
+    assert_equal(rewrite(s5 + x), ctx.series(x, 0, 0, z, w, unit_index=0))
 
     assert_equal(s1 * s2, ctx.multiply(s1, s2))
-    assert_equal(rewrite(s1 * s2), ctx.series(0, x * z, x * w + y * z, y * w))
-    assert_equal(rewrite(s2 * s2), ctx.series(0, z * z, z * w + w * z, w * w))
-    assert_equal(rewrite(s3 * s3), ctx.series(2, z * z, z * w + w * z, w * w))
-    assert_equal(rewrite(s4 * s4), ctx.series(-2, z * z, z * w + w * z, w * w))
-    assert_equal(rewrite(s4 * s5), ctx.series(-4, z * z, z * w + w * z, w * w))
+    assert_equal(rewrite(s1 * s2), ctx.series(x * z, x * w + y * z, y * w, unit_index=0))
+    assert_equal(rewrite(s2 * s2), ctx.series(z * z, z * w + w * z, w * w, unit_index=0))
+    assert_equal(rewrite(s3 * s3), ctx.series(z * z, z * w + w * z, w * w, unit_index=2))
+    assert_equal(rewrite(s4 * s4), ctx.series(z * z, z * w + w * z, w * w, unit_index=-2))
+    assert_equal(rewrite(s4 * s5), ctx.series(z * z, z * w + w * z, w * w, unit_index=-4))
 
-    assert_equal(rewrite(s3**2), ctx.series(2, z * z, z * w + w * z, w * w))
+    assert_equal(rewrite(s3**2), ctx.series(z * z, z * w + w * z, w * w, unit_index=2))
 
-    assert_equal(rewrite(s1 * z), ctx.series(0, x * z, y * z))
-    assert_equal(rewrite(z * s1), ctx.series(0, z * x, z * y))
-    assert_equal(rewrite(s2 * x), ctx.series(0, z * x, w * x))
-    assert_equal(rewrite(s3 * x), ctx.series(1, z * x, w * x))
-    assert_equal(rewrite(s5 * x), ctx.series(-3, z * x, w * x))
+    assert_equal(rewrite(s1 * z), ctx.series(x * z, y * z, unit_index=0))
+    assert_equal(rewrite(z * s1), ctx.series(z * x, z * y, unit_index=0))
+    assert_equal(rewrite(s2 * x), ctx.series(z * x, w * x, unit_index=0))
+    assert_equal(rewrite(s3 * x), ctx.series(z * x, w * x, unit_index=1))
+    assert_equal(rewrite(s5 * x), ctx.series(z * x, w * x, unit_index=-3))
 
-    assert_equal(rewrite(s1 / z), ctx.series(0, x / z, y / z))
+    assert_equal(rewrite(s1 / z), ctx.series(x / z, y / z, unit_index=0))
 
-    assert_equal(rewrite(-s1), ctx.series(0, -x, -y))
-    assert_equal(rewrite(-s3), ctx.series(1, -z, -w))
+    assert_equal(rewrite(-s1), ctx.series(-x, -y, unit_index=0))
+    assert_equal(rewrite(-s3), ctx.series(-z, -w, unit_index=1))