improve performance of stream cauchy invert

src/sage/data_structures/stream.py

@@ -102,7 +102,7 @@
 from sage.combinat.integer_vector_weighted import iterator_fast as wt_int_vec_iter
 from sage.combinat.sf.sfa import _variables_recursive, _raise_variables
 from sage.categories.hopf_algebras_with_basis import HopfAlgebrasWithBasis
-
+from sage.misc.cachefunc import cached_method
 
 class Stream():
     """
@@ -722,7 +722,7 @@ def __getitem__(self, n):
         if n >= self._degree:
             return self._constant
         i = n - self._approximate_order
-        if i < 0 or i >= len(self._initial_coefficients):
+        if i < 0 or i >= len(self._initial_coefficients): # TODO: i < 0 should not happen
             return ZZ.zero()
         return self._initial_coefficients[i]
 
@@ -1543,13 +1543,10 @@ def get_coefficient(self, n):
             sage: [h.get_coefficient(i) for i in range(10)]
             [0, 0, 1, 6, 20, 50, 105, 196, 336, 540]
         """
-        c = ZZ.zero()
-        for k in range(self._left._approximate_order,
-                       n - self._right._approximate_order + 1):
-            val = self._left[k]
-            if val:
-                c += val * self._right[n-k]
-        return c
+        return sum(l * self._right[n - k]
+                   for k in range(self._left._approximate_order,
+                                  n - self._right._approximate_order + 1)
+                   if (l := self._left[k]))
 
     def is_nonzero(self):
         r"""
@@ -1642,15 +1639,10 @@ def get_coefficient(self, n):
             sage: [h[i] for i in range(1, 10)]
             [1, 3, 4, 7, 6, 12, 8, 15, 13]
         """
