python - 通过第三个值绘制线条颜色并添加颜色条

Question

我想在桁架系统 (f) 中绘制力图，并根据它们的大小给它们颜色。所以我根据它们的节点坐标绘制这些线，并使用地震颜色图给它们一个颜色。虽然我想出了如何分配相应的颜色，但我无法在侧面绘制颜色条。我怎样才能添加这个？

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.cm as cm

def elmCoord(elements,nodes,i):
    idx1=df_index(nodes,elements['n1'][i],'name')
    idx2=df_index(nodes,elements['n2'][i],'name')
    x1=nodes['x'][idx1]
    y1=nodes['y'][idx1]
    x2=nodes['x'][idx2]
    y2=nodes['y'][idx2]
    return x1,x2,y1,y2

def df_index(df,val,col_ID): return df.index[df[col_ID] == val].tolist()[0]

nodes=[['node 1', 3, 0],['node 2', 0, 0],['node 3', 0, 4]]
elements=[['element 1','node 2','node 3'],
          ['element 2','node 1','node 3'],
          ['element 3','node 1','node 2']]
nodes=pd.DataFrame(nodes,columns=['name','x','y'])
nodes=nodes.astype({'x':'float64','y':'float64'})
elements=pd.DataFrame(elements,columns=['name','n1','n2'])

f=[-1000,2000,3000]
fig = plt.figure(figsize = (6,4))
ax = fig.add_subplot(111)
c = cm.seismic((f-np.min(f))/(np.max(f)-np.min(f)))
for i in range(len(elements)):
    x1,x2,y1,y2 = elmCoord(elements,nodes,i)
    ax.plot([x1,x2],[y1,y2],'-', linewidth=2, markersize=5, c=c[i])
fig.colorbar(f,ax=ax)
plt.show()

标签： pythonmatplotlib

Answer 1

要创建自定义 colorbar，ScalarMappable应使用相同的颜色图和规范创建 a：

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.cm as cm

def elmCoord(elements, nodes, i):
    idx1 = df_index(nodes, elements['n1'][i], 'name')
    idx2 = df_index(nodes, elements['n2'][i], 'name')
    x1 = nodes['x'][idx1]
    y1 = nodes['y'][idx1]
    x2 = nodes['x'][idx2]
    y2 = nodes['y'][idx2]
    return x1, x2, y1, y2

def df_index(df, val, col_ID):
    return df.index[df[col_ID] == val].tolist()[0]

nodes = [['node 1', 3, 0], ['node 2', 0, 0], ['node 3', 0, 4]]
elements = [['element 1', 'node 2', 'node 3'],
            ['element 2', 'node 1', 'node 3'],
            ['element 3', 'node 1', 'node 2']]
nodes = pd.DataFrame(nodes, columns=['name', 'x', 'y'])
nodes = nodes.astype({'x': 'float64', 'y': 'float64'})
elements = pd.DataFrame(elements, columns=['name', 'n1', 'n2'])

f = [-1000, 2000, 3000]
fig = plt.figure(figsize=(6, 4))
ax = fig.add_subplot(111)
norm = plt.Normalize(np.min(f), np.max(f))
cmap = plt.get_cmap('seismic')
c = cmap(norm(f))
for i in range(len(elements)):
    x1, x2, y1, y2 = elmCoord(elements, nodes, i)
    ax.plot([x1, x2], [y1, y2], '-', linewidth=2, markersize=5, c=c[i])
fig.colorbar(cm.ScalarMappable(norm=norm, cmap=cmap), ax=ax, ticks=f)
plt.show()

PS：您可以选择从f:设置刻度fig.colorbar(..., ticks=f)。