Search code examples
pythonmatplotlibcomplex-numberscolor-mapping

Plot complex roots of unity as arrow-vectors on complex plane


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)


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