-        c = ZZ.zero()
-        for k in divisors(n):
-            if k < self._left._approximate_order or n // k < self._right._approximate_order:
-                continue
-            val = self._left[k]
-            if val:
-                c += val * self._right[n//k]
-        return c
-
+        return sum(l * self._right[n//k] for k in divisors(n)
+                   if (k >= self._left._approximate_order
+                       and n // k >= self._right._approximate_order
+                       and (l := self._left[k])))
 
 class Stream_dirichlet_invert(Stream_unary):
     r"""
@@ -1752,12 +1744,10 @@ def get_coefficient(self, n):
         """
         if n == 1:
             return self._ainv
-        c = self._zero
-        for k in divisors(n):
-            if k < n:
-                val = self._series[n//k]
-                if val:
-                    c += self[k] * val
+        # TODO: isn't self[k] * l and l * self[k] the same here?
+        c = sum(self[k] * l for k in divisors(n)
+                if (k < n
+                    and (l := self._series[n // k])))
         return -c * self._ainv
 
 
@@ -1858,19 +1848,23 @@ def get_coefficient(self, n):
         fv = self._left._approximate_order
         gv = self._right._approximate_order
         if n < 0:
-            return sum(f_coeff_i * self._neg_powers[-i][n]
-                       for i in range(fv, n // gv + 1)
-                       if (f_coeff_i := self._left[i]))
+            return sum(l * self._neg_powers[-k][n]
+                       for k in range(fv, n // gv + 1)
+                       if (l := self._left[k]))
         # n > 0
         while len(self._pos_powers) <= n // gv:
             # TODO: possibly we always want a dense cache here?
-            self._pos_powers.append(Stream_cauchy_mul(self._pos_powers[-1], self._right, self._is_sparse))
-        ret = sum(f_coeff_i * self._neg_powers[-i][n] for i in range(fv, 0)
-                  if (f_coeff_i := self._left[i]))
-        if n == 0:
+            self._pos_powers.append(Stream_cauchy_mul(self._pos_powers[-1],
+                                                      self._right,
+                                                      self._is_sparse))
+        ret = sum(l * self._neg_powers[-k][n] for k in range(fv, 0)
+                  if (l := self._left[k]))
+
+        if not n:
             ret += self._left[0]
-        return ret + sum(f_coeff_i * self._pos_powers[i][n] for i in range(1, n // gv + 1)
-                         if (f_coeff_i := self._left[i]))
+
+        return ret + sum(l * self._pos_powers[k][n] for k in range(1, n // gv + 1)
+                         if (l := self._left[k]))
 
 
 class Stream_plethysm(Stream_binary):
@@ -1987,7 +1981,7 @@ def __init__(self, f, g, is_sparse, p, ring=None, include=None, exclude=None):
             self._basis = ring
         self._p = p
         g = Stream_map_coefficients(g, lambda x: p(x), is_sparse)
-        self._powers = [g]  # a cache for the powers of g
+        self._powers = [g]  # a cache for the powers of g in the powersum basis
         R = self._basis.base_ring()
         self._degree_one = _variables_recursive(R, include=include, exclude=exclude)
 
@@ -2056,20 +2050,22 @@ def get_coefficient(self, n):
 
                 return sum((c * self.compute_product(n, la)
                             for k in range(self._left._approximate_order, self._degree_f)
-                            if self._left[k]
+                            if self._left[k] # necessary, because it might be int(0)
                             for la, c in self._left[k]),
                            self._basis.zero())
 
+        # this also needs to be executed if not n and not
+        # self._right[0]
         res = sum((c * self.compute_product(n, la)
                    for k in range(self._left._approximate_order, n+1)
-                   if self._left[k]
+                   if self._left[k] # necessary, because it might be int(0)
                    for la, c in self._left[k]),
                   self._basis.zero())
         return res
 
     def compute_product(self, n, la):
         r"""
-        Compute the product ``c * p[la](self._right)`` in degree ``n``.
+        Compute the product ``p[la](self._right)`` in degree ``n``.
 
         EXAMPLES::
 
@@ -2119,17 +2115,28 @@ def compute_product(self, n, la):
         for k in wt_int_vec_iter(n - ret_approx_order, wgt):
             # TODO: it may make a big difference here if the
             #   approximate order would be updated.
-            # The test below is based on not removing the fixed block
-            #if any(d < self._right._approximate_order * m
-            #       for m, d in zip(exp, k)):
-            #    continue
-            ret += prod(self.stretched_power_restrict_degree(i, m, rao * m + d)
-                        for i, m, d in zip(wgt, exp, k))
+
+            # prod does not short-cut zero, therefore
+            # ret += prod(self.stretched_power_restrict_degree(i, m, rao * m + d)
+            #             for i, m, d in zip(wgt, exp, k))
+            # is expensive
+            lf = []
+            for i, m, d in zip(wgt, exp, k):
+                f = self.stretched_power_restrict_degree(i, m, rao * m + d)
+                if f:
+                    lf.append(f)
+                else:
+                    break
+            else:
+                ret += prod(lf)
+
         return ret
 
+    @cached_method
     def stretched_power_restrict_degree(self, i, m, d):
         r"""
-        Return the degree ``d*i`` part of ``p([i]*m)(g)``.
+        Return the degree ``d*i`` part of ``p([i]*m)(g)`` in
+        terms of ``self._basis``.
 
         EXAMPLES::
 
@@ -2152,24 +2159,28 @@ def stretched_power_restrict_degree(self, i, m, d):
             sage: B = p2.element_class(p2, {m: c for m, c in B if sum(mu.size() for mu in m) == 12})  # long time
             sage: A == B  # long time
             True
+
         """
+        # TODO: we should do lazy binary powering here
         while len(self._powers) < m:
             # TODO: possibly we always want a dense cache here?
-            self._powers.append(Stream_cauchy_mul(self._powers[-1], self._powers[0], self._is_sparse))
+            self._powers.append(Stream_cauchy_mul(self._powers[-1],
+                                                  self._powers[0],
+                                                  self._is_sparse))
         power_d = self._powers[m-1][d]
         # we have to check power_d for zero because it might be an
         # integer and not a symmetric function
         if power_d:
             if self._tensor_power is None:
                 terms = {mon.stretch(i): raised_c for mon, c in power_d
-                         if (raised_c := _raise_variables(c, i, self._degree_one))}
+                         if (raised_c := _raise_variables(c, i, self._degree_one))}  # TODO: could this ever be 0?
             else:
                 terms = {tuple((mu.stretch(i) for mu in mon)): raised_c
                          for mon, c in power_d
-                         if (raised_c := _raise_variables(c, i, self._degree_one))}
-            return self._p.element_class(self._p, terms)
+                         if (raised_c := _raise_variables(c, i, self._degree_one))}  # TODO: could this ever be 0?
+            return self._basis(self._p.element_class(self._p, terms))
 
-        return self._p.zero()
+        return self._basis.zero()
 
 
 #####################################################################
@@ -2293,7 +2304,7 @@ class Stream_rmul(Stream_scalar):
     INPUT:
 
     - ``series`` -- a :class:`Stream`
-    - ``scalar`` -- a scalar
+    - ``scalar`` -- a non-zero scalar
 
     EXAMPLES::
 
@@ -2334,7 +2345,7 @@ class Stream_lmul(Stream_scalar):
     INPUT:
 
     - ``series`` -- a :class:`Stream`
-    - ``scalar`` -- a scalar
+    - ``scalar`` -- a non-zero scalar
 
     EXAMPLES::
 
@@ -2541,6 +2552,46 @@ def _ainv(self):
         except TypeError:
             return self._series[v].inverse_of_unit()
 
+    def get_coefficient(self, n):
+        """
+        Return the ``n``-th coefficient of ``self``.
+
+        INPUT:
+
+        - ``n`` -- integer; the degree for the coefficient
+
+        EXAMPLES::
+
+            sage: from sage.data_structures.stream import (Stream_cauchy_invert, Stream_function)
+            sage: f = Stream_function(lambda n: n, True, 1)
+            sage: g = Stream_cauchy_invert(f)
+            sage: g.get_coefficient(5)
+            0
+            sage: [g.get_coefficient(i) for i in range(10)]
+            [-2, 1, 0, 0, 0, 0, 0, 0, 0, 0]
+        """
+        if not self._series._true_order:
+            self._ainv  # this computes the true order of ``self``
+        v = self._approximate_order
+        if n < v:
+            return ZZ.zero()
+        if n == v:
+            return self._ainv
+
+        # if self._series is exact with self._series._constant == 0
+        # then self._series[n - v - k] == 0 unless
+        # k > n - v - self._series._degree
+        # in general, we have self._series[n - v - k] == 0 unless
+        # n - v - self._series._approximate_order < k
+        if isinstance(self._series, Stream_exact) and not self._series._constant:
+            a = max(v, n - v - self._series._degree + 1)
+        else:
+            a = v
+        b = min(n, n - v - self._series._approximate_order)
+        c = sum(l * self._series[n - v - k] for k in range(a, b)
+                if (l := self[k]))
+        return -c * self._ainv
+
     def iterate_coefficients(self):
         """
         A generator for the coefficients of ``self``.
@@ -2788,7 +2839,7 @@ def __getitem__(self, n):
             sage: [M[i] for i in range(6)]
             [0, 0, 0, 1, 2, 3]
         """
-        return self._series[n-self._shift]
+        return self._series[n - self._shift]
 
     def __hash__(self):
         """
@@ -2825,7 +2876,8 @@ def __eq__(self, other):
             sage: M2 == Stream_shift(F, 2)
             True
         """
-        return (isinstance(other, type(self)) and self._shift == other._shift
+        return (isinstance(other, type(self))
+                and self._shift == other._shift
                 and self._series == other._series)
 
     def is_nonzero(self):
@@ -3161,7 +3213,7 @@ def __getitem__(self, n):
             sage: [f2[i] for i in range(-1, 4)]
             [0, 2, 6, 12, 20]
         """
-        return (prod(n+k for k in range(1, self._shift + 1))
+        return (prod(n + k for k in range(1, self._shift + 1))
                 * self._series[n + self._shift])
 
     def __hash__(self):
@@ -3203,7 +3255,8 @@ def __eq__(self, other):
             sage: f == Stream_derivative(a, 1, True)
             True
         """
-        return (isinstance(other, type(self)) and self._shift == other._shift
+        return (isinstance(other, type(self))
+                and self._shift == other._shift
                 and self._series == other._series)
 
     def is_nonzero(self):