Mercurial > repos > muon-spectroscopy-computational-project > larch_plot

comparison common.py @ 5:3584db5902b5 draft

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
planemo upload for repository https://github.com/MaterialsGalaxy/larch-tools/tree/main/larch_plot commit 4814f53888643f1d3667789050914675fffb7d59
author muon-spectroscopy-computational-project
date Fri, 23 Aug 2024 14:10:50 +0000
parents 35d24102cefd
children 0339eb694129
comparison
equal deleted inserted replaced
4:35d24102cefd 5:3584db5902b5
1 import re
1 from typing import Iterable 2 from typing import Iterable
2 3
3 from larch.io import extract_athenagroup, read_athena 4 from larch.io import extract_athenagroup, read_athena
4 from larch.io.athena_project import AthenaGroup 5 from larch.io.athena_project import AthenaGroup
5 from larch.symboltable import Group 6 from larch.symboltable import Group
8 9
9 def get_group(athena_group: AthenaGroup, key: str = None) -> Group: 10 def get_group(athena_group: AthenaGroup, key: str = None) -> Group:
10 group_keys = list(athena_group.keys()) 11 group_keys = list(athena_group.keys())
11 if key is None: 12 if key is None:
12 key = group_keys[0] 13 key = group_keys[0]
13 else:
14 key = key.replace("-", "_")
15 14
16 try: 15 try:
17 return extract_athenagroup(athena_group.groups[key]) 16 return extract_athenagroup(athena_group.groups[key])
18 except KeyError as e: 17 except KeyError as e:
19 raise KeyError(f"{key} not in {group_keys}") from e 18 raise KeyError(f"{key} not in {group_keys}") from e
20 19
21 20
22 def read_all_groups(dat_file: str, key: str = None) -> "dict[str, Group]": 21 def read_all_groups(dat_file: str) -> "dict[str, Group]":
23 # Cannot rely on do_ABC as _larch is None 22 # Cannot rely on do_ABC as _larch is None
24 athena_group = read_athena( 23 athena_group = read_athena(
25 dat_file, 24 dat_file,
26 do_preedge=False, 25 do_preedge=False,
27 do_bkg=False, 26 do_bkg=False,
37 36
38 return all_groups 37 return all_groups
39 38
40 39
41 def read_group(dat_file: str, key: str = None): 40 def read_group(dat_file: str, key: str = None):
41 if key:
42 match_ = key.replace(" ", "_").replace("-", "_").replace(".", "_")
43 else:
44 match_ = None
45
42 # Cannot rely on do_ABC as _larch is None 46 # Cannot rely on do_ABC as _larch is None
43 athena_group = read_athena( 47 athena_group = read_athena(
44 dat_file, 48 dat_file,
49 match=match_,
45 do_preedge=False, 50 do_preedge=False,
46 do_bkg=False, 51 do_bkg=False,
47 do_fft=False, 52 do_fft=False,
48 ) 53 )
49 group = get_group(athena_group, key) 54 group = get_group(athena_group, match_)
50 pre_edge_with_defaults(group=group) 55 pre_edge_with_defaults(group=group)
51 xftf_with_defaults(group=group) 56 xftf_with_defaults(group=group)
52 return group 57 return group
53 58
54 59
166 171
167 172
168 def read_groups(dat_files: "list[str]", key: str = None) -> Iterable[Group]: 173 def read_groups(dat_files: "list[str]", key: str = None) -> Iterable[Group]:
169 for dat_file in dat_files: 174 for dat_file in dat_files:
170 yield read_group(dat_file=dat_file, key=key) 175 yield read_group(dat_file=dat_file, key=key)
176
177
178 def sorting_key(filename: str) -> str:
179 return re.findall(r"\d+", filename)[-1]