Reputation: 659
I need to create a plot as close to this picture as possible (given the generated dataframe code below):
And here's the output plot of my code:
What I am having problems with is:
Here's my code:
import matplotlib.pyplot as plt
import matplotlib.colors as colors
import matplotlib.cm as cm
import numpy as np
import pandas as pd
ncols = 10
figsize = (20, 5)
fontsize = 14
dti = pd.date_range('2013-01-01', '2020-12-31', freq='2W')
probabilities_in_time = np.random.random((ncols, len(dti)))
probabilities_in_time = probabilities_in_time / \
probabilities_in_time.sum(axis=0)
probabilities_in_time = pd.DataFrame(probabilities_in_time).T
probabilities_in_time.index = dti
cm_subsection = np.linspace(0, 1, ncols)
colors = [cm.coolwarm(x) for x in cm_subsection]
def plot_time_probabilities(probabilities_in_time, figsize):
plt.figure(figsize=figsize)
plt.yticks(np.arange(0, 1.2, 0.2), fontsize=fontsize)
plt.xticks(fontsize=fontsize)
draw_stack_plot(colors, probabilities_in_time)
set_grid()
set_legend()
plt.show()
def draw_stack_plot(colors, probabilities_in_time):
for i, color in enumerate(colors):
if i == 0:
plt.plot(probabilities_in_time[i], color=color)
plt.fill_between(probabilities_in_time.index,
probabilities_in_time[0], color=color)
else:
probabilities_in_time[i] += probabilities_in_time[i-1]
plt.fill_between(probabilities_in_time.index,
probabilities_in_time[i], probabilities_in_time[i-1],
color=color)
plt.plot(probabilities_in_time[i], label=' Probability: {}'.format(
i), color=color)
def set_grid():
ax = plt.gca()
ax.set_axisbelow(False)
ax.xaxis.grid(True, linestyle='-', lw=1)
def set_legend():
leg = plt.legend(loc='lower left', fontsize=14, handlelength=1.3)
for i in leg.legendHandles:
i.set_linewidth(12)
plot_time_probabilities(probabilities_in_time, figsize)
Upvotes: 0
Views: 623
Reputation: 80329
To set the legend in the center, you can set loc='center'
, or you can put the legend outside. To avoid that the legend handles grow to larger, you can leave out .set_linewidth(12)
(this sets a very wide edge width of 12 points).
Shifting the colors by one position can help to show the fill borders more pronounced. To still have a correct legend, the label should then be added to fill_between
.
The code below also tries to simplify part of the calls:
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import numpy as np
import pandas as pd
ncols = 10
figsize = (20, 5)
fontsize = 14
dti = pd.date_range('2013-01-01', '2020-12-31', freq='2W')
probabilities_in_time = np.random.random((ncols, len(dti)))
probabilities_in_time = probabilities_in_time / probabilities_in_time.sum(axis=0)
probabilities_in_time = pd.DataFrame(probabilities_in_time).T
probabilities_in_time.index = dti
cm_subsection = np.linspace(0, 1, ncols)
colors = cm.coolwarm(cm_subsection)
def plot_time_probabilities(probabilities_in_time, figsize):
plt.figure(figsize=figsize)
plt.yticks(np.arange(0, 1.2, 0.2), fontsize=fontsize)
plt.xticks(fontsize=fontsize)
draw_stack_plot(colors, probabilities_in_time)
set_grid()
set_legend()
# plt.margins(x=0, y=0)
plt.margins(x=0.02)
plt.tight_layout()
plt.show()
def draw_stack_plot(colors, probabilities_in_time):
current_probabilities = 0
for i, color in enumerate(colors):
plt.fill_between(probabilities_in_time.index,
probabilities_in_time[i] + current_probabilities, current_probabilities,
color=color, label=f' Probability: {i}')
current_probabilities += probabilities_in_time[i]
plt.plot(current_probabilities,
color=colors[0] if i <= 1 else colors[-1] if i >= 8 else colors[i - 1] if i < 5 else colors[i + 1])
def set_grid():
ax = plt.gca()
ax.set_axisbelow(False)
ax.xaxis.grid(True, linestyle='-', lw=1)
def set_legend():
leg = plt.legend(loc='lower left', fontsize=14, handlelength=1.3)
# leg = plt.legend(loc='upper left', bbox_to_anchor=(1.01, 1), fontsize=14, handlelength=1.3)
# for i in leg.legendHandles:
# i.set_linewidth(12)
plot_time_probabilities(probabilities_in_time, figsize)
Upvotes: 1