Search code examples
pythonpandasstreamlitaltair

Altair: Apply condition to axis label color based on label membership in a list


I have a scheduling plot where the Y axis is categorical (each Y-entry corresponds to a piece of equipment and the x axis displays time).

I have a list of equipment that contains scheduling clashes. For each equipment item in the list of equipment with clashes, I want to highlight the equipment name in red on the plot.

I can't find a way to apply a test that checks for membership in a list.

A similar problem was posted here: Color some x-labels in altair plot? - but this solution involves encoding the test in a string which is interpreted by altair I guess. The "is in" test doesn't seem to work in this sort of construction.

Here is a minimum working example:

import streamlit as st
import altair as alt
import pandas as pd

data = [
    {"task": "Task 1", "start": 1, "finish": 10, "equipment": "XXX-101"},
    {"task": "Task 2", "start": 9, "finish": 20, "equipment": "XXX-101"},
    {"task": "Task 3", "start": 6, "finish": 8, "equipment": "XXX-102"},
    {"task": "Task 4", "start": 9, "finish": 12, "equipment": "XXX-102"},
    {"task": "Task 5", "start": 13, "finish": 18, "equipment": "XXX-102"},
    {"task": "Task 6", "start": 5, "finish": 15, "equipment": "XXX-103"},
    {"task": "Task 7", "start": 16, "finish": 18, "equipment": "XXX-103"},
    {"task": "Task 8", "start": 6, "finish": 8, "equipment": "XXX-104"},
    {"task": "Task 9", "start": 4, "finish": 12, "equipment": "XXX-104"},
    {"task": "Task 10", "start": 11, "finish": 16, "equipment": "XXX-104"},
]
dataframe = pd.DataFrame(data)

equipment_has_clashes = ["XXX-101", "XXX-104"]  # determined by another algorithm


chart = (
    alt.Chart(dataframe)
    .mark_bar(height=20)
    .encode(
        x=alt.X("start").title("time"),
        x2=("finish"),
        y=alt.Y(
            "equipment",
            axis=alt.Axis(
                labelColor=alt.condition(
                    # Want to apply test equivalent to "equipment is in equipment_has_clashes"
                    predicate="datum",
                    if_true=alt.value("red"),
                    if_false=alt.value("black"),
                )
            ),
        ),
        tooltip=[alt.Tooltip(t) for t in ["task", "equipment", "start", "finish"]],
        color=alt.Color("task", type="nominal").scale(scheme="category20"),
    )
    .properties(height=alt.Step(30))
    .configure_view(strokeWidth=1)
    .configure_axis(domain=True, grid=True)
    .interactive()
)

st.altair_chart(chart, use_container_width=True)


Solution

  • Maybe not the most elegant solution but you can pass each condition with an "OR" operator in the solution here.

    cond = " || ".join([f'datum.value == "{x}"'  for x in equipment_has_clashes])
    

    and pass it in your labelColor:

                axis=alt.Axis(
                    labelColor=alt.condition(cond, alt.value('red'), alt.value('blue')) 
                    ),
    

    Output:

    enter image description here

    Another solution might be using something like this, but I was not able to implement it.