使用itertools和配套列表跳过元素模式

Question

我有一段代码运行得很慢（上次测试时需要30到60分钟），我需要对它进行优化。这段代码是用来从Abaqus中提取数据的，主要用于结构工程模型。代码中最慢的部分是一个循环，它首先按帧（也就是模拟时间的历史记录）遍历对象模型数据库，然后在这个基础上再遍历每一个节点。问题是有大约10万个“节点”，但实际上只有大约2万个是有用的。不过幸运的是，这些节点的顺序是固定的，这样我就不需要每次都查找节点的唯一标签。我可以先在一个单独的循环中处理一次，然后在最后筛选出需要的结果。因此，我把所有的节点放进一个列表中，然后去掉重复的节点。但是从代码来看：

    timeValues = []
    peeqValues = []
    for frame in frames: #760 loops
        setValues = frame.fieldOutputs['@@@fieldOutputType'].getSubset
                    region=abaqusSet, position=ELEMENT_NODAL).values
        timeValues.append(frame.frameValue)
        for value in setValues: # 100k loops
            peeqValues.append(value.data)

它仍然需要不必要地调用value.data大约8万次。如果有人对Abaqus的odb（对象数据库）对象熟悉的话，就会知道在Python中它们的速度非常慢。而且它们只能在单线程下运行，Abaqus使用的是自己的Python版本（2.6.x）和一些特定的包（比如numpy可以用，但pandas就不行）。还有一个麻烦的地方是，你可以通过位置来访问对象，比如frames[-1]可以得到最后一帧，但你不能切片，也就是说你不能这样做：for frame in frames[0:10]: # 遍历前10个元素。

我对itertools没有经验，但我想给它提供一个节点ID的列表（或者一个真/假的列表）来映射到setValues上。每760帧的setValues的长度和模式都是一样的。也许可以这样做：

    for frame in frames: #still 760 calls
        setValues = frame.fieldOutputs['@@@fieldOutputType'].getSubset(
                    region=abaqusSet, position=ELEMENT_NODAL).values
        timeValues.append(frame.frameValue)
        # nodeSet_IDs_TF = [True, True, False, False, False, ...] same length as       
        # setValues
        filteredSetValues = ifilter(nodeSet_IDs_TF, setValues)  
        for value in filteredSetValues: # only 20k calls
            peeqValues.append(value.data)

如果有其他建议也很感谢，之后我想通过把.append()从循环中去掉，整个放到一个函数里来看看是否能提高效率。整个脚本现在已经在1.5小时内完成（之前是6小时，有时候甚至需要21小时），但一旦开始优化，就没法停下来了。

关于内存的考虑也很重要，我是在一个集群上运行这些代码的，我记得有一次我用80GB的内存也能运行。脚本在160GB的内存下肯定能正常工作，问题是如何分配到这些资源。

我搜索过解决方案，但可能是用错了关键词，我相信在循环中遇到这个问题并不罕见。

编辑 1

这是我最终使用的代码：

    # there is no compress under 2.6.x ... so use the equivalent recipe:
    from itertools import izip
    def compress(data, selectors):
        # compress('ABCDEF', [1,0,1,0,1,1]) --> ACEF
        return (d for d, s in izip(data, selectors) if s)
    def iterateOdb(frames, selectors): # minor speed up
        peeqValues = []
        timeValues = []
        append = peeqValues.append # minor speed up
        for frame in frames:
            setValues = frame.fieldOutputs['@@@fieldOutputType'].getSubset(
            region=abaqusSet, position=ELEMENT_NODAL).values
            timeValues.append(frame.frameValue)
            for value in compress(setValues, selectors): # massive speed up
                append(value.data)
        return peeqValues, timeValues

    peeqValues, timeValues = iterateOdb(frames, selectors)

最大的改进来自于使用compress(values, selectors)方法（整个脚本，包括odb部分，从大约1小时30分钟缩短到25分钟）。另外，通过append = peeqValues.append和把所有内容放在def iterateOdb(frames, selectors):中也有一些小的改进。

我参考了以下链接的建议：https://wiki.python.org/moin/PythonSpeed/PerformanceTips

感谢大家的回答和帮助！

Answer 1

除了的解决方案，还有另一种变体：

for value, selector in zip(setValues, selectors):
    if selector:
        peeqValue.append(value.data)

如果你想让输出列表的长度和setValue的长度保持一致，那么可以加一个else分支：

for value, selector in zip(setValues, selectors):
    if selector:
        peeqValue.append(value.data)
    else:
        peeqValue.append(None)

这里的selector是一个包含True和False的向量，它的长度和setValues是一样的。

在这种情况下，选择哪个其实是个人喜好。如果要遍历7600万个节点（760 x 100 000）需要30分钟，那么时间并不是花在Python的循环上。

我试过这个：

def loopit(a):  
    for i in range(760):
        for j in range(100000):
             a = a + 1
    return a

IPython的%timeit显示循环的时间是3.54秒。所以，循环可能只占总时间的0.1%。

Answer 2

如果你对 itertools 还不太熟悉，可以先在你的循环中使用 if 语句。

比如：

for index, item in enumerate(values):
    if not selectors[index]:
        continue
    ...
# where selectors is a truth array like nodeSet_IDs_TF

这样你可以更确定你得到的结果是正确的，同时也能获得使用 itertools 时大部分的性能提升。

itertools 中对应的功能是 compress。

for item in compress(values, selectors):
    ...

我对 abaqus 不是很了解，但你可以尝试找到方法把选择器提供给 abaqus，这样它就不用浪费时间去创建每个值，最后又把它们丢掉。如果 abaqus 是用来处理大数组的数据，那么这种情况很可能会发生。