One Color per Group

Part of the Data and Plots section of Coddy's Machine Learning journey. Lesson 49 of 60.

Do swimmers and runners differ in build? Give each sport its own color and see:

Eight athletes as dots, height across and weight up: four swimmers in blue at the upper right, four runners in orange at the lower left, with a legend

.unique() lists each value of a column once: here swim, then run. Loop over them, keep each group's rows, and give each group its own ax.scatter:

athletes = pd.DataFrame({
    "sport": ["swim", "run", "swim", "run", "run", "swim", "run", "swim"],
    "height": [188, 172, 191, 168, 175, 185, 170, 194],
    "weight": [84, 60, 88, 57, 63, 80, 59, 90],
})
fig, ax = plt.subplots()
for sport in athletes["sport"].unique():
    group = athletes[athletes["sport"] == sport]
    ax.scatter(group["height"], group["weight"], label=sport)
ax.set_xlabel("Height (cm)")
ax.set_ylabel("Weight (kg)")
ax.legend()
plt.show()

Each call takes the next color, and label= names it for ax.legend(). Groups that sit apart, like these, are what a model learns to tell apart.

Challenge

Medium

A greenhouse measured 18 tomatoes in tomatoes.csv: variety (cherry, plum or beefsteak), diameter_cm and weight_g. Do the varieties sit apart? The starter sets up the chart, the legend and plt.show().

Write the loop that draws one scatter per variety, diameter_cm across and weight_g up, each labelled with its variety.

Yours should look like this one, from last year's harvest:

A scatter titled Tomato Harvest, diameter across and weight up, with three colored groups and a legend: cherry small at the lower left, plum in the middle, beefsteak large at the upper right

Coddy checks your chart and describes it in the output.

Try it yourself

import pandas as pd
import matplotlib.pyplot as plt

tomatoes = pd.read_csv("tomatoes.csv")

fig, ax = plt.subplots()
ax.set_title("Tomato Harvest")
ax.set_xlabel("Diameter (cm)")
ax.set_ylabel("Weight (g)")

# Write your code below

# Don't change the lines below
ax.legend()
plt.show()