Search code examples
pythonmatplotlibpolar-coordinates

How to create a 3D polar plot


I have a 2D matrix which contains amplitudes of a function. Instead of usual Cartesian geometry, this matrix contains these amplitudes at polar points. For example, the first row contains the amplitudes at fixed radius but increasing angle. Columns have radial values at fixed angle.

I am using following python code to plot a 3D graph.

import numpy as np
import matplotlib.pyplot as plt
from os.path import exists
from os import sys
import matplotlib as mpl
from mpl_toolkits.mplot3d import Axes3D

N=int(sys.argv[1]);
filling=int(sys.argv[2]);
trial=int(sys.argv[3]);

filename = 'density_N%d_filling%d_trial%d'  % (N,filling,trial);
filename2 = 'init_%d_%d_%d'                     % (N,filling,trial);

if exists(filename2):
   a=np.loadtxt(filename2)
   NUMBINS=int(np.round(a[0]));
   dR2=a[1]
   dtheta=a[2]
   rmax = a[3];


def read_file(filename):
    with open(filename,'r') as data:
       x = []
       y = []
       for line in data:
          p = line.split()
          for i in range(NUMBINS):
              x.append(float(p[i]))
          y.append(np.array(x))
          x=[]
    return y

density = np.array(read_file(filename))

density = density*NUMBINS/(np.pi*dR2)
localfillingfraction=2*np.pi*density;
r = np.sqrt(np.arange(0,rmax*rmax,dR2) )[:NUMBINS]
theta = np.array([dtheta for point in r])
mpl.rcParams['legend.fontsize'] = 10

sdensity = np.transpose(localfillingfraction)
fig = plt.figure()
ax = fig.gca(projection='3d')

for n, armonic in enumerate(sdensity):
   ax.plot(r,theta*n,armonic)       
plt.savefig("N%d_FILLING%d_trial%d.pdf" % (N,filling,trial));
plt.title('filling fraction, v = 1/' + str(filling))
plt.xlabel('Radius (units of lB)')
plt.ylabel('angle(Rad)')
plt.show()

So filename2 contains four parameters I need for the calculation. filename is the 2D matrix with polar data. read_file() is a function I made to read this 2D matrix into a list. Though the matrix contains radial data as columns, the code was plotting it the other way round. That is why I had to transpose the matrix. In the end, I am getting following plot :Filling fraction as function of R, theta

Even though this plot does the job, but the radius and angle are still plotted like Cartesian coordinates. Can I make it look polar?


Solution

  • So I have tried this variant of the code, which uses a heat-map type polar plot which is really simple and clear to show the angular+radial distribution.

    #! /usr/bin/env python
    import numpy as np
    import matplotlib.pyplot as plt
    from os.path import exists
    from os import sys
    import matplotlib as mpl
    from matplotlib import rc
    from mpl_toolkits.mplot3d import Axes3D
    #=========================================
    
    N=int(sys.argv[1]);
    filling=int(sys.argv[2]);
    trial=int(sys.argv[3]);
    
    #=====================================================================
    
    filename = 'density_N%d_filling%d_trial%d'  % (N,filling,trial);
    filename2 = 'init_%d_%d_%d'                     % (N,filling,trial);
    
    if exists(filename2):
        a=np.loadtxt(filename2)
        NUMBINS=int(np.round(a[0]));
        print(NUMBINS)
        dR2=a[1]
        dtheta=a[2]
        rmax = a[3];
    
    #===========================================================================
    '''This part reads the 2D array of NUMBIN*NUMBIN, row by row.
       Thus, the returned list y consists of a linear array of dimension
       [NUMBIN,0] with each element is itself a list containing each row.'''
    
    def read_file(filename):
        with open(filename,'r') as data:
            x = []
            y = []
            for line in data:
                p = line.split()
                for i in range(NUMBINS):
                    x.append(float(p[i]))
                y.append(np.array(x))
                x=[]
        return y
    
    density = np.array(read_file(filename))
    #==========================================================================
    
    density = density*NUMBINS/(np.pi*dR2)
    localfillingfraction=2*np.pi*density;
    r = np.sqrt(np.arange(0,rmax*rmax,dR2) )[:NUMBINS]
    theta = np.linspace(0,2*np.pi,NUMBINS)
    
    mpl.rcParams['legend.fontsize'] = 10
    
    #===============================================================
    '''the matrix needed to be trasposed, to get the radial and angular
       data to be read correctly'''
    
    sdensity = np.transpose(localfillingfraction)
    
    fig = plt.figure()
    ax = Axes3D(fig)
    rad, th = np.meshgrid(r,theta)
    plt.subplot(projection="polar")
    plt.pcolormesh(th,rad,sdensity)
    plt.thetagrids([theta * 15 for theta in range(360//15)])
    plt.plot(theta,rad,color='k',ls='none')
    plt.grid()
    plt.savefig("N%d_FILLING%d_trial%d.jpg" % (N,filling,trial));
    plt.show()
    

    The plot is as follows:

    enter image description here