从文件构造结构化的numpy数组？

Question

我的任务是在 python 中编写一个简单的 MD-Simulation 程序，它不使用 python 内置类型（dict、list），而只使用 numpy 数组。据我了解，这允许编译代码以更快地运行。在我的代码中的某一时刻，我想从一个类似字典的对象中获取原子的质量，该对象可以通过元素名称（如MassDict['N'] = 14.0067）进行切片。

根据我的阅读，我需要使用结构化的 numpy 数组。我现在要做的是使用以下表格打开我的文件：

H 1.008
He 4.003
Li 6.941

然后构造一个结构化的numpy数组，可以使用第一列的元素名称进行切片。

我尝试制作两个 numpy 数组，然后将它们连接起来，但这似乎不是我需要的。我的代码一开始看起来并不那么好。那么如何创建一个 numpy 对象，该对象可以最佳地从文本文件中按字符串切片？ 这是我的代码：

import numpy as np
import re

def mass_el(file):

    with open(file) as inf:

        for i, line in enumerate(inf):
            pass

        elements = np.empty((i+1),dtype='S2')
        masses = np.empty((i+1),dtype=np.float32)


    with open(file) as inf:

        for i, line in enumerate(inf):
            elements[i] = re.search('[a-zA-Z]+',line).group()
            masses[i] = re.search('\d+[.]\d+',line).group()

    #???


mass_el('elements.txt')

Answer 1

使用 numpy

您可以使用numpy loadtxt 函数，如果文件中的数据字段定期分隔（CSV 文件），您可以直接从文件中读取和格式化数据。

m_els = np.loadtxt('elements.txt', dtype={'names':('element', 'mass'), 'formats':('U2', 'f')})

使用您的 3 行文件，您将获得以下 m_els 数组：

array([('H', 1.008), ('He', 4.003), ('Li', 6.941)],
    dtype=[('element', '<U2'), ('mass', '<f4')])

我猜这就是你想要的。要获得一个元素，比如氢，请执行m_els[0] 以获得第一个元组。 m_els[0][1] 获取氢气质量。

使用熊猫

使用pandas 更容易，而且速度也很快，因为它建立在numpy 之上。

import pandas as pd
m_els = pd.read_csv('elements.txt', sep='\s+', header=None, names=['element', 'mass'])

在这种情况下，m_els 是一个数据框，其中元素名称用作索引：

  element   mass
0       H  1.008
1      He  4.003
2      Li  6.941

要获取与某个元素（例如氢）对应的行，请执行 m_els.iloc[0]。要获得氢的质量，请执行 m_els.loc[0, 'mass'] 或使用元素名称 m_els['mass'].loc[m_els['element'] == 'H']。