Speeding up evaluation of power series, Python I want to plot (on the complex plane, with mpmath) the power series whose general term is the number of partitions of an integer. 

The more points I want (and the larger the degree is), the longer the plot takes.

I first wrote a naive version, added memoization of the sequence, changed the naive evaluation of the polynomial with the Horner algorithm. All these modifications led to speed ups, and I want to know if there are more optimizations I could do.

import mpmath

MAX_DEGREE = 250
NB_POINTS = 40000


def memoize(function):
    memo = {}

    def wrapper(*args):
        if args in memo:
            return memo[args]
        else:
            rv = function(*args)
            memo[args] = rv
            return rv

    return wrapper


@memoize
def p_k(k, n):
    if k == 1 or k == n:
        return 1
    if k &gt; n:
        return 0

    return p_k(k - 1, n - 1) + p_k(k, n - k)


@memoize
def p(n):
    res = 1
    for k in range(1, n):
        res = res + p_k(k, n)
    return res


def mgf(x):
    if abs(x) &gt;= 1:
        return 0
    res = 0
    x_i = x
    for i in range(1, MAX_DEGREE):
        res = res + p(i) * x_i
        x_i = x_i * x
    return res


def mgf_horner(x):
    if abs(x) &gt;= 1:
        return 0
    res = p(MAX_DEGREE)
    for i in reversed(range(1, MAX_DEGREE)):
        res = res * x + p(i)
    return res * x


if __name__ == "__main__":
    assert(p(1) == 1)
    assert(p(2) == 2)
    assert(p(5) == 7)
    assert(p(7) == 15)
    assert(round(mgf_horner(0.3), 8) == round(mgf(0.3), 8))
    assert(round(mgf_horner(0.5), 8) == round(mgf(0.5), 8))

    print("Tests passed!")

    mpmath.cplot(mgf_horner, points=NB_POINTS, re=[-1, 1], im=[-1, 1])
I think most savings are already achieved by the memoization. If you deconstruct the algorithms, your mgf and mgf_horner don't seem to do a lot different. The main difference is the direction of the summation, which affects the float rounding errors.

some small changes:


instead of reversed(range(1, MAX_DEGREE)) you can use range(MAX_DEGREE - 1, 0, -1)
you can use the built-in sum or math.fsum if you rewrite the formula's a bit


something like this:

@lru_cache(maxsize=MAX_DEGREE)
def p2(n):
    return sum((p_k(k, n) for k in range(1, n)), 1)

def mgf_fsum(x):
    if abs(x) &gt;= 1:
        return 0
    return math.fsum(p2(i) * x ** i for i in range(1, MAX_DEGREE))

def mgf_horner_fsum(x):
    if abs(x) &gt;= 1:
        return 0
    return math.fsum(p2(i) * x ** i for i in range(MAX_DEGREE, 0, -1))

def mgf_sum(x):
    if abs(x) &gt;= 1:
        return 0
    return sum(p2(i) * x ** i for i in range(1, MAX_DEGREE))

def mgf_horner_sum(x):
    if abs(x) &gt;= 1:
        return 0
    return sum(p2(i) * x ** i for i in range(MAX_DEGREE, 0, -1))


but this rounding has an effect on the results:

steps = 15
for x in range(1, steps, 1):
    print(x/steps, mgf(x/steps), mgf(x/steps)/mgf(x/steps), mgf_sum(x/steps)/mgf(x/steps), mgf_fsum(x/steps)/mgf(x/steps), mgf_horner(x/steps)/mgf(x/steps), mgf_horner_sum(x/steps)/mgf(x/steps), mgf_horner_fsum(x/steps)/mgf(x/steps))



x mgf_value mgf mgf_sum mgf_fsum mgf_horner mgf_horner_sum mgf_horner_fsum
0.06666666666666667 0.07655349092586666 1.0 1.0 1.0 1.0 1.0 1.0
0.13333333333333333 0.17795095671821595 1.0 1.0 1.0 1.0 1.0 1.0
0.2 0.3152135557353453 1.0 1.0 0.9999999999999998 0.9999999999999998 0.9999999999999998 0.9999999999999998
0.26666666666666666 0.5067809825345801 1.0 1.0 1.0000000000000007 1.0000000000000007 1.0000000000000007 1.0000000000000007
0.3333333333333333 0.7853123419985341 1.0 0.9999999999999997 1.0 0.9999999999999999 1.0000000000000002 1.0
0.4 1.2130721483139313 1.0 1.0 1.0 0.9999999999999998 1.0 1.0
0.4666666666666667 1.9206394911829312 1.0 1.0 1.0000000000000004 1.0000000000000004 1.0000000000000004 1.0000000000000004
0.5333333333333333 3.2187202179248864 1.0 1.0 1.0000000000000004 1.000000000000099 1.000000000000099 1.000000000000099
0.6 5.9870712939827015 1.0 1.0 0.9999999999999999 1.000000000201781 1.0000000002017808 1.000000000201781
0.6666666666666666 13.435823754250293 1.0 0.9999999999999999 1.0000000000000002 1.0000001435010422 1.0000001435010424 1.0000001435010424
0.7333333333333333 43.087235122138615 1.0 1.0000000000000002 1.0 1.0000353393606511 1.0000353393606507 1.000035339360651
0.8 288.2244586544328 1.0 1.0000000000000004 1.0 1.0023338650736258 1.002333865073626 1.002333865073626
0.8666666666666667 6672.576615237916 1.0 0.9999999999999998 0.9999999999999994 1.0273438993019854 1.0273438993019859 1.0273438993019859
0.9333333333333333 379389.61410647386 1.0 0.9999999999999997 0.9999999999999999 1.0860986641621089 1.086098664162109 1.086098664162109



for these tests I limited MAX_DEGREE to 70. 

As you can see, the 3 methods with increasing iteration give a different result than the 3 methods with decreasing iterations, especially with values closer to 1

Once the p(x) are calculated and in memory, this goes pretty fast, and no noticeable difference between the algorithms