read contcar file, which is used to add CO mole to the slab

read_contcar.py

def read_contcar(contcar_path):
    '''
    read data from contcar file
    return coordiante_array and atom_list
    Target: 需要读取 原子种类, 原子坐标, 晶胞参数
    return:
            lattice (a. b. c)
            atom, x, y, z, tf_x, tf_y, tf_z
    '''
    with open(contcar_path, 'r') as f:
        info = f.readlines()
        # 获取晶胞参数信息
        _lattice = info[2:5]
        lattice = np.zeros((3, 3), dtype=np.float)
        for i in range(3):
            _abc = _lattice[i].strip()
            _abc = re.sub('\s+', ' ', _abc)
            lattice[i] = [float(_) for _ in _abc.split(' ')]

        # 获取元素种类, 数目信息
        atom_type = info[5].strip().split("    ")
        atom_number = [int(_) for _ in info[6].strip().split("    ")]
        total_number = sum(atom_number)
        atom = dict(zip(atom_type, atom_number))
        atom_array = [atom_type[0]] * atom[atom_type[0]] + \
            [atom_type[1]] * atom[atom_type[1]]

        mytype = np.dtype({
            'names': ['atom', 'x', 'y', 'z', 'ftx', 'fty', 'ftz'],
            'formats': ['S30', 'f', 'f', 'f', 'S30', 'S30', 'S30']}, align=True)

        # 获取坐标信息
        coordinate = np.zeros((total_number, 1), dtype=mytype)

        _coordinates = info[9: 9 + total_number]
        for i in range(total_number):
            coorline = _coordinates[i].strip()
            coorline = re.sub('\s+', ' ', coorline)
            coorline_split = coorline.split(' ')
            xyz, ft = [float(_)
                       for _ in coorline_split[0:3]], coorline_split[3:]
            new_cor_line = [atom_array[i]]
            new_cor_line.extend(xyz)
            new_cor_line.extend(ft)
            new_cor_line = tuple(new_cor_line)
            coordinate[i] = new_cor_line
    return lattice, coordinate