Coverage for larch/io/save_restore.py: 10%

200 statements  

« prev     ^ index     » next       coverage.py v7.6.0, created at 2024-10-16 21:04 +0000

1import json 

2import time 

3import numpy as np 

4import uuid, socket, platform 

5from collections import namedtuple 

6 

7from gzip import GzipFile 

8 

9from lmfit import Parameter, Parameters 

10# from lmfit.model import Model, ModelResult 

11# from lmfit.minimizer import Minimizer, MinimizerResult 

12 

13from larch import Group, isgroup, __date__, __version__, __release_version__ 

14from ..utils import (isotime, bytes2str, str2bytes, fix_varname, 

15 read_textfile, unique_name, format_exception, unixpath) 

16from ..utils.jsonutils import encode4js, decode4js 

17 

18SessionStore = namedtuple('SessionStore', ('config', 'command_history', 'symbols')) 

19 

20EMPTY_FEFFCACHE = {'paths': {}, 'runs': {}} 

21 

22def invert_dict(d): 

23 "invert a dictionary {k: v} -> {v: k}" 

24 return {v: k for k, v in d.items()} 

25 

26def get_machineid(): 

27 "machine id / MAC address, independent of hostname" 

28 return hex(uuid.getnode())[2:] 

29 

30def is_larch_session_file(fname): 

31 return read_textfile(fname, size=64).startswith('##LARIX:') 

32 

33def save_groups(fname, grouplist): 

34 """save a list of groups (and other supported datatypes) to file 

35 

36 This is a simplified and minimal version of save_session() 

37 

38 Use 'read_groups()' to read data saved from this function 

39 """ 

40 buff = ["##LARCH GROUPLIST"] 

41 for dat in grouplist: 

42 buff.append(json.dumps(encode4js(dat))) 

43 

44 buff.append("") 

45 

46 fh = GzipFile(unixpath(fname), "w") 

47 fh.write(str2bytes("\n".join(buff))) 

48 fh.close() 

49 

50def read_groups(fname): 

51 """read a list of groups (and other supported datatypes) 

52 from a file saved with 'save_groups()' 

53 

54 Returns a list of objects 

55 """ 

56 text = read_textfile(fname) 

57 lines = text.split('\n') 

58 line0 = lines.pop(0) 

59 if not line0.startswith('##LARCH GROUPLIST'): 

60 raise ValueError(f"Invalid Larch group file: '{fname:s}'") 

61 

62 out = [] 

63 for line in lines: 

64 if len(line) > 1: 

65 out.append(decode4js(json.loads(line))) 

66 return out 

67 

68 

69def save_session(fname=None, symbols=None, histbuff=None, 

70 auto_xasgroups=False, _larch=None): 

71 """save all groups and data into a Larch Save File (.larix) 

72 A portable compressed json file, that can be loaded with `read_session()` 

73 

74 Arguments: 

75 fname (str): name of output save file. 

76 symbols [list of str or None]: names of symbols to save. [None] 

77 saving all non-core (user-generated) objects. 

78 histbuff [list of str or None]: command history, [None] 

79 saving the full history of the current session. 

80 auto_xasgroups [bool]: whether to automatically generate the 

81 `_xasgroups` dictionary for "XAS Groups" as used by 

82 Larix, which will include all symbols that are Groups 

83 and have both 'filename' and 'groupname' attributes. 

84 

85 Notes: 

86 1. if `symbols` is `None` (default), all variables outside of the 

87 core groups will be saved: this effectively saves "the whole session". 

88 A limited list of symbol names can also be provided, saving part of 

89 a project (say, one or two data sets). 

90 2. if `histbuff` is `None` (default), the full list of commands in 

91 the session will be saved. 

92 3. auto_xasgroups will generate an `_xasgroups` dictionary (used by 

93 Larix to decide what groups to present as "Data Groups") from the 

94 supplied or available symbols: every Group with a "groupname" and 

95 "filename" will be included. 

96 

97 See Also: 

98 read_session, load_session, clear_sessio 

99 

100 """ 

101 if fname is None: 

102 fname = time.strftime('%Y%b%d_%H%M') 

103 if not fname.endswith('.larix'): 

104 fname = fname + '.larix' 

105 

106 if _larch is None: 

107 raise ValueError('_larch not defined') 

108 symtab = _larch.symtable 

109 

