How to draw axis with arrows the same in Python

Question

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)?

Answer 1

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: