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
« 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
7from gzip import GzipFile
9from lmfit import Parameter, Parameters
10# from lmfit.model import Model, ModelResult
11# from lmfit.minimizer import Minimizer, MinimizerResult
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
18SessionStore = namedtuple('SessionStore', ('config', 'command_history', 'symbols'))
20EMPTY_FEFFCACHE = {'paths': {}, 'runs': {}}
22def invert_dict(d):
23 "invert a dictionary {k: v} -> {v: k}"
24 return {v: k for k, v in d.items()}
26def get_machineid():
27 "machine id / MAC address, independent of hostname"
28 return hex(uuid.getnode())[2:]
30def is_larch_session_file(fname):
31 return read_textfile(fname, size=64).startswith('##LARIX:')
33def save_groups(fname, grouplist):
34 """save a list of groups (and other supported datatypes) to file
36 This is a simplified and minimal version of save_session()
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)))
44 buff.append("")
46 fh = GzipFile(fname, "w")
47 fh.write(str2bytes("\n".join(buff)))
48 fh.close()
50def read_groups(fname):
51 """read a list of groups (and other supported datatypes)
52 from a file saved with 'save_groups()'
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}'")
62 out = []
63 for line in lines:
64 if len(line) > 1:
65 out.append(decode4js(json.loads(line)))
66 return out
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()`
74 Arguments:
75 fname (str): name of output save file.
77 See Also:
78 read_session, load_session, clear_session
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'
86 if _larch is None:
87 raise ValueError('_larch not defined')
88 symtab = _larch.symtable
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 ]
106 core_groups = symtab._sys.core_groups
107 buff.append('##Larch Core Groups: %s' % (json.dumps(core_groups)))
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>")
114 try:
115 histbuff = _larch.input.history.get(session_only=True)
116 except:
117 histbuff = None
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>")
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))
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))))
137 buff.append("##</Symbols>")
138 buff.append("")
140 fh = GzipFile(fname, "w")
141 fh.write(str2bytes("\n".join(buff)))
142 fh.close()
144def clear_session(_larch=None):
145 """clear user-definded data in a session
147 Example:
148 >>> save_session('foo.larix')
149 >>> clear_session()
151 will effectively save and then reset the existing session.
152 """
153 if _larch is None:
154 raise ValueError('_larch not defined')
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)
162def read_session(fname):
163 """read Larch Session File, returning data into new data in the
164 current session
166 Arguments:
167 fname (str): name of save file
169 Returns:
170 Tuple
171 A tuple wih entries:
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
177 See Also:
178 load_session
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}'")
188 version = line0.split()[1]
190 symbols = {}
191 config = {'Larix Version': version}
192 cmd_history = []
193 nsyms = nsym_expected = 0
194 section = symname = '_unknown_'
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)
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])
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)
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)
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
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`
249 2. to avoid name clashes, group and file names in the `_xasgroups` dictionary
250 may be modified on loading
252 """
253 if _larch is None:
254 raise ValueError('load session needs a larch session')
256 session = read_session(fname)
258 if ignore_groups is None:
259 ignore_groups = []
260 if include_xasgroups is None:
261 include_xasgroups = []
263 # special groups to merge into existing session:
264 # _feffpaths, _feffcache, _xasgroups
265 s_symbols = session.symbols
266 s_xasgroups = s_symbols.pop('_xasgroups', {})
268 s_xasg_inv = invert_dict(s_xasgroups)
270 s_feffpaths = s_symbols.pop('_feffpaths', {})
271 s_feffcache = s_symbols.pop('_feffcache', EMPTY_FEFFCACHE)
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
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
288 c_xas_gnames = list(symtab._xasgroups.values())
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)
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
307 if verbose and hasattr(symtab, sym):
308 print(f"warning overwriting '{sym}'")
309 setattr(symtab, sym, val)
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])