This gist contains the code used to generate all the graphs and data used in my 2013 Mathematics IA project ("Patterns from Complex Numbers"). utilities.py contains various helper functions and classes. diagrams.py contains the code used to render the drawings and run test data. This code has been confirmed to work using Python 2.7.3, matplotlib…

diagrams.py

#!/usr/bin/env python
'''
Contains code to generate the graphs and data used in this IA
'''


from __future__ import division

import cmath
import utilities


# Convenience functions

def find_product(nums):
    '''Finds the product of the diagonals by manually finding and 
    multiplying the distance between roots. Assumes that the first
    root in the list is the radiating point.'''
    product = 1
    for vertex in nums[1:]:
        product *= utilities.find_distance(nums[0], vertex)
    return product
    
def find_product_proof(n, r):
    '''Finds the product based on the general equation derived in 
    the paper.'''
    return n * r**((n - 1)/n)
    
def plot_diagonals(plot, nums):
    '''Given a plot, draws red lines to indicate the diagonals.
    Assumes that the first root in the list is the radiating point.'''
    for vertex in nums:
        plot.add_line(nums[0], vertex, color='red')
            

# Graphs
    
def draw_basic_complex_plane():
    '''Draws a basic complex plane used in the first section of this
    paper. (figure 1.1)'''
    plot = utilities.Plot((-1.2, 1.2), (-1.2, 1.2), 1, 1, figsize=(4, 4))
    point = utilities.S(0.5 + utilities.sympy.sin(utilities.sympy.pi/3)*1j)
    
    plot.add_points([point])
    plot.add_circle(1)
    
    plot.add_line(0, point)
    plot.add_line(0, utilities.S(0.5), color='blue')
    plot.add_line(utilities.S(0.5), point, color='blue')
    
    plot.add_annotation(r'$\theta$', (0.09, 0.03), color="purple")
    plot.add_annotation('$r$', (0.18, 0.42), color="red")
    plot.add_annotation('$a + bi$', (0.55, 0.85))
    
    plot.add_annotation(r'$a = r \cdot \cos{(\theta)}$', 
        (0.15, -0.15), color="blue")
    plot.add_annotation(r'$b = r \cdot \sin{(\theta)}i$', 
        (0.55, 0.35), color="blue")
    
    utilities.save(utilities.render(plot), 'basic-complex-plane')
    
def draw_de_moivre_3():
    '''Creates figure 2.1 a and b'''
    nums = utilities.de_moivre(3, 1, 0)
    p1 = utilities.Plot((-1.2, 1.2), (-1.2, 1.2), 0.5, 0.5, figsize=(5, 5))
    p1.add_points(nums)
    p1.add_circle(1)
    utilities.save(utilities.render(p1), 'de_moivre_3-simple')
    
    p2 = p1.copy()
    plot_diagonals(p2, nums)
    av1 = (nums[0] + nums[1]) * 3 / 2
    av2 = (nums[0] + nums[2]) * 3 / 2
    d1 = float(utilities.find_distance(nums[0], nums[1]))
    d2 = float(utilities.find_distance(nums[0], nums[2]))
    p2.add_annotation(r'$distance \approx {0:.4f}$'.format(d1), (0.1, 0.6))
    p2.add_annotation(r'$distance \approx {0:.4f}$'.format(d2), (0.1, -0.6))
    utilities.save(utilities.render(p2), 'de_moivre_3-lines')
    
def draw_diagonal_lines_3_to_10():
    '''Draws graphs for z^n - 1 = 0 with diagonals.'''
    for n in xrange(3, 11):
        nums = utilities.de_moivre(n, 1, 0)
        
        plot = utilities.Plot((-1.2, 1.2), (-1.2, 1.2), 1, 1, figsize=(3, 3))
        plot.add_points(nums)
        plot.add_circle(1)
        plot_diagonals(plot, nums)
        for index, vertex in enumerate(nums):
            plot.add_annotation('$z_{0}$'.format(index), 
                utilities.tup(vertex), offset=(5, 5))
            
        utilities.save(utilities.render(plot), 'de_moivre_{0}'.format(n))
        
def draw_3_when_theta_is_i():
    '''Draws roots and lines when z^3 - i = 0'''
    nums = utilities.de_moivre(3, 1, utilities.sympy.pi / 2)
    plot = utilities.Plot((-1.2, 1.2), (-1.2, 1.2), 1, 1, figsize=(5, 5))
    plot.add_points(nums)
    plot.add_circle(1)
    utilities.save(utilities.render(plot), 'de_moivre_3_i-simple')
    
    plot2 = plot.copy()
    plot_diagonals(plot2, nums)
    av1 = (nums[0] + nums[1]) * 3 / 2
    av2 = (nums[0] + nums[2]) * 3 / 2
    d1 = float(utilities.find_distance(nums[0], nums[1]))
    d2 = float(utilities.find_distance(nums[0], nums[2]))
    plot2.add_annotation(r'$distance \approx {0:.4f}$'.format(d1), (0.1, 0.55))
    plot2.add_annotation(r'$distance \approx {0:.4f}$'.format(d2), (0.4, -0.5))
    utilities.save(utilities.render(plot2), 'de_moivre_3_i-lines')
        
