mggg / ecological-inference

Ecological inference, in Python
MIT License
28 stars 11 forks source link

add legend to precinct_level_plot() #38

Closed gabeschoenbach closed 3 years ago

gabeschoenbach commented 3 years ago

Added a legend to the ei.precinct_level_plot() method, which calls plot_utils.plot_precincts().

We first have to pass the list of demographic group names into plot_precincts(). Since there are really 2 x num_precincts matplotlib Artists in the figure, I opted to create proxy handles for the legend. In order to ensure the legend and the KDE's draw the same colors, I put a colors list into plot_precincts() function body to which both the plot_single_ridgeplots() helper and the legend can point. Lastly, there needs to be space for the legend itself! Some experimentation showed me that creating legend_space = 5 extra lines at the top were enough to comfortably fit the legend box, and this seems to look fine when tested on 10, 20, and 42 precincts.

Miscellaneous thoughts:

karink520 commented 3 years ago

Awesome. Looks great!

And, yes, definitely, to your point about including more generalizations of the plotting functions for r x c -- we absolutely should do this! There are also places where there's currently redundancy in the plotting utilities that we can get rid of (e.g. there's a plot_kde function that is just for two groups, and plot_kdes that handles r x c). There are also non-plot-related opportunities for more refactoring with 2x2 and rxc, which I was going to hold off on until later, but hopefully not too much later.