improve performance of stream cauchy invert

src/sage/data_structures/stream.py

@@ -102,6 +102,7 @@
 from sage.misc.lazy_import import lazy_import
 from sage.combinat.integer_vector_weighted import iterator_fast as wt_int_vec_iter
 from sage.categories.hopf_algebras_with_basis import HopfAlgebrasWithBasis
+from sage.misc.cachefunc import cached_method
 
 lazy_import('sage.combinat.sf.sfa', ['_variables_recursive', '_raise_variables'])
 
@@ -1545,13 +1546,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"""
@@ -1644,15 +1642,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"""
@@ -1754,12 +1747,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
 
 
@@ -1860,19 +1851,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):
@@ -1989,7 +1984,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)
 
@@ -2054,23 +2049,21 @@ def get_coefficient(self, n):
         if not n:  # special case of 0
             if self._right[0]:
                 assert self._degree_f is not None, "the plethysm with a lazy symmetric function of valuation 0 is defined only for symmetric functions of finite support"
+                K = self._degree_f
+            else:
+                K = 1
+        else:
+            K = n + 1
 
-                return sum((c * self.compute_product(n, la)
-                            for k in range(self._left._approximate_order, self._degree_f)
-                            if self._left[k]
-                            for la, c in self._left[k]),
-                           self._basis.zero())
-
-        res = sum((c * self.compute_product(n, la)
-                   for k in range(self._left._approximate_order, n+1)
-                   if self._left[k]
-                   for la, c in self._left[k]),
-                  self._basis.zero())
-        return res
+        return sum((c * self.compute_product(n, la)
+                    for k in range(self._left._approximate_order, K)
+                    if self._left[k] # necessary, because it might be int(0)
+                    for la, c in self._left[k]),
+                   self._basis.zero())
 
     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::
 
@@ -2118,19 +2111,31 @@ def compute_product(self, n, la):
         wgt.reverse()
         exp.reverse()
         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 not f:
+                    break
+                lf.append(f)
+            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``.
+
+        INPUT:
+
+        - ``i``, ``m`` -- positive integers
+        - ``d`` -- integer
 
         EXAMPLES::
 
@@ -2153,24 +2158,31 @@ 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:
+            # _raise_variables(c, i, self._degree_one) cannot vanish
+            # because i is positive and c is non-zero
             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))}
+                terms = {mon.stretch(i):
+                         _raise_variables(c, i, self._degree_one)
+                         for mon, c in power_d}
             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)
+                terms = {tuple((mu.stretch(i) for mu in mon)):
+                         _raise_variables(c, i, self._degree_one)
+                         for mon, c in power_d}
+            return self._basis(self._p.element_class(self._p, terms))
 
-        return self._p.zero()
+        return self._basis.zero()
 
 
 #####################################################################
@@ -2294,7 +2306,7 @@ class Stream_rmul(Stream_scalar):
     INPUT:
 
     - ``series`` -- a :class:`Stream`
-    - ``scalar`` -- a scalar
+    - ``scalar`` -- a non-zero scalar
 
     EXAMPLES::
 
@@ -2335,7 +2347,7 @@ class Stream_lmul(Stream_scalar):
     INPUT:
 
     - ``series`` -- a :class:`Stream`
-    - ``scalar`` -- a scalar
+    - ``scalar`` -- a non-zero scalar
 
     EXAMPLES::
 
@@ -2470,7 +2482,8 @@ class Stream_cauchy_invert(Stream_unary):
     - ``approximate_order`` -- ``None``, or a lower bound on the
       order of the resulting stream
 
-    Instances of this class are always dense.
+    Instances of this class are always dense, because of mathematical
+    necessities.
 
     EXAMPLES::
 
@@ -2789,7 +2802,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):
         """
@@ -2826,7 +2839,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):
@@ -3162,7 +3176,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):
@@ -3204,7 +3218,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):