In [14]:
import matplotlib.pyplot as plt
import numpy as np
from scipy import signal
from sympy import *

def conv(k,n,pmf1=None,grid=None):
    f = lambda x,a,b : 1/(b-a)*(a<x)*(x<b)
    c,d = 0,1
    
    def firstconv(k,f,c,d,n) :
        a,b = 0,1
        big_grid = np.linspace(start=-b-d*k,stop=b+d*k,num=n)
        delta = big_grid[1]-big_grid[0]
        f1, f2 = f(big_grid,a,b),  f(big_grid,c,d)
        pmf1, pmf2 = f1*delta, f2*delta
        conv_pmf = signal.fftconvolve(pmf1,pmf2,'same')
        pdf1, pdf2 = pmf1/delta, pmf2/delta
        conv_pdf = conv_pmf/delta
        plt.xlim(min(a,c),b+d*1)
        plt.plot(big_grid,pdf1, label='pdf')
        plt.plot(big_grid,pdf2, label='$U('+str(c)+','+str(d)+')$')
        plt.plot(big_grid,conv_pdf, label='convoluted')
        plt.legend(loc='best'), plt.suptitle('Convolution')
        plt.show()
        print("Sum of convoluted pmf: "+str(sum(conv_pmf)))     
        return(big_grid,conv_pdf,min(a,c),b+d*1)
    
    grid, h1, a, b = firstconv(k,f,c,d,n)
    
    if k==1 :
        return(h1)
    elif k>1 :
        h = h1
        delta = grid[1]-grid[0]
        f2 = f(grid,c,d)
        for i in range(1,k) :
            pmf1, pmf2 = h*delta, f2*delta # Default : pdf of U(0,1)
            conv_pmf = signal.fftconvolve(pmf1,pmf2,'same') # fast fourier transform convolution
            pdf1, pdf2 = pmf1/delta, pmf2/delta
            conv_pdf = conv_pmf/delta
            h = conv_pdf
            plt.xlim(c, b+d*i)
            plt.plot(grid, pdf1, label='pdf')
            plt.plot(grid, pdf2, label='$U('+str(c)+','+str(d)+')$')
            plt.plot(grid, conv_pdf, label='convoluted')
            plt.legend(loc='best'), plt.suptitle('Convolution')
            plt.show()
            print("Sum of convoluted pmf: "+str(sum(conv_pmf)))
        return(conv_pdf)
 
k, n = 5, 10**3
h = conv(k,n)
Sum of convoluted pmf: 0.9940030120210587
Sum of convoluted pmf: 0.9910180179909798
Sum of convoluted pmf: 0.9880419879069375
Sum of convoluted pmf: 0.985074894850173
Sum of convoluted pmf: 0.9821167119827701