将多维numpy数组写入csv

时间:2015-05-26 08:26:28

标签: python csv numpy

我有一个包含函数值的多维numpy数组,我想把它写成一个很长的csv。我怎么能干净利落地做到这一点?我无法找到一个numpy函数,但也许我在谷歌搜索错误的条款。一个例子:

#!/usr/bin/python

import csv
import numpy as np

x = np.array([1, 2, 3, 4])
y = np.array([50, 51])
z = np.array([99, 100, 101])
f = np.arange(24).reshape((4, 2, 3))  # Contains f(x, y, z)
assert f.shape == (x.size, y.size, z.size)

## I'd like to create a csv file whose columns are x, y, z, f
## How can I do that?

## np.savetxt("test.csv", a, delimiter=",")
## TypeError: float argument required, not numpy.ndarray

## Works, but does numpy already have a function that does this?
with open("test.csv", "wb") as csvfile:
    writer = csv.writer(csvfile, delimiter=",", quotechar="'", quoting=csv.QUOTE_MINIMAL)
    writer.writerow(["x", "y", "z", "f"])
    for x_index in range(x.size):
        for y_index in range(y.size):
            for z_index in range(z.size):
                writer.writerow([x[x_index], y[y_index], z[z_index],
                                 f[x_index, y_index, z_index]])

我有三个向量x,y,z和一个包含函数值f(x,y,z)的X-by-Y-by-Z数组。换句话说,f [i,j,k]包含对应于x [i],y [j]和z [k]的函数值f。有没有更简洁的方法来写一个长的csv与列x,y,z,f?

这里有头test.csv:

x,y,z,f
1,50,99,0
1,50,100,1
1,50,101,2
1,51,99,3
1,51,100,4
1,51,101,5
2,50,99,6
2,50,100,7
2,50,101,8

编辑:这似乎也有效:

x_y_z = np.array([x for x in itertools.product(x, y, z)])
assert x_y_z.shape[0] == f.size
output_array = np.hstack((x_y_z, f.flatten().reshape((f.size, 1)))
np.savetxt("test2.csv", output_array, comments="", delimiter=",", fmt="%i",
           header="x,y,z,f")

我是否重新发明轮子?

1 个答案:

答案 0 :(得分:0)

事实上,是的,它比应该的更复杂。

给出3个列表x,y和z

import numpy as np

x = [1,2,3]
y = [4,5]
z = [6,7,8]

您需要修改此列表才能获得所有可能的组合,请以这种方式使用numpy.repeat

new_x = np.array(x).repeat(len(y)*len(z))
print new_x

>> [1 1 1 1 1 1 2 2 2 2 2 2 3 3 3 3 3 3]

new_y = np.array([y]).repeat(len(z),axis=0).repeat(len(x),axis=1)

print new_y

>> [4 4 4 5 5 5 4 4 4 5 5 5 4 4 4 5 5 5]

new_z = np.array([z]).repeat(len(x)*len(y),axis=0)
print new_z

>> [6 7 8 6 7 8 6 7 8 6 7 8 6 7 8 6 7 8]

# reshape y and z just like new_x

new_y = new_y.reshape(new_x.shape)
new_z = new_z.reshape(new_x.shape)

只是将它们连接起来!

# suppose that your vector f
f = np.array(range(len(x)*len(y)*len(z)))

matrix = np.array([new_x,new_y,new_z,f]).T
# or matrix = np.concatenate((np.concatenate((new_x,new_y),axis=1),np.concatenate((new_z,f),axis=1)),axis=1).T
print matrix

>>
[[ 1  4  6  0]
 [ 1  4  7  1]
 [ 1  4  8  2]
 [ 1  5  6  3]
 [ 1  5  7  4]
 [ 1  5  8  5]
 [ 2  4  6  6]
 [ 2  4  7  7]
 [ 2  4  8  8]
 [ 2  5  6  9]
 [ 2  5  7 10]
 [ 2  5  8 11]
 [ 3  4  6 12]
 [ 3  4  7 13]
 [ 3  4  8 14]
 [ 3  5  6 15]
 [ 3  5  7 16]
 [ 3  5  8 17]]

最后,将数组保存为csv

np.savetxt('file_name.csv',matrix)