diff --git a/orcasong/from_toml.py b/orcasong/from_toml.py new file mode 100644 index 0000000000000000000000000000000000000000..2e24ed79fb3ee95fcdf38c073751a0100d76a6dc --- /dev/null +++ b/orcasong/from_toml.py @@ -0,0 +1,62 @@ +import os +import toml +import orcasong.core +import orcasong.extractors as extractors + +# available extractors. First argument has to be the input filename +EXTRACTORS = { + "neutrino_mc": extractors.get_neutrino_mc_info_extr, + "neutrino_data": extractors.get_real_data_info_extr, +} + + +def _add_args(parser): + parser.add_argument('infile', type=str, help="Aanet file in h5 format.") + parser.add_argument('toml_file', type=str, help="Orcasong configuration in toml format.") + parser.add_argument('--detx_file', type=str, default=None, help=( + "Optional detx file to calibrate on the fly. Can not be used if a " + "detx_file has also been given in the toml file.")) + parser.add_argument('--outfile', type=str, default=None, help=( + "Path to output file. Default: Save with auto ogenerated name in cwd.")) + + +def add_parser_filegraph(subparsers): + parser = subparsers.add_parser( + "graph", + description='Produce a graph dl file from an aanet file.') + _add_args(parser) + parser.set_defaults(func=get_run_orcasong(orcasong.core.FileGraph)) + + +def add_parser_filebinner(subparsers): + parser = subparsers.add_parser( + "image", + description='Produce an image dl file from an aanet file.') + _add_args(parser) + parser.set_defaults(func=get_run_orcasong(orcasong.core.FileBinner)) + + +def get_run_orcasong(processor): + def run_orcasong(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(toml_file) + if "detx_file" in cfg: + if detx_file is not None: + raise ValueError("detx_file passed to function AND defined in toml") + detx_file = cfg.pop("detx_file") + + if "extractor" in cfg: + extractor_name = cfg.pop("extractor") + extractor_cfg = cfg.pop("extractor_config", {}) + extractor = EXTRACTORS[extractor_name](infile, **extractor_cfg) + else: + extractor = None + + processor( + det_file=detx_file, + extractor=extractor, + **cfg, + ).run(infile=infile, outfile=outfile) + return run_orcasong diff --git a/orcasong/parser.py b/orcasong/parser.py index a95bc7b53566887dceaf59eee469f29baa145d8d..6ff14d4aa7a95df60b793cc34657563eaefd3b7a 100644 --- a/orcasong/parser.py +++ b/orcasong/parser.py @@ -6,6 +6,7 @@ import argparse from orcasong.tools.concatenate import concatenate from orcasong.tools.postproc import postproc_file from orcasong.tools.shuffle2 import h5shuffle2 +import orcasong.from_toml as from_toml def _add_parser_concatenate(subparsers): @@ -119,6 +120,8 @@ def main(): ) subparsers = parser.add_subparsers() + from_toml.add_parser_filegraph(subparsers) + from_toml.add_parser_filebinner(subparsers) _add_parser_concatenate(subparsers) _add_parser_h5shuffle(subparsers) _add_parser_h5shuffle2(subparsers) diff --git a/orcasong/toml.py b/orcasong/toml.py deleted file mode 100644 index 70112a2c1739e724f3e8c0984c30a42d062ad63e..0000000000000000000000000000000000000000 --- a/orcasong/toml.py +++ /dev/null @@ -1,46 +0,0 @@ -import os -import toml -from orcasong.core import FileGraph -import orcasong.extractors as extractors - - -EXTRACTORS = { - "neutrino_mc": extractors.get_neutrino_mc_info_extr, - "neutrino_data": extractors.get_real_data_info_extr, -} - - -def add_parser(subparsers): - parser = subparsers.add_parser( - "graph", - description='Produce graph dl file from aanet file.') - parser.add_argument('infile', 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, 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(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, - extractor=extractor, - **cfg, - ).run(infile=infile, outfile=outfile) diff --git a/requirements.txt b/requirements.txt index 847f4fb56e2fab2cf9d6c0b49eea2886c20c85b3..1491cac90564e68347e26be90dba93bc1dd58708 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,7 @@ numpy h5py matplotlib +toml km3pipe>=9 psutil setuptools_scm