¿Guardar el objeto KDTree en Python?

Question

Estoy usando la implementación de KDTree de Scipy para leer un archivo grande de 300 MB. Ahora, ¿hay alguna manera de guardar la estructura de datos en el disco y volver a cargarla o me quedo atascado con los puntos sin procesar del archivo y construyendo la estructura de datos cada vez que inicio mi programa? Estoy construyendo el KDTree de la siguiente manera:

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!" 

¿Alguna sugerencia por favor?

Answer 1

KDtree utiliza clases anidadas para definir sus tipos de nodo (innernode, leafnode). Pepinillo sólo funciona en las definiciones de clase de nivel de módulo, por lo que una clase anidada dispara hacia arriba:

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 

Sin embargo, hay una (hacky) solución de mono-parchear las definiciones de clases en el scipy.spatial.kdtree en el ámbito de módulo por lo que el pickler puede encontrarlos. Si todo el código que lee y escribe objetos en escabeche KDtree instala estos parches, este truco debería funcionar bien:

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]] 

Salida:

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