Skip to content

Commit

Permalink
gh-35338: improve performance of stream cauchy invert
Browse files Browse the repository at this point in the history
    
<!-- ^^^^^
Please provide a concise, informative and self-explanatory title.
Don't put issue numbers in there, do this in the PR body below.
For example, instead of "Fixes #1234" use "Introduce new method to
calculate 1+1"
-->
### 📚 Description
Replace loops by sums with lazy assignments and provide a dedicated
`get_coefficient` for `Stream_cauchy_invert`

Fixes #34553

### 📝 Checklist

<!-- Put an `x` in all the boxes that apply. -->
<!-- If your change requires a documentation PR, please link it
appropriately -->
<!-- If you're unsure about any of these, don't hesitate to ask. We're
here to help! -->

- [X] I have made sure that the title is self-explanatory and the
description concisely explains the PR.
- [X] I have linked an issue or discussion.
- [X] I have created tests covering the changes.
- [X] I have updated the documentation accordingly.

### ⌛ Dependencies
<!-- List all open pull requests that this PR logically depends on -->
<!--
- #xyz: short description why this is a dependency
- #abc: ...
-->
    
URL: #35338
Reported by: Martin Rubey
Reviewer(s): Matthias Köppe, Travis Scrimshaw
  • Loading branch information
Release Manager committed Aug 31, 2023
2 parents 6695bec + 1d52328 commit d0759e8
Showing 1 changed file with 84 additions and 69 deletions.
153 changes: 84 additions & 69 deletions src/sage/data_structures/stream.py
Original file line number Diff line number Diff line change
Expand Up @@ -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'])

Expand Down Expand Up @@ -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"""
Expand Down Expand Up @@ -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"""
Expand Down Expand Up @@ -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


Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -1996,7 +1991,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)

Expand Down Expand Up @@ -2063,23 +2058,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::
Expand Down Expand Up @@ -2133,19 +2126,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::
Expand All @@ -2172,24 +2177,31 @@ def stretched_power_restrict_degree(self, i, m, d):
....: if sum(mu.size() for mu in m) == 12})
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()


#####################################################################
Expand Down Expand Up @@ -2313,7 +2325,7 @@ class Stream_rmul(Stream_scalar):
INPUT:
- ``series`` -- a :class:`Stream`
- ``scalar`` -- a scalar
- ``scalar`` -- a non-zero scalar
EXAMPLES::
Expand Down Expand Up @@ -2355,7 +2367,7 @@ class Stream_lmul(Stream_scalar):
INPUT:
- ``series`` -- a :class:`Stream`
- ``scalar`` -- a scalar
- ``scalar`` -- a non-zero scalar
EXAMPLES::
Expand Down Expand Up @@ -2491,7 +2503,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::
Expand Down Expand Up @@ -2810,7 +2823,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):
"""
Expand Down Expand Up @@ -2847,7 +2860,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):
Expand Down Expand Up @@ -3183,7 +3197,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):
Expand Down Expand Up @@ -3225,7 +3239,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):
Expand Down

0 comments on commit d0759e8

Please sign in to comment.