Altair: geo_shape doesn't work with selection - altair

Question

I am working on a Choropleth map that will display the similarities between different states. So when you select a state from a dropdown, the map will show the similarity it has to other states.
For this, I am using 2 datasets:
DatasetA: a long-form dataframe, with 3 columns: State 1, State 2, and the similarity between them.
DatasetB: a GeoDataFrame that contains the geometry of each state.
When I try to plot this without the selection, then it works:
alt.Chart(gdf).mark_geoshape(
).encode(
color='Similarity:O',
tooltip = ['Similarity:Q']
).properties(
projection={'type': 'albersUsa'},
width=700,
height=400
).transform_lookup(
lookup='State',
from_=alt.LookupData(source, 'State', source.columns.values)
)
But once I add the selection, then it only works when I select Wyoming (the last state in Dataset A). When I select other states, the plot disappears.
input_dropdown = alt.binding_select(options=source.State.unique())
selection = alt.selection_single(fields=['Similarity_to'], bind=input_dropdown ,init={'Similarity_to': 'New York'})
alt.Chart(gdf).mark_geoshape(
).encode(
color='Similarity:Q',
tooltip = ['Similarity:Q']
).properties(
projection={'type': 'albersUsa'},
width=700,
height=400
).transform_lookup(
lookup='State',
from_=alt.LookupData(source, 'State', source.columns.values)
).transform_filter(
selection
).add_selection(
selection
)
Here's a clip that demonstrates it: https://www.loom.com/share/292e8b1a80344cf5a998a54f453ece2c

Answer 1

I got this to work by using a transform_fold. The problem is that the transform_lookup only matches once, so if there are multiple matches in the dataset, it ignores them. So you have to use a wide-form dataset and then use the transform_fold to convert it back into long-form.
input_dropdown = alt.binding_select(options=source.State.unique())
selection = alt.selection_single(fields=['Similarity_to'], bind=input_dropdown ,init={'Similarity_to': 'New York'})
alt.Chart(gdf).mark_geoshape(
stroke='black'
).encode(
color='Similarity:Q',
tooltip = ['Similarity:Q']
).properties(
projection={'type': 'albersUsa'},
width=700,
height=400
).transform_lookup(
lookup='State',
from_=alt.LookupData(source, 'State', source.columns.values)
).transform_fold(
source.drop('State',axis=1).columns.values, # Preserve the State column, fold the rest
['Similarity_to','Similarity']
).transform_filter(
selection
).add_selection(
selection
)
I actually tried this before asking the question. But it turns out I did the transformations in the wrong order. The order of transformations matters!
You can find the full code here: https://deepnote.com/project/altairstackoverflow--zl7Wx2tQQ22U3D1cLcodA/%2Fnotebooks%2Fstates.ipynb#00011-ed9a9249-2c34-412c-a091-d87d0ddb457d