def draw_when_theta_is_i():
    '''Draws multiple graphs when z^n - 1 = 0'''
    for n in xrange(3, 6):
        nums = utilities.de_moivre(n, 1, utilities.sympy.pi / 2)
        plot = utilities.Plot((-1.2, 1.2), (-1.2, 1.2), 1, 1, figsize=(5, 5))
        plot.add_points(nums)
        plot.add_circle(1)
        plot_diagonals(plot, nums)
        utilities.save(utilities.render(plot), 'de_moivre_{0}_i'.format(n))
    
    
# Tests
    
def test_product_of_diagonals():
    '''Tests that the product of the diagonals equals n for a small amount of n'''
    for n in xrange(3, 11):
        nums = utilities.de_moivre(n, 1, 0)
        
        print 'n =', n
        for vertex in nums:
            re, im = utilities.tup(vertex)
            print '   ', float(re), ',', float(im)
        print '   product = ', float(find_product(nums))
        for i in xrange(n):
            print '   @', r'(z - \cis{{\frac{{2\pi \cdot {0}}}{{{1}}})'.format(i, n)
        for vertex in nums:
            print '   >', float(utilities.find_distance(nums[0], vertex))
            
def test_increasing_n():
    '''Tests the effects of increasing n from 3 to 1000 given z^n - 1 = 0'''
    for n in xrange(3, 1000):
        nums = utilities.de_moivre(n, 1, 0)
        prod = 1
        for vertex in nums[1:]:
            prod *= utilities.find_distance(nums[0], vertex)
        a = utilities.R(n)
        b = prod.evalf(10)
        if a != b:
            print a, b
    
def test_when_theta_is_i():
    '''Determines what happens when z^n - i = 0'''
    for n in xrange(3, 6):
        nums = utilities.de_moivre(n, 1, utilities.sympy.pi / 2)
        
        print 'n = {0}'.format(n)
        print '    prod = ', find_product(nums).evalf(10)
        for num in nums:
            print '   >', num
        for num in nums:
            print '   @', num.evalf(4)
        for num in nums:
            print '   ?', utilities.find_distance(nums[0], num).evalf(5)
        
def test_varying_n_r_and_theta():
    '''Determines what happens when n, r, and theta are varied, and
    if the general equation works.'''
    for n in xrange(3, 5):
        for r in xrange(1, 3):
            for theta in xrange(0, 8):
                theta = utilities.sympy.pi * theta / 4
                nums = utilities.de_moivre(n, r, theta)
                direct_product = find_product(nums)
                proof_product = utilities.S(find_product_proof(n, r))
                
                print '{0}: {1}: {2}'.format(n, r, theta)
                print '   ', direct_product.evalf(8)
                print '   ', proof_product.evalf(8)

def test_increasing_radius():
    '''Tests the effects of increasing the radius.'''
    plot = utilities.Plot((-2, 2), (-2, 2), 0.5, 0.5, figsize=(8, 8))
    for r in xrange(1, 10):
        nums = utilities.de_moivre(4, r, 0)
        
        prod = find_product(nums)
        print r, prod, prod.evalf(10)
        print find_product_proof(4, r)
        
        plot.add_points(nums)
        plot.add_circle(r)
    #utilities.render(plot).show()
        
    
    
def graphs():
    draw_basic_complex_plane()
    draw_de_moivre_3()
    draw_diagonal_lines_3_to_10()
    draw_3_when_theta_is_i()
    draw_when_theta_is_i()
    
def tests():
    test_product_of_diagonals()
    test_increasing_n()
    test_when_theta_is_i()
    test_varying_n_r_and_theta()
    test_increasing_radius()
    
if __name__ == '__main__':
    utilities.setup()
    graphs()
    tests()

utilities.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import division

import copy

# Mathematical libraries
import sympy
import matplotlib
import matplotlib.pyplot as pyplot

S = sympy.sympify
R = sympy.Rational


# Functions to help set up the mathematical libraries
def setup():
    sympy.init_printing(use_unicode=False, wrap_line=False, no_global=True)
    matplotlib.rc('text', usetex=True)

def force_sympy(func):
    '''An annotation which forces all input to be converted into Sympy objects,
    which gives the mathematical calculations greater precision.'''
    def wrapper(*args, **kwargs):
        new_args = []
        new_kwargs = {}
        for arg in args:
            if isinstance(arg, (int, long, float, complex)):
                new_args.append(R(arg))
            else:
                new_args.append(arg)
        for key, item in kwargs.items():
            if isinstance(item, (int, long, float, complex)):
                new_kwargs[key] = R(item)
            else:
                new_kwargs[key] = item
        return func(*new_args, **new_kwargs)
    wrapper.__name__ = func.__name__
    wrapper.__doc__ = func.__doc__
    return wrapper

@force_sympy
def cis(theta):
    # Note that in Python, complex numbers are represented using the 'j' 
    # character instead of the 'i' character, for whatever reason.
    return sympy.cos(theta) + sympy.sin(theta)*1j
    
