Python将numpy数组写入CSV文件时如何添加列标题

原代码的问题很明确：

• 写入时给每个numpy数组额外套了一层列表，zip之后只生成1行内容，每个单元格存的是整个数组的字符串，根本没把数组元素逐行展开
• 写数据之前没先写入列标题行

直接用下面的代码就能实现需求，列名可以按需调整：

import numpy as np
import csv 

A = np.array([[1,2,3],[4,5,6],[7,8,9]])
B = np.array([[11,12,13],[14,15,16],[17,18,19]])  # 对应pressure列
C = np.array([[21,22,23],[24,25,26],[27,28,29]])  # 对应velocity列
D = np.array([[31,32,33],[34,35,36],[37,38,39]])

# 先把多维数组展平成一维，才能逐行配对四个数组的对应元素
a_col = A.flatten()
b_col = B.flatten()
c_col = C.flatten()
d_col = D.flatten()

with open('Data.csv', 'w', newline='', encoding='utf-8') as f:
    writer = csv.writer(f)
    # 第一行写列标题
    writer.writerow(['A', 'pressure', 'velocity', 'D'])
    # 逐行写入数据
    writer.writerows(zip(a_col, b_col, c_col, d_col))

如果想写得更简洁，直接用numpy自带的保存方法就行，不用额外导入csv模块：

import numpy as np

A = np.array([[1,2,3],[4,5,6],[7,8,9]])
B = np.array([[11,12,13],[14,15,16],[17,18,19]])
C = np.array([[21,22,23],[24,25,26],[27,28,29]])
D = np.array([[31,32,33],[34,35,36],[37,38,39]])

# 按列拼接四个展平后的数组
data = np.column_stack([A.flat, B.flat, C.flat, D.flat])
np.savetxt(
    'Data.csv',
    data,
    delimiter=',',
    header='A,pressure,velocity,D',
    comments='', # 去掉numpy默认加在表头前的#号
    fmt='%d' # 按整数格式保存，避免出现小数
)

小提示：csv模块写入时打开文件加newline=''是官方推荐写法，可以避免Windows系统下生成的CSV出现多余空行。

内容的提问来源于stack exchange，提问作者Wiz123