I know this exact question has been asked here, however the current solution does nothing for me. I can't seem to generate a legend that has a different color for each label. I have tried the current documentation on Matplotlib to no avail. I keep getting the error that my PathCollection object has no attribute legend_elements
EDIT: Also, I want my legend to be just the Years
, unique years for the plot not how it is right now with is that each data point is mapped to my legend.
Here's what I have
import numpy as np # linear algebra
import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)
import matplotlib.pyplot as plt
from matplotlib.pyplot import figure
from matplotlib.pyplot import legend
import os
%config InlineBackend.figure_format = 'retina'
path = None
for dirname, _, filenames in os.walk('/kaggle/input'):
for filename in filenames:
path = os.path.join(dirname, filename)
# Indexes to be removed
early_demo_dividend = 13
high_income = 24
lower_middle_income = 40
north_america = 46
members = 50
post_demo = 56
_removals = [early_demo_dividend, high_income, lower_middle_income, north_america, members, post_demo]
#Read in data
df = pd.read_csv(path)
#Get the rows we want
df = df.loc[df['1960'] > 1]
df = df.drop(columns=["Code", "Type", "Indicator Name"])
#Remove the odd rows
for i in _removals:
df = df.drop(df.index[i])
#Format the dataframe
df = df.melt('Name', var_name='Year', value_name='Budget')
#Plot setup
plt.figure().set_size_inches(16,6)
plt.xticks(rotation=90)
plt.grid(True)
#Plot labels
plt.title('Military Spending of Countries')
plt.xlabel('Countries')
plt.ylabel('Budget in Billions')
#Plot data
new_year = df['Year'].astype(int)
scatter = plt.scatter(df['Name'], df['Budget'], c=(new_year / 10000) , label=new_year)
#Legend setup produce a legend with the unique colors from the scatter
legend1 = plt.legend(*scatter.legend_elements(),
loc="lower left", title="Years")
plt.add_artist(legend1)
plt.show()