Matplotlib: correct placement of subplot annotations?

Question

I want to create subplots with Matplotlib by looping over my data. However, I don't get the annotations into the correct position, apparently not even into the correct subplot. Also, the common x- and y-axis labels don't work.

My real data is more complex but here is an example that reproduces the error:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
import seaborn as sns

# create data
distributions = []
first_values = []
second_values = []
for i in range(4):
    distributions.append(np.random.normal(0, 0.5, 100))
    first_values.append(np.random.uniform(0.7, 1))
    second_values.append(np.random.uniform(0.7, 1))

# create subplot
fig, axes = plt.subplots(2, 2, figsize = (15, 10))
legend_elements = [Line2D([0], [0], color = '#76A29F', lw = 2, label = 'distribution'),
                   Line2D([0], [0], color = '#FEB302', lw = 2, label = '1st value', linestyle = '--'),
                   Line2D([0], [0], color = '#FF5D3E', lw = 2, label = '2nd value')]

# loop over data and create subplots    
for data in range(4):
    if data == 0:
        position = axes[0, 0]
    if data == 1:
        position = axes[0, 1]
    if data == 2:
        position = axes[1, 0]
    if data == 3:
        position = axes[1, 1]
        
    dist = distributions[data]
    first = first_values[data]
    second = second_values[data]
        
    sns.histplot(dist, alpha = 0.5, kde = True, stat = 'density', bins = 20, color = '#76A29F', ax = position)
    sns.rugplot(dist, alpha = 0.5, color = '#76A29F', ax = position)
    position.annotate(f'{np.mean(dist):.2f}', (np.mean(dist), 0.825), xycoords = ('data', 'figure fraction'), color = '#76A29F')
    position.axvline(first, 0, 0.75, linestyle = '--', alpha = 0.75, color = '#FEB302')
    position.axvline(second, 0, 0.75, linestyle = '-', alpha = 0.75, color = '#FF5D3E')
    position.annotate(f'{first:.2f}', (first, 0.8), xycoords = ('data', 'figure fraction'), color = '#FEB302')
    position.annotate(f'{second:.2f}', (second, 0.85), xycoords = ('data', 'figure fraction'), color = '#FF5D3E')
    position.set_xticks(np.arange(round(min(dist), 1) - 0.1, round(max(max(dist), max([first]), max([second])), 1) + 0.1, 0.1))

plt.xlabel("x-axis name")
plt.ylabel("y-axis name")
plt.legend(handles = legend_elements, bbox_to_anchor = (1.5, 0.5))
plt.show()

The resulting plot looks like this:

What I want is to have

1. the annotations in the correct subplot next to the vertical lines / the mean of the distribution
2. shared x- and y-labels for all subplot or at least for each row / column

Any help is highly appreciated!

Answer 1

If you use the function to make the subplot a single array (axes.flatten()) and modify it to draw the graph sequentially, you can draw the graph. The colors of the annotations have been partially changed for testing purposes.

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
import seaborn as sns

np.random.seed(202000104)
# create data
distributions = []
first_values = []
second_values = []
for i in range(4):
    distributions.append(np.random.normal(0, 0.5, 100))
    first_values.append(np.random.uniform(0.7, 1))
    second_values.append(np.random.uniform(0.7, 1))
    
fig, axes = plt.subplots(2, 2, figsize=(15, 10))

legend_elements = [Line2D([0], [0], color = '#76A29F', lw = 2, label = 'distribution'),
                   Line2D([0], [0], color = '#FEB302', lw = 2, label = '1st value', linestyle = '--'),
                   Line2D([0], [0], color = '#FF5D3E', lw = 2, label = '2nd value')]

for i,ax in enumerate(axes.flatten()):
    sns.histplot(distributions[i], alpha=0.5, kde=True, stat='density', bins=20, color='#76A29F', ax=ax)
    sns.rugplot(distributions[i], alpha=0.5, color='#76A29F', ax=ax)
    ax.annotate(f'{np.mean(distributions[i]):.2f}', (np.mean(distributions[i]), 0.825), xycoords='data', color='red')
    ax.axvline(first_values[i], 0, 0.75, linestyle = '--', alpha = 0.75, color = '#FEB302')
    ax.axvline(second_values[i], 0, 0.75, linestyle = '-', alpha = 0.75, color = '#FF5D3E')
    ax.annotate(f'{first_values[i]:.2f}', (first_values[i], 0.8), xycoords='data', color='#FEB302')
    ax.annotate(f'{second_values[i]:.2f}', (second_values[i], 0.85), xycoords='data', color = '#FF5D3E')
    ax.set_xticks(np.arange(round(min(distributions[i]), 1) - 0.1, round(max(max(distributions[i]), max([first_values[i]]), max([second_values[i]])), 1) + 0.1, 0.1))

plt.xlabel("x-axis name")
plt.ylabel("y-axis name")
plt.legend(handles = legend_elements, bbox_to_anchor = (1.35, 0.5))

plt.show()