Update plot.py

[milescsmith]

Fixes KeyError: [...] not in index when running celltypist.plot.dotplot()

[ChuanXu1]

@milescsmith, can you provide more details on why this change is needed? Thank you!

[milescsmith]

Sure. The problem is that if prediction_order is not passed to celltypist.plot.dotplot() (and thus to celltypist.plot._get_fraction_prob_df()), prediction_order is set to all possible categories in the passed AnnotationResult.predicted_labels; however, the dot_color_df is derived from the sliced pred = predictions.predicted_labels[use_as_prediction], so that by the time you get to dot_color_df = dot_color_df.loc[prediction_order], the values in prediction_order may be larger than the possible values in dot_color_df and thus _get_fraction_prob_df() and dotplot() fail due to a KeyError. The bit of code I added ensures the values match or that prediction_order is a subset.

[ChuanXu1]

@milescsmith, I think dot_color_df is not sliced, so their values should be matched. Do you any code to produce the scenario you mentioned?