para mate4

gistfile1.py

from pylab import *
import cmath
ion()

def zroots(k, z):
	roots = []
	theta = arctan2(z.imag, z.real)
	modroot = abs(z)**(1.0/k)
	for i in range(k):
		real = modroot * cos(float(theta+2*i*pi)/k)
		imag = modroot * sin(float(theta+2*i*pi)/k)
		roots.append(complex(real, imag))
	return roots


def polarplot(zs):
	polar()
	for z in zs:
		p = cmath.polar(z)
		plot((0,p[1]), (0,p[0]))


def rootandplot(k, z):
	clf()
	zs = zroots(k, z)
	zs.insert(0, z)
	polarplot(zs)
	titles = ["k=" + str(i) for i in range(k)]
	titles.insert(0, str(z))
	legend(titles)

def funcplot(func, args, *opts, **kopts):
	ys = []
	for arg in args:
		ys.append(func(arg))
	plot(args, ys, *opts, **kopts)