使用生成器来填充创建一个不需要for循环的numpy数组。

Question

我有一个简单的生成器来执行随机行走，我使用一个生成器来执行随机行走，因为我想运行它一些步骤，分析它，然后再运行一些。另一个很好的原因是，因为我想把数学模型和模拟隔离开来。

我希望使用一个 numpy array 存储来自生成器的数据（生成器返回一个包含有 time, position, velocity, acceleration). 为了达到这个目的，我之前创建了numpy数组，然后在for循环中（用枚举法），我填满了所有位置。

我的问题是，有什么方法可以避免出现这样的 for loop 完全没有？

代码如下。

import numpy as np

def random_walk(initial_position = 0, accel = 0, pa=0.25, pb=.5):
    """Initial position (often 0), acceleration, 0 < pa < pb < 1"""
    # Time, x-position, Velocity, Acceleration
    t, x, v, a = 0, initial_position, 0, accel
    yield (t, x, v, a)

    while True:        
        # Roll the dices
        god_wishes = np.random.random()

        if god_wishes <= pa:
                # Increase acceleration
                a += .1
        elif god_wishes <= pb:
                # Reduce acceleration
                a -= .1

        # Lets avoid too much acceleration
        a = np.clip(a, -.5, +.5)

        # How much time has passed, since last update?
        dt = np.random.random()
        v += dt*a
        x += dt*v
        t += dt

        yield (t, x, v, a)

然后，我创建了一个模拟与

rw = random_walk(0, .2, .25, .5)

我做了个实验

experiment = np.zeros((1000,4))
for i in range(1000):
    experiment[i] = rw.__next__()

这就是 for 循环，我真的很想避免。最欢迎一个性能高效的方法。

有什么提示吗？

Answer 1

你可以大大加快 random_walk 发电机 不 使用 np.clip.

这个基准显示了速度的提升。

import numpy as np

def clip(value, lower, upper):
    return lower if value < lower else upper if value > upper else value

def random_walk(initial_position = 0, accel = 0, pa=0.25, pb=.5, clip_function=np.clip):
    """Initial position (often 0), acceleration, 0 < pa < pb < 1"""
    # Time, x-position, Velocity, Acceleration
    t, x, v, a = 0, initial_position, 0, accel

    yield (t, x, v, a)

    while True:
        # # Roll the dices
        god_wishes = np.random.random()

        if god_wishes <= pa:
                # Increase acceleration
                a += .1
        elif god_wishes <= pb:
                # Reduce acceleration
                a -= .1

        # # Lets avoid too much acceleration
        a = clip_function(a, -.5, +.5)

        # # How much time has passed, since last update?
        dt = np.random.random()
        v += dt*a
        x += dt*v
        t += dt

        yield (t, x, v, a)

from timeit import timeit

def f1():
    rw = random_walk(0, .2, .25, .5)
    experiment = np.zeros((1000,4))
    for i in range(1000):
        experiment[i] = rw.__next__()
    return experiment

def f2():
    rw = random_walk(0, .2, .25, .5, clip_function=clip)
    return np.array([next(rw) for _ in range(1000)])

t1 = timeit(lambda: f1(), number=10)
t2 = timeit(lambda: f2(), number=10)

print(t1)
print(t2)

在我的机器上打印(AMD 2400G, numPy 1.18.3)。

0.26705767000385094
0.01964799000415951