Coverage for /builds/kinetik161/ase/ase/io/fortranfile.py: 35.37%

1# Copyright 2008-2010 Neil Martinsen-Burrell 

2# 

3# Permission is hereby granted, free of charge, to any person obtaining a copy 

4# of this software and associated documentation files (the "Software"), to deal 

5# in the Software without restriction, including without limitation the rights 

6# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 

7# copies of the Software, and to permit persons to whom the Software is 

8# furnished to do so, subject to the following conditions: 

9 

10# The above copyright notice and this permission notice shall be included in 

11# all copies or substantial portions of the Software. 

12 

13# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 

14# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 

15# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 

16# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 

17# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 

18# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 

19# THE SOFTWARE. 

20 

21"""Defines a file-derived class to read/write Fortran unformatted files. 

22 

23The assumption is that a Fortran unformatted file is being written by 

24the Fortran runtime as a sequence of records. Each record consists of 

25an integer (of the default size [usually 32 or 64 bits]) giving the 

26length of the following data in bytes, then the data itself, then the 

27same integer as before. 

28 

29Examples 

30-------- 

31 

32To use the default endian and precision settings, one can just do: 

33 

34>>> f = FortranFile('filename') 

35>>> x = f.readReals() 

36 

37One can read arrays with varying precisions: 

38 

39>>> f = FortranFile('filename') 

40>>> x = f.readInts('h') 

41>>> y = f.readInts('q') 

42>>> z = f.readReals('f') 

43 

44Where the format codes are those used by Python's struct module. 

45 

46One can change the default endian-ness and header precision: 

47 

48>>> f = FortranFile('filename', endian='>', header_prec='l') 

49 

50for a file with little-endian data whose record headers are long 

51integers. 

52""" 

53 

54import numpy 

55 

56try: 

57 file 

58except NameError: 

59 # For python3 compatibility 

60 from io import FileIO as file 

61 

62 

63class FortranFile(file): 

64 

65 """File with methods for dealing with fortran unformatted data files""" 

66 

67 def _get_header_length(self): 

68 return numpy.dtype(self._header_prec).itemsize 

69 _header_length = property(fget=_get_header_length) 

70 

71 def _set_endian(self, c): 

72 """Set endian to big (c='>') or little (c='<') or native (c='=') 

73 

74 :Parameters: 

75 `c` : string 

76 The endian-ness to use when reading from this file. 

77 """ 

78 if c in '<>@=': 

79 if c == '@': 

80 c = '=' 

81 self._endian = c 

82 else: 

83 raise ValueError('Cannot set endian-ness') 

84 

85 def _get_endian(self): 

86 return self._endian 

87 ENDIAN = property( 

88 fset=_set_endian, 

89 fget=_get_endian, 

90 doc="Possible endian values are '<', '>', '@', '='", 

91 ) 

92 

93 def _set_header_prec(self, prec): 

94 if prec in 'hilq': 

95 self._header_prec = prec 

96 else: 

97 raise ValueError('Cannot set header precision') 

98 

99 def _get_header_prec(self): 

100 return self._header_prec 

101 HEADER_PREC = property( 

102 fset=_set_header_prec, 

103 fget=_get_header_prec, 

104 doc="Possible header precisions are 'h', 'i', 'l', 'q'", 

105 ) 

106 

107 def __init__(self, fname, endian='@', header_prec='i', *args, **kwargs): 

108 """Open a Fortran unformatted file for writing. 

109 

110 Parameters 

111 ---------- 

112 endian : character, optional 

113 Specify the endian-ness of the file. Possible values are 

114 '>', '<', '@' and '='. See the documentation of Python's 

115 struct module for their meanings. The default is '>' (native 

116 byte order) 

117 header_prec : character, optional 

118 Specify the precision used for the record headers. Possible 

119 values are 'h', 'i', 'l' and 'q' with their meanings from 

120 Python's struct module. The default is 'i' (the system's 

121 default integer). 

122 

123 """ 

124 file.__init__(self, fname, *args, **kwargs) 

125 self.ENDIAN = endian 

