以类似字典的方式将新项添加到某个结构化数组中

时间:2013-04-17 00:25:11

标签: python numpy

我想在numpy中扩展结构化数组对象,以便我可以轻松添加新元素。

例如,对于一个简单的结构化数组

>>> import numpy as np
>>> x=np.ndarray((2,),dtype={'names':['A','B'],'formats':['f8','f8']})
>>> x['A']=[1,2]
>>> x['B']=[3,4]

我想轻松添加新元素x['C']=[5,6],但是会出现与未定义名称'C'相关联的错误。

只需向np.ndarray添加新方法即可:

import numpy as np
class sndarray(np.ndarray):
    def column_stack(self,i,x):
        formats=['f8']*len(self.dtype.names)
        new=sndarray(shape=self.shape,dtype={'names':list(self.dtype.names)+[i],'formats':formats+['f8']})
        for key in self.dtype.names:
            new[key]=self[key]

        new[i]=x

        return new

然后,

>>> x=sndarray((2,),dtype={'names':['A','B'],'formats':['f8','f8']})
>>> x['A']=[1,2]
>>> x['B']=[3,4]
>>> x=x.column_stack('C',[4,4])
>>> x
sndarray([(1.0, 3.0, 4.0), (2.0, 4.0, 4.0)], 
  dtype=[('A', '<f8'), ('B', '<f8'), ('C', '<f8')])

是否有任何方式可以像字典一样添加新元素?例如

>>> x['C']=[4,4]
>>> x
sndarray([(1.0, 3.0, 4.0), (2.0, 4.0, 4.0)], 
  dtype=[('A', '<f8'), ('B', '<f8'), ('C', '<f8')])

更新

使用__setitem__我离理想的解决方案还有一步之遥,因为我不知道如何:

  

更改self

中引用的对象
import numpy as np

class sdarray(np.ndarray):
    def __setitem__(self, i,x):
    if i in self.dtype.names:
        super(sdarray, self).__setitem__(i,x)
    else:
        formats=['f8']*len(self.dtype.names)
        new=sdarray(shape=self.shape,dtype={'names':list(self.dtype.names)+[i],'formats':formats+['f8']})
        for key in self.dtype.names:
           new[key]=self[key]

        new[i]=x

        self.with_new_column=new

然后

>>> x=sndarray((2,),dtype={'names':['A','B'],'formats':['f8','f8']})
>>> x['A']=[1,2]
>>> x['B']=[3,4]
>>> x['C']=[4,4]
>>> x=x.with_new_column #extra uggly step!
>>> x
sndarray([(1.0, 3.0, 4.0), (2.0, 4.0, 4.0)], 
  dtype=[('A', '<f8'), ('B', '<f8'), ('C', '<f8')])

更新2 在所选答案中正确实施后,我发现问题已由pandas DataFrame对象解决:

>>> import numpy as np
>>> import pandas as pd
>>> x=np.ndarray((2,),dtype={'names':['A','B'],'formats':['f8','f8']})
>>> x=pd.DataFrame(x)
>>> x['A']=[1,2]
>>> x['B']=[3,4]
>>> x['C']=[4,4]
>>> x
   A  B  C
0  1  3  4
1  2  4  4
>>> 

1 个答案:

答案 0 :(得分:3)

使用numpy.recarray代替,在我的numpy 1.6.1中,您获得了一个额外的方法field,当您从numpy.ndarray继承时,该方法不存在。

This questionthis one (if using numpy 1.3)还讨论了向structured array添加字段的问题。从那里你会看到使用:

import numpy.lib.recfunctions as rf
rf.append_fields( ... )

可以大大简化您的生活。乍一看,我认为这个函数将附加到原始数组,但它会创建一个新实例。下面显示的class使用的是__setitem__()的解决方案,效果非常好。

您发现导致丑陋的解决方案的问题是reported in another question。问题在于,当您执行self=...时,您只是将new对象存储在变量中,但实体sdarray未被更新。也许可以从其方法中直接销毁和重建class,但根据that讨论,可以创建以下class,其中ndarray不是子类,但内部存储和调用。添加了一些其他方法以使其工作,看起来您正在直接使用ndarray。我没有详细测试它。

自动调整good solution has been presented here的大小。您也可以在代码中加入。

import numpy as np

class sdarray(object):
    def __init__(self, *args, **kwargs):
        self.recarray =  np.recarray( *args, **kwargs)

    def __getattr__(self,attr):
        if hasattr( self.recarray, attr ):
            return getattr( self.recarray, attr )
        else:
            return getattr( self, attr )

    def __len__(self):
        return self.recarray.__len__()

    def __add__(self,other):
        return self.recarray.__add__(other)

    def __sub__(self,other):
        return self.recarray.__sub__(other)

    def __mul__(self,other):
        return self.recarray.__mul__(other)

    def __rmul__(self,other):
        return self.recarray.__rmul__(other)

    def __getitem__(self,i):
        return self.recarray.__getitem__(i)

    def __str__(self):
        return self.recarray.__str__()

    def __repr__(self):
        return self.recarray.__repr__()

    def __setitem__(self, i, x):
        keys = []
        formats = []
        if i in self.dtype.names:
            self.recarray.__setitem__(i,x)
        else:
            for name, t in self.dtype.fields.iteritems():
                keys.append(name)
                formats.append(t[0])
            keys.append( i )
            formats.append( formats[-1] )
            new = np.recarray( shape = self.shape,
                              dtype = {'names'  : keys,
                                       'formats': formats} )
            for k in keys[:-1]:
                new[k] = self[k]
            new[i] = x
            self.recarray = new