Question

I would like to plot n roots of unity using matplotlib, with each one as a different coloured arrow.

It should look like a star shape, with the arrows equally spaced pointing outwards onto the unit circle.

matplotlib has a function for drawing an arrow, but is there any way to do this using complex numbers, or do I have to convert to real cartesians?

Also, does there exist an array of stock colors, so that regardless of how many roots I wish to display, it will give me an array of distinct colors? (rather than say seven almost identical shades of red)

Was it helpful?

Solution

import numpy as np
import pylab as plt
import itertools

n = 13
roots = np.roots( [1,] + [0,]*(n-1) + [-1,] )
colors = itertools.cycle(['r', 'g', 'b', 'y'])

plt.figure(figsize=(6,6))

for root in roots:
    plt.arrow(0,0,root.real,root.imag,ec=colors.next())


plt.xlim(-1.5,1.5)
plt.ylim(-1.5,1.5)
plt.show()

enter image description here

The roots of unity are calculated in a manner similar to this answer.

Update: If you want to use seaborn, you can get unique colors quite easily:

import numpy as np
import pylab as plt
import itertools

import seaborn as sns
n = 13
colors = sns.color_palette("hls", n)
roots = np.roots( [1,] + [0,]*(n-1) + [-1,] )

# Sorted by angle
idx = np.argsort([np.angle(x) for x in roots])
roots = roots[idx]

plt.figure(figsize=(6,6))

for root,c in zip(roots,colors):
    plt.arrow(0,0,root.real,root.imag,ec=c,lw=3)

plt.xlim(-1.25,1.25)
plt.ylim(-1.25,1.25)
plt.show()

enter image description here

Licensed under: CC-BY-SA with attribution
Not affiliated with StackOverflow
scroll top