Python中numpy矩阵赋值错误：原子坐标提取结果异常如何解决？

问题描述 #

我编写了如下Python代码，试图从输入文件中提取原子的三维坐标并存储到shape为(6,3)的numpy矩阵atoms中：

#! /usr/bin/env python
import sys
import numpy
import math

inputfile = open(sys.argv[1], 'r')

atoms = numpy.zeros((6, 3), dtype='float64')
for line in inputfile:
    sline = line.split()
    header = sline[0]
    for i in range(1, 4, 1):
        atoms[i] = float(sline[i])
        # print(float(sline[i]))
print(atoms)

输入文件内容：

O5    -2.978    -1.342     0.528
 C1    -4.475    -1.062     0.491
 C2    -4.710    -0.049    -0.519
 C3    -4.193     1.266    -0.144
 C4    -3.045     1.113     0.856
 C5    -2.268    -0.070     0.555

运行后得到错误结果：

[[ 0.     0.     0.   ]
 [-2.268 -2.268 -2.268]
 [-0.07  -0.07  -0.07 ]
 [ 0.555  0.555  0.555]
 [ 0.     0.     0.   ]
 [ 0.     0.     0.   ]]

结果与预期不符，需排查代码问题。

问题分析 #

代码存在两个关键错误：

• 行索引缺失：外层循环遍历每一行原子数据，但没有记录当前处理的是第几个原子，导致每一行的坐标都在覆盖atoms的第1、2、3行，最终只有最后一行的三个数值分别填充到这三行的所有列，其余行保持初始0值。
• 赋值逻辑错误：atoms[i] = float(sline[i])中，atoms[i]是numpy数组的整行，将单个浮点数赋值给整行，会让该行所有列都被设置为该数值，完全违背了提取三维坐标（每行对应一个原子的x、y、z值）的需求。

修复后的代码 #

#! /usr/bin/env python
import sys
import numpy

inputfile = open(sys.argv[1], 'r')

atoms = numpy.zeros((6, 3), dtype='float64')
row_idx = 0  # 跟踪当前处理的原子行索引
for line in inputfile:
    sline = line.split()
    if not sline:  # 跳过空行，增强鲁棒性
        continue
    # 将当前原子的三个坐标赋值给对应行
    atoms[row_idx] = [float(sline[1]), float(sline[2]), float(sline[3])]
    row_idx += 1
    if row_idx >= 6:  # 处理完6个原子后提前终止
        break
print(atoms)

修复说明 #

• 新增row_idx变量，确保每一行原子数据对应atoms矩阵的正确行。
• 直接将当前行的三个坐标组成列表赋值给atoms[row_idx]，numpy会自动将列表转换为匹配的行向量，无需逐列赋值。
• 加入空行判断和行数限制，避免处理多余内容导致索引越界。

内容的提问来源于stack exchange，提问作者Another.Chemist