利用NumPy高效修改二维数组的2x2分块

import numpy as np
import itertools

# 假设 transitions 是一个查找表，将2x2块的字节表示映射到新的布尔数组
# transitions = {b'\x00\x00\x00\x00': np.array([True, True, True, True]), …}
# grid 是一个二维NumPy数组
# yshift, xshift, ny, nx 定义了迭代范围

# 传统的Python循环方式
# for y, x in itertools.product(range(yshift, ny, 2), range(xshift, nx, 2)):
#     grid[y:y+2, x:x+2].flat = transitions[grid[y:y+2, x:x+2].tobytes()]

import numpy as np

# 示例：创建一个10x10的0/1随机数组
A = np.random.randint(0, 2, (10, 10))
print("原始数组 A:\n", A)

# 计算步长
# A.strides[0] 是行步长（移动一行所需的字节数）
# A.strides[1] 是列步长（移动一列所需的字节数）
# 对于2x2块，外部视图的行步长是 A.strides[0]*2
# 外部视图的列步长是 A.strides[1]*2
# 内部2x2块的行步长是 A.strides[0]
# 内部2x2块的列步长是 A.strides[1]
Av = np.lib.stride_tricks.as_strided(A, 
                                     shape=(A.shape[0]//2, A.shape[1]//2, 2, 2), 
                                     strides=(A.strides[0]*2, A.strides[1]*2, A.strides[0], A.strides[1]))

print("\nAv 的形状:", Av.shape) # (5, 5, 2, 2)
print("Av 是 A 的视图吗？", np.shares_memory(A, Av)) # True

# 创建一个多维查找表 lut
# lut 的形状 (2,2,2,2, 2,2) 意味着：
# 前四个维度对应2x2块的四个布尔值 (0或1)
# 后两个维度是替换后的2x2结果块
lut = np.zeros((2, 2, 2, 2, 2, 2), dtype=A.dtype)

# 填充查找表（根据实际转换规则）
# 例如，如果块是 [[0,0],[0,0]]，替换为 [[1,1],[1,1]]
lut[0, 0, 0, 0] = [[1, 1], [1, 1]]
# 如果块是 [[0,0],[0,1]]，替换为 [[1,1],[1,0]]
lut[0, 0, 0, 1] = [[1, 1], [1, 0]]
# 如果块是 [[0,0],[1,0]]，替换为 [[1,1],[0,1]]
lut[0, 0, 1, 0] = [[1, 1], [0, 1]]
# 更多规则...
lut[1, 1, 0, 0] = [[1, 1], [1, 1]] # 示例规则

# 使用高级索引进行赋值
# Av[...,0,0] 提取所有2x2块的左上角元素
# Av[...,0,1] 提取所有2x2块的右上角元素
# 依此类推
Av[:] = lut[Av[..., 0, 0], Av[..., 0, 1], Av[..., 1, 0], Av[..., 1, 1]]

print("\n修改后的数组 A (方法一):\n", A)

# 创建一个扁平化索引的查找表 lut2
# lut2 的形状 (16, 2, 2) 意味着：
# 第一个维度是0-15的索引
# 后两个维度是替换后的2x2结果块
lut2 = np.zeros((16, 2, 2), dtype=A.dtype)

# 填充查找表（根据实际转换规则）
# 索引 0 (0000b) 对应 [[0,0],[0,0]]
lut2[0] = [[1, 1], [1, 1]]
# 索引 1 (0001b) 对应 [[0,0],[0,1]]
lut2[1] = [[1, 1], [1, 0]]
# 索引 2 (0010b) 对应 [[0,0],[1,0]]
lut2[2] = [[1, 1], [0, 1]]
# 索引 12 (1100b) 对应 [[1,1],[0,0]]
lut2[12] = [[1, 1], [1, 1]] # 示例规则

# 重新生成一个随机数组 A 进行演示
A_new = np.random.randint(0, 2, (10, 10))
Av_new = np.lib.stride_tricks.as_strided(A_new, 
                                        shape=(A_new.shape[0]//2, A_new.shape[1]//2, 2, 2), 
                                        strides=(A_new.strides[0]*2, A_new.strides[1]*2, A_new.strides[0], A_new.strides[1]))

# 计算每个2x2块的扁平化索引
# [[8,4],[2,1]] 是权重矩阵
idx = (Av_new * [[8, 4], [2, 1]]).sum(axis=(2, 3))

# 使用扁平化索引进行赋值
Av_new[:] = lut2[idx]

print("\n修改后的数组 A (方法二):\n", A_new)

# 假设 A 和 lut2 已经定义
# 重新生成一个随机数组 A 进行演示
A_partial = np.random.randint(0, 2, (10, 10))
Av_partial = np.lib.stride_tricks.as_strided(A_partial, 
                                             shape=(A_partial.shape[0]//2, A_partial.shape[1]//2, 2, 2), 
                                             strides=(A_partial.strides[0]*2, A_partial.strides[1]*2, A_partial.strides[0], A_partial.strides[1]))

# 计算局部区域的索引
idx_partial = (Av_partial[2:4, 2:4] * [[8, 4], [2, 1]]).sum(axis=(2, 3))

# 对 Av 的局部区域进行赋值
Av_partial[2:4, 2:4] = lut2[idx_partial]

print("\n局部修改后的数组 A (2:4, 2:4 块):\n", A_partial)