如何在 Python matplotlib 中的曲线下填充彩虹色

Question

I want to fill rainbow color under a curve. Actually the function matplotlib.pyplot.fill_between can fill area under a curve with a single color.

import matplotlib.pyplot as plt
import numpy as np
x = np.linspace(0, 100, 50) 
y = -(x-50)**2 + 2500
plt.plot(x,y)
plt.fill_between(x,y, color='green')
plt.show()

Is there a knob I can tweak the color to be rainbow? Thanks.

Answer 1

This is pretty easy to hack if you want "fill" with a series of rectangles:

import numpy as np
import pylab as plt

def rect(x,y,w,h,c):
    ax = plt.gca()
    polygon = plt.Rectangle((x,y),w,h,color=c)
    ax.add_patch(polygon)

def rainbow_fill(X,Y, cmap=plt.get_cmap("jet")):
    plt.plot(X,Y,lw=0)  # Plot so the axes scale correctly

    dx = X[1]-X[0]
    N  = float(X.size)

    for n, (x,y) in enumerate(zip(X,Y)):
        color = cmap(n/N)
        rect(x,0,dx,y,color)

# Test data    
X = np.linspace(0,10,100)
Y = .25*X**2 - X
rainbow_fill(X,Y)
plt.show()

You can smooth out the jagged edges by making the rectangles smaller (ie use more points). Additionally you could use a trapezoid (or even an interpolated polynomial) to refine the "rectangles".

Answer 2

If you mean giving some clever argument to "color=" I'm afraid this doesn't exist to the best of my knowledge. You could do this manually by setting a quadratic line for each color and varying the offset. Filling between them with the correct colors will give a rainbowish This makes a fun project to learn some python but if you don't feel like trying here is an example:

import matplotlib.pyplot as plt
import numpy as np

x = np.linspace(0, 100, 50) 

y_old = -(x-50)**2 + 2500
for delta, color in zip([2250, 2000, 1750, 1500, 1250, 1000], ["r", "orange", "g", "b", "indigo", "violet"] ):
    y_new = -(x-50)**2 + delta
    plt.plot(x, y, "-k")
    plt.fill_between(x, y_old, y_new, color=color)
    y_old = y_new

plt.ylim(0, 2500)
plt.show()

As you will notice this does not look like a rainbow. This is because the function we are using is a quadratic, in actual fact a rainbow is made of circles with different radii (there is also a fun maths project here!). This is also plotable by matplotlib, I would try this and make it so you can plot more than the 7 colors in the rainbow eg plot 1000 colors spanning the entire spectrum to make it really look like a rainbow!

Answer 3

Here is a modified solution of the accepted answer that uses trapezoids instead of rectangles.

import numpy as np
import pylab as plt

# a solution that uses rectangles
def rect(x,y,w,h,c):
    ax = plt.gca()
    polygon = plt.Rectangle((x,y),w,h,color=c)
    ax.add_patch(polygon)

# a solution that uses trapezoids
def polygon(x1,y1,x2,y2,c):
    ax = plt.gca()
    polygon = plt.Polygon( [ (x1,y1), (x2,y2), (x2,0), (x1,0) ], color=c )
    ax.add_patch(polygon)

def rainbow_fill(X,Y, cmap=plt.get_cmap("jet")):
    plt.plot(X,Y,lw=0)  # Plot so the axes scale correctly

    dx = X[1]-X[0]
    N  = float(X.size)

    for n, (x,y) in enumerate(zip(X,Y)):
        color = cmap(n/N)
        # uncomment to use rectangles
        # rect(x,0,dx,y,color)
        # uncomment to use trapezoids
        if n+1 == N: continue
        polygon(x,y,X[n+1],Y[n+1],color)

# Test data    
X = np.linspace(0,10,100)
Y = .25*X**2 - X
rainbow_fill(X,Y)
plt.show()