# Simple program to modify a lammps data file
from pylab import *
import re
infile = 'rst1.data'
outfile = 'modrst1.data'
datain = open(infile,'r')
dataout = open(outfile,'w')
nfound = 0
re_xlo = re.compile('xlo')
re_ylo = re.compile('ylo')
re_atoms = re.compile('atoms')
re_Atoms = re.compile('Atoms')
while (nfound==0):
    id = datain.readline().rstrip()
    dataout.write('%s\n' % id)
    px = id.find('xlo')
    py = id.find('ylo')
    pnatoms = id.find('atoms')
    if (px>=0):
        xx = map(float,id[:px-1].split(' '))
    if (py>=0):
        yy = map(float,id[:py-1].split(' '))
    if (pnatoms>=0):
        natoms = int(id[:pnatoms])
    if (re_Atoms.search(id)>=0):
        nfound=1
xcm = average(xx)
ycm = average(yy)
rad = (xx[1]-xx[0])*0.33
rad2=rad*rad
# Now read atom data and modify output
id = datain.readline().rstrip()
dataout.write('%s\n' % id)
for i in range(natoms):
    id = datain.readline().rstrip()
    dta = map(float,id.split(' '))
    x = dta[2]
    y = dta[3]
    z = dta[4]
    dx = x-xcm
    dy = y-ycm
    rr = dx*dx+dy*dy
    if (rr<rad2):
        dta[1]=1
    else:
        dta[1]=2
    dataout.write('%d %d %f %f %f %d %d %d\n' % (dta[0],dta[1],dta[2],dta[3],dta[4],dta[5],dta[6],dta[7]))
# Copy the rest of the file
for line in datain.readlines():
    dataout.write(line)
# Close files
datain.close()
dataout.close()