Changes to prod() from discussion and initial implementation of infin…

…ite sums.

src/sage/data_structures/stream.py

@@ -3244,15 +3244,15 @@ def is_nonzero(self):
         """
         return self._series.is_nonzero()
 
-class Stream_infinite_product(Stream):
+
+class Stream_infinite_operator(Stream):
     r"""
-    Stream defined by an infinite product.
+    Stream defined by applying an operator an infinite number of times.
 
-    The ``iterator`` returns elements either `1 + p_i` to compute the
-    product `\prod_{i \in I} (1 + p_i)`. The valuation of `p_i` is weakly
-    increasing as we iterate over `I` and there are only finitely many
-    terms with any fixed valuation. In particular, this *assumes* the
-    product is nonzero.
+    The ``iterator`` returns elements `s_i` to compute an infinite operator.
+    The valuation of `s_i` is weakly increasing as we iterate over `I` and
+    there are only finitely many terms with any fixed valuation with
+    `s_i \neq 0`. In particular, this *assumes* the result is nonzero.
 
     .. WARNING::
 
@@ -3261,20 +3261,18 @@ class Stream_infinite_product(Stream):
     INPUT:
 
     - ``iterator`` -- the iterator for the factors
-    - ``constant_order`` -- the order corresponding to the constant term
     """
-    def __init__(self, iterator, constant_order):
+    def __init__(self, iterator):
         """
         Initialize ``self``.
 
         EXAMPLES::
 
-            sage: from sage.data_structures.stream import Stream_infinite_product
+            sage: from sage.data_structures.stream import Stream_infinite_sum
         """
-        self._prod_iter = iterator
+        self._op_iter = iterator
         self._cur = None
         self._cur_order = -infinity
-        self._constant_order = constant_order
         super().__init__(False)
 
     @lazy_attribute
@@ -3300,24 +3298,27 @@ def _advance(self):
 
         EXAMPLES::
 
-            sage: from sage.data_structures.stream import Stream_infinite_product
+            sage: from sage.data_structures.stream import Stream_infinite_sum
         """
         if self._cur is None:
-            temp = next(self._prod_iter)
-            self._cur = temp
+            temp = next(self._op_iter)
+            self.initial(temp)
             self._cur_order = temp._coeff_stream._approximate_order
         order = self._cur_order
         while order == self._cur_order:
             try:
-                next_factor = next(self._prod_iter)
+                next_factor = next(self._op_iter)
             except StopIteration:
                 self._cur_order = infinity
                 return
-            self._cur *= next_factor
+            coeff_stream = next_factor._coeff_stream
+            while coeff_stream._approximate_order < order:
+                # This check also updates the next_factor._approximate_order
+                if coeff_stream[coeff_stream._approximate_order]:
+                    order = coeff_stream._approximate_order
+                    raise ValueError(f"invalid product computation with invalid order {order} < {self._cur_order}")
+            self.apply_operator(next_factor)
             # nonzero checks are safer than equality checks (i.e. in lazy series)
-            if order == self._constant_order and not (next_factor._coeff_stream[order] - 1):
-                order += 1
-                break
             while not next_factor._coeff_stream[order]:
                 order += 1
         self._cur_order = order
@@ -3328,7 +3329,7 @@ def __getitem__(self, n):
 
         EXAMPLES::
 
-            sage: from sage.data_structures.stream import Stream_infinite_product
+            sage: from sage.data_structures.stream import Stream_infinite_sum
         """
         while n >= self._cur_order:
             self._advance()
@@ -3341,7 +3342,7 @@ def order(self):
 
         EXAMPLES::
 
-            sage: from sage.data_structures.stream import Stream_infinite_product
+            sage: from sage.data_structures.stream import Stream_infinite_sum
         """
         return self._approximate_order
 
@@ -3351,9 +3352,9 @@ def __hash__(self):
 
         EXAMPLES::
 
-            sage: from sage.data_structures.stream import Stream_infinite_product
+            sage: from sage.data_structures.stream import Stream_infinite_sum
         """
-        return hash((type(self), self._ring, self._prod_iter))
+        return hash((type(self), self._ring, self._op_iter))
 
     def __ne__(self, other):
         """
