# Aitken's delta squared convergence acceleration
# Pramode, "Python generator tricks"
# http://linuxgazette.net/100/pramode.html

import itertools

def aitken(seq):
    """ 
    apply Aitken's method to sequence seq with terms {s_i}
     s_i' = s_i - (s_i+1 - s_i)^2/(s_i+2 - 2 s_i+1 + s_i)
    """
    s_i   = next(seq)  
    s_ip1 = next(seq)  
    s_ip2 = next(seq) 
    while True:
        yield s_i - (s_ip1 - s_i)**2/(s_ip2 - 2*s_ip1 + s_i)
        s_i, s_ip1, s_ip2 = s_ip1, s_ip2, next(seq)


def leibnitz_pi():
    """
    Leibnitz formula
     pi/4 = sum_k=0^inf (-1)**k/(2k+1)
    the partial sum is 
     s_N = sum_k=0^N (-1)**k/(2k+1)
    """
    sN = 0
    j = 1
    for i in itertools.cycle([1,-1]):
        yield sN
        sN += i/j
        j  += 2

def pi2_over_six():
    """
     pi^2/6 = sum_k=1^inf 1/k^2
    the partial sum is 
     s_N = sum_k=1^N 1/k^2
    """
    sN = 1
    k = 1
    while True:
        yield sN
        k += 1
        sN += 1/k**2       

if __name__=='__main__':

    import numpy as np
    N = 50
    
    print('computing pi/4 from Leibnitz series')
    res_lei = list(itertools.islice(leibnitz_pi(),N))
    res_ait = list(itertools.islice(aitken(leibnitz_pi()),N))

    print('computing pi^2/6 from series')
    res_pi6  = list(itertools.islice(pi2_over_six(),N))
    res_ait6 = list(itertools.islice(aitken(pi2_over_six()),N))

    
      
    import matplotlib.pyplot as plt
    plt.rc('font', family='serif')
    fig, ax = plt.subplots(2, sharex=True)
    ax1,ax2 = ax[0],ax[1]
    x   = range(N)
    xx  = [x[0],x[-1]]
    yy  = [np.pi/4,np.pi/4]    
    ax1.plot(x,res_lei,'ro-',label=r"$\sum_{k=0}^N \frac{(-1)^k}{2k+1}$"+' (Leibnitz)')
    ax1.plot(x,res_ait,'go-',label='Aitken accel.')
    ax1.plot(xx,yy,'b-',label=r"$\pi/4$")
    eps = 0.1
    ax1.set_ylim([np.pi/4-eps,np.pi/4+3*eps])
    ax1.legend()

    
    ax1.set_title('Aitken '+r"$\Delta^2$"+'-accelerated partial sums\n (Aitken shifted left by 2 steps)')
    x   = range(N)
    xx  = [x[0],x[-1]]
    yy  = [np.pi**2/6,np.pi**2/6]
    ax2.plot(x,res_pi6,'ro-',label=r"$\sum_{k=1}^N \frac{1}{k^2}$")
    ax2.plot(x,res_ait6,'go-',label='Aitken accel.')
    ax2.plot(xx,yy,'b-',label=r"$\pi^2/6$")
    eps = 0.4
    ax2.set_ylim([np.pi**2/6-eps,np.pi**2/6+3*eps])
    ax2.legend()
    
    plt.xlabel('N')
    plt.xlim(-1,N)
    plt.show()