快捷導(dǎo)航

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

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

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

Matplotlib繪制漸變的彩色曲線

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

Matplotlib的使用非常簡(jiǎn)單，以繪制一條單周期的正弦曲線為例，只需要寥寥數(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個(gè)點(diǎn)，并計(jì)算它們的正弦值，再調(diào)用plot函數(shù)將這200個(gè)點(diǎn)的正弦值連成一條曲線，通過(guò)plot函數(shù)的color參數(shù)（也可以簡(jiǎn)寫為c）設(shè)置正弦曲線的顏色為綠色——當(dāng)然，也可以設(shè)置為你想要的任意一種顏色。

在這里插入圖片描述

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

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

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

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

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

這個(gè)方法就是Matplotlib的LineCollection對(duì)象。

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

下面這個(gè)函數(shù)實(shí)現(xiàn)可數(shù)值到顏色的映射，給定一組數(shù)據(jù)和一個(gè)顏色映射方案，即返回每個(gè)數(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]

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

如果需要繪制Colorbar的話，還需要為L(zhǎng)ineCollection指定和數(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繪圖時(shí)，axis不會(huì)自動(dòng)在x軸和y軸兩端留白，通常需要調(diào)用set_xlim和set_ylim兩個(gè)函數(shù)設(shè)置坐標(biāo)軸的范圍。最終的顯示效果如下圖所示。

在這里插入圖片描述