@@ -3365,7 +3366,7 @@ def __ne__(self, other):
 
         EXAMPLES::
 
-            sage: from sage.data_structures.stream import Stream_infinite_product
+            sage: from sage.data_structures.stream import Stream_infinite_sum
         """
         if not (isinstance(other, type(self)) and other._ring == self._ring):
             return True
@@ -3381,6 +3382,54 @@ def is_nonzero(self):
 
         EXAMPLES::
 
-            sage: from sage.data_structures.stream import Stream_infinite_product
+            sage: from sage.data_structures.stream import Stream_infinite_sum
         """
         return True
+
+class Stream_infinite_sum(Stream_infinite_operator):
+    r"""
+    Stream defined by an infinite sum.
+
+    The ``iterator`` returns elements `s_i` to compute the product
+    `\sum_{i \in I} s_i`. See :class:`Stream_infinite_operator`
+    for restrictions on the `s_i`.
+
+    INPUT:
+
+    - ``iterator`` -- the iterator for the factors
+    """
+    def initial(self, obj):
+        """
+        Set the initial data.
+        """
+        self._cur = obj
+
+    def apply_operator(self, next_obj):
+        """
+        Apply the operator to ``next_obj``.
+        """
+        self._cur += next_obj
+
+class Stream_infinite_product(Stream_infinite_operator):
+    r"""
+    Stream defined by an infinite product.
+
+    The ``iterator`` returns elements `p_i` to compute the product
+    `\prod_{i \in I} (1 + p_i)`. See :class:`Stream_infinite_operator`
+    for restrictions on the `p_i`.
+
+    INPUT:
+
+    - ``iterator`` -- the iterator for the factors
+    """
+    def initial(self, obj):
+        """
+        Set the initial data.
+        """
+        self._cur = obj + 1
+
+    def apply_operator(self, next_obj):
+        """
+        Apply the operator to ``next_obj``.
+        """
+        self._cur = self._cur + self._cur * next_obj

src/sage/rings/lazy_series_ring.py

@@ -55,6 +55,7 @@
 from sage.misc.lazy_attribute import lazy_attribute
 
 from sage.rings.integer_ring import ZZ
