在Python中保存KDTree对象？

Question

我正在使用Scipy的KDTree实现来读取一个300MB的大文件。现在，我是否可以将数据结构保存到磁盘上并重新加载它，还是只能每次启动程序时从文件中读取原始点并构建数据结构？我按以下方式构建KDTree：

def buildKDTree(self):
        self.kdpoints = numpy.fromfile("All", sep=' ')
        self.kdpoints.shape = self.kdpoints.size / self.NDIM, NDIM
        self.kdtree = KDTree(self.kdpoints, leafsize = self.kdpoints.shape[0]+1)
        print "Preparing KDTree... Ready!"

请问有什么建议吗？

- Legend

---

你尝试过使用pickle吗？ - helloworld922

当我尝试在KDTree对象上使用cPickle时，我的机器出现了错误。 - JoshAdel

Answer 1

KDtree使用嵌套类来定义其节点类型（innernode，leafnode）。Pickle仅适用于模块级别的类定义，因此嵌套类会使其失效：

import cPickle

class Foo(object):
    class Bar(object):
        pass

obj = Foo.Bar()
print obj.__class__
cPickle.dumps(obj)

<class '__main__.Bar'>
cPickle.PicklingError: Can't pickle <class '__main__.Bar'>: attribute lookup __main__.Bar failed

然而，有一个（hacky的）解决方法，即在模块范围内将类定义打补丁到scipy.spatial.kdtree中，以便pickler可以找到它们。 如果读取和写入pickle KDtree对象的所有代码都安装了这些补丁，则此hack应该可以正常工作：

import cPickle
import numpy
from scipy.spatial import kdtree

# patch module-level attribute to enable pickle to work
kdtree.node = kdtree.KDTree.node
kdtree.leafnode = kdtree.KDTree.leafnode
kdtree.innernode = kdtree.KDTree.innernode

x, y = numpy.mgrid[0:5, 2:8]
t1 = kdtree.KDTree(zip(x.ravel(), y.ravel()))
r1 = t1.query([3.4, 4.1])
raw = cPickle.dumps(t1)

# read in the pickled tree
t2 = cPickle.loads(raw)
r2 = t2.query([3.4, 4.1])
print t1.tree.__class__
print repr(raw)[:70]
print t1.data[r1[1]], t2.data[r2[1]]

输出：

<class 'scipy.spatial.kdtree.innernode'>
"ccopy_reg\n_reconstructor\np1\n(cscipy.spatial.kdtree\nKDTree\np2\nc_
[3 4] [3 4]

- samplebias

---

你有Cython cKDTree的补丁吗？ - denis

很抱歉，@Denis，我没有cKDTree的补丁。应该可以实现某种形式的保存/加载方法，但由于cKDTree节点是malloc'd结构体而不是类，所以需要更多的自定义操作。 - samplebias

不幸的是，我遇到了错误：“在调用Python对象时超过了最大递归深度”。公平地说，我的树是基于一个包含100万个5D坐标的列表计算出来的，由于从该数组中计算只需要几分钟（数组本身可以通过numpy保存和加载），所以我认为我必须接受这种情况。 - CastleH

非常好的解决方案，这个方法对其它类型也适用吗？你如何研究哪些类型可以以这种方式公开？ - raam86