Home > Enterprise >  Change legend label based on dataframe value
Change legend label based on dataframe value

Time:12-04

I currently have an X value df["long"] and a Y value df["lat"]. Each X and Y point represents a house and I'm wanting the scatter points to appear blue if df["waterfront"] = 0 or red if df["waterfront"] = 1.

My attempt is shown below but I just can't figure out how to insert both labels and legends.

import pandas as pd
import matplotlib.pyplot as plt

d = {'long': [1, 2], 'lat': [3, 4], 'waterfront': [0, 1]}
df = pd.DataFrame(data=d)

colors = {0:'tab:blue', 1:'tab:red'}

plt.scatter(df['long'], df['lat'], 
            c=[colors[i] for i in df['waterfront']], 
            zorder=3, label=("Without Waterfront","With Waterfront"))
plt.xlabel ("Longitude")
plt.ylabel ("Latitude")
plt.title ("Longitude versus Latitude")
plt.grid(zorder=0)
plt.legend()
plt.show()

Actual Output Actual Output

Expected Output Expected Output

CodePudding user response:

You can use seaborn in order to do this simply:

import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

# data 
d = {'long': [1, 2], 'lat': [3, 4], 'waterfront': [0, 1]}
df = pd.DataFrame(data=d)

# plot colors
colors = {0:'tab:blue', 1:'tab:red'}

# create scatter plot
g = sns.scatterplot(data=df, x="long", y="lat", hue="waterfront", palette=colors)

# set legend outside plot
g.legend(loc='center left', bbox_to_anchor=(1, 0.5), ncol=1)

# change legend labels
for t, l in zip(g.legend_.texts, ["Without Waterfront", "With Waterfront"]):
    t.set_text(l)

# set labels and plot
plt.xlabel ("Longitude")
plt.ylabel ("Latitude")
plt.title ("Longitude versus Latitude")
plt.show()

Output:

output1

Solution without seaborn:

import pandas as pd
import matplotlib.pyplot as plt

# data 
d = {'long': [1, 2], 'lat': [3, 4], 'waterfront': [0, 1]}
df = pd.DataFrame(data=d)

# create scatter plot
plt.scatter(df.loc[df['waterfront'] == 0,'long'], df.loc[df['waterfront'] == 0,'lat'], 
            color='tab:blue', label="Without Waterfront")
plt.scatter(df.loc[df['waterfront'] == 1,'long'], df.loc[df['waterfront'] == 1,'lat'], 
            color='tab:red', label="With Waterfront")

# set labels and plot
plt.xlabel ("Longitude")
plt.ylabel ("Latitude")
plt.title ("Longitude versus Latitude")
plt.legend(loc='center left', bbox_to_anchor=(1, 0.5), ncol=1)
plt.show()

Output:

output2

  • Related