Matplotlib 3维图中的箭头设置

Matplotlib是一个使用Python语言开发的数据可视化库，可用于绘制各种静态、动态、交互式的图形。在 Matplotlib 3D 绘图中，设置箭头是一个常见的需求，本文将介绍如何在 Matplotlib 3D 图中添加箭头。

3D 矢量绘制

在 Matplotlib 中绘制 3D 矢量需要使用 mplot3d 轴。在这里，我们将讨论如何绘制一个线段与两个箭头。

我们需要使用三个关键的库：

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

首先，我们可以生成一些随机数据，以便在 3D 空间中绘制。

# 准备数据
X, Y, Z = np.array([[0, 0, 0], [1, 2, 3], [4, 5, 6]]).T
U, V, W = np.array([[1, 2, 3], [2, 3, 4], [1, 0, 1]]).T

然后，创建一个关于 3D 绘图的坐标轴。

# 创建一个 3d 绘图坐标轴
fig = plt.figure()
ax = fig.gca(projection='3d')

最后，将数据传递给 Axes3D.quiver() 函数并进行绘制。

# 绘制 3D 矢量
ax.quiver(X, Y, Z, U, V, W, length=1, arrow_length_ratio=0.3, normalize=False)

其中，X、Y、Z 为矢量的空间位置，U、V、W 为矢量的方向和大小，length 为矢量的标度长度，arrow_length_ratio 为箭头长度和标度长度的比例，normalize 为矢量是否应规划化（即使得所有箭头的长度相同，但方向不一定相同）。

我们可以使用以下代码查看整个过程有多么简单：

# 准备数据
X, Y, Z = np.array([[0, 0, 0], [1, 2, 3], [4, 5, 6]]).T
U, V, W = np.array([[1, 2, 3], [2, 3, 4], [1, 0, 1]]).T

# 创建一个 3d 绘图坐标轴
fig = plt.figure()
ax = fig.gca(projection='3d')

# 绘制 3D 矢量
ax.quiver(X, Y, Z, U, V, W, length=1, arrow_length_ratio=0.3, normalize=False)

# 显示图像
plt.show()

在 Matplotlib 3D 图中添加箭头

绘制 3D 矢量并不总是足够的。Matplotlib 还提供了一些方法来为我们绘制 3D 箭头，具体方法如下：

def plotArrow3D(ax, x, y, z, dx, dy, dz, arrowLen=1.0, color=None, mutation_scale=10, lw=1):
    """`ax` is a matplotlib `Axes3D` object.
        `x`, `y`, `z` is the start point of the arrow.
        `dx`, `dy`, `dz` is the vector of the arrow.
        `arrowLen` is the length of the arrow head.
        `color` is the color of the arrow.
        `mutation_scale` is the scale factor of the arrow head.
        `lw` is the line width of the arrow.
    """
    # 计算箭头的长度和角度
    length = arrowLen * np.sqrt(dx**2 + dy**2 + dz**2)
    alpha = np.arccos(dz/length)
    beta = np.arctan2(dy,dx)

    # 绘制箭头线
    ax.plot([x, x+dx], [y, y+dy], [z, z+dz], color=color, lw=lw)

    # 绘制箭头尖端
    ax.plot([x+dx], [y+dy], [z+dz], marker='o', color=color)

    # 绘制箭头头部
    X,Y,Z = arrowHead(x+dx, y+dy, z+dz, alpha, beta, length, mutation_scale)
    ax.plot(X, Y, Z, color=color, lw=lw)

def arrowHead(x, y, z, alpha, beta, length, mutation_scale):
    """计算箭头头部坐标"""
    X,Y,Z = [],[],[]

    # 计算箭头尖端的坐标
    X.append(x)
    Y.append(y)
    Z.append(z)

    X.append(x - 0.5*mutation_scale*length*np.sin(alpha)*np.cos(beta))
    Y.append(y - 0.5*mutation_scale*length*np.sin(alpha)*np.sin(beta))
    Z.append(z - 0.5*mutation_scale*length*np.cos(alpha))

    X.append(x - 0.5*mutation_scale*length*np.sin(alpha)*np.cos(beta)
             + 0.1*length*np.cos(alpha)*np.cos(beta))
    Y.append(y - 0.5*mutation_scale*length*np.sin(alpha)*np.sin(beta)
             + 0.1*length*np.cos(alpha)*np.sin(beta))
    Z.append(z - 0.5*mutation_scale*length*np.cos(alpha) - 0.1*length*np.sin(alpha))

    X.append(x - 0.5*mutation_scale*length*np.sin(alpha)*np.cos(beta)
             - 0.1*length*np.cos(alpha)*np.cos(beta))
    Y.append(y - 0.5*mutation_scale*length*np.sin(alpha)*np.sin(beta)
             - 0.1*length*np.cos(alpha)*np.sin(beta))
    Z.append(z - 0.5*mutation_scale*length*np.cos(alpha) - 0.1*length*np.sin(alpha))

    return X,Y,Z

在上述代码中，plotArrow3D() 函数用于在 3D 坐标系中绘制箭头，arrowHead() 函数用于计算箭头头部的坐标。我们可以使用以下代码进行演示：

# 准备数据
X, Y, Z = np.array([[0, 0, 0], [1, 2, 3], [4, 5, 6]]).T
U, V, W = np.array([[1, 2, 3], [2, 3, 4], [1, 0, 1]]).T

# 创建一个 3d 绘图坐标轴
fig = plt.figure()
ax = fig.gca(projection='3d')

# 绘制 3D 矢量
ax.quiver(X, Y, Z, U, V, W, length=1, arrow_length_ratio=0.3, normalize=False)

# 在第一个点上添加箭头
plotArrow3D(ax, X[0], Y[0], Z[0], U[0], V[0], W[0], arrowLen=0.4, mutation_scale=20, color='red')

# 显示图像
plt.show()