高效的numpy数组创建

5

给定x,我想要产生一个numpy数组x, log(x),其中x的形状为s,结果的形状为(*s, 2)。最简单的方法是什么?如果x只是一个浮点数,则我希望结果的形状为(2,)

一个丑陋的方法是:

import numpy as np

x = np.asarray(x)
result = np.empty((*x.shape, 2))
result[..., 0] = x
result[..., 1] = np.log(x)
1个回答

5
重要的是要将美学与性能分开。有时候丑陋的代码运行速度很快。实际上,在这里就是这种情况。尽管创建一个空数组,然后给切片赋值可能看起来不太美观,但它确实很快。
import numpy as np
import timeit 
import itertools as IT
import pandas as pd

def using_empty(x):
    x = np.asarray(x)
    result = np.empty(x.shape + (2,))
    result[..., 0] = x
    result[..., 1] = np.log(x)
    return result

def using_concat(x):
    x = np.asarray(x)
    return np.concatenate([x, np.log(x)], axis=-1).reshape(x.shape+(2,), order='F')

def using_stack(x):
    x = np.asarray(x)
    return np.stack([x, np.log(x)], axis=x.ndim)

def using_ufunc(x):
    return np.array([x, np.log(x)])
using_ufunc = np.vectorize(using_ufunc, otypes=[np.ndarray])

tests = [np.arange(600),
         np.arange(600).reshape(20,30),
         np.arange(960).reshape(8,15,8)]

# check that all implementations return the same result
for x in tests:
    assert np.allclose(using_empty(x), using_concat(x))
    assert np.allclose(using_empty(x), using_stack(x))


timing = []
funcs = ['using_empty', 'using_concat', 'using_stack', 'using_ufunc']
for test, func in IT.product(tests, funcs):
    timing.append(timeit.timeit(
        '{}(test)'.format(func), 
        setup='from __main__ import test, {}'.format(func), number=1000))

timing = pd.DataFrame(np.array(timing).reshape(-1, len(funcs)), columns=funcs)
print(timing)

在我的机器上运行,得到如下时间测试结果:
   using_empty  using_concat  using_stack  using_ufunc
0     0.024754      0.025182     0.030244     2.414580
1     0.025766      0.027692     0.031970     2.408344
2     0.037502      0.039644     0.044032     3.907487

所以在应用于tests时,using_empty是最快的选择。

请注意,np.stack正好符合你的要求,因此

np.stack([x, np.log(x)], axis=x.ndim)

看起来相当不错,但它也是测试的三个选项中最慢的。


请注意,除了速度慢之外,using_ufunc 还返回一个对象 dtype 的数组:

In [236]: x = np.arange(6)

In [237]: using_ufunc(x)
Out[237]: 
array([array([  0., -inf]), array([ 1.,  0.]),
       array([ 2.        ,  0.69314718]),
       array([ 3.        ,  1.09861229]),
       array([ 4.        ,  1.38629436]), array([ 5.        ,  1.60943791])], dtype=object)

这与所期望的结果不同:

In [240]: using_empty(x)
Out[240]: 
array([[ 0.        ,        -inf],
       [ 1.        ,  0.        ],
       [ 2.        ,  0.69314718],
       [ 3.        ,  1.09861229],
       [ 4.        ,  1.38629436],
       [ 5.        ,  1.60943791]])

In [238]: using_ufunc(x).shape
Out[238]: (6,)

In [239]: using_empty(x).shape
Out[239]: (6, 2)

我实际上最喜欢stack。我通常不会为了速度优化我的Python代码 :) - Neil G
您介意添加@ufunc装饰器作为另一种实现方式吗? - Neil G
我不熟悉@ufunc装饰器。你是指这个吗? - unutbu
抱歉,我的意思是这个:http://docs.scipy.org/doc/numpy-1.10.1/reference/generated/numpy.vectorize.html - Neil G
嗯,我不确定那该怎么做。我的天真尝试是 def using_ufunc(x): return x, np.log(x); using_ufunc = np.vectorize(using_ufunc); using_ufunc(np.arange(6)) 返回的是一个数组元组,而不是期望的结果。你能看出如何使用 np.vectorize 来实现吗? - unutbu
显示剩余4条评论

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接