快捷導(dǎo)航

使用Python的Matplotlib繪制漸變的彩色曲線

更新時間：2023年08月18日 10:14:23 作者：天元浪子

這篇文章主要介紹了使用Python的Matplotlib繪制漸變的彩色曲線,Matplotlib是Python生態(tài)圈中應(yīng)用最為廣泛的2D繪圖庫,通過擴展工具包mplot3d,也可以繪制簡單的3D模型,需要的朋友可以參考下

Matplotlib繪制漸變的彩色曲線

Matplotlib是Python生態(tài)圈中應(yīng)用最為廣泛的2D繪圖庫，通過擴展工具包mplot3d，也可以繪制簡單的3D模型。

Matplotlib的使用非常簡單，以繪制一條單周期的正弦曲線為例，只需要寥寥數(shù)行代碼即可完成。

import numpy as np
import matplotlib.pyplot as plt
x = np.linspace(0, 2*np.pi, 200)
y = np.sin(x)
plt.plot(x, y, color='green')
plt.show()

上面的代碼在 [ 0 , 2 π ] [0, 2\pi] [0,2π]之間以均勻間隔選取了200個點，并計算它們的正弦值，再調(diào)用plot函數(shù)將這200個點的正弦值連成一條曲線，通過plot函數(shù)的color參數(shù)（也可以簡寫為c）設(shè)置正弦曲線的顏色為綠色——當(dāng)然，也可以設(shè)置為你想要的任意一種顏色。

在這里插入圖片描述

不過，略有遺憾的是，Matplotlib沒有提供直接的方法為這條曲線上的每一個點設(shè)置不同的顏色，而這一需求在數(shù)據(jù)可視化方面卻是非常普遍的。

比如，在一條變化復(fù)雜的曲線上，如果每個點的數(shù)值越大顏色越趨近于紅色，數(shù)值越小顏色越趨近于藍(lán)色，就更容易展示數(shù)據(jù)的變化規(guī)律。

既然沒有直接的方法可用，那就變通一下，將這200個點拆分成199條線段，每條線段設(shè)置不同的顏色，不就可以實現(xiàn)上面的需求了嗎？

這樣做的確可行，不過需要調(diào)用plot函數(shù)199次，過程漫長到令人無法忍受，并且更新顯示（比如縮放）時也會重復(fù)這個漫長的調(diào)用過程。

幸運的是，Matplotlib提供了另外一種方法，基本思路也是將這200個點拆分成199條線段，每條線段設(shè)置不同的顏色，只是不需要反復(fù)調(diào)用plot函數(shù)，因此速度和單次調(diào)用plot函數(shù)相差無幾。

這個方法就是Matplotlib的LineCollection對象。

在使用LineCollection之前，還需要生成199條線段的顏色。

下面這個函數(shù)實現(xiàn)可數(shù)值到顏色的映射，給定一組數(shù)據(jù)和一個顏色映射方案，即返回每個數(shù)據(jù)的映射顏色。

import numpy as np
import matplotlib.pyplot as plt
def color_map(data, cmap):
    """數(shù)值映射為顏色"""
    dmin, dmax = np.nanmin(data), np.nanmax(data)
    cmo = plt.cm.get_cmap(cmap)
    cs, k = list(), 256/cmo.N
    for i in range(cmo.N):
        c = cmo(i)
        for j in range(int(i*k), int((i+1)*k)):
            cs.append(c)
    cs = np.array(cs)
    data = np.uint8(255*(data-dmin)/(dmax-dmin))
    return cs[data]

實例化LineCollection時，需要提供多條線段的列表（每條線段可以有多個點組成），以及每條線段的顏色組成的列表，也可以設(shè)置線寬和線型。

如果需要繪制Colorbar的話，還需要為LineCollection指定和數(shù)據(jù)映射相同的顏色映射方案。完整代碼如下。

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
def color_map(data, cmap):
    """數(shù)值映射為顏色"""
    dmin, dmax = np.nanmin(data), np.nanmax(data)
    cmo = plt.cm.get_cmap(cmap)
    cs, k = list(), 256/cmo.N
    for i in range(cmo.N):
        c = cmo(i)
        for j in range(int(i*k), int((i+1)*k)):
            cs.append(c)
    cs = np.array(cs)
    data = np.uint8(255*(data-dmin)/(dmax-dmin))
    return cs[data]
x = np.linspace(0, 2*np.pi, 200)
y = np.sin(x)
ps = np.stack((x,y), axis=1)
segments = np.stack((ps[:-1], ps[1:]), axis=1)
cmap = 'viridis' # jet, hsv等也是常用的顏色映射方案
colors = color_map(np.cos(x)[:-1], cmap)
colors = color_map(y[:-1], cmap)
line_segments = LineCollection(segments, colors=colors, linewidths=3, linestyles='solid', cmap=cmap)
fig, ax = plt.subplots()
ax.set_xlim(np.min(x)-0.1, np.max(x)+0.1)
ax.set_ylim(np.min(y)-0.1, np.max(y)+0.1)
ax.add_collection(line_segments)
cb = fig.colorbar(line_segments, cmap='jet')
plt.show()

使用LineCollection繪圖時，axis不會自動在x軸和y軸兩端留白，通常需要調(diào)用set_xlim和set_ylim兩個函數(shù)設(shè)置坐標(biāo)軸的范圍。最終的顯示效果如下圖所示。

在這里插入圖片描述