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)

Était-ce utile?

La 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

Licencié sous: CC-BY-SA avec attribution
Non affilié à StackOverflow
scroll top