+from sage.rings.infinity import infinity
 from sage.rings.polynomial.laurent_polynomial_ring import LaurentPolynomialRing
 from sage.rings.polynomial.polynomial_ring_constructor import PolynomialRing
 from sage.rings.lazy_series import (LazyModuleElement,
@@ -890,28 +891,40 @@ def is_exact(self):
         """
         return self.base_ring().is_exact()
 
-    def prod(self, args, index_set=None):
+    def prod(self, f, a=None, b=infinity, add_one=False):
         r"""
         The product of elements of ``self``.
 
         INPUT:
 
-        - ``args`` -- a list (or iterable) of elements of ``self`` or when
-          ``index_set`` is given, a function that returns elements of ``self``
-        - ``index_set`` -- (optional) an indexing set for the product or
-          or ``True``
+        - ``f`` -- a list (or iterable) of elements of ``self``
+        - ``a``, ``b`` -- optional arguments
+        - ``add_one`` -- (default: ``False``); if ``True``, then converts a
+          lazy series `p_i` from ``args`` into `1 + p_i` for the product
 
-        If ``index_set`` is an iterable, then ``args`` should be a function
-        that takes in an index and returns an element `1 + p_i` of ``self`` to
-        compute the product `\prod_{i \in I} (1 + p_i)`. The valuation of `p_i`
-        is weakly increasing as we iterate over `I` and there are only
-        finitely many terms with any fixed valuation. If ``index_set=True``,
-        then this will treat ``args`` as an iterator for a potentially
-        infinite product. In particular, this assumes the product is nonzero.
+        If ``a`` and ``b`` are both integers, then this returns the product
+        `\prod_{i=a}^b f(i)`, where `f(i) = p_i` if ``add_one=False`` or
+        `f(i) = 1 + p_i` otherwise. If ``b`` is not specified, then we consider
+        `b = \infty`. Note this corresponds to the Python `range(a, b+1)`.
+
+        If `a` is any other iterable, then this returns the product
+        `\prod_{i \in a} f(i)`, where `f(i) = p_i` if ``add_one=False`` or
+        `f(i) = 1 + p_i`.
+
+        .. NOTE::
+
+            For infinite products, it is faster to use ``add_one=True`` since
+            the implementation is based on `p_i` in `\prod_i (1 + p_i)`.
+
+        .. WARNING::
+
+            When ``f`` is an infinite generator, then the first argument
+            ``a`` must be ``True``. Otherwise this will loop forever.
 
         .. WARNING::
 
-            If invalid input is given, this may loop forever.
+            For an *infinite* product of the form `\prod_i (1 + p_i)`,
+            if `p_i = 0`, then this will loop forever.
 
         EXAMPLES::
 
@@ -923,30 +936,135 @@ def prod(self, args, index_set=None):
             1 + t + 2*t^2 + 3*t^3 + 5*t^4 + 7*t^5 + 11*t^6 + O(t^7)
             sage: euler - L.euler()
             O(t^7)
+            sage: L.prod(lambda n: -t^n, 1, add_one=True)
+            1 - t - t^2 + t^5 + O(t^7)
 
-            sage: L.prod((1 - t^n for n in PositiveIntegers()), index_set=True)
+            sage: L.prod((1 - t^n for n in PositiveIntegers()), True)
+            1 - t - t^2 + t^5 + O(t^7)
+            sage: L.prod((-t^n for n in PositiveIntegers()), True, add_one=True)
             1 - t - t^2 + t^5 + O(t^7)
 
-            sage: L.prod((1 + t^(n-3) for n in PositiveIntegers()), index_set=True)
+            sage: L.prod((1 + t^(n-3) for n in PositiveIntegers()), True)
             2*t^-3 + 4*t^-2 + 4*t^-1 + 4 + 6*t + 10*t^2 + 16*t^3 + O(t^4)
 
+            sage: L.prod(lambda n: 2 + t^n, -3, 5)
+            96*t^-6 + 240*t^-5 + 336*t^-4 + 840*t^-3 + 984*t^-2 + 1248*t^-1
+             + 1980 + 1668*t + 1824*t^2 + 1872*t^3 + 1782*t^4 + 1710*t^5
+             + 1314*t^6 + 1122*t^7 + 858*t^8 + 711*t^9 + 438*t^10 + 282*t^11
+             + 210*t^12 + 84*t^13 + 60*t^14 + 24*t^15
+            sage: L.prod(lambda n: t^n / (1 + abs(n)), -2, 2, add_one=True)
+            1/3*t^-3 + 5/6*t^-2 + 13/9*t^-1 + 25/9 + 13/9*t + 5/6*t^2 + 1/3*t^3
+            sage: L.prod(lambda n: t^-2 + t^n / n, -4, -2)
+            1/24*t^-9 - 1/8*t^-8 - 1/6*t^-7 + 1/2*t^-6
+
             sage: D = LazyDirichletSeriesRing(QQ, "s")
-            sage: ret = D.prod(lambda p: (1+D(1, valuation=p)).inverse(), index_set=Primes())
-            sage: ret
+            sage: D.prod(lambda p: (1+D(1, valuation=p)).inverse(), Primes())
             1 - 1/(2^s) - 1/(3^s) + 1/(4^s) - 1/(5^s) + 1/(6^s) - 1/(7^s) + O(1/(8^s))
-        """
-        if index_set is None:
-            return super().prod(args)
-        if index_set is True:
-            it = args
+
+            sage: D.prod(lambda p: D(1, valuation=p), Primes(), add_one=True)
+            1 + 1/(2^s) + 1/(3^s) + 1/(5^s) + 1/(6^s) + 1/(7^s) + O(1/(8^s))
+        """
+        if a is None:
+            if add_one:
+                return super().prod(self.one() + g for g in f)
+            return super().prod(f)
+
+        if a is True:
+            it = f
+        elif a in ZZ:
+            if b != infinity:
+                if add_one:
+                    return super().prod(1 + f(i) for i in range(a, b+1))
+                return super().prod(f(i) for i in range(a, b+1))
+            from sage.sets.non_negative_integers import NonNegativeIntegers
+            it = (f(i+a) for i in NonNegativeIntegers())
         else:
-            it = (args(i) for i in index_set)
+            it = (f(i) for i in a)
+
+        # NOTE: We must have a new variable name for each new iterator
+        if not add_one:
+            data = (g - 1 for g in it)
+        else:
+            data = it
+
         from sage.data_structures.stream import Stream_infinite_product
-        if self._minimal_valuation is not None and self._minimal_valuation > 0:
-            # Only for Dirichlet series (currently)
-            coeff_stream = Stream_infinite_product(it, self._minimal_valuation)
+        coeff_stream = Stream_infinite_product(data)
+        return self.element_class(self, coeff_stream)
+
+    def sum(self, f, a=None, b=infinity):
+        r"""
+        The sum of elements of ``self``.
+
+        INPUT:
+
+        - ``f`` -- a list (or iterable or function) of elements of ``self``
+        - ``a``, ``b`` -- optional arguments
+
+        If ``a`` and ``b`` are both integers, then this returns the sum
+        `\sum_{i=a}^b f(i)`. If ``b`` is not specified, then we consider
+        `b = \infty`. Note this corresponds to the Python `range(a, b+1)`.
+
+        If `a` is any other iterable, then this returns the sum
+        `\sum{i \in a} f(i)`.
+
+        .. WARNING::
+
+            When ``f`` is an infinite generator, then the first argument
+            ``a`` must be ``True``. Otherwise this will loop forever.
+
+        .. WARNING::
+
+            For an *infinite* sum of the form `\sum_i s_i`,
+            if `s_i = 0`, then this will loop forever.
+
+        EXAMPLES::
+
+            sage: L.<t> = LazyLaurentSeriesRing(QQ)
+            sage: L.sum(lambda n: t^n / (n+1), PositiveIntegers())
+            1/2*t + 1/3*t^2 + 1/4*t^3 + 1/5*t^4 + 1/6*t^5 + 1/7*t^6 + 1/8*t^7 + O(t^8)
+
+        We verify the Rogers-Ramanujan identities up to degree 100::
+
+            sage: L.<q> = LazyPowerSeriesRing(QQ)
+            sage: Gpi = L.prod(lambda k: -q^(1+5*k), 0, oo, add_one=True)
+            sage: Gpi *= L.prod(lambda k: -q^(4+5*k), 0, oo, add_one=True)
+            sage: Gp = 1 / Gpi
+            sage: G = L.sum(lambda n: q^(n^2) / prod(1 - q^(k+1) for k in range(n)), 0, oo)
+            sage: G - Gp
+            O(q^7)
+            sage: all(G[k] == Gp[k] for k in range(100))
+            True
+
+            sage: Hpi = L.prod(lambda k: -q^(2+5*k), 0, oo, add_one=True)
+            sage: Hpi *= L.prod(lambda k: -q^(3+5*k), 0, oo, add_one=True)
+            sage: Hp = 1 / Hpi
+            sage: H = L.sum(lambda n: q^(n^2+n) / prod(1 - q^(k+1) for k in range(n)), 0, oo)
+            sage: H - Hp
+            O(q^7)
+            sage: all(H[k] == Hp[k] for k in range(100))
+            True
+
+        ::
+
+            sage: D = LazyDirichletSeriesRing(QQ, "s")
+            sage: D.sum(lambda p: D(1, valuation=p), Primes())
+            1/(2^s) + 1/(3^s) + 1/(5^s) + 1/(7^s) + O(1/(9^s))
+        """
+        if a is None:
+            return super().sum(f)
+
+        if a is True:
+            it = f
+        elif a in ZZ:
+            if b != infinity:
+                return super().sum(f(i) for i in range(a, b+1))
+            from sage.sets.non_negative_integers import NonNegativeIntegers
+            it = (f(i+a) for i in NonNegativeIntegers())
         else:
-            coeff_stream = Stream_infinite_product(it, 0)
+            it = (f(i) for i in a)
+
+        from sage.data_structures.stream import Stream_infinite_sum
+        coeff_stream = Stream_infinite_sum(it)
         return self.element_class(self, coeff_stream)
 
     def _test_invert(self, **options):