[FIXED] How to draw axis with arrows the same in Python

Issue

I plot the function, and write code for plotting graph of this function:

import seaborn as sns
import matplotlib.pyplot as plt
import numpy as np
from matplotlib import ticker
from matplotlib import rc


rc('text', usetex=True)


def fct(x):
    if -2 <= x < -1:
        y = 1.0
    elif -1 <= x < 0:
        y = -1.0
    elif 0 <= x < 0.5:
        y = 2.0
    elif 0.5 <= x < 1:
        y = -2
    else:
        y = 0
    return y


x = np.arange(-2.1, 1.2, 0.003)
yf = [fct(i) for i in x]

plt.style.use('science')  
sns.set(font_scale=2)
sns.set_style("whitegrid")
fig, ax = plt.subplots(figsize=(14, 12))
tick_spacing = 0.5
ax.xaxis.set_major_locator(ticker.MultipleLocator(tick_spacing))
ax.set_ylabel('$f(x)$', rotation='vertical')
ax.set_xlabel('$x$', labelpad=5)
ax.set_ylabel(ax.get_ylabel(), rotation=0, ha='right')
sns.lineplot(x=x, y=yf, color='black', linestyle='-', linewidth=1.5, label='$f(x)$')
ax.legend(loc="upper left")
ax.set_xlim(-2.005, 1.005)

fig.savefig('graph1', dpi=600, bbox_inches='tight')  # save

I want to draw axis and arrow of each axis. And my question is how to draw axis and arrows at the same style (from image)?

Solution

To plot the axis with arrows, you can use the function matplotlib.pyplot.arrow. I have shown you one possible implementation in the following function plot_arrows.

import seaborn as sns
import matplotlib.pyplot as plt
import numpy as np
from matplotlib import ticker
from matplotlib import rc

rc('text', usetex=True)


def plot_arrows(
    fig, ax, 
    width:float=5e-4, 
    head_length:float=8e-3, 
    head_width_factor:float=20, 
    color:str='black'
    ) -> None:
    figsize_x, figsize_y = fig.get_size_inches()
    xlim_min, xlim_max = ax.get_xlim()
    ylim_min, ylim_max = ax.get_ylim()
    span_x = xlim_max - xlim_min
    span_y = ylim_max - ylim_min
    widthx = width*figsize_x/span_x
    widthy = width*figsize_y/span_y
    head_lengthy = head_length*figsize_x/span_x
    head_lengthx = head_length*figsize_y/span_y

    plt.arrow(xlim_min, 0, xlim_max - xlim_min, 0, width=widthx, color=color, length_includes_head=True, head_width=head_width_factor*widthx, head_length=head_lengthx)
    plt.arrow(0, ylim_min, 0, ylim_max-ylim_min, width=widthy, color=color, length_includes_head=True, head_width=head_width_factor*widthy, head_length=head_lengthy)

def fct(x):
    if -2 <= x < -1:
        y = 1.0
    elif -1 <= x < 0:
        y = -1.0
    elif 0 <= x < 0.5:
        y = 2.0
    elif 0.5 <= x < 1:
        y = -2
    else:
        y = 0
    return y


x = np.arange(-2.1, 1.2, 0.003)
yf = [fct(i) for i in x]

plt.style.use('science')  
sns.set(font_scale=2)
sns.set_style("whitegrid")
fig, ax = plt.subplots(figsize=(14, 12))
tick_spacing = 0.5
ax.xaxis.set_major_locator(ticker.MultipleLocator(tick_spacing))
ax.set_ylabel('$f(x)$', rotation='vertical')
ax.set_xlabel('$x$', labelpad=5)
ax.set_ylabel(ax.get_ylabel(), rotation=0, ha='right')
sns.lineplot(x=x, y=yf, color='black', linestyle='-', linewidth=1.5, label='$f(x)$')
ax.legend(loc="upper left")
ax.set_xlim(-2.005, 1.005)

fig.savefig('graph1', dpi=600, bbox_inches='tight')  # save

Output:

Answered By - francois-marie