Search code examples
pandasmatplotlibxticks

How to set values of a vertical stem plot as xticks labels?


I would like to reverse a grouped data and use group name as xtick label to draw it side by side. below demo mostly good but the label position not as expected.

import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

def main():
  data = [['AAAAAA',8],['AAAAAA',9],['AAAAAA',10],['BBBBBB',5],['BBBBBB',6],['BBBBBB',7],['CCCCCC',1],['CCCCCC',2],['CCCCCC',3],['CCCCCC',4]]
  df = pd.DataFrame(data,columns=['name','value'])

  dfg = df.groupby('name')
  fig, ax = plt.subplots(figsize=(8, 4))
  i = 0
  ymin = df['value'].min()
  c1='#ececec'
  c2='#bcbcbc'
  color=c1  
  for ix, row in reversed(tuple(dfg)):
    print(ix,row)
    n = len(row['name'])
    x = np.linspace(i,i + n,n)
    ax.stem(x,row['value'])

    font_dict = {'family':'serif','color':'darkred', 'size':8}
    ax.text(i + n/2,ymin,ix,ha='right',va='top',rotation=90, fontdict=font_dict)
    if color == c1:
        color = c2
    else:
        color = c1
    
    plt.axvspan(i, i+n, facecolor=color, alpha=0.5)    
    i += len(row)    
    
  ax.xaxis.set_ticks_position('none') 
  plt.setp( ax.get_xticklabels(), visible=False)
  ax.grid(axis='y',color='gray', linestyle='dashed', alpha=1)    
  ax.spines[["top", "right"]].set_visible(False)
  
  fig.tight_layout() 
    
  plt.show()
  return

main()

Output: enter image description here

Welcome to comment any other proper way to do this, or how to improve the xticks down, use ymin properly not good way to do it.


Solution

  • If my understanding of what you are trying to achieve is correct, here is one way to do it:

    import numpy as np
    import pandas as pd
    from matplotlib import pyplot as plt
    
    df = pd.DataFrame(
        [
            ["AAAAAA", 8],
            ["AAAAAA", 9],
            ["AAAAAA", 10],
            ["BBBBBB", 5],
            ["BBBBBB", 6],
            ["BBBBBB", 7],
            ["CCCCCC", 1],
            ["CCCCCC", 2],
            ["CCCCCC", 3],
            ["CCCCCC", 4],
        ],
        columns=["name", "value"],
    )
    
    fig, ax = plt.subplots(figsize=(8, 4))
    i = 0
    c1 = "#ececec"
    c2 = "#bcbcbc"
    color = c1
    ticks = {}
    
    for ix, row in reversed(tuple(df.groupby("name"))):
        # Create stem plot
        n = len(row["name"])
        x = np.linspace(i, i + n, n)
        ax.stem(x, row["value"])
    
        # Create axvspan plot 
        if color == c1:
            color = c2
        else:
            color = c1
        ax.axvspan(i, i + n, facecolor=color, alpha=0.5)
    
        # Save positions and names in a dict
        for key, name in zip(x, row["name"]):
            if key not in ticks.keys():
                ticks[key] = name
            else:
                # Deal with multiple names for same tick
                ticks[key] += f"\n{name}"
        i += len(row)
    
    # Add ticks and ticks labels
    ax.set_xticks(ticks=list(ticks.keys()))
    ax.set_xticklabels(list(ticks.values()), fontsize=10, rotation="vertical")
    
    # In Jupyter notebook
    fig
    

    Output:

    enter image description here

    And to avoid repeating the labels, you can, for instance, do:

    ax.set_xticklabels(
        [
            "",
            "CCCCCC",
            "",
            "CCCCCC\nBBBBBB",
            "BBBBBB",
            "BBBBBB\nAAAAAA",
            " " * 20 + "AAAAAA",
            "",
        ],
        fontsize=10,
    )
    
    # In Jupyter notebook
    fig
    

    Output:

    enter image description here