NumPy中生成带条件依赖范围的meshgrid：高级技巧与应用

import numpy as np

# 设定目标网格维度 n
n = 3

# 定义 x 和 z 的范围，使用 n 个点
x = np.linspace(0, 1, n)
z = np.linspace(0, 1, n)

# 定义 y 的范围，使用 2*n - 1 个点
# 这里的 2*n - 1 是关键，它保证了在后续筛选后能得到 n^3 个点
y = np.linspace(0, 1, 2 * n - 1) # 对于 n=3，y 将有 5 个点

# 生成初始的超集网格
X_full, Y_full, Z_full = np.meshgrid(x, y, z, indexing='ij')

# 找到满足条件 X <= Y 的所有点的索引
indices = np.nonzero(X_full <= Y_full)

# 使用这些索引来筛选 X, Y, Z 数组
X_filtered = X_full[indices]
Y_filtered = Y_full[indices]
Z_filtered = Z_full[indices]

# 将筛选后的数组重塑为目标的三维网格形状
X = X_filtered.reshape([n, n, n])
Y = Y_filtered.reshape([n, n, n])
Z = Z_filtered.reshape([n, n, n])

import numpy as np

def generate_conditional_meshgrid(n: int):
    """
    生成一个 n x n x n 的三维网格，满足 x <= y 的条件。

    参数:
        n (int): 目标网格的维度大小。

    返回:
        tuple: (X, Y, Z) 三个 n x n x n 的 NumPy 数组，代表生成的网格坐标。
    """
    if not isinstance(n, int) or n <= 0:
        raise ValueError("n 必须是正整数。")

    # 1. 定义 x 和 z 的范围，使用 n 个点
    x_coords = np.linspace(0, 1, n)
    z_coords = np.linspace(0, 1, n)

    # 2. 定义 y 的范围，使用 2*n - 1 个点
    # 这是确保筛选后能得到 n^3 个点的关键
    y_coords = np.linspace(0, 1, 2 * n - 1)

    # 3. 生成初始的超集网格
    # 使用 'ij' 索引模式，使 X 对应 x_coords 的行，Y 对应 y_coords 的列，Z 对应 z_coords 的深度
    X_full, Y_full, Z_full = np.meshgrid(x_coords, y_coords, z_coords, indexing='ij')

    # 4. 找到满足条件 X <= Y 的所有点的索引
    indices = np.nonzero(X_full <= Y_full)

    # 5. 使用这些索引来筛选 X, Y, Z 数组
    X_filtered = X_full[indices]
    Y_filtered = Y_full[indices]
    Z_filtered = Z_full[indices]

    # 6. 将筛选后的数组重塑为目标的三维网格形状
    # 经过上述步骤，X_filtered, Y_filtered, Z_filtered 的长度都恰好是 n*n*n
    X = X_filtered.reshape([n, n, n])
    Y = Y_filtered.reshape([n, n, n])
    Z = Z_filtered.reshape([n, n, n])

    return X, Y, Z

# 示例使用
n_dim = 3
X_mesh, Y_mesh, Z_mesh = generate_conditional_meshgrid(n_dim)

print(f"生成的 X 维度: {X_mesh.shape}")
print(f"生成的 Y 维度: {Y_mesh.shape}")
print(f"生成的 Z 维度: {Z_mesh.shape}")

# 验证条件是否满足 (例如，检查第一个切片)
# print("\nX_mesh[:, 0, 0]:", X_mesh[:, 0, 0])
# print("Y_mesh[:, 0, 0]:", Y_mesh[:, 0, 0])
# print("Z_mesh[:, 0, 0]:", Z_mesh[:, 0, 0])

# 随机检查几个点是否满足 X <= Y
# for _ in range(5):
#     i, j, k = np.random.randint(0, n_dim, size=3)
#     print(f"X[{i},{j},{k}]={X_mesh[i,j,k]}, Y[{i},{j},{k}]={Y_mesh[i,j,k]}, Z[{i},{j},{k}]={Z_mesh[i,j,k]} -> X <= Y: {X_mesh[i,j,k] <= Y_mesh[i,j,k]}")

# 确保所有点都满足条件
assert np.all(X_mesh <= Y_mesh)
print("\n所有网格点都满足 X <= Y 条件。")