Redefine x axis in a heatmap seaborn

Question

I have a Dataframe with some cities. This DataFrame has 5 columns (city, date, and 3 columns about values).

The date value goes from 01/01/2015 to 31/12/2019.

I did a code that plots it. But I need to improve two things: The x axis needs to be monthly, and I also need a black line to delimit theses monthly date.

total_matrix = total.pivot_table(index="Estacao", columns="Date", values="Media_Horaria_MP10")

fig, ax = plt.subplots(figsize=(18, 9), dpi=150) 

ax = sns.heatmap(total_matrix, cmap = 'crest')
for i in range(total_matrix.shape[1]+1):
    ax.axhline(i, color='black', lw=2)

plt.title('MP10')
plt.show()

As you can see, I already have delimited the different cities ("Estacao"). How can I improve it?

Answer 1

I am giving my answer without testing as data for total wasn't provided. Here is how you can do it

import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt

# Assuming total is your DataFrame
# Ensure 'Date' is in datetime format
total['Date'] = pd.to_datetime(total['Date'])

# Pivot the data
total_matrix = total.pivot_table(
    index="Estacao", columns="Date", values="Media_Horaria_MP10")

# Create the plot
fig, ax = plt.subplots(figsize=(18, 9), dpi=150)

# Plot the heatmap
ax = sns.heatmap(total_matrix, cmap='crest', cbar=True)


# This iterates over the columns (dates), adds a vertical line for the start of each month
for date in total_matrix.columns:
    if date.day == 1: 
        ax.axvline(total_matrix.columns.get_loc(date), color='green', lw=2)

# Set x-axis with monthly labels
ax.set_xticks([total_matrix.columns.get_loc(date)
              for date in total_matrix.columns if date.day == 1])
ax.set_xticklabels([date.strftime('%b-%Y')
                   for date in total_matrix.columns if date.day == 1], rotation=45)

plt.title('MP10')
plt.show()