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('{0:5}{1:12.6f}{2:12.6f}{3:12.6f}\n' 

53 .format(len(atoms), *origin)) 

54 

55 for i in range(3): 

56 n = data.shape[i] 

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

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

59 

60 positions = atoms.positions / Bohr 

61 numbers = atoms.numbers 

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

63 fileobj.write('{0:5}{1:12.6f}{2:12.6f}{3:12.6f}{4:12.6f}\n' 

64 .format(Z, 0.0, x, y, z)) 

65 

66 data.tofile(fileobj, sep='\n', format='%e') 

67 

68 

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

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

71 

72 fileobj : str or file 

73 Location to the cubefile. 

74 read_data : boolean 

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

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

77 and the dimensions of the corresponding voxels are read. 

78 program: str 

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

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

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

82 to catch castep files from the comment lines. 

83 verbose : bool 

84 Print some more information to stdout. 

85 

86 Returns a dict with the following keys: 

87  

88 * 'atoms': Atoms object 

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

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

91 """ 

92 

93 readline = fileobj.readline 

94 line = readline() # the first comment line 

95 line = readline() # the second comment line 

96 

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

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

99 axes = [] 

100 if 'OUTER LOOP' in line.upper(): 

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

102 if not axes: 

103 axes = [0, 1, 2] 

104 

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

106 if 'castep2cube' in line: 

107 program = 'castep' 

108 if verbose: 

109 print('read_cube identified program: castep') 

110 

111 # Third line contains actual system information: 

112 line = readline().split() 

113 natoms = int(line[0]) 

114 

115 # Origin around which the volumetric data is centered 

116 # (at least in FHI aims): 

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

118 

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

120 shape = [] 

121 

122 # the upcoming three lines contain the cell information 

123 for i in range(3): 

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

125 shape.append(int(n)) 

126 

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

128 # identical to the first one 

129 if program == 'castep': 

130 n -= 1 

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

132 

133 numbers = np.empty(natoms, int) 

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

135 for i in range(natoms): 

136 line = readline().split() 

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

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

139 

140 positions *= Bohr 

141 

142 atoms = Atoms(numbers=numbers, positions=positions, cell=cell) 

143 

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

145 # contain this kind of information 

146 if program == 'castep': 

147 atoms.pbc = True 

148 

149 dct = {'atoms': atoms} 

150 

151 if read_data: 

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

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

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

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

156 

157 if program == 'castep': 

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

159 # dimension equals the very first one. 

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

161 

162 dct['data'] = data 

163 dct['origin'] = origin 

164 

165 return dct 

166 

167 

168def read_cube_data(filename): 

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

170 but also the contained volumetric data. 

171 """ 

172 dct = read(filename, format='cube', read_data=True, full_output=True) 

173 return dct['data'], dct['atoms']