Elisabeth Shevtsova
Elisabeth Shevtsova

Reputation: 659

Matplotlib fill_between edge

I need to create a plot as close to this picture as possible (given the generated dataframe code below): enter image description here

And here's the output plot of my code: enter image description here

What I am having problems with is:

  1. The edge of fill_between is not sharp as in the picture. What I have is some kind of white shadow. How do I change the line between the fillings to match a target picture?
  2. How do I aling legend color lines to the center, but not to the left border which my code does?

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

Answers (1)

JohanC
JohanC

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)

example plot

Upvotes: 1

Related Questions