2011/08/11

python tips: draw binary pdf graph with matplotlib



this is a code for drawing a graph of binary probability distribution with matplotlib.

#!/usr/bin/env python                                                           
#coding: utf-8                                                                  
                                                                                
import pylab as pl                                                              
                                                                                
def combination(n, r):                                                          
    if (r == 0) or (n == r):                                                    
        return 1.0                                                              
                                                                                
    denom = reduce(lambda x, y: x*y, [i for i in range(1, n+1)])                
    mol_l = reduce(lambda x, y: x*y, [i for i in range(1, n-r+1)])              
    mol_r = reduce(lambda x, y: x*y, [i for i in range(1, r+1)])                
    return float(denom / (mol_l * mol_r))                                       
                                                                                
def get_probability(p, n, r):                                                   
    q = 1 - p                                                                   
    return combination(n, r) * (p ** r) * (q ** (n-r))                          
                                                                                
def plot_bin_distrib(p, n, lt, lab):                                            
    xl = [nr for nr in range(0, n+1, 1)]                                        
    yl = [get_probability(p, n, nr) for nr in range(0, n+1, 1)]                 
    pl.plot(xl, yl, lt, label=lab)                                              
                                                                                
def main():                                                                     
    p = 1.0/3.0                                                                 
    n = 10                                                                      
    plot_bin_distrib(p, n, 'b-', 'Bin(%s,%f)' %(n,p))                           
    n = 20                                                                      
    plot_bin_distrib(p, n, 'r-', 'Bin(%s,%f)' %(n,p))                           
    n = 30                                                                      
    plot_bin_distrib(p, n, 'g-', 'Bin(%s,%f)' %(n,p))                           
    n = 40                                                                      
    plot_bin_distrib(p, n, 'y-', 'Bin(%s,%f)' %(n,p))                           
                                                                                
    pl.ylabel('p(x)')                                                           
    pl.xlabel('x')                                                              
    pl.legend()                                                                 
    pl.show()                                                                   
                                                                                
if __name__ == "__main__":                                                      
    main()

and here is output from above code.


No comments:

Post a Comment

100