126 self.HEADER_PREC = header_prec 

127 

128 def _read_exactly(self, num_bytes): 

129 """Read in exactly num_bytes, raising an error if it can't be done.""" 

130 data = b'' 

131 while True: 

132 L = len(data) 

133 if L == num_bytes: 

134 return data 

135 else: 

136 read_data = self.read(num_bytes - L) 

137 if read_data == b'': 

138 raise OSError('Could not read enough data.' 

139 ' Wanted %d bytes, got %d.' % (num_bytes, L)) 

140 data += read_data 

141 

142 def _read_check(self): 

143 return numpy.frombuffer( 

144 self._read_exactly(self._header_length), 

145 dtype=self.ENDIAN + self.HEADER_PREC, 

146 )[0] 

147 

148 def _write_check(self, number_of_bytes): 

149 """Write the header for the given number of bytes""" 

150 self.write(numpy.array( 

151 number_of_bytes, dtype=self.ENDIAN + self.HEADER_PREC, 

152 ).tostring()) 

153 

154 def readRecord(self): 

155 """Read a single fortran record""" 

156 L = self._read_check() 

157 data_str = self._read_exactly(L) 

158 check_size = self._read_check() 

159 if check_size != L: 

160 raise OSError('Error reading record from data file') 

161 return data_str 

162 

163 def writeRecord(self, s): 

164 """Write a record with the given bytes. 

165 

166 Parameters 

167 ---------- 

168 s : the string to write 

169 

170 """ 

171 length_bytes = len(s) 

172 self._write_check(length_bytes) 

173 self.write(s) 

174 self._write_check(length_bytes) 

175 

176 def readString(self): 

177 """Read a string.""" 

178 return self.readRecord() 

179 

180 def writeString(self, s): 

181 """Write a string 

182 

183 Parameters 

184 ---------- 

185 s : the string to write 

186 

187 """ 

188 self.writeRecord(s) 

189 

190 _real_precisions = 'df' 

191 

192 def readReals(self, prec='f'): 

193 """Read in an array of real numbers. 

194 

195 Parameters 

196 ---------- 

197 prec : character, optional 

198 Specify the precision of the array using character codes from 

199 Python's struct module. Possible values are 'd' and 'f'. 

200 

201 """ 

202 

203 if prec not in self._real_precisions: 

204 raise ValueError('Not an appropriate precision') 

205 

206 data_str = self.readRecord() 

207 return numpy.frombuffer(data_str, dtype=self.ENDIAN + prec) 

208 

209 def writeReals(self, reals, prec='f'): 

210 """Write an array of floats in given precision 

211 

212 Parameters 

213 ---------- 

214 reals : array 

215 Data to write 

216 prec` : string 

217 Character code for the precision to use in writing 

218 """ 

219 if prec not in self._real_precisions: 

220 raise ValueError('Not an appropriate precision') 

221 

222 nums = numpy.array(reals, dtype=self.ENDIAN + prec) 

223 self.writeRecord(nums.tostring()) 

224 

225 _int_precisions = 'hilq' 

226 

227 def readInts(self, prec='i'): 

228 """Read an array of integers. 

229 

230 Parameters 

231 ---------- 

232 prec : character, optional 

233 Specify the precision of the data to be read using 

234 character codes from Python's struct module. Possible 

235 values are 'h', 'i', 'l' and 'q' 

236 

237 """ 

238 if prec not in self._int_precisions: 

239 raise ValueError('Not an appropriate precision') 

240 

241 data_str = self.readRecord() 

242 return numpy.frombuffer(data_str, dtype=self.ENDIAN + prec) 

243 

244 def writeInts(self, ints, prec='i'): 

245 """Write an array of integers in given precision 

246 

247 Parameters 

248 ---------- 

249 reals : array 

250 Data to write 

251 prec : string 

252 Character code for the precision to use in writing 

253 """ 

254 if prec not in self._int_precisions: 

255 raise ValueError('Not an appropriate precision') 

256 

257 nums = numpy.array(ints, dtype=self.ENDIAN + prec) 

258 self.writeRecord(nums.tostring())