Colored bar plots with confidence intervals

Question

I needed to create a bar plot that show:

• the mean from some series;

• them 95% confidence interval; and,

• bars might be colored blue if they are definitely above this value (given the confidence interval), red if they are definitely below this value, or white if they contain this value.

Well, I coded the following lines, but considering I am learning Python, I am always thinking: is this a pythonic way?

You can see this code on here.

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import norm
from matplotlib.cm import ScalarMappable

np.random.seed(12345)

df = pd.DataFrame([np.random.normal(32000,200000,3650), 
                   np.random.normal(43000,100000,3650), 
                   np.random.normal(43500,140000,3650), 
                   np.random.normal(48000,70000,3650)], 
                  index=[1992,1993,1994,1995])

# Y select by user
y_user = 40000

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

# Calculating the probabilities
# I definitely believe this could be more readable
colorlist = [norm.sf(40000, loc=df[df.index == year].mean(axis = 1), scale=df[df.index == year].sem(axis = 1)).item(0) for year in df.index]

# Selecting the pallete
my_cmap = plt.cm.get_cmap('RdBu')
colors = my_cmap(colorlist)

# Creating the bars
ax.bar(df.index, df.mean(axis = 1), yerr=df.sem(axis = 1) * norm.ppf(0.975), color=colors, capsize=10)

# Ploting the y choose by user
plt.axhline(y = y_user, color = 'r', linestyle = '--')

# Controling the xticks
plt.xticks(df.index)

# Creating the colorbar 
sm = ScalarMappable(cmap=my_cmap)
sm.set_array([])

cbar = plt.colorbar(sm)
cbar.set_label('Probability to contain the value choose ({})'.format(y_user), rotation=270, labelpad=25)
cbar.set_ticks(np.arange(0, 1.1 , 0.1))

fig.savefig("GregOliveira_Week3.png")

plt.show()

Answer 1

Sort your imports.

Move your code into functions instead of leaving it in the global namespace.

Don't use np.random; it's deprecated in favour of RNG instances. In my suggested code I use a seed producing data vaguely similar but not equal to yours.

Don't call normal four times; call it once and specify a shape whose outer dimension is 4.

Don't make your dataframe with four rows and 3650 columns; this makes everything more difficult. Your dataframe should be transposed so that the columns correspond to years.

Is y_user actually selected by the user? It seems not. You're using a notebook. Notebooks have input utilities; use one of these (I have not shown this).

I definitely believe this could be more readable

You're absolutely right! Don't use a list comprehension at all; make a single vectorised call to sf().

Avoid plt.axhline when you have an axis reference to use instead. Same for xticks.

Delete set_array(); you don't need it.

Grammar and spelling: Y select by user -> Y selected by user; pallete -> palette; Ploting -> Plotting; Controling -> Controlling; Probability to contain the value choose -> Probability to contain the chosen value.

Delete nearly all of your comments since they're obvious to the average reader. Consider adding comments for non-obvious things, such as spelling out the acronyms for ppf and sem.

I consider the use of arange for your ticks to be more legible as a call to linspace, since the endpoint is included.

mean() and sem() are called multiple times; call these only once and keep a variable.

You need a chart title and axis labels. Other than the obvious "Year" I have not shown these.

Suggested

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.cm import ScalarMappable
from numpy.random import default_rng
from scipy.stats import norm


def make_data(
    n_years: int = 4, first_year: int = 1992, n_samples: int = 3650,
) -> pd.DataFrame:
    rand = default_rng(seed=18)
    return pd.DataFrame(
        rand.normal(
            loc  =( 32e3,  43e3,  43.5e3, 48e3),
            scale=(200e3, 100e3, 140.0e3, 70e3),
            size=(n_samples, n_years),
        ),
        columns=range(first_year, first_year + n_years),
    )


def make_plot(df: pd.DataFrame, y_user: float) -> plt.Figure:
    fig, ax = plt.subplots()

    mean = df.mean()
    sem = df.sem()  # Standard error of mean

    color_list = norm.sf(x=y_user, loc=mean, scale=sem)  # survival function
    my_cmap = plt.cm.get_cmap('RdBu')
    colors = my_cmap(color_list)

    ax.bar(
        df.columns, mean, color=colors, capsize=10,
        yerr=sem * norm.ppf(0.975),  # percent-point function
    )
    ax.axhline(y=y_user, color='r', linestyle='--')
    ax.set_xticks(df.columns)
    ax.set_xlabel('Year')

    cbar = plt.colorbar(ScalarMappable(cmap=my_cmap))
    cbar.set_label(
        f'Probability to contain chosen value ({y_user:.0f})',
        rotation=270, labelpad=25,
    )
    cbar.set_ticks(np.linspace(0, 1, 11))

    return fig


def main() -> None:
    y_user = 40e3  # selected by user (eventually?)
    df = make_data()
    make_plot(df, y_user)
    plt.show()


if __name__ == '__main__':
    main()

Output