3

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)

P i
  • 29,020
  • 36
  • 159
  • 267

1 Answers1

5
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

Community
  • 1
  • 1
Hooked
  • 84,485
  • 43
  • 192
  • 261
  • 1
    +1 - For the random distinct color portion, it's often handy to borrow from `pandas`. E.g. `pandas.tools.plotting._get_standard_colors(num, color_type='random')` – Joe Kington Apr 25 '14 at 17:57
  • 2
    @JoeKington Yes, I agree. Typically, I'll pull a brewer color scheme from `seaborn`. I wasn't sure how complicated I wanted to make the answer, since the heart of the question is about the complex plotting itself - not the color mapping. – Hooked Apr 25 '14 at 17:59
  • That's amazing -- so few lines of code! Yes, I'm definitely asking about the colour mapping as well :) – P i Apr 25 '14 at 18:59
  • 1
    @Pi As Joe and I mentioned, look into either the seaborn library or pandas for this functionality "out-of-the-box". http://www.stanford.edu/~mwaskom/software/seaborn/ and http://pandas.pydata.org/ are good places to get started. – Hooked Apr 25 '14 at 19:01
  • 1
    @Pi I've added an example using seaborn - it's really quite easy! – Hooked Apr 25 '14 at 19:08
  • I'd recommend the `husl` palette if you're using seaborn, that way the different arrows will have the same perceptual brightness. – mwaskom Apr 25 '14 at 23:43
  • 1
    You could also put the drawing loop under a statement that says `with sns.color_palette("husl", 10):` and then you don't need to deal with getting the list yourself or passing it to the `arrow` function. – mwaskom Apr 25 '14 at 23:44
  • @mwaskom Awesome! I didn't know `sns.color_palette` can act as a context manager - you've simplified a ton of code for me! – Hooked Apr 26 '14 at 01:45
  • Indeed, `color_palette`, `axes_style`, and `plotting_context` are all context managers for temporarily changing the look of plots. – mwaskom Apr 26 '14 at 02:22
  • python3 threw an error for the first plot: `AttributeError: 'itertools.cycle' object has no attribute 'next'` but seaborn code works. – jsky Apr 13 '20 at 13:01