如何从扁平化的二维NumPy数组中提取数据块？

针对你求解二维泊松方程时，需要将二维网格操作转换为Fortran序扁平化向量操作的需求（适配Jax迭代求解器的向量输入要求），以下是无需显式循环、高效直接的实现方案：

核心逻辑：利用Fortran序索引映射 #

对于形状为(nx, nz)的二维数组，按**Fortran序（列优先）**扁平化后，二维坐标(i, j)（i为行索引，j为列索引）对应的扁平化索引公式为：

flat_idx = i + j * nx

基于这个映射关系，我们可以直接计算内部点及其上下左右邻点的扁平化索引，无需先转回二维数组。

高效实现代码 #

import numpy as np

nx = 5
nz = 7
numGPs = nx * nz

# 生成原二维网格（Fortran序）及扁平化数组
GPs_matrix = np.arange(numGPs).reshape((nx, nz), order='F')
GPs_flat = GPs_matrix.reshape(-1, order='F')

# 1. 定义内部点的行、列索引范围
inner_rows = np.arange(1, nx - 1)  # 排除首尾行
inner_cols = np.arange(1, nz - 1)  # 排除首尾列

# 2. 利用广播生成所有内部点及邻点的扁平化索引
# 内部点索引
cor_idx = inner_rows[:, None] + inner_cols * nx
# 上方邻点（同列，行相同，列+1）
top_idx = inner_rows[:, None] + (inner_cols + 1) * nx
# 下方邻点（同列，行相同，列-1）
btm_idx = inner_rows[:, None] + (inner_cols - 1) * nx
# 右方邻点（同行，列相同，行+1）
rgt_idx = (inner_rows + 1)[:, None] + inner_cols * nx
# 左方邻点（同行，列相同，行-1）
lft_idx = (inner_rows - 1)[:, None] + inner_cols * nx

# 3. 初始化结果并计算均值
av_flat = np.zeros_like(GPs_flat)
av_flat[cor_idx.flatten()] = (GPs_flat[top_idx.flatten()] + GPs_flat[btm_idx.flatten()] +
                              GPs_flat[rgt_idx.flatten()] + GPs_flat[lft_idx.flatten()]) / 4

# 验证与原二维方法结果一致
av_original = np.zeros_like(GPs_matrix)
av_original[1:-1, 1:-1] = (GPs_matrix[1:-1, 2:] + GPs_matrix[1:-1, :-2] +
                           GPs_matrix[2:, 1:-1] + GPs_matrix[:-2, 1:-1]) / 4
print(np.array_equal(av_flat, av_original.reshape(-1, order='F')))  # 输出True

方案优势 #

1. 无显式循环：利用NumPy广播机制生成索引，避免了低效的Python循环，适配Jax的向量化优化需求。
2. 无冗余操作：直接通过索引映射计算，无需反复reshape转换二维/扁平化数组，内存和计算效率更高。
3. Jax兼容性：将代码中的np替换为jax.numpy即可直接迁移到Jax环境，完美适配bicgstab等迭代求解器的向量输入要求。

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