@force_sympy
def find_distance(a, b):
    '''Finds the distance between two points'''
    ax, ay = sympy.re(a), sympy.im(a)
    bx, by = sympy.re(b), sympy.im(b)
    return sympy.sqrt((ax - bx)**R(2) + (ay - by)**R(2))
    
@force_sympy
def de_moivre(n, radius, theta):
    '''Returns all the roots using de Moivre's theorem'''
    magnitude = radius**(1 / n)
    return [magnitude * cis((theta + k * 2 * sympy.pi) / n) for k in range(n)]
    
@force_sympy
def tup(number):
    '''Converts a complex number into a Cartesian coordinate.'''
    return (sympy.re(number), sympy.im(number))
    


class Plot(object):
    '''An object representing a graph.'''
    def __init__(self, xrange, yrange, xinterval, yinterval, figsize=(8,8)):
        self.xrange = xrange
        self.yrange = yrange
        self.xinterval = xinterval
        self.yinterval = yinterval
        self.figsize = figsize
        self.points = []
        self.circles = []
        self.lines = []
        self.annotations = []
        
    def add_points(self, numbers, color=None, **kwargs):
        '''Adds a series of complex numbers in a list to the plot.'''
        x = [sympy.re(num) for num in numbers]
        y = [sympy.im(num) for num in numbers]
        
        args = (x, y)
        kwargs['color'] = color
        
        self.points.append((args, kwargs))
        
    def add_circle(self, radius, origin=(0, 0), color='lightgreen', 
                   facecolor='none', **kwargs):
        args = (origin,)
        kwargs['radius'] = radius
        kwargs['edgecolor'] = color
        kwargs['facecolor'] = facecolor
        self.circles.append((args, kwargs))
        
    def add_line(self, a, b, color='red', **kwargs):
        x1, y1 = tup(a)
        x2, y2 = tup(b)
        args = ((x1, x2), (y1, y2))
        kwargs['color'] = color
        kwargs['zorder'] = 3
        self.lines.append((args, kwargs))
        
    def add_annotation(self, string, coords, offset=(0, 0), **kwargs):
        '''Adds text.'''
        args = (string,)
        kwargs['xy'] = coords
        kwargs['xycoords'] = 'data'
        kwargs['xytext'] = offset
        kwargs['textcoords'] = 'offset points'
        kwargs['zorder'] = 4
        self.annotations.append((args, kwargs))
        
    def copy(self):
        '''Returns a copy of this object.'''
        return copy.deepcopy(self)
        
        
def center_origin(axis):
    '''Given a matplotlib axis, centers them in the middle of the image.'''
    axis.spines['right'].set_color('none')
    axis.spines['top'].set_color('none')
    axis.xaxis.set_ticks_position('bottom')
    axis.spines['bottom'].set_position(('data',0))
    axis.yaxis.set_ticks_position('left')
    axis.spines['left'].set_position(('data',0))
    
def add_axis_labels(axis, xrange, yrange, xinterval, yinterval):
    '''Correctly formats the tick labels based on the script object'''
    def fix(string):
        if string.startswith('.'):
            string = '0' + string
        if string.endswith('.'):
            string = string + '0'
        return string
        
    def x_format_func(x, pos):
        return '$' + fix(('%F' % x).strip('0')) + '$'
    def y_format_func(y, pos):
        return '$' + fix(('%F' % y).strip('0')) + 'i$' if y != 0 else ''
        
    x_formatter = matplotlib.ticker.FuncFormatter(x_format_func)
    y_formatter = matplotlib.ticker.FuncFormatter(y_format_func)
    
    x_locator = matplotlib.ticker.MultipleLocator(xinterval)
    y_locator = matplotlib.ticker.MultipleLocator(yinterval)
    
    axis.xaxis.set_major_formatter(x_formatter)
    axis.xaxis.set_major_locator(x_locator)
    axis.yaxis.set_major_formatter(y_formatter)
    axis.yaxis.set_major_locator(y_locator)
    
    axis.set_xlim(*xrange)
    axis.set_ylim(*yrange)
    
def render(plot):
    '''Converts a plot into a matplotlib figure which can then be turned
    into a png'''
    figure = pyplot.figure(figsize=plot.figsize)
    axis = figure.add_subplot(1, 1, 1)
    
    for args, kwargs in plot.points:
        axis.scatter(*args, **kwargs)
    for args, kwargs in plot.circles:
        axis.add_patch(pyplot.Circle(*args, **kwargs))
    for args, kwargs in plot.lines:
        axis.add_line(pyplot.Line2D(*args, **kwargs))
    for args, kwargs in plot.annotations:
        axis.annotate(*args, **kwargs)
        
    
    center_origin(axis)
    add_axis_labels(axis, plot.xrange, plot.yrange, 
        plot.xinterval, plot.yinterval)
    
    return figure
    
def save(figure, name):
    '''Saves a rendered plot to a folder.'''
    path = ''.join([
        './images/',
        name,
        '.png'])
    figure.savefig(path, bbox_inches='tight')
    

def block():
    '''Halts execution until user input.'''
    raw_input('Continue? >>>')