Creating Horizontal Bars with Itertools and For Loops (Python)

Question

I've been playing around with Matplotlib and created a horizontal bar using the following algorithm (Full code and junk data provided at the bottom of this 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'))

which produces the following graphical output:

So I've been trying to make the code more efficient by introducing itertools

I managed to simplify the above code into a version 2:

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

Great, this also produces the above bar in the same order with the same colours.

The Problem

But I'm struggling with my next objective which is to replace facecolors=('tab:brown', 'tab:blue', 'tab:green') with a function that uses a for loop. This function ideally selects the correct colour for each bar based on the thickness. All 3 bars return a brown colour as the function continuously returns the value associated with the else statement (see image below).

I've attempted substituting next(cycle_thick) in place of the variable cycle_think in the function, but then only one of the colours is correct again.

The colour_checker() function is as follows:

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

Any hints or suggestions welcomed!

Full Code and Junk Data

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

question from:https://stackoverflow.com/questions/65946823/creating-horizontal-bars-with-itertools-and-for-loops-python

Answer 1

Since the intention of the outcome was ambiguous, I have created examples for all three versions I can imagine.

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

Sample output: