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

Question

I've been given the task of writing a simple MD-Simulation program in python which does not utilize python builtin types (dict, list) but only numpy arrays. From what I understood, this allows the code to be compiled to run faster. At one point in my code, I want to get the mass of an atom from a dictionary-like object, that can be sliced by the element name like MassDict['N'] = 14.0067 .

From what I've read, I'd need to use a structured numpy array. What I want to do right now is open my file with the following form:

H 1.008
He 4.003
Li 6.941

and then construct a structured numpy array, which can be sliced by using the element names in the first column.

I tried making two numpy arrays and then concatenating them, but that doesn't seem to be what I need. My code looks not that great to begin with. So how do I create a numpy object which can be sliced by string from a text file optimally? Here's my code:

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

Using numpy

You could use the numpy loadtxt function, which allows you to read and format directly from a file your data, if the data fields in the file are regularly separated (a CSV file).

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

Using your 3-lines file, you get the following m_els array:

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

Which is what you want I guess. To get an element, say hydrogen, do m_els[0] to get the first tuple. m_els[0][1] to get the hydrogen mass.

Using pandas

It's even easier using pandas which is also fast since it's builtd on top of numpy .

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

In this case m_els is a dataframe, where element names are used as indexes:

  element   mass
0       H  1.008
1      He  4.003
2      Li  6.941

To get the row corresponding to an element, say hydrogen, do m_els.iloc[0] . To get the hydrogen mass, do m_els.loc[0, 'mass'] or, using the element name, m_els['mass'].loc[m_els['element'] == 'H'] .