From 7452033ae3dd9321a47a3bdb2fb36fb844147c2d Mon Sep 17 00:00:00 2001
From: Stefan Reck <stefan.reck@fau.de>
Date: Fri, 16 Apr 2021 12:21:55 +0200
Subject: [PATCH] polish and add to main parser
---
orcasong/from_toml.py | 62 +++++++++++++++++++++++++++++++++++++++++++
orcasong/parser.py | 3 +++
orcasong/toml.py | 46 --------------------------------
requirements.txt | 1 +
4 files changed, 66 insertions(+), 46 deletions(-)
create mode 100644 orcasong/from_toml.py
delete mode 100644 orcasong/toml.py
diff --git a/orcasong/from_toml.py b/orcasong/from_toml.py
new file mode 100644
index 0000000..2e24ed7
--- /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 a95bc7b..6ff14d4 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 70112a2..0000000
--- 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 847f4fb..1491cac 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,6 +1,7 @@
numpy
h5py
matplotlib
+toml
km3pipe>=9
psutil
setuptools_scm
--
GitLab