Coverage for /Users/Newville/Codes/xraylarch/larch/io/save_restore.py: 12%

170 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2023-11-09 10:08 -0600

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, is_gzip, 

15 read_textfile, unique_name) 

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(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, _larch=None): 

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

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

72 

73 

74 Arguments: 

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

76 

77 See Also: 

78 read_session, load_session, clear_session 

79 

80 """ 

81 if fname is None: 

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

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

84 fname = fname + '.larix' 

85 

86 if _larch is None: 

87 raise ValueError('_larch not defined') 

88 symtab = _larch.symtable 

89 

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

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

92 "##<CONFIG>", 

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

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

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

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

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

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

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

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

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

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

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

104 ] 

105 

106 core_groups = symtab._sys.core_groups 

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

108 

109 config = symtab._sys.config 

110 for attr in dir(config): 

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

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

113 

114 try: 

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

116 except: 

117 histbuff = None 

118 

119 if histbuff is not None: 

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

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

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

123 

124 syms = [] 

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

126 if attr in core_groups: 

127 continue 

128 syms.append(attr) 

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

130 

131 for attr in dir(symtab): 

132 if attr in core_groups: 

133 continue 

134 buff.append('<:%s:>' % attr) 

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

136 

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

138 buff.append("") 

139 

140 fh = GzipFile(fname, "w") 

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

142 fh.close() 

143 

144def clear_session(_larch=None): 

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

146 

147 Example: 

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

149 >>> clear_session() 

150 

151 will effectively save and then reset the existing session. 

152 """ 

153 if _larch is None: 

154 raise ValueError('_larch not defined') 

155 

156 core_groups = _larch.symtable._sys.core_groups 

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

158 if attr not in core_groups: 

159 delattr(_larch.symtable, attr) 

160 

161 

162def read_session(fname): 

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

164 current session 

165 

166 Arguments: 

167 fname (str): name of save file 

168 

169 Returns: 

170 Tuple 

171 A tuple wih entries: 

172 

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

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

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

176 

177 See Also: 

178 load_session 

179 

180 

181 """ 

182 text = read_textfile(fname) 

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

184 line0 = lines.pop(0) 

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

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

187 

188 version = line0.split()[1] 

189 

190 symbols = {} 

191 config = {'Larix Version': version} 

192 cmd_history = [] 

193 nsyms = nsym_expected = 0 

194 section = symname = '_unknown_' 

195 

196 for line in lines: 

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

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

199 if ':' in section: 

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

201 if section.startswith('/'): 

202 section = '_unknown_' 

203 elif section == 'session commands': 

204 cmd_history.append(line) 

205 

206 elif section == 'symbols': 

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

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

209 else: 

210 try: 

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

212 except: 

213 print("decode failed:: ", symname, line[:150]) 

214 

215 else: 

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

217 line = line[2:] 

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

219 key = key.strip() 

220 val = val.strip() 

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

222 try: 

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

224 except: 

225 print("decode failed @## ", val[:150]) 

226 config[key] = val 

227 return SessionStore(config, cmd_history, symbols) 

228 

229 

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

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

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

233 

234 Arguments: 

235 fname (str): name of session file 

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

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

238 even if not expicitly set in `_xasgroups` 

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

240 Returns: 

241 None 

242 

243 Notes: 

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

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

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

247 `_xasgroups` : dict mapping "File Name" and "Group Name", used in `XAS Viewer` 

248 

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

250 may be modified on loading 

251 

252 """ 

253 if _larch is None: 

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

255 

256 session = read_session(fname) 

257 

258 if ignore_groups is None: 

259 ignore_groups = [] 

260 if include_xasgroups is None: 

261 include_xasgroups = [] 

262 

263 # special groups to merge into existing session: 

264 # _feffpaths, _feffcache, _xasgroups 

265 s_symbols = session.symbols 

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

267 

268 s_xasg_inv = invert_dict(s_xasgroups) 

269 

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

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

272 

273 symtab = _larch.symtable 

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

275 symtab._xasgroups = {} 

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

277 symtab._feffpaths = {} 

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

279 symtab._feffcache = EMPTY_FEFFCACHE 

280 

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

282 symtab._sys.restored_sessions = {} 

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

284 'config': session.config, 

285 'command_history': session.command_history} 

286 symtab._sys.restored_sessions[fname] = restore_data 

287 

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

289 

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

291 if sym in ignore_groups: 

292 if sym in s_xasgroups.values(): 

293 s_key = s_xasg_inv[sym] 

294 s_xasgroups.pop(s_key) 

295 s_xasg_inv = invert_dict(s_xasgroups) 

296 

297 continue 

298 if sym in c_xas_gnames or sym in include_xasgroups: 

299 newsym = unique_name(sym, c_xas_gnames) 

300 c_xas_gnames.append(newsym) 

301 if sym in s_xasgroups.values(): 

302 s_key = s_xasg_inv[sym] 

303 s_xasgroups[s_key] = newsym 

304 s_xasg_inv = invert_dict(s_xasgroups) 

305 sym = newsym 

306 

307 if verbose and hasattr(symtab, sym): 

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

309 setattr(symtab, sym, val) 

310 

311 symtab._xasgroups.update(s_xasgroups) 

312 symtab._feffpaths.update(s_feffpaths) 

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

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