Coverage for /builds/kinetik161/ase/ase/io/xyz.py: 100.00%

1"""Reference implementation of reader and writer for standard XYZ files. 

2 

3See https://en.wikipedia.org/wiki/XYZ_file_format 

4 

5Note that the .xyz files are handled by the extxyz module by default. 

6""" 

7from ase.atoms import Atoms 

8 

9 

10def read_xyz(fileobj, index): 

11 # This function reads first all atoms and then yields based on the index. 

12 # Perfomance could be improved, but this serves as a simple reference. 

13 # It'd require more code to estimate the total number of images 

14 # without reading through the whole file (note: the number of atoms 

15 # can differ for every image). 

16 lines = fileobj.readlines() 

17 images = [] 

18 while len(lines) > 0: 

19 symbols = [] 

20 positions = [] 

21 natoms = int(lines.pop(0)) 

22 lines.pop(0) # Comment line; ignored 

23 for _ in range(natoms): 

24 line = lines.pop(0) 

25 symbol, x, y, z = line.split()[:4] 

26 symbol = symbol.lower().capitalize() 

27 symbols.append(symbol) 

28 positions.append([float(x), float(y), float(z)]) 

29 images.append(Atoms(symbols=symbols, positions=positions)) 

30 yield from images[index] 

31 

32 

33def write_xyz(fileobj, images, comment='', fmt='%22.15f'): 

34 comment = comment.rstrip() 

35 if '\n' in comment: 

36 raise ValueError('Comment line should not have line breaks.') 

37 for atoms in images: 

38 natoms = len(atoms) 

39 fileobj.write('%d\n%s\n' % (natoms, comment)) 

40 for s, (x, y, z) in zip(atoms.symbols, atoms.positions): 

41 fileobj.write('%-2s %s %s %s\n' % (s, fmt % x, fmt % y, fmt % z)) 

42 

43 

44# Compatibility with older releases 

45simple_read_xyz = read_xyz 

46simple_write_xyz = write_xyz