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

Question

我想扩展 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__ 我离理想的解决方案还有一步，因为我不知道如何:

change the object referenced at 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
>>> 

Answer 1

改为使用 numpy.recarray，在我的 numpy 1.6.1 中，您会得到一个额外的方法 field，当您从numpy.ndarray.

This question或 this one (if using numpy 1.3)还讨论了将字段添加到结构化数组。从那里你会看到使用:

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

可以大大简化您的生活。乍一看，我以为这个函数会附加到原始数组，但它会创建一个新实例。下面显示的类正在使用您的__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