优化基于规则集的NumPy分形数组生成函数的高效实现问询

看起来你的分形生成代码因为嵌套的Python循环拖慢了速度——这是numpy使用中很常见的性能瓶颈，因为Python级别的循环开销远大于底层C实现的向量化操作。下面是针对你的场景设计的更高效、更Pythonic的改写方案：

核心优化思路 #

你的原始代码用了四层嵌套循环（通道循环→迭代循环→y轴循环→x轴循环），其中最影响性能的是内层的y/x循环：每次手动切片赋值new_matrix[y*3:...]的操作不仅冗余，还会触发多次数组内存分配和复制。

我们可以利用numpy的数组索引广播和形状重塑来替代这些循环，把核心操作完全交给numpy的底层优化逻辑处理。

改写后的高效实现 #

import numpy as np

def decode(rules, fractal_iterations, seed):
    # 从规则集自动获取单规则的尺寸（这里是3），避免硬编码
    rule_size = rules.shape[-1]
    num_channels = rules.shape[1]  # 对应你代码中的dimensions，应该是3
    
    # 初始化种子矩阵：将单种子扩展为多通道的初始矩阵 (num_channels, 1, 1)
    seed_arr = np.asarray(seed)
    seed_matrices = np.tile(seed_arr[np.newaxis, ...], (num_channels, 1, 1))
    
    for _ in range(fractal_iterations):
        current_shape = seed_matrices.shape[1:]  # (M, N)
        M, N = current_shape
        
        # 向量化获取所有元素对应的规则块：形状变为 (num_channels, M, N, rule_size, rule_size)
        # 这里假设rules的形状是 (256, num_channels, rule_size, rule_size)，和你原始代码一致
        blocks = np.array([rules[seed_matrices[i], i] for i in range(num_channels)])
        
        # 重塑为目标尺寸：先交换轴把规则块的行维度和种子的列维度对齐，再合并
        seed_matrices = blocks.swapaxes(2, 3).reshape(num_channels, M * rule_size, N * rule_size)
    
    # 把通道轴从第一个位置移到最后，匹配原始输出格式
    return np.moveaxis(seed_matrices, 0, -1)

进一步提速：完全向量化通道处理 #

如果可以调整规则集rules的形状为(num_channels, 256, rule_size, rule_size)（把通道维度放在最前面），我们可以去掉最后一层Python循环，实现完全向量化：

import numpy as np

def decode_fully_vectorized(rules, fractal_iterations, seed):
    rule_size = rules.shape[-1]
    num_channels = rules.shape[0]
    
    seed_arr = np.asarray(seed)
    seed_matrices = np.tile(seed_arr[np.newaxis, ...], (num_channels, 1, 1))
    
    for _ in range(fractal_iterations):
        M, N = seed_matrices.shape[1:]
        
        # 利用numpy的广播索引，一次性获取所有通道的规则块
        # rules形状为(num_channels, 256, rule_size, rule_size)
        blocks = rules[np.arange(num_channels)[:, None, None], seed_matrices]
        
        # 直接重塑，无需循环
        seed_matrices = blocks.swapaxes(2, 3).reshape(num_channels, M * rule_size, N * rule_size)
    
    return np.moveaxis(seed_matrices, 0, -1)

为什么这能提速？ #

1. 去掉内层循环：原始代码中y/x的嵌套循环被numpy的向量化索引和形状操作替代，这些操作在C层面执行，速度比Python循环快10~100倍（取决于迭代次数和矩阵大小）。
2. 减少内存操作：通过一次性生成所有规则块再重塑，避免了多次小切片的内存分配和复制，降低了开销。
3. 更简洁的代码：去掉冗余的手动维度计算和切片赋值，代码可读性更强，也更符合numpy的使用习惯。

验证正确性 #

你可以用示例规则测试：

• 规则0对应[[0,0,0],[0,1,0],[0,0,0]]，规则1对应[[1,1,1],[0,0,0],[1,1,1]]
• 种子[[0]]迭代1次后，输出的每个通道都应该是规则0对应的3x3矩阵，和原始代码结果一致。

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