在 Matplotlib 等高线图中绘制坐标轴线或原点

背景介绍

Matplotlib 是 Python 中应用最广泛的绘图库之一，它提供了丰富的图形展示功能和大量的适用数据类型。在 Matplotlib 中，等高线图是一种常用的图形展示方式。在绘制等高线图的过程中，有时需要增加坐标轴线或将原点移动到指定位置，这样的高级绘图方法绝对是展示数据的很好选择。

绘制坐标轴线

假设我们有一组数据。

import numpy as np

x = np.linspace(-2, 2, 21)
y = np.linspace(-2, 2, 21)
X, Y = np.meshgrid(x, y)
Z = np.sqrt(X**2 + Y**2)

我们可以按照如下方法简单绘制等高线图。

import matplotlib.pyplot as plt

plt.contour(X, Y, Z)
plt.show()

这段代码先定义了数据 x, y，然后用 meshgrid() 函数建立网格。最后计算 Z 值，绘制等高线图。执行以上代码将得到与下图类似的图形。

默认情况下，等高线图中是没有坐标轴的，我们可以使用如下代码增加 x 轴和 y 轴坐标轴线。

# 增加坐标轴线
ax = plt.gca()  # 获取当前的坐标轴
ax.spines['right'].set_color('none') 
ax.spines['top'].set_color('none')
ax.spines['bottom'].set_position(('data', 0))  
ax.spines['left'].set_position(('data', 0))
plt.contour(X, Y, Z) 
plt.show()

在这个例子中，我们首先通过 gca() 函数获取当前的坐标轴。然后使用 spines 属性，设置上、右侧坐标轴为无色。接着将 bottom 和 left 坐标轴设置为从 data (即实际值)上开始，也就是在 x=0 和 y=0 上。

绘制新的原点

除了增加坐标轴线还有一个需求就是移动原点到其它位置。例如我们要在 x 轴方向上建立新的原点位置，坐标值为 1.5。这段代码可以得到相应的效果。

# 在 x 轴方向上建立新的原点位置
ax = plt.gca()
ax.spines['right'].set_color('none')  
ax.spines['top'].set_color('none') 
ax.spines['bottom'].set_position(('data', 1.5))  
ax.spines['left'].set_position(('data', 0))
plt.contour(X, Y, Z) 
plt.show()

这里先获取坐标轴对象，然后将所有的 axis.spines 设定成无色。接着，通过 spines[‘bottom’].set_position() 函数，令 x 轴从 x=1.5 开始。

以上代码也可以应用于 y 轴方向上的新原点位置。只需要将第 4 行代码改为以下代码即可。

ax.spines['left'].set_position(('data', 1.5))

将代码封装成函数

如果想要多次使用这个绘图方法，可以将代码封装成一个函数进行调用。以下是相应的函数。

def contour_with_axis(data, x_axis_pos=0, y_axis_pos=0):
    """
    绘制带有坐标轴线的等高线图，并设置新的原点位置。

    Parameters
    ----------
    data : numpy.array
        数据。
    x_axis_pos : float, optional
        新 x 轴原点坐标值，默认为 0。
    y_axis_pos : float, optional
        新 y 轴原点坐标值，默认为 0。

    Returns
    -------
    None.

    """
    x = np.linspace(-2, 2, 21)
    y = np.linspace(-2, 2, 21)
    X, Y = np.meshgrid(x, y)
    Z = data(X, Y)

    fig, ax = plt.subplots()
    ax.spines['right'].set_color('none')
    ax.spines['top'].set_color('none')
    ax.spines['bottom'].set_position(('data', x_axis_pos))
    ax.spines['left'].set_position(('data', y_axis_pos))
    plt.contour(X, Y, Z)
    plt.show()

这个函数接受三个参数，其中 data 是数据，x_axis_pos 是新的 x 轴原点位置，y_axis_pos 是新的 y 轴原点位置。如果这两个参数不传递，则默认都是 0。在这个函数中，我们首先生成 X, Y 轴网格，然后计算出 Z 值并绘制等高线图。接着设置 x,y 轴坐标轴线和原点位置，并最终画出等高线图。