2011/07/23

python tips: creating polar graph using matplotlib

#!/usr/bin/env python                                                                                                                                                                                      
                                                                                                                                                                                                           
import sys                                                                                                                                                                                                 
import csv                                                                                                                                                                                                 
import codecs                                                                                                                                                                                              
                                                                                                                                                                                                           
import matplotlib                                                                                                                                                                                          
import numpy as np                                                                                                                                                                                         
from matplotlib.pyplot import figure, show, rc, grid                                                                                                                                                       
                                                                                                                                                                                                           
def main():                                                                                                                                                                                                
    argvs = sys.argv                                                                                                                                                                                       
    argc = len(argvs)                                                                                                                                                                                      
                                                                                                                                                                                                           
    if (argc != 4):                                                                                                                                                                                        
        print "Usage: python %s input_file output_file title" %(argvs[0])                                                                                                                                  
        sys.exit(0)                                                                                                                                                                                        
                                                                                                                                                                                                           
    in_file = argvs[1]                                                                                                                                                                                     
    out_file = argvs[2]                                                                                                                                                                                    
    title = argvs[3]                                                                                                                                                                                       
                                                                                                                                                                                                           
    csvfile = codecs.open(in_file, 'r', 'utf-8')                                                                                                                                                           
                                                                                                                                                                                                           
    list_x = []                                                                                                                                                                                            
    list_y = []                                                                                                                                                                                            
                                                                                                                                                                                                           
    for row in csv.reader(csvfile):                                                                                                                                                                        
        radian = 0                                                                                                                                                                                         
                                                                                                                                                                                                           
        if float(row[0]) < 0:                                                                                                                                                                              
            radian = (float(row[0])+360.0)/180                                                                                                                                                             
        else:                                                                                                                                                                                              
            radian = float(row[0])/180                                                                                                                                                                     
                                                                                                                                                                                                           
        list_x.append(np.pi*radian)                                                                                                                                                                        
        list_y.append(float(row[1]))                                                                                                                                                                       
                                                                                                                                                                                                           
    init_graph_settings()                                                                                                                                                                                  
    draw_polar_graph(list_x, list_y, out_file, title)                                                                                                                                                      
    csvfile.close()                                                                                                                                                                                        
                                                                                                                                                                                                           
def init_graph_settings():                                                                                                                                                                                 
    # radar green, solid grid lines                                                                                                                                                                        
    rc('grid', color='#316931', linewidth=0.5, linestyle='--')                                                                                                                                             
    rc('xtick', labelsize=15)                                                                                                                                                                              
    rc('ytick', labelsize=10)                                        

def draw_polar_graph(list_x, list_y, out_file, title):                                                                                                                                                     
    """                                                                                                                                                                                                    
    draws polar graph                                                                                                                                                                                      
                                                                                                                                                                                                           
    :param list_x: list contains x values                                                                                                                                                                  
    :param list_y: list contains y values                                                                                                                                                                  
    :param title: title for a graph                                                                                                                                                                        
    :type list_x: list contains float values                                                                                                                                                               
    :type list_y: list contains float values                                                                                                                                                               
    :type title: string.                                                                                                                                                                                   
    """                                                                                                                                                                                                    
                                                                                                                                                                                                           
    # force square figure and square axes looks better for polar, IMO                                                                                                                                      
    width, height = matplotlib.rcParams['figure.figsize']                                                                                                                                                  
    size = min(width, height)                                                                                                                                                                              
                                                                                                                                                                                                           
    # make a square figure                                                                                                                                                                                 
    fig = figure(figsize=(size, size))                                                                                                                                                                     
                                                                                                                                                                                                           
    ax = fig.add_axes([0.1, 0.1, 0.8, 0.8], polar=True, axisbg='#d5de9c')                                                                                                                                  
    ax.plot(list_x, list_y, 'r.', markersize=2)                                                                                                                                                            
    ax.set_rmax(10000.0)                                                                                                                                                                                   
    ax.set_title(title, fontsize=20)                                                                                                                                                                       
    grid(True)                                                                                                                                                                                             
    fig.savefig(out_file, dpi=100)                                                                                                                                                                         
                                                                                                                                                                                                           
if __name__ == "__main__":                                                                                                                                                                                 
    main()

No comments:

Post a Comment

100