110 buff = ["##LARIX: 1.0 Larch Session File", 

111 "##Date Saved: %s" % time.strftime('%Y-%m-%d %H:%M:%S'), 

112 "##<CONFIG>", 

113 "##Machine Platform: %s" % platform.system(), 

114 "##Machine Name: %s" % socket.gethostname(), 

115 "##Machine MACID: %s" % get_machineid(), 

116 "##Machine Version: %s" % platform.version(), 

117 "##Machine Processor: %s" % platform.machine(), 

118 "##Machine Architecture: %s" % ':'.join(platform.architecture()), 

119 "##Python Version: %s" % platform.python_version(), 

120 "##Python Compiler: %s" % platform.python_compiler(), 

121 "##Python Implementation: %s" % platform.python_implementation(), 

122 "##Larch Release Version: %s" % __release_version__, 

123 "##Larch Release Date: %s" % __date__, 

124 ] 

125 

126 core_groups = symtab._sys.core_groups 

127 buff.append('##Larch Core Groups: %s' % (json.dumps(core_groups))) 

128 

129 config = symtab._sys.config 

130 for attr in dir(config): 

131 buff.append('##Larch %s: %s' % (attr, json.dumps(getattr(config, attr, None)))) 

132 buff.append("##</CONFIG>") 

133 

134 

135 if histbuff is None: 

136 try: 

137 histbuff = _larch.input.history.get(session_only=True) 

138 except: 

139 histbuff = None 

140 

141 if histbuff is not None: 

142 buff.append("##<Session Commands>") 

143 buff.extend(["%s" % l for l in histbuff]) 

144 buff.append("##</Session Commands>") 

145 

146 if symbols is None: 

147 symbols = [] 

148 for attr in symtab.__dir__(): # insert order, not alphabetical order 

149 if attr not in core_groups: 

150 symbols.append(attr) 

151 nsyms = len(symbols) 

152 

153 _xasgroups = None 

154 if '_xasgroups' not in symbols and auto_xasgroups: 

155 nsyms +=1 

156 _xasgroups = {} 

157 for sname in symbols: 

158 obj = getattr(symtab, sname, None) 

159 if isgroup(obj): 

160 gname = getattr(obj, 'groupname', None) 

161 fname = getattr(obj, 'filename', None) 

162 if gname is not None and fname is not None: 

163 _xasgroups[fname] = gname 

164 

165 buff.append("##<Symbols: count=%d>" % len(symbols)) 

166 if _xasgroups is not None: 

167 buff.append('<:_xasgroups:>') 

168 buff.append(json.dumps(encode4js(_xasgroups))) 

169 

170 for attr in symbols: 

171 if attr not in core_groups: 

172 buff.append(f'<:{attr}:>') 

173 buff.append(json.dumps(encode4js(getattr(symtab, attr)))) 

174 

175 buff.append("##</Symbols>") 

176 buff.append("") 

177 

178 fh = GzipFile(unixpath(fname), "w") 

179 fh.write(str2bytes("\n".join(buff))) 

180 fh.close() 

181 

182def clear_session(_larch=None): 

183 """clear user-definded data in a session 

184 

185 Example: 

186 >>> save_session('foo.larix') 

187 >>> clear_session() 

188 

189 will effectively save and then reset the existing session. 

190 """ 

191 if _larch is None: 

192 raise ValueError('_larch not defined') 

193 

194 core_groups = _larch.symtable._sys.core_groups 

195 for attr in _larch.symtable.__dir__(): 

196 if attr not in core_groups: 

197 delattr(_larch.symtable, attr) 

198 

199 

200def read_session(fname, clean_xasgroups=True): 

201 """read Larch Session File, returning data into new data in the 

202 current session 

203 

204 Arguments: 

205 fname (str): name of save file 

206 

207 Returns: 

208 Tuple 

209 A tuple wih entries: 

210 

211 | configuration - a dict of configuration for the saved session. 

212 | command_history - a list of commands in the saved session. 

213 | symbols - a dict of Larch/Python symbols, groups, etc 

214 

215 See Also: 

216 load_session 

217 

218 

219 """ 

220 text = read_textfile(fname) 

221 lines = text.split('\n') 

222 line0 = lines.pop(0) 

223 if not line0.startswith('##LARIX:'): 

224 raise ValueError(f"Invalid Larch session file: '{fname:s}'") 

225 

226 version = line0.split()[1] 

227 

228 symbols = {} 

229 config = {'Larix Version': version} 

230 cmd_history = [] 

231 nsyms = nsym_expected = 0 

232 section = symname = '_unknown_' 

233 for line in lines: 

234 if line.startswith("##<"): 

235 section = line.replace('##<','').replace('>', '').strip().lower() 

236 if ':' in section: 

237 section, options = section.split(':', 1) 

238 if section.startswith('/'): 

239 section = '_unknown_' 

240 elif section == 'session commands': 

241 cmd_history.append(line) 

242 

243 elif section == 'symbols': 

244 if line.startswith('<:') and line.endswith(':>'): 

245 symname = line.replace('<:', '').replace(':>', '') 

246 else: 

247 try: 

248 symbols[symname] = decode4js(json.loads(line)) 

249 except: 

250 print(''.join(format_exception())) 

