Coverage for /builds/ase/ase/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 

10 

11 

12def read_xyz(fileobj, index): 

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

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

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

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

17 # can differ for every image). 

18 lines = fileobj.readlines() 

19 images = [] 

20 while len(lines) > 0: 

21 symbols = [] 

22 positions = [] 

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

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

25 for _ in range(natoms): 

26 line = lines.pop(0) 

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

28 symbol = symbol.lower().capitalize() 

29 symbols.append(symbol) 

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

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

32 yield from images[index] 

33 

34 

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

36 comment = comment.rstrip() 

37 if '\n' in comment: 

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

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