如何在循环中绘制水平条形图根据另一列中的值更改条形图上的颜色

Question

我需要在matplotlib中绘制时间(时间戳)与空间(IntersectionId)单个水平条形图。条形图的颜色将根据currState中的另一列以时间间隔进行更改。颜色将是红色、绿色、黄色。我试图创建一个颜色和值的字典，但不确定如何在循环中使用它们来根据值更改颜色。我已经附上了一个样本csv下面连同一个代码，我试图实现什么，我已经写到现在。

category_colors = { 'red' : [2,3] , 'yellow' : [5,6] , 'green' : [7,8]}
date_test =  df_sample['timestamp']
y_test = ['123456']
data = np.array(list(df_sample.currState))
fig, ax = plt.subplots(figsize=(10, 1))
ax = plt.barh(y_test,date_test,label="trafficsignal")
data_cum = data.cumsum
plt.xlabel('timestamp') 
plt.ylabel('space')
plt.title('TimeSpace')
plt.legend()
plt.show()

timestamp                                       currState          IntersectionId    
2020-02-26 16:12:13.131484                        3                    12345
2020-02-26 16:12:14.131484                        3                    12345
2020-02-26 16:12:15.131484                        3                    12345
2020-02-26 16:12:16.131484                        5                    12345
2020-02-26 16:12:17.131484                        5                    12345
2020-02-26 16:12:18.131484                        5                    12345
2020-02-26 16:12:19.131484                        6                    12345
2020-02-26 16:12:20.131484                        6                    12345
2020-02-26 16:12:21.131484                        6                    12345

当前绘图：

​

​

所需的绘图：

​

​

Answer 1

我不知道有什么绘图包可以让您根据示例表的结构以一种简单的方式创建此图。一种选择是计算一个start和一个end变量，然后像this question的答案中那样创建图，例如使用this answer中的Altair Gantt chart。

这里，我提供了两个使用matplotlib的解决方案。通过查看matplotlib gallery，我偶然发现了broken_barh绘图函数，该函数提供了一种创建所需绘图的方法。使用它时有两个主要的障碍需要克服：

1. 决定x轴使用的单位，并计算xranges参数accordingly;
2. Creating和格式化x刻度和刻度标签。

让我首先创建一个与您的数据集相似的示例数据集，请注意，您需要根据您的代码调整color_dict：

import numpy as np                 # v 1.19.2
import pandas as pd                # v 1.1.3
import matplotlib.pyplot as plt    # v 3.3.2
import matplotlib.dates as mdates

## Create sample dataset

# Light color codes
gre = 1
yel_to_red = 2
red = 3
yel_to_gre = 4
color_dict = {1: 'green', 2: 'yellow', 3: 'red', 4: 'yellow'}

# Light color duration in seconds
sec_g = 45
sec_yr = 3
sec_r = 90
sec_yg = 1

# Light cycle
light_cycle = [gre, yel_to_red, red, yel_to_gre]
sec_cycle = [sec_g, sec_yr, sec_r, sec_yg]

ncycles = 3
sec_total = ncycles*sum(sec_cycle)

# Create variables and store them in a pandas dataframe with the datetime as index
IntersectionId = 12345
currState = np.repeat(ncycles*light_cycle, repeats=ncycles*sec_cycle)
time_sec = pd.date_range(start='2021-01-04 08:00:00', freq='S', periods=sec_total)
df = pd.DataFrame(dict(IntersectionId = np.repeat(12345, repeats=ncycles*sum(sec_cycle)),
                       currState = currState),
                  index = time_sec)

​

broken_barh函数接受元组格式的数据，对于组成水平条的每个彩色矩形，您需要提供左下角的xy坐标以及沿每个轴的长度，如下所示：

xranges=[(x1_start, x1_length), (x2_start, x2_length), ... ], yranges=(y_all_start, y_all_width)

请注意，yranges适用于所有矩形。为x轴选择的单位决定了必须如何处理数据以及如何创建x刻度和刻度标签。这里有两个替代方案。

​

​

Matplotlib broken_barh，matplotlib日期编号为x轴刻度

在这种方法中，提取光线发生变化的行的时间戳，然后将其转换为matplotlib date numbers。这使得使用matplotlib date tick locator和formatter成为可能。这种使用x轴值的matplotlib日期来简化刻度格式化的方法受到了this answer by ImportanceOfBeingErnest的启发。

对于这个解决方案和下一个解决方案，由于this Gist by alimanfoo提供的一般思想，用于获取光照变化指数和计算周期长度的代码都基于this answer by Jaime。

## Compute variables needed to define the plotting function arguments

states = np.array(df['currState'])

# Create a list of indices of the rows where the light changes
# (i.e. where a new currState code section starts)
starts_indices = np.where(np.concatenate(([True], states[:-1] != states[1:])))

