lammps--将模型文件只保留原子坐标信息Python实现

#!/usr/bin/env python

filename = 'impact_AL2O3.data'
atoms = ['O','Al']
ni = len(atoms)  # number of atom types
id = []
id.append(['1','3'])   # atom types included are lumped as one same atom type
id.append(['2'])   # atom types included are lumped as one same atom type

f1 = open(filename, 'r')
lines = f1.readlines()
switch1 = False
switch2 = False
cube = []
massid = []
coordinates = []
index = 0

for i, line in enumerate(lines):

    if 'Bonds' in line:
        break

    if 'atoms' in line:
        data = line.split()
        na = int(data[0])    # number of atoms in file1

    if 'xlo' in line or 'ylo' in line or 'zlo' in line:
        cube.append(line)

    if 'Masses' in line:
        index = i
        switch1 = True
        continue

    for j in range(ni):
        if switch1 == True and i == index+1+int(id[j][0]):
            massid.append(line)

    if 'Atoms' in line:
        switch2 = True
        continue

    if switch2 == True and line != '\n':
        data = line.split()
        for j in range(ni):
            if data[2] in id[j]:
                data[2] = str(j+1)
        data.remove(data[1])
        for k in range(2):
            data.remove(data[9])
        line = ' '.join(data) + '\n'
        coordinates.append(line)

f1.close()


f2 = open('impact_Al2O3.txt', 'w')

f2.write('LAMMPS data file.\n')
f2.write('\n')
f2.write(str(na)+' atoms\n')
f2.write(str(ni)+' atom types\n')
f2.write('\n')
for line in cube:
    f2.write(line)
f2.write('\n')
f2.write('Masses\n')
f2.write('\n')
for line in massid:
    f2.write(line)
f2.write('\n')
f2.write('Atoms # charge\n')
f2.write('\n')
for line in coordinates:
    f2.write(line)
f2.write('\n')

f2.close()