Example code of drawing graphs and such

utilities.py

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

from __future__ import division

import sympy
import matplotlib
import matplotlib.pyplot as pyplot

S = sympy.sympify
R = sympy.Rational

def setup_sympy():
    sympy.init_printing(use_unicode=False, wrap_line=False, no_global=True)

def force_sympy(func):
    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):
    return sympy.cos(theta) + sympy.sin(theta)*1j
    
@force_sympy
def find_distance(a, b):
    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):
    magnitude = radius**(1 / n)
    return [magnitude * cis((theta + k * 2 * sympy.pi) / n) for k in range(n)]
    
def center_origin(axis):
    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, magnitude, interval):
    def frange(x, y, jump):
        while x < y:
            if -0.0001 < x < 0.0001:
                yield 0
            else:
                yield x
            x += jump
    x = list(frange(-magnitude - interval, magnitude + interval*2, interval))
    y = [n for n in x if n != 0]
    
    x_labels = ['${0}$'.format(n) for n in x]
    y_labels = ['${0}i$'.format(n) for n in y]
    
    axis.set_xticks(x)
    axis.set_xticklabels(x_labels)
    axis.set_yticks(y)
    axis.set_yticklabels(y_labels)
    
    return axis
    
def plot(numbers):
    x = [sympy.re(num) for num in numbers]
    y = [sympy.im(num) for num in numbers] 
    
    figure = pyplot.figure(figsize=(8,8))
    axis = figure.add_subplot(1, 1, 1)
    
    axis.add_patch(pyplot.Circle((0, 0), radius=1.00, edgecolor='g', facecolor='none'))
    axis.scatter(x, y)
    
    center_origin(axis)
    add_axis_labels(axis, 1.0, 0.2)
                         
    axis.set_xlim(-1.2, 1.2)
    axis.set_ylim(-1.2, 1.2)
    
    figure.show()
    figure.savefig('test1.png')
    
    
def test():
    print cis(0)
    print cis(2 * sympy.pi / 3)
    a1, a2, a3 = de_moivre(3, 1, 0)
    print [a1, a2, a3]
    plot(de_moivre(5, 1, 0))
    print find_distance(a1, a2)
    
def block():
    raw_input('Continue? >>>')
    
if __name__ == '__main__':
    setup_sympy()
    test()
    block()