PageRenderTime 45ms CodeModel.GetById 20ms RepoModel.GetById 0ms app.codeStats 0ms

/ase/io/trajectory.py

https://gitlab.com/keeeto/ase
Python | 380 lines | 369 code | 4 blank | 7 comment | 1 complexity | 62bae661fdc0e0bd3f2d3402f95718ca MD5 | raw file
  1. from __future__ import print_function
  2. import warnings
  3. from ase.calculators.singlepoint import SinglePointCalculator, all_properties
  4. from ase.constraints import dict2constraint
  5. from ase.atoms import Atoms
  6. from ase.io.aff import affopen, DummyWriter, InvalidAFFError
  7. from ase.io.jsonio import encode, decode
  8. from ase.io.pickletrajectory import PickleTrajectory
  9. from ase.parallel import world
  10. __all__ = ['Trajectory', 'PickleTrajectory']
  11. def Trajectory(filename, mode='r', atoms=None, properties=None, master=None):
  12. """A Trajectory can be created in read, write or append mode.
  13. Parameters:
  14. filename: str
  15. The name of the file. Traditionally ends in .traj.
  16. mode: str
  17. The mode. 'r' is read mode, the file should already exist, and
  18. no atoms argument should be specified.
  19. 'w' is write mode. The atoms argument specifies the Atoms
  20. object to be written to the file, if not given it must instead
  21. be given as an argument to the write() method.
  22. 'a' is append mode. It acts as write mode, except that
  23. data is appended to a preexisting file.
  24. atoms: Atoms object
  25. The Atoms object to be written in write or append mode.
  26. properties: list of str
  27. If specified, these calculator properties are saved in the
  28. trajectory. If not specified, all supported quantities are
  29. saved. Possible values: energy, forces, stress, dipole,
  30. charges, magmom and magmoms.
  31. master: bool
  32. Controls which process does the actual writing. The
  33. default is that process number 0 does this. If this
  34. argument is given, processes where it is True will write.
  35. The atoms, properties and master arguments are ignores in read mode.
  36. """
  37. if mode == 'r':
  38. return TrajectoryReader(filename)
  39. return TrajectoryWriter(filename, mode, atoms, properties, master=master)
  40. class TrajectoryWriter:
  41. """Writes Atoms objects to a .traj file."""
  42. def __init__(self, filename, mode='w', atoms=None, properties=None,
  43. extra=[], master=None):
  44. """A Trajectory writer, in write or append mode.
  45. Parameters:
  46. filename: str
  47. The name of the file. Traditionally ends in .traj.
  48. mode: str
  49. The mode. 'r' is read mode, the file should already exist, and
  50. no atoms argument should be specified.
  51. 'w' is write mode. The atoms argument specifies the Atoms
  52. object to be written to the file, if not given it must instead
  53. be given as an argument to the write() method.
  54. 'a' is append mode. It acts as write mode, except that
  55. data is appended to a preexisting file.
  56. atoms: Atoms object
  57. The Atoms object to be written in write or append mode.
  58. properties: list of str
  59. If specified, these calculator properties are saved in the
  60. trajectory. If not specified, all supported quantities are
  61. saved. Possible values: energy, forces, stress, dipole,
  62. charges, magmom and magmoms.
  63. master: bool
  64. Controls which process does the actual writing. The
  65. default is that process number 0 does this. If this
  66. argument is given, processes where it is True will write.
  67. """
  68. if master is None:
  69. master = (world.rank == 0)
  70. self.master = master
  71. self.atoms = atoms
  72. self.properties = properties
  73. self.numbers = None
  74. self.pbc = None
  75. self.masses = None
  76. self._open(filename, mode)
  77. def _open(self, filename, mode):
  78. if mode not in 'aw':
  79. raise ValueError('mode must be "w" or "a".')
  80. if self.master:
  81. self.backend = affopen(filename, mode, tag='ASE-Trajectory')
  82. if len(self.backend) > 0:
  83. r = affopen(filename)
  84. self.numbers = r.numbers
  85. self.pbc = r.pbc
  86. else:
  87. self.backend = DummyWriter()
  88. def write(self, atoms=None, **kwargs):
  89. """Write the atoms to the file.
  90. If the atoms argument is not given, the atoms object specified
  91. when creating the trajectory object is used.
  92. Use keyword arguments to add extra properties::
  93. writer.write(atoms, energy=117, dipole=[0, 0, 1.0])
  94. """
  95. b = self.backend
  96. if atoms is None:
  97. atoms = self.atoms
  98. if hasattr(atoms, 'interpolate'):
  99. # seems to be a NEB
  100. neb = atoms
  101. assert not neb.parallel or world.size == 1
  102. for image in neb.images:
  103. self.write(image)
  104. return
  105. while hasattr(atoms, 'atoms_for_saving'):
  106. # Seems to be a Filter or similar, instructing us to
  107. # save the original atoms.
  108. atoms = atoms.atoms_for_saving
  109. if len(b) == 0:
  110. b.write(version=1)
  111. # Atomic numbers and periodic boundary conditions are only
  112. # written once - in the header. Store them here so that we can
  113. # check that they are the same for all images:
  114. self.numbers = atoms.get_atomic_numbers()
  115. self.pbc = atoms.get_pbc()
  116. else:
  117. if (atoms.pbc != self.pbc).any():
  118. raise ValueError('Bad periodic boundary conditions!')
  119. elif len(atoms) != len(self.numbers):
  120. raise ValueError('Bad number of atoms!')
  121. elif (atoms.numbers != self.numbers).any():
  122. raise ValueError('Bad atomic numbers!')
  123. write_atoms(b, atoms, write_header=(len(b) == 0))
  124. calc = atoms.get_calculator()
  125. if calc is None and len(kwargs) > 0:
  126. calc = SinglePointCalculator(atoms)
  127. if calc is not None:
  128. if not hasattr(calc, 'get_property'):
  129. calc = OldCalculatorWrapper(calc)
  130. c = b.child('calculator')
  131. c.write(name=calc.name)
  132. for prop in all_properties:
  133. if prop in kwargs:
  134. x = kwargs[prop]
  135. else:
  136. if self.properties is not None:
  137. if prop in self.properties:
  138. x = calc.get_property(prop, atoms)
  139. else:
  140. x = None
  141. else:
  142. try:
  143. x = calc.get_property(prop, atoms,
  144. allow_calculation=False)
  145. except (NotImplementedError, KeyError):
  146. # KeyError is needed for Jacapo.
  147. x = None
  148. if x is not None:
  149. if prop in ['stress', 'dipole']:
  150. x = x.tolist()
  151. c.write(prop, x)
  152. info = {}
  153. for key, value in atoms.info.items():
  154. try:
  155. encode(value)
  156. except TypeError:
  157. warnings.warn('Skipping "{0}" info.'.format(key))
  158. else:
  159. info[key] = value
  160. if info:
  161. b.write(info=info)
  162. b.sync()
  163. def close(self):
  164. """Close the trajectory file."""
  165. self.backend.close()
  166. def __len__(self):
  167. return world.sum(len(self.backend))
  168. class TrajectoryReader:
  169. """Reads Atoms objects from a .traj file."""
  170. def __init__(self, filename):
  171. """A Trajectory in read mode.
  172. The filename traditionally ends in .traj.
  173. """
  174. self.numbers = None
  175. self.pbc = None
  176. self.masses = None
  177. self._open(filename)
  178. def _open(self, filename):
  179. try:
  180. self.backend = affopen(filename, 'r')
  181. except InvalidAFFError:
  182. raise RuntimeError('This is not a valid ASE trajectory file. '
  183. 'If this is an old-format (version <3.9) '
  184. 'PickleTrajectory file you can convert it '
  185. 'with ase.io.trajectory.convert("%s") '
  186. 'or:\n\n $ python -m ase.io.trajectory %s'
  187. % (filename, filename))
  188. self._read_header()
  189. def _read_header(self):
  190. b = self.backend
  191. if b.get_tag() != 'ASE-Trajectory':
  192. raise IOError('This is not a trajectory file!')
  193. if len(b) > 0:
  194. self.pbc = b.pbc
  195. self.numbers = b.numbers
  196. self.masses = b.get('masses')
  197. self.constraints = b.get('constraints', '[]')
  198. def close(self):
  199. """Close the trajectory file."""
  200. self.backend.close()
  201. def __getitem__(self, i=-1):
  202. b = self.backend[i]
  203. atoms = read_atoms(b, header=[self.pbc, self.numbers, self.masses,
  204. self.constraints])
  205. if 'calculator' in b:
  206. results = {}
  207. c = b.calculator
  208. for prop in all_properties:
  209. if prop in c:
  210. results[prop] = c.get(prop)
  211. calc = SinglePointCalculator(atoms, **results)
  212. calc.name = b.calculator.name
  213. atoms.set_calculator(calc)
  214. return atoms
  215. def __len__(self):
  216. return len(self.backend)
  217. def __iter__(self):
  218. for i in range(len(self)):
  219. yield self[i]
  220. def read_atoms(backend, header=None):
  221. b = backend
  222. if header:
  223. pbc, numbers, masses, constraints = header
  224. else:
  225. pbc = b.pbc
  226. numbers = b.numbers
  227. masses = b.get('masses')
  228. constraints = b.get('constraints', '[]')
  229. atoms = Atoms(positions=b.positions,
  230. numbers=numbers,
  231. cell=b.cell,
  232. masses=masses,
  233. pbc=pbc,
  234. info=b.get('info'),
  235. constraint=[dict2constraint(d)
  236. for d in decode(constraints)],
  237. momenta=b.get('momenta'),
  238. magmoms=b.get('magmoms'),
  239. charges=b.get('charges'),
  240. tags=b.get('tags'))
  241. return atoms
  242. def write_atoms(backend, atoms, write_header=True):
  243. b = backend
  244. if write_header:
  245. b.write(pbc=atoms.pbc.tolist(),
  246. numbers=atoms.numbers)
  247. if atoms.constraints:
  248. if all(hasattr(c, 'todict') for c in atoms.constraints):
  249. b.write(constraints=encode(atoms.constraints))
  250. if atoms.has('masses'):
  251. b.write(masses=atoms.get_masses())
  252. b.write(positions=atoms.get_positions(),
  253. cell=atoms.get_cell().tolist())
  254. if atoms.has('tags'):
  255. b.write(tags=atoms.get_tags())
  256. if atoms.has('momenta'):
  257. b.write(momenta=atoms.get_momenta())
  258. if atoms.has('magmoms'):
  259. b.write(magmoms=atoms.get_initial_magnetic_moments())
  260. if atoms.has('charges'):
  261. b.write(charges=atoms.get_initial_charges())
  262. def read_traj(filename, index):
  263. trj = TrajectoryReader(filename)
  264. for i in range(*index.indices(len(trj))):
  265. yield trj[i]
  266. def write_traj(filename, images):
  267. """Write image(s) to trajectory."""
  268. trj = TrajectoryWriter(filename, mode='w')
  269. if isinstance(images, Atoms):
  270. images = [images]
  271. for atoms in images:
  272. trj.write(atoms)
  273. trj.close()
  274. class OldCalculatorWrapper:
  275. def __init__(self, calc):
  276. self.calc = calc
  277. try:
  278. self.name = calc.name
  279. except AttributeError:
  280. self.name = calc.__class__.__name__.lower()
  281. def get_property(self, prop, atoms, allow_calculation=True):
  282. try:
  283. if (not allow_calculation and
  284. self.calc.calculation_required(atoms, [prop])):
  285. return None
  286. except AttributeError:
  287. pass
  288. method = 'get_' + {'energy': 'potential_energy',
  289. 'magmom': 'magnetic_moment',
  290. 'magmoms': 'magnetic_moments',
  291. 'dipole': 'dipole_moment'}.get(prop, prop)
  292. try:
  293. result = getattr(self.calc, method)(atoms)
  294. except AttributeError:
  295. raise NotImplementedError
  296. return result
  297. def convert(name):
  298. import os
  299. t = TrajectoryWriter(name + '.new')
  300. for atoms in PickleTrajectory(name, _warn=False):
  301. t.write(atoms)
  302. t.close()
  303. os.rename(name, name + '.old')
  304. os.rename(name + '.new', name)
  305. def main():
  306. import optparse
  307. parser = optparse.OptionParser(usage='python -m ase.io.trajectory '
  308. 'a1.traj [a2.traj ...]',
  309. description='Convert old trajectory '
  310. 'file(s) to new format. '
  311. 'The old file is kept as a1.traj.old.')
  312. opts, args = parser.parse_args()
  313. for name in args:
  314. convert(name)
  315. if __name__ == '__main__':
  316. main()