# Append the last index to be able to compute the duration of the last
# light color period recorded in the dataset
starts_end_indices = np.append(starts_indices, states.size-1)

# Get the timestamps of those rows and convert them to python datetime format
starts_end_pydt = df.index[starts_end_indices].to_pydatetime()

# Convert the python timestamps to matplotlib date number that is used as the
# x-axis unit, this makes it easier to format the tick labels
starts_end_x = mdates.date2num(starts_end_pydt)

# Get the duration of each light color in matplotlib date number units
lengths = np.diff(starts_end_x)

# Add one second (computed in python datetime units) to the duration of
# the last light to make the bar chart left and right inclusive instead
# of just left inclusive
pydt_second = (max(starts_end_x) - min(starts_end_x))/starts_end_indices[-1]
lengths[-1] = lengths[-1] + pydt_second

# Compute the arguments for the broken_barh plotting function
xranges = [(start, length) for start, length in zip(starts_end_x, lengths)]
yranges = (0.75, 0.5)
colors = df['currState'][starts_end_indices[:-1]].map(color_dict)


## Create horizontal bar with colors by using the broken_barh function
## and format ticks and tick labels

fig, ax = plt.subplots(figsize=(10,2))
ax.broken_barh(xranges, yranges, facecolors=colors, zorder=2)

# Create and format x ticks and tick labels
loc = mdates.AutoDateLocator()
ax.xaxis.set_major_locator(loc)
formatter = mdates.AutoDateFormatter(loc)
formatter.scaled[1/(24.*60.)] = '%H:%M:%S' # adjust this according to time range
ax.xaxis.set_major_formatter(formatter)

# Format y-axis and create y tick and tick label
ax.set_ylim(0, 2)
ax.set_yticks([1])
ax.set_yticklabels([df['IntersectionId'][0]])

plt.grid(axis='x', alpha=0.5, zorder=1)

plt.show()

​

​

​

​

以秒为x轴刻度的Matplotlib broken_barh

这种方法利用了表的索引可以用来以秒为单位计算灯光持续时间的事实。缺点是这次必须从头开始创建x ticks和tick标签。代码的编写使得标签根据数据集覆盖的总持续时间自动具有良好的格式。唯一需要调整的是刻度数，因为这取决于图形的宽度。

用于在刻度之间自动选择适当时间步长的代码基于this answer by kennytm。datetime字符串格式代码在here中列出。

## Compute the variables needed for the plotting function arguments
## using the currState variable

states = np.array(df['currState'])

# Create list of indices indicating the rows where the currState code
# changes: note the comma to unpack the tuple
starts_indices, = np.where(np.concatenate(([True], states[:-1] != states[1:])))
# Compute durations of each light in seconds
lengths = np.diff(starts_indices, append=states.size)


## Compute the arguments for the plotting function

xranges = [(start, length) for start, length in zip(starts_indices, lengths)]
yranges = (0.75, 0.5)
colors = df['currState'][starts_indices].map(color_dict)


## Create horizontal bar with colors using the broken_barh function

fig, ax = plt.subplots(figsize=(10,2))
ax.broken_barh(xranges, yranges, facecolors=colors, zorder=2)


## Create appropriate x ticks and tick labels

# Define time variable and parameters needed for computations
time = pd.DatetimeIndex(df.index).asi8 // 10**9 # time is in seconds
tmin = min(time)
tmax = max(time)
trange = tmax-tmin

# Choose the approximate number of ticks, the exact number depends on
# the automatically selected time step
approx_nticks = 6 # low number selected because figure width is only 10 inches
round_time_steps = [15, 30, 60, 120, 180, 240, 300, 600, 900, 1800, 3600, 7200, 14400]
time_step = min(round_time_steps, key=lambda x: abs(x - trange//approx_nticks))

# Create list of x ticks including the right boundary of the last time point
# in the dataset regardless of whether not it is aligned with the time step
timestamps = np.append(np.arange(tmin, tmax, time_step), tmax+1)
xticks = timestamps-tmin
ax.set_xticks(xticks)

# Create x tick labels with format depending on time step
fmt_time = '%H:%M:%S' if time_step <= 60 else '%H:%M'
xticklabels = [pd.to_datetime(ts, unit='s').strftime(fmt_time) for ts in timestamps]
ax.set_xticklabels(xticklabels)


## Format y-axis limits, tick and tick label

ax.set_ylim(0, 2)
ax.set_yticks([1])
ax.set_yticklabels([df['IntersectionId'][0]])


plt.grid(axis='x', alpha=0.5, zorder=1)

plt.show()

​

​

​

更多文档：to_datetime、to_pydatetime、strftime