使用 Itertools 和 For 循环创建单杠 (Python)

Question

我一直在玩弄 Matplotlib 并使用以下算法创建了一个单杠（完整代码和垃圾数据在本 post 底部提供）。

# 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)))

太棒了，这也以相同的顺序用相同的颜色生成了上面的条。

问题

但我正在为我的下一个 objective 苦苦挣扎，即将 facecolors=('tab:brown', 'tab:blue', 'tab:green') 替换为使用 for loop 的 function。此功能理想地根据厚度为每个条选择正确的颜色。所有 3 条 return 棕色作为函数连续 returns 与 else 语句关联的值（见下图）。

我尝试用 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()

Answer 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()

示例输出：

使用 Itertools 和 For 循环创建单杠 (Python)

Creating Horizontal Bars with Itertools and For Loops (Python)

python

loops

for-loop

matplotlib

itertools