Coverage for orcasong/from_toml.py: 62%

1import os 

2from pathlib import Path 

3import toml 

4import orcasong.core 

5import orcasong.extractors as extractors 

6 

7# built-in extractors. First argument has to be the input filename, 

8# other parameters can be set via '[extractor_config]' in the toml 

9EXTRACTORS = { 

10 "nu_chain_neutrino": extractors.get_neutrino_mc_info_extr, 

11 "nu_chain_muon": extractors.get_muon_mc_info_extr, 

12 "nu_chain_noise": extractors.get_random_noise_mc_info_extr, 

13 "nu_chain_data": extractors.get_real_data_info_extr, 

14 "bundle_mc": extractors.BundleMCExtractor, 

15 "bundle_data": extractors.BundleDataExtractor, 

16} 

17 

18MODES = { 

19 "graph": orcasong.core.FileGraph, 

20 "image": orcasong.core.FileBinner, 

21} 

22 

23 

24def add_parser_run(subparsers): 

25 parser = subparsers.add_parser( 

26 "run", description="Produce a dl file from an aanet file." 

27 ) 

28 parser.add_argument("infile", type=str, help="Aanet file in h5 format.") 

29 parser.add_argument( 

30 "config", 

31 type=str, 

32 help=( 

33 "Orcasong configuration in toml format. Use prefix 'orcasong:' to load " 

34 "a toml from OrcaSong/configs." 

35 ), 

36 ) 

37 parser.add_argument( 

38 "--detx", 

39 type=str, 

40 default=None, 

41 help=("Optional detx file to calibrate on the fly."), 

42 ) 

43 parser.add_argument( 

44 "--outfile", 

45 type=str, 

46 default=None, 

47 help=("Path to output file. Default: Save with auto generated name in cwd."), 

48 ) 

49 parser.set_defaults(func=run_orcasong) 

50 

51 

52def run_orcasong(infile, config, detx=None, outfile=None): 

53 setup_processor(infile, config, detx).run(infile=infile, outfile=outfile) 

54 

55 

56def setup_processor(infile, toml_file, detx_file=None): 

57 if toml_file.startswith("orcasong:"): 

58 toml_file = _get_config(toml_file[9:]) 

59 cfg = toml.load(toml_file) 

60 processor = _get_verbose(cfg.pop("mode"), MODES) 

61 

62 if "extractor" in cfg: 

63 extractor_cfg = cfg.pop("extractor_config", {}) 

64 extractor = _get_verbose(cfg.pop("extractor"), EXTRACTORS)( 

65 infile, **extractor_cfg 

66 ) 

67 else: 

68 extractor = None 

69 

70 return processor( 

71 det_file=detx_file, 

72 extractor=extractor, 

73 **cfg, 

74 ) 

75 

76 

77def _get_verbose(key, d): 

78 if key not in d: 

79 raise KeyError(f"Unknown key '{key}' (available: {list(d.keys())})") 

80 return d[key] 

81 

82 

83def _get_config(filename): 

84 direc = os.path.join(Path(orcasong.core.__file__).parents[1], "configs") 

85 files = {file: os.path.join(direc, file) for file in os.listdir(direc)} 

86 return _get_verbose(filename, files)