python - 使用 matplotlib 绘制大量点并耗尽内存

Question

我有一个大的 (~6GB) 格式简单的文本文件

x1 y1 z1
x2 y2 z2
...

由于我可能会多次加载此数据，出于效率原因，我创建了一个 np.memmap 文件:

X,Y,Z = np.memmap(f_np_mmap,dtype='float32',mode='r',shape=shape).T

我想做的是情节:

plt.scatter(X, Y, 
           color=custom_colorfunction(Z), 
           alpha=.01, s=.001, marker='s', linewidth=0)

这非常适用于较小的数据集。但是，对于这个更大的数据集，我的内存不足。我检查过 plt.scatter 是否占用了所有内存；我可以单步执行 X,Y,Z 就好了。有没有办法“栅格化” Canvas ，这样我就不会用完内存？我不需要缩放和平移图像，它直接进入磁盘。我意识到我可以对数据进行分箱并绘制它，但我不确定如何使用自定义颜色图和 alpha 值来做到这一点。

Answer 1

@tcaswell 关于覆盖 Axes.draw 方法的建议绝对是最灵活的方法。

但是，您可以使用/滥用 blitting 来执行此操作而无需子类化 Axes。每次都使用 draw_artist 而不恢复 Canvas 。

还有一个额外的技巧:我们需要一个特殊的save 方法，因为所有其他方法都在保存之前绘制 Canvas ，这将清除我们之前在其上绘制的所有内容。

此外，正如 tcaswell 指出的那样，为每个项目调用 draw_artist 相当慢，因此对于大量的点，您需要对输入数据进行分 block 。分 block 会显着加快速度，但这种方法总是比绘制单个 PathCollection 慢。

无论如何，这些答案中的任何一个都应该可以缓解您的内存问题。这是一个简单的例子。

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import _png
from itertools import izip

def main():
    # We'll be saving the figure's background, so let's make it transparent.
    fig, ax = plt.subplots(facecolor='none')

    # You'll have to know the extent of the input beforehand with this method.
    ax.axis([0, 10, 0, 10])

    # We need to draw the canvas before we start adding points.
    fig.canvas.draw()

    # This won't actually ever be drawn. We just need an artist to update.
    col = ax.scatter([5], [5], color=[0.1, 0.1, 0.1], alpha=0.3)

    for xy, color in datastream(int(1e6), chunksize=int(1e4)):
        col.set_offsets(xy)
        col.set_color(color)
        ax.draw_artist(col)

    save(fig, 'test.png')

def datastream(n, chunksize=1):
    """Returns a generator over "n" random xy positions and rgb colors."""
    for _ in xrange(n//chunksize):
        xy = 10 * np.random.random((chunksize, 2))
        color = np.random.random((chunksize, 3))
        yield xy, color

def save(fig, filename):
    """We have to work around `fig.canvas.print_png`, etc calling `draw`."""
    renderer = fig.canvas.renderer
    with open(filename, 'w') as outfile:
        _png.write_png(renderer._renderer.buffer_rgba(),
                       renderer.width, renderer.height,
                       outfile, fig.dpi)

main()

此外，您可能会注意到顶部和左侧的书脊被拉长了。您可以通过在保存之前重新绘制这两个脊椎(ax.draw_artist(ax.spines['top']) 等)来解决这个问题。