使用 Itertools 和 For 循环创建水平条 (Python)

Posted 2023-02-25

【中文标题】使用 Itertools 和 For 循环创建水平条 (Python)

【英文标题】：Creating Horizontal Bars with Itertools and For Loops (Python)

【发布时间】：2021-01-28 23:42:38

【问题描述】：

我一直在使用 Matplotlib，并使用以下算法创建了一个水平条（本文底部提供了完整的代码和垃圾数据）。

# Version 1
ax.broken_barh([(depth_start[0], thick[0]), (depth_start[1], thick[1]), (depth_start[2], thick[2])], (25, 0.8),
               facecolors=('tab:brown', 'tab:blue', 'tab:green'))

产生以下图形输出：

所以我一直试图通过引入itertools 来提高代码效率

我设法将上述代码简化为版本 2：

# Version 2
for i in thick:
    ax.broken_barh([(next(cycle_depth), next(cycle_thick))], (15, 0.8), facecolors=(next(cycle_colour)))

太好了，这也会以相同的顺序以相同的颜色生成上面的条形图。

问题

但我正在努力实现我的下一个目标，即将facecolors=('tab:brown', 'tab:blue', 'tab:green') 替换为使用for loop 的function。此功能理想地根据粗细为每个条形选择正确的颜色。由于函数不断返回与 else 语句关联的值，所有 3 个条形都返回棕色（见下图）。

我尝试用 next(cycle_thick) 代替函数中的变量 cycle_think，但后来只有一种颜色再次正确。

colour_checker() 函数如下：

def colour_checker():
    if cycle_thick == 10:
        return 'tab:green'
    elif cycle_thick == 20:
        return 'tab:blue'
    else:
        return 'tab:brown'

# Version 3
for i in thick:
    ax.broken_barh([(next(cycle_depth), next(cycle_thick))], (10, 0.8), facecolors=colour_checker())

欢迎任何提示或建议！

完整代码和垃圾数据

import itertools
import matplotlib.pyplot as plt

# Junk data in the form of lists
depth_start = [90, 70, 40]  # top of lithology
thick = [30, 20, 10]  # thickness for each lithology
colour = ('tab:brown', 'tab:blue', 'tab:green')

# Lists to be cycled through
cycle_colour = itertools.cycle(colour)
cycle_depth = itertools.cycle(depth_start)
cycle_thick = itertools.cycle(thick)

#setting up the plot
fig, ax = plt.subplots()

def colour_checker():
    if cycle_thick == [0]:
        return 'tab:green'
    elif cycle_thick == [1]:
        return 'tab:blue'
    else:
        return 'tab:brown'


# Version 1
ax.broken_barh([(depth_start[0], thick[0]), (depth_start[1], thick[1]), (depth_start[2], thick[2])], (25, 0.8),
               facecolors=('tab:brown', 'tab:blue', 'tab:green'))

# Version 2
for i in thick:
    ax.broken_barh([(next(cycle_depth), next(cycle_thick))], (15, 0.8), facecolors=(next(cycle_colour)))

# Version 3
for i in thick:
    ax.broken_barh([(next(cycle_depth), next(cycle_thick))], (10, 0.8), facecolors=colour_checker())

ax.set_ylabel('X_UTM Position')
ax.set_xlabel('MAMSL')

plt.show()

【问题讨论】：

您坚持使用循环仪吗？任务也不完全清楚。您想在条形图中组合三个列表中的每个索引元素吗？无论其他两个列表生成为条形，您是否要循环显示三种颜色？如果条形图具有特定长度（例如，10 是绿色），您是否想要特定颜色？您的问题的所有三种解释都是可能的，并且处理方式不同。

您好，T 先生，感谢您的反馈，感谢您抽出宝贵时间解决我的困惑问题。我必须努力使我的问题变得清晰。是的，我的目标是为特定长度的 barh 分配特定的颜色。我没有必要使用循环。很像您使用 zip() 函数的方式。将使用更多。

唷。我担心你实际上想要第四个不同的版本。很高兴这解决了。 zip 确实是摆脱笨拙的索引任务的绝佳工具。尽情探索 matplotlib。

【参考方案1】：

由于结果的意图不明确，我为我能想象的所有三个版本都创建了示例。

import matplotlib.pyplot as plt

# Junk data in the form of lists
depth_start = [90, 70, 40, 200, 170, 140]  # top of lithology
thick = [30, 20, 10, 20, 10, 30]  # thickness for each lithology
colour = ('tab:brown', 'tab:blue', 'tab:green')

#setting up the plot
fig, ax = plt.subplots()

#Version 1: using zip to chain all three lists
for start, length, color in zip(depth_start, thick, colour+colour[::-1]):
    ax.broken_barh([(start, length)], (-0.4, 0.8), facecolors=color)
    
#Version 2: color cycler repetitive color assignments 
from itertools import cycle
cycle_colour = cycle(colour)

for start, length in zip(depth_start, thick):
    ax.broken_barh([(start, length)], (0.6, 0.8), facecolors=next(cycle_colour))
    
#Version 3: lookup table to color bars of a specific length with a certain color
color_dic = 30: 'tab:brown', 20: 'tab:blue', 10: 'tab:green'

for start, length in zip(depth_start, thick):
    ax.broken_barh([(start, length)], (1.6, 0.8), facecolors=color_dic[length])

ax.set_yticks(range(3)) 
ax.set_yticklabels(["Version 1", "Version 2", "Version 3"])   

plt.show()

示例输出：