Customizing Plots with Python Matplotlib | by Carolina Bento | Towards Data Science

A central part of Data Science and Data Analysis is how you visualize the data. How you make use of visualizations tools has an important role in defining how you communicate insights.
My language of choice to explore and visualize data is Python.
In this article, I want to walk you through my framework for going from visualizing raw data to having a beautiful plot that is not just eye-catching but emphases the core insights you want to convey.
In this example I'm going to be using a dataset of workout sessions used in a previous article. It looks like this
A bare bones scatter plot would look like this
Which you can replicate with the following code
import pandas as pdimport matplotlib.pyplot as plt#loading datasetdf = pd.read_csv(‘workout_log.csv’)df.columns = [‘date’, ‘distance_km’, ‘duration_min’, ‘delta_last_workout’, ‘day_category’]def scatterplot(df, x_dim, y_dim):  x = df[x_dim]  y = df[y_dim]  fig, ax = plt.subplots(figsize=(10, 5))  ax.scatter(x, y)  plt.show()scatterplot(df, ‘distance_km’, ‘duration_min’)
The usual next step for me is to label the axes and add a title so each plot is appropriately labeled.
The code change is minimal, but definitely makes a difference.
import pandas as pdimport matplotlib.pyplot as plt#loading datasetdf = pd.read_csv(‘workout_log.csv’)df.columns = [‘date’, ‘distance_km’, ‘duration_min’, ‘delta_last_workout’, ‘day_category’]def scatterplot(df, x_dim, y_dim):  x = df[x_dim]  y = df[y_dim]  fig, ax = plt.subplots(figsize=(10, 5))  ax.scatter(x, y)   #adds a title and axes labels  ax.set_title('Distance vs Workout Duration')  ax.set_xlabel('Distance (Km)')  ax.set_ylabel('Workout Duration (min)')  plt.show()scatterplot(df, ‘distance_km’, ‘duration_min’)
In order to change the default box around the plot, we have to actually remove some of the plot's borders.
import pandas as pdimport matplotlib.pyplot as plt#loading datasetdf = pd.read_csv(‘workout_log.csv’)df.columns = [‘date’, ‘distance_km’, ‘duration_min’, ‘delta_last_workout’, ‘day_category’]def scatterplot(df, x_dim, y_dim):  x = df[x_dim]  y = df[y_dim]  fig, ax = plt.subplots(figsize=(10, 5))  ax.scatter(x, y)   #adds a title and axes labels  ax.set_title('Distance vs Workout Duration')  ax.set_xlabel('Distance (Km)')  ax.set_ylabel('Workout Duration (min)')   #removing top and right borders  ax.spines['top'].set_visible(False)  ax.spines['right'].set_visible(False)  plt.show()scatterplot(df, ‘distance_km’, ‘duration_min’)
Something that I usually like to add to my plots are major gridlines. It helps with readability by reducing the amount of white background. You can play around with the its width linewidth and transparency alpha.
import pandas as pdimport matplotlib.pyplot as plt#loading datasetdf = pd.read_csv(‘workout_log.csv’)df.columns = [‘date’, ‘distance_km’, ‘duration_min’, ‘delta_last_workout’, ‘day_category’]def scatterplot(df, x_dim, y_dim):  x = df[x_dim]  y = df[y_dim]  fig, ax = plt.subplots(figsize=(10, 5))  ax.scatter(x, y)   #adds a title and axes labels  ax.set_title('Distance vs Workout Duration')  ax.set_xlabel('Distance (Km)')  ax.set_ylabel('Workout Duration (min)')   #removing top and right borders  ax.spines['top'].set_visible(False)  ax.spines['right'].set_visible(False)  #adds major gridlines  ax.grid(color='grey', linestyle='-', linewidth=0.25, alpha=0.5)  plt.show()scatterplot(df, ‘distance_km’, ‘duration_min’)
You can see that some of the dots in the plot overlap. To improve readability even more, we can adjust the dots' transparency — alpha.
import pandas as pdimport matplotlib.pyplot as plt#loading datasetdf = pd.read_csv(‘workout_log.csv’)df.columns = [‘date’, ‘distance_km’, ‘duration_min’, ‘delta_last_workout’, ‘day_category’]def scatterplot(df, x_dim, y_dim):  x = df[x_dim]  y = df[y_dim]  fig, ax = plt.subplots(figsize=(10, 5))  #customizes alpha for each dot in the scatter plot  ax.scatter(x, y, alpha=0.70)   #adds a title and axes labels  ax.set_title('Distance vs Workout Duration')  ax.set_xlabel('Distance (Km)')  ax.set_ylabel('Workout Duration (min)')   #removing top and right borders  ax.spines['top'].set_visible(False)  ax.spines['right'].set_visible(False)  #adds major gridlines  ax.grid(color='grey', linestyle='-', linewidth=0.25, alpha=0.5)  plt.show()scatterplot(df, ‘distance_km’, ‘duration_min’)
There is still a bit of overlap, but at least the transparency improved the readability of the majority of the dots.
Since we have the day category we can also try identifying each dot in our plot with a different color.
For that you can choose from two different approaches:
Pick the colors yourself using tools like Adobe Kuler’s color wheel
Use Python's color maps
#1 Defining your own color palette
import pandas as pdimport matplotlib.pyplot as plt#loading datasetdf = pd.read_csv(‘workout_log.csv’)df.columns = [‘date’, ‘distance_km’, ‘duration_min’, ‘delta_last_workout’, ‘day_category’]def scatterplot(df, x_dim, y_dim):  x = df[x_dim]  y = df[y_dim]fig, ax = plt.subplots(figsize=(10, 5))   #defining an array of colors    colors = ['#2300A8', '#00A658']  #assigns a color to each data point  ax.scatter(x, y, alpha=0.70, color=colors)   #adds a title and axes labels  ax.set_title('Distance vs Workout Duration')  ax.set_xlabel('Distance (Km)')  ax.set_ylabel('Workout Duration (min)')   #removing top and right borders  ax.spines['top'].set_visible(False)  ax.spines['right'].set_visible(False)#adds major gridlines  ax.grid(color='grey', linestyle='-', linewidth=0.25, alpha=0.5)plt.show()scatterplot(df, ‘distance_km’, ‘duration_min’)
#2 Using Python Color Maps
To paint each dot according to its day category I need to introduce a few new components in the code
Import the color map library
Take the day category as a parameter, so the corresponding color can be mapped
Use parameterc from the scatter method to assign the color sequence
Use parameter cmap to assign the color map to be used. I'm going to use the brg color map
import pandas as pdimport matplotlib.cm as cmimport matplotlib.pyplot as plt#loading datasetdf = pd.read_csv(‘workout_log.csv’)df.columns = [‘date’, ‘distance_km’, ‘duration_min’, ‘delta_last_workout’, ‘day_category’]def scatterplot(df, x_dim, y_dim, category):  x = df[x_dim]  y = df[y_dim]  fig, ax = plt.subplots(figsize=(10, 5))  #applies the custom color map along with the color sequence  ax.scatter(x, y, alpha=0.70, c= df[category], cmap=cm.brg)   #adds a title and axes labels  ax.set_title('Distance vs Workout Duration')  ax.set_xlabel('Distance (Km)')  ax.set_ylabel('Workout Duration (min)')   #removing top and right borders  ax.spines['top'].set_visible(False)  ax.spines['right'].set_visible(False)  #adds major gridlines  ax.grid(color='grey', linestyle='-', linewidth=0.25, alpha=0.5)  plt.show()scatterplot(df, ‘distance_km’, ‘duration_min’, ‘day_category’)
So far, we've been using the native scatter method to plot each data point. In order to add a legend, we'll have to change the code a little bit.
We'll have to
Take the day category as a parameter, so we have our labels
Convert the numerical (0,1) labels into categorical labels (weekday, weekend)
Iterate through the dataset in order to assign a label to each data point
import pandas as pdimport matplotlib.cm as cmimport matplotlib.pyplot as plt#loading datasetdf = pd.read_csv(‘workout_log.csv’)df.columns = [‘date’, ‘distance_km’, ‘duration_min’, ‘delta_last_workout’, ‘day_category’]def scatterplot(df, x_dim, y_dim, category):   x = df[x_dim]   y = df[y_dim]   #converting original (numerical) labels into categorical labels   categories = df[category].apply(lambda x: 'weekday' if x == 0 else 'weekend')   fig, ax = plt.subplots(figsize=(10, 5))   #assigns a color to each data point   colors = ['#2300A8', '#00A658']   #iterates through the dataset plotting each data point and assigning it its corresponding color and label   for i in range(len(df)):     ax.scatter(x.ix[i], y.ix[i], alpha=0.70, color = colors[i%len(colors)], label=categories.ix[i])   #adds title and axes labels   ax.set_title('Distance vs Workout Duration')   ax.set_xlabel('Distance (Km)')   ax.set_ylabel('Workout Duration (min)')   #removing top and right borders   ax.spines['top'].set_visible(False)   ax.spines['right'].set_visible(False)   #adds major gridlines   ax.grid(color='grey', linestyle='-', linewidth=0.25, alpha=0.5)   #adds legend   ax.legend(categories.unique())   plt.show()scatterplot(df, 'distance_km', 'duration_min', 'day_category')
And there you have it! A customized scatter plot from which it's now easier to understand the data and draw some insights.
Thanks for reading!