1# -*- coding: utf-8 -*- 

2# 

3# Copyright (c) 2017, the cclib development team 

4# 

5# This file is part of cclib (http://cclib.github.io) and is distributed under 

6# the terms of the BSD 3-Clause License. 

7 

8"""A writer for XYZ (Cartesian coordinate) files.""" 

9 

10from cclib.io import filewriter 

11 

12 

13class XYZ(filewriter.Writer): 

14 """A writer for XYZ (Cartesian coordinate) files.""" 

15 

16 def __init__(self, ccdata, splitfiles=False, 

17 firstgeom=False, lastgeom=False, allgeom=False, 

18 *args, **kwargs): 

19 """Initialize the XYZ writer object. 

20 

21 Inputs: 

22 ccdata - An instance of ccData, parse from a logfile. 

23 splitfiles - Boolean to write multiple files if multiple files are requested. [TODO] 

24 firstgeom - Boolean to write the first available geometry from the logfile. 

25 lastgeom - Boolean to write the last available geometry from the logfile. 

26 allgeom - Boolean to write all available geometries from the logfile. 

27 """ 

28 

29 self.required_attrs = ('natom', 'atomcoords', 'atomnos') 

30 

31 # Call the __init__ method of the superclass 

32 super(XYZ, self).__init__(ccdata, *args, **kwargs) 

33 

34 self.do_firstgeom = firstgeom 

35 self.do_lastgeom = lastgeom 

36 self.do_allgeom = allgeom 

37 

38 self.natom = str(self.ccdata.natom) 

39 self.element_list = [self.pt.element[Z] for Z in self.ccdata.atomnos] 

40 

41 def generate_repr(self): 

42 """Generate the XYZ representation of the logfile data.""" 

43 

44 # Options for output (to a single file): 

45 # 1. Write all geometries from an optimization, which programs like VMD 

46 # can read in like a trajectory. 

47 # 2. Write the final converged geometry, which for any job other than 

48 # a geometry optimization would be the single/only geometry. 

49 # 3. Write the very first geometry, which for any job other than a 

50 # geometry optimization would be the single/only geometry. 

51 # 4. Write the first and last geometries from a geometry optimization. 

52 # 5. Write arbitrary structures via zero-based indexing. 

53 # TODO: Options for output (to multiple files) 

54 

55 xyzblock = [] 

56 

57 lencoords = len(self.ccdata.atomcoords) 

58 

59 # Collect the indices. 

60 if lencoords == 1 or self.do_firstgeom: 

61 self.indices.add(0) 

62 if self.do_lastgeom: 

63 self.indices.add(lencoords - 1) 

64 if self.do_allgeom: 

65 for i in range(lencoords): 

66 self.indices.add(i) 

67 

68 # Generate the XYZ string for each index. 

69 indices = sorted(self.indices) 

70 if not indices: 

71 indices = [-1] 

72 for i in indices: 

73 xyzblock.append(self._xyz_from_ccdata(i)) 

74 

75 # Ensure an extra newline at the very end. 

76 xyzblock.append('') 

77 

78 return '\n'.join(xyzblock) 

79 

80 def _xyz_from_ccdata(self, index): 

81 """Create an XYZ file of the geometry at the given index.""" 

82 

83 atomcoords = self.ccdata.atomcoords[index] 

84 existing_comment = "" if "comments" not in self.ccdata.metadata \ 

85 else self.ccdata.metadata["comments"][index] 

86 

87 # Create a comment derived from the filename and the index. 

88 if index == -1: 

89 geometry_num = len(self.ccdata.atomcoords) 

90 else: 

91 geometry_num = index + 1 

92 if self.jobfilename is not None: 

93 comment = "{}: Geometry {}".format(self.jobfilename, geometry_num) 

94 else: 

95 comment = "Geometry {}".format(geometry_num) 

96 # Wrap the geometry number part of the comment in square brackets, 

97 # prefixing it with one previously parsed if it existed. 

98 if existing_comment: 

99 comment = "{} [{}]".format(existing_comment, comment) 

100 else: 

101 comment = "[{}]".format(comment) 

102 

103 atom_template = '{:3s} {:15.10f} {:15.10f} {:15.10f}' 

104 block = [] 

105 block.append(self.natom) 

106 block.append(comment) 

107 for element, (x, y, z) in zip(self.element_list, atomcoords): 

108 block.append(atom_template.format(element, x, y, z)) 

109 return '\n'.join(block)