如何修改Python代码绘制抛物面两个限定区域的公共部分？

现有如下Python代码，可生成抛物面并分别绘制其高于下限或低于上限的区域，现希望修改代码，绘制这两个限定区域的公共部分（即同时满足r_min ≤ X²+Y² ≤ r_max的区域）：

import numpy as np
import matplotlib.pyplot as plt

if __name__ == "__main__":
    r: float = 1.5
    r_min: float = 0.5 * r
    r_max: float = 0.75 * r
    n_points: int = 100

    r: np.ndarray = np.linspace(0, r, n_points)
    r_pellet: np.ndarray = np.linspace(r_min, r_max, n_points)
    theta: np.ndarray = np.linspace(0, 2 * np.pi, n_points)
    R, THETA = np.meshgrid(r, theta)

    # Deal with figures
    fig = plt.figure()
    ax = fig.add_subplot(projection = "3d")

    # Create projectile surface
    X, Y = R * np.cos(THETA), R * np.sin(THETA)
    Z = (X**2 + Y**2)
    ax.plot_surface(X, Y, Z, cmap="cool", linewidth=0, antialiased=True)

    # Create pellet surface
    # Above lower limit
    X_SUP = np.where(r_min <= X**2 + Y**2, X, np.NAN)
    Y_SUP = np.where(r_min <= X**2 + Y**2, Y, np.NAN)
    Z_SUP = X_SUP**2 + Y_SUP**2

    # Under upper limit
    X_INF = np.where(X**2 + Y**2 <= r_max, X, np.NAN)
    Y_INF = np.where(X**2 + Y**2 <= r_max, Y, np.NAN)
    Z_INF = X_INF**2 + Y_INF**2

    ax.plot_surface(X, Y, Z_INF, cmap="hot", linewidth=0, antialiased=True)

    ax.set_xlabel("X")
    ax.set_ylabel("Y")
    ax.set_zlabel("Z")

    plt.show()

解决方案 #

要绘制两个区域的公共部分，只需将两个条件合并——筛选同时满足X²+Y² ≥ r_min和X²+Y² ≤ r_max的点，把不满足条件的坐标设为NaN（matplotlib会自动忽略这些点）。

修改后的代码如下：

import numpy as np
import matplotlib.pyplot as plt

if __name__ == "__main__":
    r: float = 1.5
    r_min: float = 0.5 * r
    r_max: float = 0.75 * r
    n_points: int = 100

    r: np.ndarray = np.linspace(0, r, n_points)
    theta: np.ndarray = np.linspace(0, 2 * np.pi, n_points)
    R, THETA = np.meshgrid(r, theta)

    # Deal with figures
    fig = plt.figure()
    ax = fig.add_subplot(projection = "3d")

    # Create projectile surface
    X, Y = R * np.cos(THETA), R * np.sin(THETA)
    Z = (X**2 + Y**2)
    ax.plot_surface(X, Y, Z, cmap="cool", linewidth=0, antialiased=True)

    # 绘制上下限的公共区域
    # 合并两个条件：同时满足大于等于r_min和小于等于r_max
    mask = (X**2 + Y**2 >= r_min) & (X**2 + Y**2 <= r_max)
    X_COMMON = np.where(mask, X, np.NAN)
    Y_COMMON = np.where(mask, Y, np.NAN)
    Z_COMMON = X_COMMON**2 + Y_COMMON**2

    ax.plot_surface(X_COMMON, Y_COMMON, Z_COMMON, cmap="hot", linewidth=0, antialiased=True)

    ax.set_xlabel("X")
    ax.set_ylabel("Y")
    ax.set_zlabel("Z")

    plt.show()

关键修改说明

• 用mask变量存储元素级逻辑与的结果（numpy中需用&而非and，因为要对每个数组元素逐一判断）
• 基于mask筛选出同时满足两个条件的坐标，不满足的设为NaN
• 直接绘制筛选后的公共区域即可

内容的提问来源于stack exchange，提问作者aheuchamps