Square spiral with matplotlib (2)

Some time ago I write a post showing how to draw a square spiral using matplotlib (embedded in pyqt4). By this time I was aware that an easier way to find the orientation changes musts exists. Here is my easier way (without pyqt4):

#!/usr/bin/python2.7
# encoding: utf-8

import sys
import matplotlib.pyplot as plt

def calculate_directions( num_pts ):
    '''We know we will start, always, from right to left. So
    the pattern will always be right-up-left-down. We can calculate
    the directions focusing in the repetition of this pattern.
    This function calculates the directions using the notation:
     * 1: right
     * 2: up
     * 3: left
     * 4: right
    '''
    pattern = '1234'
    times = num_pts / 4
    directions = pattern * times
    for i in range( num_pts - times * 4 ):
        directions += pattern[ i ]
    return directions

def calculate_points( num_pts ):
    '''Using directions from calculate_directions we start on 0,0
    and we start adding on x or y axis (if 1 and 3) or subtracting
    on x or y axis (if 2 and 4)to move the cursor around the plot.
    The result is a list of points.
    '''
    points, values = [ ], range( num_pts, -1, -1 )
    lx, ly = 0, 0
    for d in calculate_directions( num_pts ):
        inc = values.pop( 0 )
        points.append( ( lx, ly ) )
        if d in ( '1', '3' ):  # x-axis
            nlx = lx + inc if d == '1' else lx - inc
            lx = nlx
        else:  # y-axis
            nly = ly + inc if d == '2' else ly - inc
            ly = nly
    return points

def draw_plot( pts_lst, num_edgs ):
    '''Giving the list of points we create a plot and we draw the
    edges using two consecutive points of the list.
    '''
    fig = plt.figure()
    ax = fig.add_subplot( 111 )
    ax.set_xlim( [ -1, num_edgs + 1 ] )
    ax.set_ylim( [ -1, num_edgs ] )
    ax.plot( map( lambda x: x[ 0 ], pts_lst ), map( lambda x: x[ 1 ], pts_lst ) )
    plt.show()

def main():
    '''We check in sys.argv if we have two or more arguments. If
    there are more than one we try to use the second one as an
    integer.
    '''
    if len( sys.argv ) < 2:
        points = 18
    else:
        try:
            points = int( sys.argv[ 1 ] )
        except:
            points = 7
    lst = calculate_points( points )
    draw_plot( lst, points )

if __name__ == '__main__':
    main()

Download the complete code file.

Advertisements

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out / Change )

Twitter picture

You are commenting using your Twitter account. Log Out / Change )

Facebook photo

You are commenting using your Facebook account. Log Out / Change )

Google+ photo

You are commenting using your Google+ account. Log Out / Change )

Connecting to %s

%d bloggers like this: