Visualizing Data in Python Using plt.scatter()

Getting Started With plt.scatter()

Before you can start working with plt.scatter() , you’ll need to install Matplotlib. You can do so using Python’s standard package manger, pip, by running the following command in the console :

$ python -m pip install matplotlib

Now that you have Matplotlib installed, consider the following use case. A café sells six different types of bottled orange drinks. The owner wants to understand the relationship between the price of the drinks and how many of each one he sells, so he keeps track of how many of each drink he sells every day. You can visualize this relationship as follows:

import matplotlib.pyplot as plt

price = [2.50, 1.23, 4.02, 3.25, 5.00, 4.40]
sales_per_day = [34, 62, 49, 22, 13, 19]

plt.scatter(price, sales_per_day)
plt.show()

In this Python script, you import the pyplot submodule from Matplotlib using the alias plt. This alias is generally used by convention to shorten the module and submodule names. You then create lists with the price and average sales per day for each of the six orange drinks sold.

Finally, you create the scatter plot by using plt.scatter() with the two variables you wish to compare as input arguments. As you’re using a Python script, you also need to explicitly display the figure by using plt.show().

When you’re using an interactive environment, such as a console or a Jupyter Notebook, you don’t need to call plt.show(). In this tutorial, all the examples will be in the form of scripts and will include the call to plt.show().

Here’s the output from this code:

This plot shows that, in general, the more expensive a drink is, the fewer items are sold. However, the drink that costs $4.02 is an outlier, which may show that it’s a particularly popular product. When using scatter plots in this way, close inspection can help you explore the relationship between variables. You can then carry out further analysis, whether it’s using linear regression or other techniques.

Comparing plt.scatter() and plt.plot()

You can also produce the scatter plot shown above using another function within matplotlib.pyplot. Matplotlib’s plt.plot() is a general-purpose plotting function that will allow you to create various different line or marker plots.

You can achieve the same scatter plot as the one you obtained in the section above with the following call to plt.plot(), using the same data:

plt.plot(price, sales_per_day, "o")
plt.show()

In this case, you had to include the marker "o" as a third argument, as otherwise plt.plot() would plot a line graph. The plot you created with this code is identical to the plot you created earlier with plt.scatter().

In some instances, for the basic scatter plot you’re plotting in this example, using plt.plot() may be preferable. You can compare the efficiency of the two functions using the timeit module:

import timeit
import matplotlib.pyplot as plt

price = [2.50, 1.23, 4.02, 3.25, 5.00, 4.40]
sales_per_day = [34, 62, 49, 22, 13, 19]

print(
    "plt.scatter()",
    timeit.timeit(
        "plt.scatter(price, sales_per_day)",
        number=1000,
        globals=globals(),
    ),
)
print(
    "plt.plot()",
    timeit.timeit(
        "plt.plot(price, sales_per_day, 'o')",
        number=1000,
        globals=globals(),
    ),
)

The performance will vary on different computers, but when you run this code, you’ll find that plt.plot() is significantly more efficient than plt.scatter(). When running the example above on my system, plt.plot() was over seven times faster.

If you can create scatter plots using plt.plot(), and it’s also much faster, why should you ever use plt.scatter()? You’ll find the answer in the rest of this tutorial. Most of the customizations and advanced uses you’ll learn about in this tutorial are only possible when using plt.scatter(). Here’s a rule of thumb you can use:

• If you need a basic scatter plot, use plt.plot(), especially if you want to prioritize performance.
• If you want to customize your scatter plot by using more advanced plotting features, use plt.scatter().

In the next section, you’ll start exploring more advanced uses of plt.scatter().

Changing the Size

Let’s return to the café owner you met earlier in this tutorial. The different orange drinks he sells come from different suppliers and have different profit margins. You can show this additional information in the scatter plot by adjusting the size of the marker. The profit margin is given as a percentage in this example:

import matplotlib.pyplot as plt
import numpy as np

price = np.asarray([2.50, 1.23, 4.02, 3.25, 5.00, 4.40])
sales_per_day = np.asarray([34, 62, 49, 22, 13, 19])
profit_margin = np.asarray([20, 35, 40, 20, 27.5, 15])

plt.scatter(x=price, y=sales_per_day, s=profit_margin * 10)
plt.show()

You can notice a few changes from the first example. Instead of lists, you’re now using NumPy arrays. You can use any array-like data structure for the data, and NumPy arrays are commonly used in these types of applications since they enable element-wise operations that are performed efficiently. The NumPy module is a dependency of Matplotlib, which is why you don’t need to install it manually.

You’ve also used named parameters as input arguments in the function call. The parameters x and y are required, but all other parameters are optional.

The parameter s denotes the size of the marker. In this example, you use the profit margin as a variable to determine the size of the marker and multiply it by 10 to display the size difference more clearly.

You can see the scatter plot created by this code below:

The size of the marker indicates the profit margin for each product. The two orange drinks that sell most are also the ones that have the highest profit margin. This is good news for the café owner!

Changing the Color

Many of the customers of the café like to read the labels carefully, especially to find out the sugar content of the drinks they’re buying. The café owner wants to emphasize his selection of healthy foods in his next marketing campaign, so he categorizes the drinks based on their sugar content and uses a traffic light system to indicate low, medium, or high sugar content for the drinks.

You can add color to the markers in the scatter plot to show the sugar content of each drink:

# ...

low = (0, 1, 0)
medium = (1, 1, 0)
high = (1, 0, 0)

sugar_content = [low, high, medium, medium, high, low]

plt.scatter(
    x=price,
    y=sales_per_day,
    s=profit_margin * 10,
    c=sugar_content,
)
plt.show()

You define the variables low, medium, and high to be tuples, each containing three values that represent the red, green, and blue color components, in that order. These are RGB color values. The tuples for low, medium, and high represent green, yellow, and red, respectively.

You then defined the variable sugar_content to classify each drink. You use the optional parameter c in the function call to define the color of each marker. Here’s the scatter plot produced by this code:

The café owner has already decided to remove the most expensive drink from the menu as this doesn’t sell well and has a high sugar content. Should he also stop stocking the cheapest of the drinks to boost the health credentials of the business, even though it sells well and has a good profit margin?

Changing the Shape

The café owner has found this exercise very useful, and he wants to investigate another product. In addition to the orange drinks, you’ll now also plot similar data for the range of cereal bars available in the café:

import matplotlib.pyplot as plt
import numpy as np

low = (0, 1, 0)
medium = (1, 1, 0)
high = (1, 0, 0)

price_orange = np.asarray([2.50, 1.23, 4.02, 3.25, 5.00, 4.40])
sales_per_day_orange = np.asarray([34, 62, 49, 22, 13, 19])
profit_margin_orange = np.asarray([20, 35, 40, 20, 27.5, 15])
sugar_content_orange = [low, high, medium, medium, high, low]

price_cereal = np.asarray([1.50, 2.50, 1.15, 1.95])
sales_per_day_cereal = np.asarray([67, 34, 36, 12])
profit_margin_cereal = np.asarray([20, 42.5, 33.3, 18])
sugar_content_cereal = [low, high, medium, low]

plt.scatter(
    x=price_orange,
    y=sales_per_day_orange,
    s=profit_margin_orange * 10,
    c=sugar_content_orange,
)
plt.scatter(
    x=price_cereal,
    y=sales_per_day_cereal,
    s=profit_margin_cereal * 10,
    c=sugar_content_cereal,
)
plt.show()

In this code, you refactor the variable names to take into account that you now have data for two different products. You then plot both scatter plots in a single figure. This gives the following output:

Unfortunately, you can no longer figure out which data points belong to the orange drinks and which to the cereal bars. You can change the shape of the marker for one of the scatter plots:

import matplotlib.pyplot as plt
import numpy as np

low = (0, 1, 0)
medium = (1, 1, 0)
high = (1, 0, 0)

price_orange = np.asarray([2.50, 1.23, 4.02, 3.25, 5.00, 4.40])
sales_per_day_orange = np.asarray([34, 62, 49, 22, 13, 19])
profit_margin_orange = np.asarray([20, 35, 40, 20, 27.5, 15])
sugar_content_orange = [low, high, medium, medium, high, low]

price_cereal = np.asarray([1.50, 2.50, 1.15, 1.95])
sales_per_day_cereal = np.asarray([67, 34, 36, 12])
profit_margin_cereal = np.asarray([20, 42.5, 33.3, 18])
sugar_content_cereal = [low, high, medium, low]

plt.scatter(
    x=price_orange,
    y=sales_per_day_orange,
    s=profit_margin_orange * 10,
    c=sugar_content_orange,
)
plt.scatter(
    x=price_cereal,
    y=sales_per_day_cereal,
    s=profit_margin_cereal * 10,
    c=sugar_content_cereal,
    marker="d",
)
plt.show()

You keep the default marker shape for the orange drink data. The default marker is "o", which represents a dot. For the cereal bar data, you set the marker shape to "d", which represents a diamond marker. You can find the list of all markers you can use in the documentation page on markers. Here are the two scatter plots superimposed on the same figure:

You can now distinguish the data points for the orange drinks from those for the cereal bars. But there is one problem with the last plot you created that you’ll explore in the next section.

Changing the Transparency

One of the data points for the orange drinks has disappeared. There should be six orange drinks, but only five round markers can be seen in the figure. One of the cereal bar data points is hiding an orange drink data point.

You can fix this visualization problem by making the data points partially transparent using the alpha value:

# ...

plt.scatter(
    x=price_orange,
    y=sales_per_day_orange,
    s=profit_margin_orange * 10,
    c=sugar_content_orange,
    alpha=0.5,
)
plt.scatter(
    x=price_cereal,
    y=sales_per_day_cereal,
    s=profit_margin_cereal * 10,
    c=sugar_content_cereal,
    marker="d",
    alpha=0.5,
)

plt.title("Sales vs Prices for Orange Drinks and Cereal Bars")
plt.legend(["Orange Drinks", "Cereal Bars"])
plt.xlabel("Price (Currency Unit)")
plt.ylabel("Average weekly sales")
plt.text(
    3.2,
    55,
    "Size of marker = profit margin\n" "Color of marker = sugar content",
)

plt.show()

You’ve set the alpha value of both sets of markers to 0.5, which means they’re semitransparent. You can now see all the data points in this plot, including those that coincide:

You’ve also added a title and other labels to the plot to complete the figure with more information about what’s being displayed.