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

1# fmt: off 

2 

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

4 

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

6 

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

8""" 

9from ase.atoms import Atoms 

10from ase.io.utils import validate_comment_line 

11 

12 

13def read_xyz(fileobj, index): 

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

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

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

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

18 # can differ for every image). 

19 lines = fileobj.readlines() 

20 images = [] 

21 while len(lines) > 0: 

22 symbols = [] 

23 positions = [] 

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

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

26 for _ in range(natoms): 

27 line = lines.pop(0) 

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

29 symbol = symbol.lower().capitalize() 

30 symbols.append(symbol) 

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

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

33 yield from images[index] 

34 

35 

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

37 comment = validate_comment_line(comment) 

38 

39 for atoms in images: 

40 natoms = len(atoms) 

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

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

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

44 

45 

46# Compatibility with older releases 

47simple_read_xyz = read_xyz 

48simple_write_xyz = write_xyz