March 12, 2013

Matlab's Cylinder Command in Python

Matlab has a "cylinder" command that will give you n coordinates of a cylinder of radius R. Here R is a vector of radii, not necessarily all the same.  You can think of R as a line or curve, and the resulting cylinder comes from rotating the line around the z-axis.

Matplotlib has commands to plot certain shapes, such as circles, but not any that give you n evenly spaced coordinates approximating the shape. This admittedly has little to do with economics. I use it for an initial guess for the set of supergame equilibria in a repeated game (see the supergametools library). This is more for the random matlab to python translation project.

It's not hard to program up a method that does this, but to save you all some time I'm posting a version of the command in Python, seen below or downloadable on GitHub.

import numpy as np

def cylinder(r,n):
    '''
    Returns the unit cylinder that corresponds to the curve r.
    INPUTS:  r - a vector of radii
             n - number of coordinates to return for each element in r

    OUTPUTS: x,y,z - coordinates of points
    '''

    # ensure that r is a column vector
    r = np.atleast_2d(r)
    r_rows,r_cols = r.shape
    
    if r_cols > r_rows:
        r = r.T

    # find points along x and y axes
    points  = np.linspace(0,2*np.pi,n+1)
    x = np.cos(points)*r
    y = np.sin(points)*r

    # find points along z axis
    rpoints = np.atleast_2d(np.linspace(0,1,len(r)))
    z = np.ones((1,n+1))*rpoints.T
    
    return x,y,z

See below for an example.

import matplotlib.pyplot as plt
import mpl_toolkits.mplot3d.axes3d as p3

# set curve vector r and set n to 20
r = np.linspace(2,.5,100)
n = 20

# get points from cylinder and plot
[x,y,z] = cylinder(r,n)

# make plots
fig=plt.figure()
ax = p3.Axes3D(fig)
ax.plot_wireframe(x,y,z)
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')
plt.show()