如何更新数组中许多对象的属性值

Question

我有一个 N 个粒子类 p = [Particle, Particle, ...] 对象的 numpy 数组，其中每个对象都有一个属性 pos 和 vel，每个数组 [x,y]，都是浮点数。我想知道更新数组中所有pos 和vel 属性值的最有效方法是什么。 我想对大量粒子运行以下更新，其中acc 是一个二维数组，N 很多 [x,y]：

for i,body in enumerate(p):
    body.vel -= acc[i] * timestep
    body.pos += body.vel * timestep

我想知道是否有更有效的方法来更改 p 元素的每个位置和速度属性，而不是使用相应的加速度元素逐个元素设置每个 pos 和 vel 元素的值。似乎必须有一个pythonic方式来做到这一点。

我希望像 p[:].pos 这样的东西能够访问 pos 值的数组。

Answer 1

也许numpy.recarray 会在这里为您提供帮助。使用recarray 而不是ndarray 更能满足您的要求。 recarray 允许使用属性进行字段访问。这里有一些示例代码：

acc = numpy.array([1, 2, 3])
timestep = 1
paticle_num = len(acc)

arr = numpy.recarray(paticle_num, dtype=[('pos', float), ('vel', float)])
# initialize
arr.pos = numpy.arange(paticle_num)
arr.vel = numpy.arange(paticle_num)

# operate in the whole dimension easily, I think that is what you want.
arr.vel -= acc * timestep
arr.pos += arr.vel * timestep

我只是将pos 和vel 视为浮点数以简化问题，如果它们是像[x, y] 这样的浮点数，就这样。

Answer 2

我建议使用pandas.DataFrame。

我做了一些实现来比较速度。您现在的选择几乎与使用 Particle 对象的简单 Python 列表相同。

import random
import timeit
import numpy as np
import pandas as pd

class Particle:
    def __init__(self, pos, vel):
        self.pos = pos
        self.vel = vel

    def __repr__(self):
        return 'Particle({:.2f} | {:.2f})'.format(self.pos, self.vel)

def f1(data, acc, time_step=2):
    for i, p in enumerate(data):
        p.vel -= acc[i] * time_step
        p.pos += p.vel * time_step

    return data

def f2(data, acc, time_step=2):
    df['vel'] -= acc * time_step
    df['pos'] += df['vel'] * time_step

    return data

if __name__ == '__main__':
    for n1 in (10**1, 10**3, 10**5):
        particle_list = [
            [random.random(), random.random()]
            for _ in range(n1)]
        acceleration_arr = np.random.random((n1, ))
        acceleration_arr_2 = acceleration_arr.reshape((n1, 1))

        # option 1
        particle_list_1 = [
            Particle(pos, vel)
            for pos, vel in particle_list]

        # option 2
        df = pd.DataFrame(
            data=particle_list,
            columns=['pos', 'vel'])

        # assure results are equal
        ret_1 = f1(particle_list_1, acceleration_arr)
        ret_2 = f2(df, acceleration_arr)
        # convert to lists
        ret_1 = [(p.pos, p.vel) for p in ret_1]
        ret_2 = [(p['pos'], p['vel']) for _, p in ret_2.iterrows()]
        # print('ret_1', ret_1)
        # print('ret_2', ret_2)
        assert ret_1 == ret_2

        # compare duration
        repetitions = 100
        t1 = timeit.timeit(
            'f1(particle_list_1, acceleration_arr)',
            'from __main__ import f1, acceleration_arr, particle_list_1',
            number=repetitions)
        t2 = timeit.timeit(
            'f2(df, acceleration_arr)',
            'from __main__ import f2, acceleration_arr, df',
            number=repetitions)
        print('n={:10d} | {:30s} {:.6f}'.format(n1, 'list with for-loop', t1))
        print('n={:10d} | {:30s} {:.6f}'.format(n1, 'pandas.DataFrame', t2))
        print()

运行这段代码给了我这个输出，表明pandas 版本随着数据大小的增长要快得多：

n=        10 | list with for-loop             0.001032
n=        10 | pandas.DataFrame               0.064379     # pandas is slower

n=      1000 | list with for-loop             0.106632
n=      1000 | pandas.DataFrame               0.067613     # pandas is faster

n=    100000 | list with for-loop             9.986003
n=    100000 | pandas.DataFrame               0.115627     # pandas is a lot faster

当然，我的示例实现与您描述的不完全一样，但它表明有一种更快的方法可以完成您的目标。