I have the following code that uses 4th order Runge Kutta to solve the equations of motion of a double pendulum. I am attempting to animate the positions of the two pendulums using FuncAnimation. I tried removing the 0th element of each coordinate list after every 10th new element is plotted. I thought this would create a fading trail effect. Instead it just plotted one point. (note x1 and y1 are the coordinates of the first pendulum, and x2 and y2 are the coordinates of the second coupled pendulum) I feel like I'm over complicating things. when I remove the if i % 10 ==1: loop the code animates fine, but it does not delete any previous points. I feel this should be fairly easy to do but I could not find any straight forward examples, and the documentation was not extremely helpful.
In short I have 4 lists of coordinates that describe the position of two objects. I want to animate the current position of these two objects and not any previous positions.
import matplotlib.pyplot as plt
import numpy as np
import matplotlib.animation as animation
from matplotlib.animation import FuncAnimation
from itertools import count
from IPython import display
h=.001
t=0
l1=1
l2=1
m1=1
m2=1
theta1= 0
theta2 = np.pi/2
thetadot1=0
thetadot2=0
g=9.8
x1=[]
y1=[]
x2=[]
y2=[]
E=[]
T=[]
Theta_dot1=[]
Theta_dot2=[]
kinetic_energy=[]
potential_energy=[]
while(t<4):
KE = .5*m1*((thetadot1*l1*np.cos(theta1))**2+(thetadot1*l1*np.sin(theta1))**2) + .5*m2*((thetadot1*l1*np.cos(theta1)+thetadot2*l2*np.cos(theta2))**2+(thetadot1*l1*np.sin(theta1)+thetadot2*l2*np.sin(theta2))**2)
PE = (m1*g*(l1-l1*np.cos(theta1))+m2*g*(l1+l2-l1*np.cos(theta1)-l2*np.cos(theta2)))
E.append(KE+PE)
kinetic_energy.append(KE)
potential_energy.append(PE)
Theta_dot1.append(thetadot1)
Theta_dot2.append(thetadot2)
T.append(t)
x1.append(l1*np.sin(theta1))
y1.append(-l1*np.cos(theta1))
x2.append(l1*np.sin(theta1)+l2*np.sin(theta2))
y2.append(-l1*np.cos(theta1)-l2*np.cos(theta2))
P1 = (-g*(2*m1+m2)*np.sin(theta1)-m2*g*np.sin(theta1-2*theta2)-2*np.sin(theta1-theta2)*m2*(l2*thetadot2**2+l1*thetadot1**2*(np.cos(theta1-theta2))))/(l1*(2*m1+m2*(1-np.cos(2*(theta1-theta2)))))
L1 = (2*np.sin(theta1-theta2)*(l1*thetadot1**2*(m1+m2)+np.cos(theta1)*g*(m1+m2)+l2*m2*thetadot2**2*np.cos(theta1-theta2)) )/(l2*(2*m1+m2*(1-np.cos(2*(theta1-theta2)))))
G1 = thetadot1
H1 = thetadot2
thetadot1_1 = thetadot1 + (h/2)*P1
thetadot2_1 = thetadot2 + (h/2)*L1
theta1_1 = theta1 + (h/2)*G1
theta2_1 = theta2 + (h/2)*H1
P2 = (-g*(2*m1+m2)*np.sin(theta1_1)-m2*g*np.sin(theta1_1-2*theta2_1)-2*np.sin(theta1_1-theta2_1)*m2*(l2*thetadot2_1**2+l1*thetadot1_1**2*(np.cos(theta1_1-theta2_1))))/(l1*(2*m1+m2*(1-np.cos(2*(theta1_1-theta2_1)))))
L2 = (2*np.sin(theta1_1-theta2_1)*(l1*thetadot1_1**2*(m1+m2)+np.cos(theta1_1)*g*(m1+m2)+l2*m2*thetadot2_1**2*np.cos(theta1_1-theta2_1)) )/(l2*(2*m1+m2*(1-np.cos(2*(theta1_1-theta2_1)))))
G2 = thetadot1_1
H2 = thetadot2_1
thetadot1_2 = thetadot1_1 + (h/2)*P2
thetadot2_2 = thetadot2_1 + (h/2)*L2
theta1_2 = theta1_1 + (h/2)*G2
theta2_2 = theta2_1 + (h/2)*H2
P3 = (-g*(2*m1+m2)*np.sin(theta1_2)-m2*g*np.sin(theta1_2-2*theta2_2)-2*np.sin(theta1_2-theta2_2)*m2*(l2*thetadot2_2**2+l1*thetadot1_2**2*(np.cos(theta1_2-theta2_2))))/(l1*(2*m1+m2*(1-np.cos(2*(theta1_2-theta2_2)))))
L3 = (2*np.sin(theta1_2-theta2_2)*(l1*thetadot1_2**2*(m1+m2)+np.cos(theta1_2)*g*(m1+m2)+l2*m2*thetadot2_2**2*np.cos(theta1_2-theta2_2)) )/(l2*(2*m1+m2*(1-np.cos(2*(theta1_2-theta2_2)))))
G3 = thetadot1_2
H3 = thetadot2_2
thetadot1_3 = thetadot1_2 + (h/2)*P3
thetadot2_3 = thetadot2_2 + (h/2)*L3
theta1_3 = theta1_2 + (h/2)*G3
theta2_3 = theta2_2 + (h/2)*H3
P4 = (-g*(2*m1+m2)*np.sin(theta1_3)-m2*g*np.sin(theta1_3-2*theta2_3)-2*np.sin(theta1_3-theta2_3)*m2*(l2*thetadot2_3**2+l1*thetadot1_3**2*(np.cos(theta1_3-theta2_3))))/(l1*(2*m1+m2*(1-np.cos(2*(theta1_3-theta2_3)))))
L4 = (2*np.sin(theta1_3-theta2_3)*(l1*thetadot1_3**2*(m1+m2)+np.cos(theta1_3)*g*(m1+m2)+l2*m2*thetadot2_3**2*np.cos(theta1_3-theta2_3)) )/(l2*(2*m1+m2*(1-np.cos(2*(theta1_3-theta2_3)))))
G4 = thetadot1_3
H4 = thetadot2_3
thetadot1 = thetadot1 + (h/6.0) * (P1+(2.*P2)+(2.0*P3) + P4)
thetadot2 = thetadot2 + (h/6.0) * (L1+(2.*L2)+(2.0*L3) + L4)
theta1 = theta1 + (h/6.0) * (G1+(2.*G2)+(2.0*G3) + G4)
theta2 = theta2 + (h/6.0) * (H1+(2.*H2)+(2.0*H3) + H4)
t=t+h
fig, axes = plt.subplots(nrows = 1, ncols = 1, figsize = (15,5))
axes.set_ylim(-2.5, 2.5)
axes.set_xlim(-2.5, 2.5)
plt.style.use("ggplot")
x_1,y_1,x_2,y_2 = [], [], [], []
def animate(i):
if i % 10 ==1: #this is the stuff that is not working
x_1.remove(0)
y_1.remove(0)
x_2.remove(0)
y_2.remove(0)
x_1.append(x1[i*10])
y_1.append(y1[i*10])
x_2.append(x2[i*10])
y_2.append(y2[i*10])
axes.plot(x_1,y_1,'.', color="red")
axes.plot(x_2,y_2,'.', color="gray", linewidth=0.5)
anim = FuncAnimation(fig, animate, interval=.1)
If you want fading lines you have to deal with the extra complexity of LineCollection, which is not easy to master.
import matplotlib.colors as colors
from matplotlib.collections import LineCollection
from matplotlib.lines import Line2D
n_points_to_render = 500
# opacity of the segments
alphas = np.linspace(0, 1, n_points_to_render)
# create "solid" color maps with a varying opacity
red = colors.to_rgb("red") + (0.0,)
redfade = colors.LinearSegmentedColormap.from_list('my', [red, "red"])
green = colors.to_rgb("green") + (0.0,)
greenfade = colors.LinearSegmentedColormap.from_list('my', [green, "green"])
def get_segments(i):
# LineCollection requires segments
_x1 = x1[i:i+n_points_to_render]
_y1 = y1[i:i+n_points_to_render]
_x2 = x2[i:i+n_points_to_render]
_y2 = y2[i:i+n_points_to_render]
points1 = np.vstack((_x1, _y1)).T.reshape(-1, 1, 2)
segments1 = np.hstack((points1[:-1], points1[1:]))
points2 = np.vstack((_x2, _y2)).T.reshape(-1, 1, 2)
segments2 = np.hstack((points2[:-1], points2[1:]))
return segments1, segments2
fig, ax = plt.subplots(nrows = 1, ncols = 1)
ax.set_ylim(-2.5, 2.5)
ax.set_xlim(-2.5, 2.5)
plt.style.use("ggplot")
# create and add two LineCollections
segments1, segments2 = get_segments(0)
lc1 = LineCollection(segments1, array=alphas, cmap=redfade, lw=2)
lc2 = LineCollection(segments2, array=alphas, cmap=greenfade, lw=2)
line1 = ax.add_collection(lc1)
line2 = ax.add_collection(lc2)
def animate(i):
segments1, segments2 = get_segments(i)
line1.set_segments(segments1)
line2.set_segments(segments2)
# create a legend as LineCollection doesn't have any by default
l1_legend = Line2D([0, 1], [0, 1], color="r", linewidth=2)
l2_legend = Line2D([0, 1], [0, 1], color="g", linewidth=2)
ax.legend([l1_legend, l2_legend], ['Line 1', 'Line 2'])
anim = FuncAnimation(fig, animate, frames=len(x1) - n_points_to_render, interval=30)
plt.show()
Related
How can we make a figure like one attached in matplotlib? Thick horizontal lines as markers with their ends connected through dashed lines. I tried, but couldn't find this marker style.
Following is my attempt -
import numpy as np
import matplotlib.pyplot as plt
rc = [0,1,2,3,4,5]
nci = [0.0000,-0.0001,0.4325,1.2711,0.2162,0.0988]
plt.plot(rc,nci,'b--',marker="_",markersize='15')
plt.xlabel("Reaction Coordinate",fontsize=22)
plt.ylabel('Relative Energy (eV)',fontsize=22)
plt.xticks(fontsize=22)
plt.yticks(fontsize=22)
plt.show()
There are several methods to draw what you want.You can use axline, define marker size with the scale of x axis, and others.
Here I write a simple (but bad-looking) solution.
import numpy as np
import matplotlib.pyplot as plt
# Data
len_arr = 6
x = [0, 1, 2, 3, 4, 5]
y = [0.0000,-0.0001,0.4325,1.2711,0.2162,0.0988]
x_gap = 0.2
# Dashed line
x_dash = np.empty(len_arr * 2)
for idx in range(len_arr):
x_dash[idx * 2 + 0] = x[idx] - x_gap
x_dash[idx * 2 + 1] = x[idx] + x_gap
y_dash = np.repeat(y, 2)
# Streched marker
x_strech = np.empty(len_arr * 3)
for idx in range(len_arr):
x_strech[idx * 3 + 0] = x[idx] - x_gap
x_strech[idx * 3 + 1] = x[idx] + x_gap
x_strech[idx * 3 + 2] = None
y_strech = np.repeat(y, 3)
# Plot
fig = plt.figure(figsize=(8, 6))
plt.plot(x_dash, y_dash, linestyle = '--', linewidth = 1, c = 'orange')
plt.plot(x_strech, y_strech, linestyle = '-', linewidth = 3, c = 'orange')
The result is as follows.
Currently I'm trying to plot a graph showing the rank of some equipment in operation, the rank goes from 1 to 300 (1 is the best, 300 is the worst) over a few days (df columns). What I'm trying to do, is a graph similar to this:
And what I got is this:
I would like to make the lines inclined as it is on the first graph instead of vertical, but I can't figure it out how. I found the base for the first graph on this question here and I started the code from there, this is what I end up having:
import matplotlib.pyplot as plt
import matplotlib.ticker as plticker
import numpy as np
def energy_rank(data, marker_width=.5, color='blue'):
y_data = np.repeat(data, 2)
x_data = np.empty_like(y_data)
x_data[0::2] = np.arange(1, len(data)+1) - (marker_width/2)
x_data[1::2] = np.arange(1, len(data)+1) + (marker_width/2)
lines = []
lines.append(plt.Line2D(x_data, y_data, lw=0.8, linestyle='dashed', color=color,alpha=1,marker='.'))
for x in range(0,len(data)*2, 2):
lines.append(plt.Line2D(x_data[x:x+2], y_data[x:x+2], lw=2, linestyle='solid', color=color))
return lines
head = 8
dfPlot = vazio.sort_values(dia, ascending = True).head(head)
data = dfPlot.to_numpy()
colorsHEX=('#FE5815','#001A70','#2F5C22','#B01338','#00030D','#2DE1FC','#2E020C','#B81D8C')
artists = []
for row, color in zip(data, colorsHEX):
artists.extend(energy_rank(row, color=color))
eixoXDatas = pd.to_datetime(list(vazio.columns),format='%d/%m/%y').strftime('%d/%b')
fig, ax = plt.subplots()
plt.xticks(np.arange(len(vazio.columns)),
eixoXDatas,
rotation = 35,
fontsize = 14)
plt.yticks(fontsize = 14)
plt.xlabel('Dias', fontsize=18)
plt.ylabel('Ranking', fontsize=18)
fig = plt.gcf()
fig.set_size_inches(16, 8)
for artist in artists:
ax.add_artist(artist)
ax.set_ybound([0,15])
ax.set_ylim(ax.get_ylim()[::-1])
ax.set_xbound([-0.1,float(len(vazio.columns))+2.5])
plt.yticks(np.arange(1,16,step=1))
ax.grid(axis='y',alpha=0.5)
lastDay = vazio.sort_values(vazio.iloc[:,-1:].columns.values[0], ascending = True).iloc[:,-1:]
lastDay = lastDay.head(head)
for inverter, pos in lastDay.iterrows():
ax.annotate(inverter, xy =(plt.gca().get_xlim()[1]-2.4, pos), color=colorsHEX[int(pos)-1])
I tried implementing on energy_rank function, removing the +/- parts on x_data but I only could end up with inclined lines with dots instead of the horizontal lines. Can anyone help me out how can I mantain the horziontal lines and instead of vertical dashed lines, implement inclined lines as the example above?
I imagine that is vertical because the points change on top of the x ticks. If you observe the 1st image, the horizontal bars are centralized on each x tick, so the lines "have some room" to be inclined.
vazio dataframe is as follows (contains the rank of each equipment):
Equipment 21-03-27 21-03-28 21-03-29 21-03-30 21-03-31 21-04-01 21-04-02
P01-INV-1-1 1 1 1 1 1 2 2
P01-INV-1-2 2 2 4 4 5 1 1
P01-INV-1-3 4 4 3 5 6 10 10
Here is an adaption of your energy_rank function creating horizontal line segments together with their connections. The line drawing part is inspired by this tutorial example. Optionally the area below the lines can be filled.
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
import numpy as np
def energy_rank(data, marker_width=.5, color='blue', ax=None, fill=False):
ax = ax or plt.gca()
y = data
x = np.arange(1, len(data) + 1)
segments1 = np.array([x - marker_width / 2, y, x + marker_width / 2, y]).T.reshape(-1, 2, 2)
lc1 = LineCollection(segments1, color=color)
lc1.set_linewidth(2)
lc1.set_linestyle('-')
lines_hor = ax.add_collection(lc1)
segments2 = np.array([x[:-1] + marker_width / 2, y[:-1], x[1:] - marker_width / 2, y[1:]]).T.reshape(-1, 2, 2)
lc2 = LineCollection(segments2, color=color)
lc2.set_linewidth(0.5)
lc2.set_linestyle('--')
lines_connect = ax.add_collection(lc2)
if fill:
ax.fill_between(segments1.reshape(-1,2)[:,0], segments1.reshape(-1,2)[:,1],
color=color, alpha=0.05)
return lines_hor, lines_connect
fig, ax = plt.subplots()
M, N = 5, 25
y = np.random.uniform(-2, 2, (M, N)).cumsum(axis=1)
y += np.random.uniform(0.5, 2, (M, 1)) - y.min(axis=1, keepdims=True)
colorsHEX = ('#FE5815', '#001A70', '#2F5C22', '#B01338', '#00030D')
for yi, color in zip(y, colorsHEX):
energy_rank(yi, ax=ax, color=color)
ax.set_xlim(0, N + 1)
ax.set_ylim(0, y.max() + 1)
plt.show()
I've created an animated plot of a wavefunction, psi:
def psi(x, t):
real = 0.4*np.cos(0.4*x - 0.08*t) + 0.6*np.cos(0.6*x - 0.18*t)
imag = 0.4*np.sin(0.4*x - 0.08*t) + 0.6*np.sin(0.6*x - 0.18*t)
square = real**2 + imag**2
return real, imag, square
I've then animated it successfully, however on adding the axvspan fill I've encountered an issue:
x = np.linspace(-10, 1000, 10000)
fig, (ax1, ax2) = plt.subplots(2,1)
line1, = ax1.plot([], [])
line2, = ax1.plot([], [])
line3, = ax2.plot([], [])
line = [line1, line2, line3]
def animate(i):
y1, y2, y3 = psi(x, t=i/2)
line1.set_data(x, y1)
line2.set_data(x, y2)
line3.set_data(x, y3)
spline = UnivariateSpline(x, y3-max(y3)/2, s=0)
r1, r2 = spline.roots()
ax2.axvspan(r1, r2, facecolor='b', alpha=0.5)
plt.legend(['Max Probability = %1.3f' % (max(y3))])
return line,
anim = animation.FuncAnimation(fig, animate, frames=600, interval = 100, blit=False, repeat=False)
It starts like
and it ends like .
Every iteration of the animation function, the fill increases across the page having started off as filling half the graph (I'd include a gif but at the moment that's a struggle I'm having with anaconda). I'm working under the assumption this is because the axes don't clear properly, however with the blit=false I assumed this wouldn't be a problem?
As asked for - the full psi function is detailed below:
n = 15
amp_scale = np.linspace(0, 0.8, n)
amp_init = norm.pdf(amp_scale, 0.4, 0.2)
#normalise wavefunction to prob=1
amp = []
for i in range(n):
amp_val = amp_init[i]/sum(amp_init)
amp.append(amp_val)
k = np.linspace(1.4, 2.6, n)
def psi (x, t=1, n=1, a = 1, k = 1, m = 1):
psi_real = 0
psi_imag = 0
for i in range(n):
a_val = a[i]
k_val = k[i]
w = (k_val**2)/(2*m)
psi_real+=a_val*np.cos(k_val*x - w*t)
psi_imag+=a_val*np.sin(k_val*x - w*t)
psi_squared = psi_real**2 + psi_imag**2
return psi_real, psi_imag, psi_squared
In your current code, a new axvspan() is continually added, never removed. You could explicitly remove the old span inside animate(). Or, similarly to what happens to the lines, update the position. A span is internally represented as a polygon, of which the coordinates can be set via .set_xy().
The function psi in the post doesn't seem to be the same as the function that generated the example plots. This also made that I couldn't get to work to calculate the spline and the roots. I replaced them by some simpler positions to show how the span can be updated during the animation.
The code also adds explicit x and y limits, as they weren't set in the question's code.
import matplotlib.pyplot as plt
from matplotlib import animation
import numpy as np
from scipy.interpolate import UnivariateSpline
def psi(x, t):
# the function from the question is adapted to more resemble the plot
real = (0.4 * np.cos(0.4 * x - 0.08 * t) + 0.6 * np.cos(0.6 * x - 0.18 * t)) * np.exp(- (x - t) ** 2 / 5000)
imag = (0.4 * np.sin(0.4 * x - 0.08 * t) + 0.6 * np.sin(0.6 * x - 0.18 * t)) * np.exp(- (x - t) ** 2 / 5000)
square = real ** 2 + imag ** 2
return real, imag, square
x = np.linspace(-10, 1000, 10000)
fig, (ax1, ax2) = plt.subplots(nrows=2, sharex=True)
line1, = ax1.plot([], [])
line2, = ax1.plot([], [])
line3, = ax2.plot([], [])
line = [line1, line2, line3]
span1 = ax2.axvspan(0, 0, facecolor='b', alpha=0.2)
ax1.set_xlim(x[0], x[-1]/2)
ax1.set_ylim(-1, 1)
ax2.set_ylim(0, 1.1)
def animate(i):
y1, y2, y3 = psi(x, t=i / 2)
line1.set_data(x, y1)
line2.set_data(x, y2)
line3.set_data(x, y3)
# this didn't work for me, spline.roots() gave me a long array of values
spline = UnivariateSpline(x, y3 - max(y3) / 2, s=0)
r1, r2 = spline.roots()[[0, -1]] # [[0, -1]] takes the first and the last
# r1, r2 = i - 50, i + 50
span1.set_xy([[r1, 0], [r1, 1], [r2, 1], [r2, 0], [r1, 0]])
plt.legend(['Max Probability = %1.3f' % (max(y3))])
return line, span1, ax2.legend_,
anim = animation.FuncAnimation(fig, animate, frames=600, interval=100, blit=False, repeat=False)
plt.plot()
The resulting end frame looks like (note that a slightly different function is used):
I am trying to color lines with a colormap so that the gradient color indicates the length of the line. My current program colors lines, depending on where they are on the x-axis.
All variables ending with _2 belong to the diagonal line.
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from matplotlib.colors import LinearSegmentedColormap
x_nom = [0, 1]
y_nom = [0, 1]
# x,y coordinates
x1 = 0
x2 = 0.5
y1 = 0
y2 = 1
x1_2 = 0
x2_2 = 1
y1_2 = 0
y2_2 = 1
# create segments from coordinates
x_coords = np.linspace(x1, x2, 100)
y_coords = np.linspace(y1, y2, 100)
x_coords_2 = np.linspace(x1_2, x2_2, 100)
y_coords_2 = np.linspace(y1_2, y2_2, 100)
# "resolution"
res_ar = np.linspace(x1, x2, len(x_coords))
res_ar_2 = np.linspace(x1_2, x2_2, len(x_coords_2))
# reshape array to fit matplotlib segmentation
points = np.array([x_coords, y_coords]).T.reshape(-1, 1, 2)
points_2 = np.array([x_coords_2, y_coords_2]).T.reshape(-1, 1, 2)
segments = np.concatenate([points[:-1], points[1:]], axis=1)
segments_2 = np.concatenate([points_2[:-1], points_2[1:]], axis=1)
fig, ax = plt.subplots(figsize=(5,5))
# create custom colormap
cvals = [0., 1.]
colors = ['black', 'white']
norm = plt.Normalize(min(cvals), max(cvals))
tuples = list(zip(map(norm, cvals), colors))
cmap = LinearSegmentedColormap.from_list('', tuples)
# plot segmented line
lc = LineCollection(segments, cmap=cmap, norm=norm)
lc_2 = LineCollection(segments_2, cmap=cmap, norm=norm)
lc.set_array((res_ar))
lc_2.set_array((res_ar_2))
ax.add_collection(lc)
ax.add_collection(lc_2)
plt.show()
As you can see, the diagonal line is plotted in the full colormap (color regions 0 to 1). It has a length of 1.41 (square root of x=1 and y=1). The upper line has a length of 1.12 (x=0.5, y=1). The upper line is about 0.8 times as long as the diagonal line. How would I use the colormap to color the line in region 0 to 0.8 instead of (as of now) 0 to 0.5?
Edit:
Thanks to tmdavison I was able to work it out. This code colors the line, depending on its length:
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from matplotlib.colors import LinearSegmentedColormap
x_nom = [0, 1]
y_nom = [0, 1]
# x,y coordinates
x1 = 0
x2 = 0.8
y1 = 0
y2 = 1
x1_2 = 0
x2_2 = 1
y1_2 = 0
y2_2 = 1
# create segments from coordinates
x_coords = np.linspace(x1, x2, 100)
y_coords = np.linspace(y1, y2, 100)
x_coords_2 = np.linspace(x1_2, x2_2, 100)
y_coords_2 = np.linspace(y1_2, y2_2, 100)
l_nom = np.sqrt(x_nom[1]**2 + y_nom[1]**2)
l = np.sqrt(x2**2 + y2**2)
correct = l/l_nom
# "resolution"
res_ar = np.linspace(x1, x2, len(x_coords)) * correct/x2
res_ar_2 = np.linspace(x1_2, x2_2, len(x_coords_2))
# reshape array to fit matplotlib segmentation
points = np.array([x_coords, y_coords]).T.reshape(-1, 1, 2)
points_2 = np.array([x_coords_2, y_coords_2]).T.reshape(-1, 1, 2)
segments = np.concatenate([points[:-1], points[1:]], axis=1)
segments_2 = np.concatenate([points_2[:-1], points_2[1:]], axis=1)
fig, ax = plt.subplots(figsize=(5,5))
# create custom colormap
cvals = [0., 1.]
colors = ['black', 'white']
norm = plt.Normalize(min(cvals), max(cvals))
tuples = list(zip(map(norm, cvals), colors))
cmap = LinearSegmentedColormap.from_list('', tuples)
# plot segmented line
lc = LineCollection(segments, cmap=cmap, norm=norm)
lc_2 = LineCollection(segments_2, cmap=cmap, norm=norm)
lc.set_array((res_ar))
lc_2.set_array((res_ar_2))
ax.add_collection(lc)
ax.add_collection(lc_2)
plt.show()
You can normalise the data array you are using to set the colours by the maximum value it will have. You could do this in two ways:
when you construct the res_ar and res_ar_2 arrays
res_ar = np.linspace(x1, x2, len(x_coords)) / x2
res_ar_2 = np.linspace(x1_2, x2_2, len(x_coords_2)) / x2_2
or, when you use those arrays later
lc.set_array((res_ar / x2))
lc_2.set_array((res_ar_2 / x2_2))
Doing one of those thing will mean the colour scale is always going from 0 to 1.
Strictly speaking you don't need to do this for you res_ar array, since that already goes from 0 to 1, but I added it here to be more general in case that line changed in the future.
I wonder if there is a more elegant way to draw the polygon in below code, or with a special plot function or parameter ?
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import norm
x = np.linspace(-4,4,150)
# plot density with shaded area showing Pr(-2 <= x <= 1)
lb = -2
ub = 1
d=norm.pdf(x)
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.plot(x, d)
### can this be done more elegantly ###
sx = np.linespace(lb,ub,100)
sd = norm.pdf(sx)
sx = [lb] + sx + [ub]
sd = [0] + list(sd) + [0]
xy = np.transpose(np.array([sx, sd]))
pgon = plt.Polygon(xy, color='b')
#######################################
ax.add_patch(pgon)
plt.show()
Perhaps you are looking for plt.fill_between:
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import norm
x = np.linspace(-4,4,150)
# plot density with shaded area showing Pr(-2 <= x <= 1)
lb = -2
ub = 1
d = norm.pdf(x)
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.plot(x, d)
idx = np.searchsorted(x,[lb,ub])
sx = x[idx[0]:idx[1]]
sd = d[idx[0]:idx[1]]
plt.fill_between(sx, sd, 0, color = 'b')
plt.show()