251 print("decode failed:: ", symname, repr(line)[:50]) 

252 else: 

253 if line.startswith('##') and ':' in line: 

254 line = line[2:] 

255 key, val = line.split(':', 1) 

256 key = key.strip() 

257 val = val.strip() 

258 if '[' in val or '{' in val: 

259 try: 

260 val = decode4js(json.loads(val)) 

261 except: 

262 print(''.join(format_exception())) 

263 print("decode failed @## ", repr(val)[:50]) 

264 config[key] = val 

265 if '_xasgroups' in symbols and clean_xasgroups: 

266 missing = [] 

267 for name, group in symbols['_xasgroups'].items(): 

268 if group not in symbols: 

269 missing.append(name) 

270 for name in missing: 

271 symbols['_xasgroups'].pop(name) 

272 

273 return SessionStore(config, cmd_history, symbols) 

274 

275 

276def load_session(fname, ignore_groups=None, include_xasgroups=None, _larch=None, verbose=False): 

277 """load all data from a Larch Session File into current larch session, 

278 merging into existing groups as appropriate (see Notes below) 

279 

280 Arguments: 

281 fname (str): name of session file 

282 ignore_groups (list of strings): list of symbols to not import 

283 include_xasgroups (list of strings): list of symbols to import as XAS spectra, 

284 even if not expicitly set in `_xasgroups` 

285 verbose (bool): whether to print warnings for overwrites [False] 

286 Returns: 

287 None 

288 

289 Notes: 

290 1. data in the following groups will be merged into existing session groups: 

291 `_feffpaths` : dict of "current feff paths" 

292 `_feffcache` : dict with cached feff paths and feff runs 

293 `_xasgroups` : dict mapping "File Name" and "Group Name", used in `Larix` 

294 

295 2. to avoid name clashes, group and file names in the `_xasgroups` dictionary 

296 may be modified on loading 

297 

298 """ 

299 if _larch is None: 

300 raise ValueError('load session needs a larch session') 

301 

302 session = read_session(fname) 

303 

304 if ignore_groups is None: 

305 ignore_groups = [] 

306 if include_xasgroups is None: 

307 include_xasgroups = [] 

308 

309 # special groups to merge into existing session: 

310 # _feffpaths, _feffcache, _xasgroups 

311 s_symbols = session.symbols 

312 s_xasgroups = s_symbols.pop('_xasgroups', {}) 

313 

314 s_xasg_inv = invert_dict(s_xasgroups) 

315 

316 s_feffpaths = s_symbols.pop('_feffpaths', {}) 

317 s_feffcache = s_symbols.pop('_feffcache', EMPTY_FEFFCACHE) 

318 

319 symtab = _larch.symtable 

320 if not hasattr(symtab, '_xasgroups'): 

321 symtab._xasgroups = {} 

322 if not hasattr(symtab, '_feffpaths'): 

323 symtab._feffpaths = {} 

324 if not hasattr(symtab, '_feffcache'): 

325 symtab._feffcache = EMPTY_FEFFCACHE 

326 

327 if not hasattr(symtab._sys, 'restored_sessions'): 

328 symtab._sys.restored_sessions = {} 

329 restore_data = {'date': isotime(), 

330 'config': session.config, 

331 'command_history': session.command_history} 

332 symtab._sys.restored_sessions[fname] = restore_data 

333 

334 c_xas_gnames = list(symtab._xasgroups.values()) 

335 

336 for sym, val in s_symbols.items(): 

337 if sym in ignore_groups: 

338 if sym in s_xasgroups.values(): 

339 s_key = s_xasg_inv[sym] 

340 s_xasgroups.pop(s_key) 

341 s_xasg_inv = invert_dict(s_xasgroups) 

342 

343 continue 

344 if sym in c_xas_gnames or sym in include_xasgroups: 

345 newsym = unique_name(sym, c_xas_gnames) 

346 c_xas_gnames.append(newsym) 

347 if sym in s_xasgroups.values(): 

348 s_key = s_xasg_inv[sym] 

349 s_xasgroups[s_key] = newsym 

350 s_xasg_inv = invert_dict(s_xasgroups) 

351 sym = newsym 

352 

353 if verbose and hasattr(symtab, sym): 

354 print(f"warning overwriting '{sym}'") 

355 setattr(symtab, sym, val) 

356 

357 symtab._feffpaths.update(s_feffpaths) 

358 

359 symtab._xasgroups.update(s_xasgroups) 

360 missing = [] 

361 for name, group in symtab._xasgroups.items(): 

362 if group not in symtab: 

363 missing.append(name) 

364 for name in missing: 

365 symtab._xasgroups.pop(name) 

366 

367 for name in ('paths', 'runs'): 

368 symtab._feffcache[name].update(s_feffcache[name])