Search code examples
pythongraphdata-visualizationvisualizationaltair

Set custom color scheme and merge legends in Altair


I want to have a single legend summarizing shape and color of a scatter plot.

And I want to choose the colors for the points myself.

I know how to do each of these things independently. But when I try to do both at once, I end up with two legends instead of one.

My code looks like:

x_axis = 'Publication date'
y_axis = 'Training compute (FLOPs)'

import altair as alt
alt.themes.enable('fivethirtyeight')
selection = alt.selection_multi(fields=['Domain'], bind='legend')

domain_to_color =  {
    'Vision' : '#6d904f', 
    'Language' : '#b96db8',   
    'Games' : '#30a2da', 
    'Drawing' : 'black', 
    'Speech' : 'black',   
    'Other' : '#e5ae38', 
    'Large Scale' : '#fc4f30',     
    'All' : '#30a2da'
    }

present_domains = [domain for domain in domain_to_color.keys() if domain in df['Domain'].unique()]
domain_scale = alt.Scale(
        domain=present_domains,
        range=[domain_to_color[domain] for domain in present_domains],
        )

## Chart with historical data
chart = alt.Chart(df, width=1100, height=600,)\
.mark_point(size=120, filled=False)\
.encode(
  x=alt.X(f'{x_axis}:{"T" if x_axis == "Publication date" else "Q"}',
          scale=alt.Scale(type='linear', 
                          domain=(df[x_axis].min(), df[x_axis].max())), 
          axis = alt.Axis(grid=True)
          ),
  y=alt.Y(f'{y_axis}:Q',
          scale=alt.Scale(type='log', 
                          domain=(df[y_axis].min(), df[y_axis].max())), 
          axis=alt.Axis(format=".1e", grid = True)
          ),
  color= alt.Color('Domain', scale=domain_scale),
  shape = 'Domain',
  opacity=alt.condition(selection, alt.value(1), alt.value(0.2)),
)

chart = chart.add_selection(selection)

chart

This produces a chart like this:

enter image description here

There are two legends! If I change the line color = alt.Color('Domain', scale=domain_scale), for color = 'Domain' then I end up with a single legend. But the colors are not the ones I want.

How can I do both things at once?


Solution

  • I think the problem is that your scale domains do not match, so there's no way to merge them. You should try making the domains explicitly match:

      color=alt.Color('Domain', scale=domain_scale),
      shape=alt.Shape('Domain', scale=alt.Scale(domain=domain_scale.domain),