pythonmatplotlibcolorbar

Changing the color of a colorbar in certain sections based on a third variable


I have the following code. I'm trying to have the values of De_below_line show up as a different color on the colorbar. The colorbar itself should have the range from De.min() to De.max(). While I can get the colors to plot differently on the scatterplot, the values don't transfer across to the colorbar easily. I've included a picture of a potential colorbar below of how I would like it to look and hopefully clarify what I'm trying to do. Any help would be greatly appreciated.

    import pandas as pd
    import numpy as np
    import matplotlib.pyplot as plt
    from matplotlib.ticker import FormatStrFormatter
    from matplotlib.colors import ListedColormap
    
    xls = pd.ExcelFile('20240109 - CPeT Export - Rev2.xlsx')
    God = pd.ExcelFile('20231213-Vs calculations.xlsx')
    n = len(xls.sheet_names)
    #------Raw data-------------------------------------------------------------------------------------#
    dfs = pd.read_excel(xls, sheet_name='SCPTU-12-CPT', skiprows=185, skipfooter=0)
    title = dfs.values[0,0]
    qt = dfs.values[:,5]                                        # Cone tip resistance (MPa)
    ov = dfs.values[:,10]                                       # Vertical total stress (kPa)
    Qtn = dfs.values[:,18]                                      # Normalised cone tip resistance with stress component (-)
    De = dfs.values[:,1]                                        # Imports depth values (m)
    Gor = pd.read_excel(God, sheet_name='SCPTU-12-CPT', skiprows=185, skipfooter=0)
    Go = Gor.values[:,8]                                        # Imports the correct small-strain shear modulus (MPa)
    D = Gor.values[:,0]                                         # Imports depth values (m)
    #------Calculations---------------------------------------------------------------------------------#
    ov = ov/1000
    qn = qt-ov                                                  # Calculates qn
    IG = Go/qn                                                  # Calculates normalised rigidty index
    #------Plotting the results-------------------------------------------------------------------------#
    fig = plt.figure()
    
    y = 2280.4*IG**-1.333
    De_below_line = De[Qtn < y]
    
    cmap = ListedColormap(['blue', 'red'])  # Two colors: below the line, above the line
    colors = np.where(Qtn < y, 'blue', 'red')  # De.min() for below the line, De.max() for above the line
    
    ax = fig.gca()
    ax.set_xscale('log')
    ax.set_yscale('log')
    ax.yaxis.set_major_formatter(FormatStrFormatter('%g'))
    ax.xaxis.set_major_formatter(FormatStrFormatter('%g'))
    plt.xlabel('Normalised Rigidity Index, $I_G$')
    plt.ylabel('Normalised Cone Tip Resistance, Qtn')
    sc = ax.scatter(IG, Qtn, s=60, label=title, c=colors, cmap=cmap)
    plt.plot([1,330],[(330/1)**(1/0.75),1], linewidth=1, color="black",linestyle='-', marker='None')
    
    cbar = plt.colorbar(sc, ticks=[De.min(), De.max()])
    cbar.ax.invert_yaxis()
    
    ax.set_xlim(1, 1000)
    ax.set_ylim(1, 1000)
    plt.show()

enter image description here


Solution

  • It looks like you don't want a regular colorbar. With imshow you can show a custom property of the DE values. The expression Qtn < y, where both parts are numpy arrays, gets converted in a new array with False and True values. When interpreted as number, this becomes 0 and 1. Changing the order of the colors in the colormap will map 1 (True, so Qtn < y to blue) and 0 to red.

    imshow needs its input as 2D, so the array needs to be reshaped. aspect='auto' avoids the default "square pixels". origin='lower' lets the image start at the bottom. extent sets the xmin, xmax, ymin, ymax as the coordinates (the x values aren't important).

    The code below creates a figure with two subplots: the scatter plot at the left, and the bar at the right. The code supposes the DE array is strictly ordered from low to high.

    import matplotlib.pyplot as plt
    from matplotlib. Colors import ListedColormap
    import numpy as np
    
    cmap = ListedColormap(['crimson', 'skyblue'])  # Two colors: above the line, below the line
    
    # create some random data for testing
    np.random.seed(1123)
    IG = np.logspace(0, 3, 20)
    Qtn = np.random.uniform(2, 999, 20)
    y = 2280.4 * IG ** -1.333
    De = np.linspace(0, 100, 20)
    
    fig, (ax, ax_DE) = plt.subplots(ncols=2, figsize=(8, 6), gridspec_kw={'width_ratios': [10, 1]})
    
    # create the scatter plot IG vs Qtn
    ax.set_xscale('log')
    ax.set_yscale('log')
    ax.scatter(IG, Qtn, s=60, label='IG vs Qtn', c=(Qtn < y), cmap=cmap)
    # ax. Plot([1, 330], [(330 / 1) ** (1 / 0.75), 1], linewidth=1, color="black", linestyle='-', marker='None')
    ax.plot(IG, y, linestyle='-', color='black', label='y')
    ax.set_xlabel('IG')
    ax.set_ylabel('Qtn')
    ax.set_title('IG vs Qtn')
    
    # create the bar for DE
    ax_DE.imshow((Qtn < y).reshape(-1, 1), aspect='auto', cmap=cmap, origin='lower', extent=(0, 1, De.min(), De.max()))
    ax_DE.set_xticks([])  # erase x ticks
    ax_DE.yaxis.tick_right()  # place y ticks at the right
    ax_DE.set_title('Depth')
    
    plt.tight_layout()
    plt.show()
    

    mimicking a colorbar with custom data