Talk:Mod:Hunt Research Group/recentre xyz.py

#!/usr/bin/python

# HOWTO: ./recentre_xyz.py <INPUT FILE> <OUTPUT FILE> <X TRANS> <Y TRANS> <Z TRANS> (<UNIT CELL LENGTH>)
#
# <INPUT FILE> = input xyz file
# <OUTPUT FILE> = output xyz file
# <X TRANS> = translate all x-coordinates by this value
# <Y TRANS> = translate all y-coordinates by this value
# <Z TRANS> = translate all z-coordinates by this value
# <UNIT CELL LENGTH> (optional) = unit cell length - note a cubic system is assumed.

import sys
def main(file,output,x_trans,y_trans,z_trans,unit_cell_len=None):
  with open(file) as in_f:
    with open(output,'w') as out_f:
      first_line = in_f.readline()
      out_f.write(first_line)
      no_atoms = int(first_line.strip())
      sec_line = in_f.readline()
      out_f.write(sec_line)
      for a in range(no_atoms):
        atom_line = in_f.readline()
        #print ('atom_line ' + str(a+1) + ' ' + atom_line)
        atom_type,x_coord,y_coord,z_coord = atom_line.split()
        x_coord = float(x_coord) + x_trans
        y_coord = float(y_coord) + y_trans
        z_coord = float(z_coord) + z_trans
        if unit_cell_len:
          if x_coord > unit_cell_len/2:
            x_coord = x_coord - unit_cell_len
          elif x_coord < -unit_cell_len/2:
            x_coord = x_coord + unit_cell_len
          
          if y_coord > unit_cell_len/2:
            y_coord = y_coord - unit_cell_len
          elif y_coord < -unit_cell_len/2:
            y_coord = y_coord + unit_cell_len

          if z_coord > unit_cell_len/2:
            z_coord = z_coord - unit_cell_len
          elif z_coord < -unit_cell_len/2:
            z_coord = z_coord + unit_cell_len
        new_atom_line = atom_type + '   ' + str(x_coord) + '   ' + str(y_coord) + '   ' + str(z_coord) + '\n'
        out_f.write(new_atom_line)
      next_line = in_f.readline()
      n_atoms_check = int(next_line.strip())
      #print ('n_atoms_check:' + str(n_atoms_check))
      while n_atoms_check == no_atoms:
        #print (next_line)
        out_f.write(next_line)
        sec_line = in_f.readline()
        #print (sec_line)
        out_f.write(sec_line)
        for a in range(no_atoms):
          atom_line = in_f.readline()
          #print ('atom_line ' + str(a+1) + ' ' + atom_line)
          atom_type,x_coord,y_coord,z_coord = atom_line.split()
          x_coord = float(x_coord) + x_trans
          y_coord = float(y_coord) + y_trans
          z_coord = float(z_coord) + z_trans
          if unit_cell_len:
            if x_coord > unit_cell_len/2:
              x_coord = x_coord - unit_cell_len
            elif x_coord < -unit_cell_len/2:
              x_coord = x_coord + unit_cell_len
          
            if y_coord > unit_cell_len/2:
              y_coord = y_coord - unit_cell_len
            elif y_coord < -unit_cell_len/2:
              y_coord = y_coord + unit_cell_len

            if z_coord > unit_cell_len/2:
              z_coord = z_coord - unit_cell_len
            elif z_coord < -unit_cell_len/2:
              z_coord = z_coord + unit_cell_len
          new_atom_line = atom_type + '   ' + str(x_coord) + '   ' + str(y_coord) + '   ' + str(z_coord) + '\n'
          out_f.write(new_atom_line)
        next_line = in_f.readline()
        #print ('next_line: ' + str(next_line))
        if len(next_line.split()) > 0:
          n_atoms_check = int(next_line.strip())
        else:
          n_atoms_check = None

  return

if __name__ == "__main__":
  if len(sys.argv) == 6:
   sys.exit(main(sys.argv[1],sys.argv[2],float(sys.argv[3]),float(sys.argv[4]),float(sys.argv[5]),None))
  elif len(sys.argv) == 7:
   sys.exit(main(sys.argv[1],sys.argv[2],float(sys.argv[3]),float(sys.argv[4]),float(sys.argv[5]),float(sys.argv[6])))
  else:
    print ('Wrong number of arguments. Please try again')
    sys.exit()