在numpy中返回类元素的向量 [英] Returning a vector of class elements in numpy

问题描述

我可以使用numpy的vectorize函数创建某个任意类的对象数组:

I can use numpy's vectorize function to create an array of objects of some arbitrary class:

import numpy as np

class Body:
    """
    Simple class to represent a point mass in 2D space, more to 
    play with numpy than anything else...
    """

    def __init__(self, position, mass, velocity):
        self.position = position
        self.mass     = mass
        self.velocity = velocity

    def __repr__(self):
        return "m = {} p = {} v = {}".format(self.mass, 
                self.position, self.velocity)

if __name__ == '__main__':

    positions  = np.array([0 + 0j, 1 + 1j, 2 + 0j])
    masses     = np.array([2,      5,      1])
    velocities = np.array([0 + 0j, 0 + 1j, 1 + 0j])

    vBody  = np.vectorize(Body)

    points = vBody(positions, masses, velocities)

现在，如果我想从points数组中检索包含(例如)velocities的向量，则可以使用普通的Python列表理解

Now, if I wanted to retrieve a vector containing (say) the velocities from the points array, I could just use an ordinary Python list comprehension

    v = [p.velocity for p in points]

但是有numpy -thonic方法吗?在大型数组上，这比使用列表理解更有效吗?

But is there a numpy-thonic way to do it? On large arrays would this be more efficient than using a list comprehension?

推荐答案

因此，我鼓励您不要将numpy数组与dtype一起使用.但是，这里实际上是一个结构，因此可以使用numpy来发挥优势，使用结构化数组.因此，首先，创建一个dtype:

So, I would encourage you not to use numpy arrays with an object dtype. However, what you have here is essentially a struct, so you could use numpy to your advantage using a structured array. So, first, create a dtype:

>>> import numpy as np
>>> bodytype = np.dtype([('position', np.complex), ('mass', np.float), ('velocity', np.complex)])

然后，初始化您的身体数组:

Then, initialize your body array:

>>> bodyarray = np.zeros((len(positions),), dtype=bodytype)
>>> bodyarray
array([(0j, 0.0, 0j), (0j, 0.0, 0j), (0j, 0.0, 0j)],
      dtype=[('position', '<c16'), ('mass', '<f8'), ('velocity', '<c16')])

现在，您可以轻松设置值:

Now, you can set your values easily:

>>> positions  = np.array([0 + 0j, 1 + 1j, 2 + 0j])
>>> masses     = np.array([2,      5,      1])
>>> velocities = np.array([0 + 0j, 0 + 1j, 1 + 0j])
>>> bodyarray['position'] = positions
>>> bodyarray['mass'] = masses
>>> bodyarray['velocity'] = velocities

现在您有了一系列实体"，它们可以充分利用numpy的优势，并允许您像这样访问属性":

And now you have an array of "bodies" that can take full advantage of numpy as well as letting you access "attributes" like this:

>>> bodyarray
array([(0j, 2.0, 0j), ((1+1j), 5.0, 1j), ((2+0j), 1.0, (1+0j))],
      dtype=[('position', '<c16'), ('mass', '<f8'), ('velocity', '<c16')])
>>> bodyarray['mass']
array([ 2.,  5.,  1.])
>>> bodyarray['velocity']
array([ 0.+0.j,  0.+1.j,  1.+0.j])
>>> bodyarray['position']
array([ 0.+0.j,  1.+1.j,  2.+0.j])
>>>

请注意此处

>>> bodyarray.shape
(3,)

这篇关于在numpy中返回类元素的向量的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！