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 

8import time 

9 

10import numpy as np 

11 

12from ase.atoms import Atoms 

13from ase.io import read 

14from ase.units import Bohr 

15 

16ATOMS = 'atoms' 

17CASTEP = 'castep' 

18DATA = 'data' 

19 

20 

21def write_cube(file_obj, atoms, data=None, origin=None, comment=None): 

22 """Function to write a cube file. 

23 

24 file_obj: str or file object 

25 File to which output is written. 

26 atoms: Atoms 

27 The Atoms object specifying the atomic configuration. 

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

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

30 origin : 3-tuple 

31 Origin of the volumetric data (units: Angstrom) 

32 comment : str, optional (default = None) 

33 Comment for the first line of the cube file. 

34 """ 

35 

36 if data is None: 

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

38 data = np.asarray(data) 

39 

40 if data.dtype == complex: 

41 data = np.abs(data) 

42 

43 if comment is None: 

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

45 else: 

46 comment = comment.strip() 

47 file_obj.write(comment) 

48 

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

50 

51 if origin is None: 

52 origin = np.zeros(3) 

53 else: 

54 origin = np.asarray(origin) / Bohr 

55 

56 file_obj.write( 

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

58 len(atoms), *origin)) 

59 

60 for i in range(3): 

61 n = data.shape[i] 

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

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

64 

65 positions = atoms.positions / Bohr 

66 numbers = atoms.numbers 

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

68 file_obj.write( 

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

70 Z, 0.0, x, y, z) 

71 ) 

72 

73 data.tofile(file_obj, sep="\n", format="%e") 

74 

75 

76def read_cube(file_obj, read_data=True, program=None, verbose=False): 

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

78 

79 file_obj : str or file 

80 Location to the cube file. 

81 read_data : boolean 

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

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

84 and the dimensions of the corresponding voxels are read. 

85 program: str 

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

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

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

89 to catch castep files from the comment lines. 

90 verbose : bool 

91 Print some more information to stdout. 

92 

93 Returns a dict with the following keys: 

94 

95 * 'atoms': Atoms object 

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

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

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

99 """ 

100 

101 readline = file_obj.readline 

102 line = readline() # the first comment line 

103 line = readline() # the second comment line 

104 

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

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

107 axes = [] 

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

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

110 if not axes: 

111 axes = [0, 1, 2] 

112 

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

114 if "castep2cube" in line: 

115 program = CASTEP 

116 if verbose: 

117 print("read_cube identified program: castep") 

118 

119 # Third line contains actual system information: 

120 line = readline().split() 

121 num_atoms = int(line[0]) 

122 

123 # num_atoms can be negative. 

124 # Negative num_atoms indicates we have extra data to parse after 

125 # the coordinate information. 

126 has_labels = num_atoms < 0 

127 num_atoms = abs(num_atoms) 

128 

129 # There is an optional last field on this line which indicates 

130 # the number of values at each point. It is typically 1 (the default) 

131 # in which case it can be omitted, but it may also be > 1, 

132 # for example if there are multiple orbitals stored in the same cube. 

133 num_val = int(line[4]) if len(line) == 5 else 1 

134 

135 # Origin around which the volumetric data is centered 

136 # (at least in FHI aims): 

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

138 

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

140 shape = [] 

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

142 

143 # The upcoming three lines contain the cell information 

144 for i in range(3): 

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

146 shape.append(int(n)) 

147 

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

149 # identical to the first one 

150 if program == CASTEP: 

151 n -= 1 

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

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

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

155 

156 numbers = np.empty(num_atoms, int) 

157 positions = np.empty((num_atoms, 3)) 

158 for i in range(num_atoms): 

159 line = readline().split() 

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

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

162 

163 positions *= Bohr 

164 

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

166 

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

168 # contain this kind of information 

169 if program == CASTEP: 

170 atoms.pbc = True 

171 

172 dct = {ATOMS: atoms} 

173 labels = [] 

174 

175 # If we originally had a negative num_atoms, parse the extra fields now. 

176 # The first field of the first line tells us how many other fields there 

177 # are to parse, but we have to guess how many rows this information is 

178 # split over. 

179 if has_labels: 

180 # Can't think of a more elegant way of doing this... 

181 fields = readline().split() 

182 nfields = int(fields[0]) 

183 labels.extend(fields[1:]) 

184 

185 while len(labels) < nfields: 

186 fields = readline().split() 

187 labels.extend(fields) 

188 

189 labels = [int(x) for x in labels] 

190 

191 if read_data: 

192 # Cube files can contain more than one density, 

193 # so we need to be a little bit careful about where one ends 

194 # and the next begins. 

195 raw_volume = [float(s) for s in file_obj.read().split()] 

196 # Split each value at each point into a separate list. 

197 raw_volumes = [np.array(raw_volume[offset::num_val]) 

198 for offset in range(0, num_val)] 

199 

200 datas = [] 

201 

202 # Adjust each volume in turn. 

203 for data in raw_volumes: 

204 data = data.reshape(shape) 

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

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

207 

208 if program == CASTEP: 

209 # Due to the PBC applied in castep2cube, the last entry 

210 # along each dimension equals the very first one. 

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

212 

213 datas.append(data) 

214 

215 datas = np.array(datas) 

216 

217 dct[DATA] = datas[0] 

218 dct["origin"] = origin 

219 dct["spacing"] = spacing 

220 dct["labels"] = labels 

221 dct["datas"] = datas 

222 

223 return dct 

224 

225 

226def read_cube_data(filename): 

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

228 but also the contained volumetric data. 

229 """ 

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

231 return dct[DATA], dct[ATOMS]