get all the frames from an xyz file

xyz.py

import numpy as np
import re


def get_frames_from_xyz(filename, use_cols):
    """
    Get all data from different frames from an xyz file

    Args:
        filename (str): the path of the xyz file to parse
        use_cols (list [int]): the indices of columnes to obtain
            sometimes it is necessary to ignore the Label / atomType columes

    Return:
        numpy.ndarray: the data in different frame, shape (n_frame, n_particle, n_column)
    """
    f = open(filename, 'r')
    frames = []
    for line in f:
        is_head = re.match(r'(\d+)\n', line)
        if is_head:
            frames.append([])
            particle_num = int(is_head.group(1))
            f.readline()  # jump through comment line
            for j in range(particle_num):
                data = re.split(r'\s', f.readline())
                data = [data[i] for i in use_cols]
                frames[-1].append(list(map(float, data)))
    f.close()
    return np.array(frames)