python的matplotlib中，如何把散点图的点画成与点运动方向一样的箭头

例如小球位置坐标（3，1），速度坐标为（0.5，0.5），如何在（3，1）处画出一个斜向45度的箭头

import numpy as np
import matplotlib.pyplot as plt

def drawArrow1(A, B):
    fig = plt.figure()
    ax = fig.add_subplot(121)
    # fc: filling color
    # ec: edge color
    ax.arrow(A[0], A[1], B[0]-A[0], B[1]-A[1],
             length_includes_head=True,# 增加的长度包含箭头部分
             head_width=0.2, head_length=0.2, fc='r', ec='b')
    ax.set_xlim(0,5)
    ax.set_ylim(0,5)
    ax.grid()
    ax.set_aspect('equal') #x轴y轴等比例
    plt.show()
    plt.tight_layout()

a = np.array([3, 1])
b = np.array([3 + 0.5, 1 + 0.5])
drawArrow1(a, b)