1""" 

2IO support for the Gaussian cube format. 

3 

4See the format specifications on: 

5http://local.wasp.uwa.edu.au/~pbourke/dataformats/cube/ 

6""" 

7 

8 

9import numpy as np 

10import time 

11from ase.atoms import Atoms 

12from ase.io import read 

13from ase.units import Bohr 

14 

15 

16def write_cube(fileobj, atoms, data=None, origin=None, comment=None): 

17 """ 

18 Function to write a cube file. 

19 

20 fileobj: str or file object 

21 File to which output is written. 

22 atoms: Atoms object 

23 Atoms object specifying the atomic configuration. 

24 data : 3dim numpy array, optional (default = None) 

25 Array containing volumetric data as e.g. electronic density 

26 origin : 3-tuple 

27 Origin of the volumetric data (units: Angstrom) 

28 comment : str, optional (default = None) 

29 Comment for the first line of the cube file. 

30 """ 

31 

32 if data is None: 

33 data = np.ones((2, 2, 2)) 

34 data = np.asarray(data) 

35 

36 if data.dtype == complex: 

37 data = np.abs(data) 

38 

39 if comment is None: 

40 comment = "Cube file from ASE, written on " + time.strftime("%c") 

41 else: 

42 comment = comment.strip() 

43 fileobj.write(comment) 

44 

45 fileobj.write("\nOUTER LOOP: X, MIDDLE LOOP: Y, INNER LOOP: Z\n") 

46 

47 if origin is None: 

48 origin = np.zeros(3) 

49 else: 

50 origin = np.asarray(origin) / Bohr 

51 

52 fileobj.write( 

53 "{0:5}{1:12.6f}{2:12.6f}{3:12.6f}\n".format( 

54 len(atoms), *origin)) 

55 

56 for i in range(3): 

57 n = data.shape[i] 

58 d = atoms.cell[i] / n / Bohr 

59 fileobj.write("{0:5}{1:12.6f}{2:12.6f}{3:12.6f}\n".format(n, *d)) 

60 

61 positions = atoms.positions / Bohr 

62 numbers = atoms.numbers 

63 for Z, (x, y, z) in zip(numbers, positions): 

64 fileobj.write( 

65 "{0:5}{1:12.6f}{2:12.6f}{3:12.6f}{4:12.6f}\n".format( 

66 Z, 0.0, x, y, z) 

67 ) 

68 

69 data.tofile(fileobj, sep="\n", format="%e") 

70 

71 

72def read_cube(fileobj, read_data=True, program=None, verbose=False): 

73 """Read atoms and data from CUBE file. 

74 

75 fileobj : str or file 

76 Location to the cubefile. 

77 read_data : boolean 

78 If set true, the actual cube file content, i.e. an array 

79 containing the electronic density (or something else )on a grid 

80 and the dimensions of the corresponding voxels are read. 

81 program: str 

82 Use program='castep' to follow the PBC convention that first and 

83 last voxel along a direction are mirror images, thus the last 

84 voxel is to be removed. If program=None, the routine will try 

85 to catch castep files from the comment lines. 

86 verbose : bool 

87 Print some more information to stdout. 

88 

89 Returns a dict with the following keys: 

90 

91 * 'atoms': Atoms object 

92 * 'data' : (Nx, Ny, Nz) ndarray 

93 * 'origin': (3,) ndarray, specifying the cube_data origin. 

94 * 'spacing': (3, 3) ndarray, representing voxel size 

95 """ 

96 

97 readline = fileobj.readline 

98 line = readline() # the first comment line 

99 line = readline() # the second comment line 

100 

101 # The second comment line *CAN* contain information on the axes 

102 # But this is by far not the case for all programs 

103 axes = [] 

104 if "OUTER LOOP" in line.upper(): 

105 axes = ["XYZ".index(s[0]) for s in line.upper().split()[2::3]] 

106 if not axes: 

107 axes = [0, 1, 2] 

108 

109 # castep2cube files have a specific comment in the second line ... 

110 if "castep2cube" in line: 

111 program = "castep" 

112 if verbose: 

113 print("read_cube identified program: castep") 

114 

115 # Third line contains actual system information: 

116 line = readline().split() 

117 natoms = int(line[0]) 

118 

119 # Origin around which the volumetric data is centered 

120 # (at least in FHI aims): 

121 origin = np.array([float(x) * Bohr for x in line[1::]]) 

122 

123 cell = np.empty((3, 3)) 

124 shape = [] 

125 spacing = np.empty((3, 3)) 

126 

127 # the upcoming three lines contain the cell information 

128 for i in range(3): 

129 n, x, y, z = [float(s) for s in readline().split()] 

130 shape.append(int(n)) 

131 

132 # different PBC treatment in castep, basically the last voxel row is 

133 # identical to the first one 

134 if program == "castep": 

135 n -= 1 

136 cell[i] = n * Bohr * np.array([x, y, z]) 

137 spacing[i] = np.array([x, y, z]) * Bohr 

138 pbc = [(v != 0).any() for v in cell] 

139 

140 numbers = np.empty(natoms, int) 

141 positions = np.empty((natoms, 3)) 

142 for i in range(natoms): 

143 line = readline().split() 

144 numbers[i] = int(line[0]) 

145 positions[i] = [float(s) for s in line[2:]] 

146 

147 positions *= Bohr 

148 

149 atoms = Atoms(numbers=numbers, positions=positions, cell=cell, pbc=pbc) 

150 

151 # CASTEP will always have PBC, although the cube format does not 

152 # contain this kind of information 

153 if program == "castep": 

154 atoms.pbc = True 

155 

156 dct = {"atoms": atoms} 

157 

158 if read_data: 

159 data = np.array([float(s) 

160 for s in fileobj.read().split()]).reshape(shape) 

161 if axes != [0, 1, 2]: 

162 data = data.transpose(axes).copy() 

163 

164 if program == "castep": 

165 # Due to the PBC applied in castep2cube, the last entry along each 

166 # dimension equals the very first one. 

167 data = data[:-1, :-1, :-1] 

168 

169 dct["data"] = data 

170 dct["origin"] = origin 

171 dct["spacing"] = spacing 

172 

173 return dct 

174 

175 

176def read_cube_data(filename): 

177 """Wrapper function to read not only the atoms information from a cube file 

178 but also the contained volumetric data. 

179 """ 

180 dct = read(filename, format="cube", read_data=True, full_output=True) 

181 return dct["data"], dct["atoms"]