如何在numpy savetxt中格式化，使零仅保存为“0”

Question

我正在把一个 numpy 的稀疏数组（密集型）保存成一个csv文件。结果是这个csv文件有3GB大。问题是里面95%的单元格都是0.0000。我使用了 fmt='%5.4f' 的格式。请问我该怎么格式化并保存，使得零只保存为0，而非零的浮点数则用 '%5.4f' 的格式保存？我相信如果能做到这一点，我可以把3GB的文件缩减到300MB。

我正在使用

np.savetxt('foo.csv', arrayDense, fmt='%5.4f', delimiter = ',')

谢谢！祝好

Answer 1

如果你只保存稀疏矩阵中非零的数值（在下面的例子中是m），这样会更好。你可以通过以下方式来实现：

fname = 'row_col_data.txt'
m = m.tocoo()
a = np.vstack((m.row, m.col, m.data)).T
header = '{0}, {1}'.format(*m.shape)
np.savetxt(fname, a, header=header, fmt=('%d', '%d', '%5.4f'))

然后可以通过以下方式重新组合稀疏矩阵：

row, col, data = np.loadtxt(fname, skiprows=1, unpack=True)
shape = map(int, open(fname).next()[1:].split(','))
m = coo_matrix((data, (row, col)), shape=shape)

Answer 2

还有一个简单的选项，可能适合你的需求，那就是使用'g'这个标识符。如果你更关心有效数字，而不是一定要看到固定数量的数字，并且不介意在科学计数法和普通浮点数之间切换，这个方法就很合适。例如：

np.savetxt("foo.csv", arrayDense, fmt='%5.4g', delimiter=',') 

如果arrayDense是这样的：

matrix([[ -5.54900000e-01,   0.00000000e+00,   0.00000000e+00],
    [  0.00000000e+00,   3.43560000e-08,   0.00000000e+00],
    [  0.00000000e+00,   0.00000000e+00,   3.43422000e+01]])

你用的方法会得到：

-0.5549,0.0000,0.0000
0.0000,0.0000,0.0000
0.0000,0.0000,34.3422

而上面的结果会变成：

-0.5549,    0,    0
0,3.436e-08,    0
0,    0,34.34

这种方法也更灵活。注意，使用'g'而不是'f'时，你不会丢失数据（比如3.4356e-08而不是0.0000）。不过，这当然还是要看你设置的精度。

Answer 3

如果你查看一下np.savetxt的源代码，你会发现，虽然里面有很多代码是用来处理参数和Python 2与Python 3之间的差异，但实际上它的核心就是一个简单的循环，逐行处理数据，把每一行格式化后写入文件。所以，如果你自己写一个这样的函数，性能上不会有太大影响。比如，下面这个简化的函数就是用来写入紧凑的零：

def savetxt_compact(fname, x, fmt="%.6g", delimiter=','):
    with open(fname, 'w') as fh:
        for row in x:
            line = delimiter.join("0" if value == 0 else fmt % value for value in row)
            fh.write(line + '\n')

举个例子：

In [70]: x
Out[70]: 
array([[ 0.        ,  0.        ,  0.        ,  0.        ,  1.2345    ],
       [ 0.        ,  9.87654321,  0.        ,  0.        ,  0.        ],
       [ 0.        ,  3.14159265,  0.        ,  0.        ,  0.        ],
       [ 0.        ,  0.        ,  0.        ,  0.        ,  0.        ],
       [ 0.        ,  0.        ,  0.        ,  0.        ,  0.        ],
       [ 0.        ,  0.        ,  0.        ,  0.        ,  0.        ]])

In [71]: savetxt_compact('foo.csv', x, fmt='%.4f')

In [72]: !cat foo.csv
0,0,0,0,1.2345
0,9.8765,0,0,0
0,3.1416,0,0,0
0,0,0,0,0
0,0,0,0,0
0,0,0,0,0

接着，既然你已经在写自己的savetxt函数了，那不如让它也能处理稀疏矩阵，这样你就不用在保存之前把它转换成（密集的）numpy数组了。（我假设这个稀疏数组是用scipy.sparse中的某种稀疏表示法实现的。）在下面的函数中，唯一的变化就是把... for value in row改成了... for value in row.A[0]。

def savetxt_sparse_compact(fname, x, fmt="%.6g", delimiter=','):
    with open(fname, 'w') as fh:
        for row in x:
            line = delimiter.join("0" if value == 0 else fmt % value for value in row.A[0])
            fh.write(line + '\n')

示例：

In [112]: a
Out[112]: 
<6x5 sparse matrix of type '<type 'numpy.float64'>'
    with 3 stored elements in Compressed Sparse Row format>

In [113]: a.A
Out[113]: 
array([[ 0.        ,  0.        ,  0.        ,  0.        ,  1.2345    ],
       [ 0.        ,  9.87654321,  0.        ,  0.        ,  0.        ],
       [ 0.        ,  3.14159265,  0.        ,  0.        ,  0.        ],
       [ 0.        ,  0.        ,  0.        ,  0.        ,  0.        ],
       [ 0.        ,  0.        ,  0.        ,  0.        ,  0.        ],
       [ 0.        ,  0.        ,  0.        ,  0.        ,  0.        ]])

In [114]: savetxt_sparse_compact('foo.csv', a, fmt='%.4f')

In [115]: !cat foo.csv
0,0,0,0,1.2345
0,9.8765,0,0,0
0,3.1416,0,0,0
0,0,0,0,0
0,0,0,0,0
0,0,0,0,0