content: # -*- coding: utf-8 -*-
  
  # from framesplit import trajectory
  # too slow using this module
  
  import matplotlib.pyplot as plt
  
  dirName = r"F:\simulations\asphaltenes\na-mont\TMBO-oil\water\373-continue/"
  xyzName = 'all.xyz'
  hetero = 'O'  # 'oh' 'N' 'sp'      'O' 'Np' 'sp'
  
  with open(dirName + xyzName, 'r') as foo:
      coords = foo.readlines()
      nAtoms = int(coords[0])
      nFrames = int(len(coords) / (nAtoms + 2))
      pos = []
      for i in range(nFrames):
          istart = i * (nAtoms + 2)
          iend = (i + 1) * (nAtoms + 2)
          pos.append(coords[istart:iend])
  
  # for i in range(200):
  #     print coords[i]
  
  heteroatom = 0
  # all of my molecules have atoms less than 200
  for i in range(200):
      x = pos[0][i].split()[0]
      if x == hetero:
          heteroatom = i
          break
  
  heteroZ = []
  for p in pos:
      # print p[heteroatom].split()[0]
      zx = float(p[heteroatom].split()[3])
      if zx < 10:
          zx = zx + 80
      heteroZ.append(zx)
  
  with open(dirName + 'heteroZ.dat', 'w') as foo:
      for i, z in enumerate(heteroZ):
          print >> foo, "%3d %8.5f" % (i, z)
  
  # energy plot
  plt.figure(0, figsize=(8, 4))
  figName = dirName + 'heteroZ.png'
  plt.title('z of heteroatom', fontsize=20)
  plt.plot(range(len(heteroZ)-1), heteroZ[1:], linewidth=2)
  plt.grid(True)
  plt.xlabel('steps')
  plt.ylabel('Z')
  plt.axis([0, len(heteroZ)*1.1, 0, max(heteroZ)*1.1])
  plt.savefig(figName, format='png', dpi=300)
  plt.close()
  