Need legend to communicate color for plotly graph go.Heatmap

Question

I'm having a bit of trouble trying to figure out how to have the legend indicate color for my plotly graph objects heatmap. The legend only needs to communicate the graph colors and what they mean. Essentially all I need is "Purple = Data Not Found" and "Yellow = Data Found". Can anyone help me with this, please?

import plotly.graph_objects as go
import numpy as np
import pandas as pd
from datetime import datetime

df = pd.read_csv('vipr_viri_all_final_df.csv')

truck = '251_QV02_KW_T680_HD_471192'

truck_df = df[df['Truck']==truck]
truck_date_df = truck_df['Date']

dates = list(truck_df['Date'])

vipr_bool = list(truck_df['VIPR Data'])
vipr_bool = [0 if item == False else item for item in vipr_bool]
vipr_bool = [1 if item == True else item for item in vipr_bool]
vipr_first_true = vipr_bool.index(1)
vipr_first_true_reversed = vipr_bool[::-1]
vipr_first_index_reversed = vipr_first_true_reversed.index(1)
vipr_last_true = len(vipr_bool)-1-vipr_first_index_reversed
vipr_low = truck_date_df[vipr_first_true:vipr_first_true+1]
vipr_low = vipr_low.iloc[0]
vipr_high = truck_date_df[vipr_last_true:vipr_last_true+1]
vipr_high = vipr_high.iloc[0]

viri_bool = list(truck_df['Viriciti Data'])
viri_bool = [0 if item == False else item for item in viri_bool]
viri_bool = [1 if item == True else item for item in viri_bool]
viri_first_true = viri_bool.index(1)
viri_first_true_reversed = viri_bool[::-1]
viri_first_index_reversed = viri_first_true_reversed.index(1)
viri_last_true = len(vipr_bool)-1-viri_first_index_reversed
viri_low = truck_date_df[viri_first_true:viri_first_true+1]
viri_low = viri_low.iloc[0]
viri_high = truck_date_df[viri_last_true:viri_last_true+1]
viri_high = viri_high.iloc[0]

if vipr_low < viri_low:
    low_range = vipr_low
else:
    low_range = viri_low

if vipr_high > viri_high:
    high_range = vipr_high
else:
    high_range = viri_high

fig = go.Figure(data=go.Heatmap(
        x=truck_date_df,
        y=['VIPR', 'Viriciti'],
        z=[vipr_bool, viri_bool],
        name=truck,
        colorscale='Viridis',
        showscale=False,
        showlegend=True),
        layout=go.Layout(title=go.layout.Title(text=truck),
                         xaxis_title='Date',
                         yaxis_title='Database',
                         xaxis_range=[low_range, high_range]
    )
)

fig.show()

Output of df.head().to_dict()

{'Unnamed: 0': {0: 0, 1: 1, 2: 2, 3: 3, 4: 4},
 'Truck': {0: '220_2_INTL_PRSTR_ET_619076',
  1: '220_2_INTL_PRSTR_ET_619076',
  2: '220_2_INTL_PRSTR_ET_619076',
  3: '220_2_INTL_PRSTR_ET_619076',
  4: '220_2_INTL_PRSTR_ET_619076'},
 'Date': {0: '2014-01-01',
  1: '2014-01-02',
  2: '2014-01-03',
  3: '2014-01-04',
  4: '2014-01-05'},
 'VIPR Data': {0: False, 1: False, 2: False, 3: False, 4: False},
 'Viriciti Data': {0: False, 1: False, 2: False, 3: False, 4: False}}

Answer 1

You can separate your boolean lists [0,0,1,...0] into a list with None as a mask: [None,None,1,...None] and [0,0,None,...0]. Then you can add these as separate heatmap traces so that a legend entry will be generated for each trace.

import plotly.graph_objects as go
import numpy as np
import pandas as pd
from datetime import datetime

# df = pd.read_csv('vipr_viri_all_final_df.csv')

# truck = '251_QV02_KW_T680_HD_471192'
# truck_df = df[df['Truck']==truck]

## recreate data with a similar structure to yours
np.random.seed(42)

truck_df = pd.DataFrame({
    'VIPR Data': np.random.choice([True, False], size=400, p=[0.1,0.9]),
    'Viriciti Data': np.random.choice([True, False], size=400, p=[0.1,0.9]),
    'Date': pd.date_range('2021-01-01',periods=400)
})

truck_date_df = truck_df['Date']

dates = list(truck_df['Date'])

vipr_bool = list(truck_df['VIPR Data'])
vipr_bool = [0 if item == False else item for item in vipr_bool]
vipr_bool = [1 if item == True else item for item in vipr_bool]
vipr_first_true = vipr_bool.index(1)
vipr_first_true_reversed = vipr_bool[::-1]
vipr_first_index_reversed = vipr_first_true_reversed.index(1)
vipr_last_true = len(vipr_bool)-1-vipr_first_index_reversed
vipr_low = truck_date_df[vipr_first_true:vipr_first_true+1]
vipr_low = vipr_low.iloc[0]
vipr_high = truck_date_df[vipr_last_true:vipr_last_true+1]
vipr_high = vipr_high.iloc[0]

viri_bool = list(truck_df['Viriciti Data'])
viri_bool = [0 if item == False else item for item in viri_bool]
viri_bool = [1 if item == True else item for item in viri_bool]
viri_first_true = viri_bool.index(1)
viri_first_true_reversed = viri_bool[::-1]
viri_first_index_reversed = viri_first_true_reversed.index(1)
viri_last_true = len(vipr_bool)-1-viri_first_index_reversed
viri_low = truck_date_df[viri_first_true:viri_first_true+1]
viri_low = viri_low.iloc[0]
viri_high = truck_date_df[viri_last_true:viri_last_true+1]
viri_high = viri_high.iloc[0]

if vipr_low < viri_low:
    low_range = vipr_low
else:
    low_range = viri_low

if vipr_high > viri_high:
    high_range = vipr_high
else:
    high_range = viri_high

fig = go.Figure()

vipr_ones = [None if v==0 else v for v in vipr_bool]
vipr_zeros = [None if v==1 else v for v in vipr_bool]
viri_ones = [None if v==0 else v for v in viri_bool]
viri_zeros = [None if v==1 else v for v in viri_bool]

fig.add_trace(go.Heatmap(
        x=truck_date_df,
        y=['VIPR', 'Viriciti'],
        z=[vipr_zeros, viri_zeros],
        name='Purple = Data Not Found',
        colorscale=[(0.00, "#440154"), (1.0, "#440154")],
        showscale=False,
        showlegend=True,
    )
)

fig.add_trace(go.Heatmap(
        x=truck_date_df,
        y=['VIPR', 'Viriciti'],
        z=[vipr_ones, viri_ones],
        name='Yellow = Data Found',
        colorscale=[(0.00, "#fde725"), (1.0, "#fde725")],
        showscale=False,
        showlegend=True,
    )
)

fig.update_layout(title=go.layout.Title(text='truck'),
                         xaxis_title='Date',
                         yaxis_title='Database',
                         xaxis_range=[low_range, high_range])

fig.show()