diff --git a/orcasong/toml.py b/orcasong/toml.py index 94e4374c2fcb9f05cb8e3bc562bb3a4f4125ab6b..70112a2c1739e724f3e8c0984c30a42d062ad63e 100644 --- a/orcasong/toml.py +++ b/orcasong/toml.py @@ -1,11 +1,12 @@ import os import toml from orcasong.core import FileGraph -from orcasong.extractors import get_neutrino_mc_info_extr +import orcasong.extractors as extractors EXTRACTORS = { - "neutrino": get_neutrino_mc_info_extr, + "neutrino_mc": extractors.get_neutrino_mc_info_extr, + "neutrino_data": extractors.get_real_data_info_extr, } @@ -14,21 +15,29 @@ def add_parser(subparsers): "graph", description='Produce graph dl file from aanet file.') parser.add_argument('infile', type=str) - parser.add_argument('config', type=str) - parser.add_argument('detx_file', type=str) + parser.add_argument('toml_file', type=str) + parser.add_argument('--detx_file', type=str, default=None) parser.add_argument('--outfile', type=str, default=None) return parser.parse_args() -def make_graph(infile, config, detx_file, outfile=None): +def make_graph(infile, toml_file, detx_file=None, outfile=None): if outfile is None: outfile = f"{os.path.splitext(os.path.basename(infile))[0]}_dl.h5" - cfg = toml.load(config) - extractor_cfg = cfg.pop("extractor") - extractor_name = extractor_cfg.pop("name") - inps = {"infile": infile, "config": config, "detx_file": detx_file, "outfile": outfile} - extractor = EXTRACTORS[extractor_name](inps, **extractor_cfg) + cfg = toml.load(toml_file) + if "detx_file" in cfg: + if detx_file is not None: + raise ValueError + detx_file = cfg.pop("detx_file") + + extractor_name = cfg.pop("extractor") + if "extractor_config" in cfg: + extractor_cfg = cfg.pop("extractor_config") + else: + extractor_cfg = {} + + extractor = EXTRACTORS[extractor_name](infile, **extractor_cfg) FileGraph( det_file=detx_file,