diff --git a/examples/orcasong_neturino_mc.toml b/examples/orcasong_example.toml similarity index 62% rename from examples/orcasong_neturino_mc.toml rename to examples/orcasong_example.toml index 4f54ca0b2ef5e10e1890efb2cf70232c23beb005..5ab2719ebb2f4939acc1231858363b836bc6fb8c 100644 --- a/examples/orcasong_neturino_mc.toml +++ b/examples/orcasong_example.toml @@ -1,8 +1,11 @@ # the mode to run orcasong in; either 'graph' or 'image' mode="graph" -# arguments for running orcasong (see FileGraph or FileBinner in orcasong.core). +# arguments for FileGraph or FileBinner (see orcasong.core) max_n_hits = 2000 time_window = [-100, 5000] +# can also give the arguments of orcasong.core.BaseProcessor, +# which are shared between modes +chunksize=16 # built-in extractor function to use extractor = "neutrino_mc" diff --git a/orcasong/from_toml.py b/orcasong/from_toml.py index 9958fd1d13343f5e8f5e336ef04499f037ff9e11..09e0a8e8345f2fdd229fec6b1f4026b8cc4b48ff 100644 --- a/orcasong/from_toml.py +++ b/orcasong/from_toml.py @@ -1,4 +1,3 @@ -import os import toml import orcasong.core import orcasong.extractors as extractors @@ -31,15 +30,17 @@ def add_parser_run(subparsers): 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" + setup_processor(infile, toml_file, detx_file).run( + infile=infile, outfile=outfile) + +def setup_processor(infile, toml_file, detx_file=None): cfg = toml.load(toml_file) - processor = _get_verbose(cfg["mode"], MODES) + processor = _get_verbose(cfg.pop("mode"), MODES) if "detx_file" in cfg: if detx_file is not None: - raise ValueError("detx_file passed to function AND defined in toml") + raise ValueError("detx_file passed to run AND defined in toml") detx_file = cfg.pop("detx_file") if "extractor" in cfg: @@ -48,11 +49,11 @@ def run_orcasong(infile, toml_file, detx_file=None, outfile=None): else: extractor = None - processor( + return processor( det_file=detx_file, extractor=extractor, **cfg, - ).run(infile=infile, outfile=outfile) + ) def _get_verbose(key, d): diff --git a/tests/test_from_toml.py b/tests/test_from_toml.py new file mode 100644 index 0000000000000000000000000000000000000000..5ee6fb619b988518461edda23407a4963f6dda65 --- /dev/null +++ b/tests/test_from_toml.py @@ -0,0 +1,36 @@ +from unittest import TestCase +import os +import orcasong +import orcasong.from_toml as from_toml + +EXAMPLES = os.path.join( + os.path.dirname(os.path.dirname(orcasong.__file__)), "examples" +) + + +def test_extr(infile): + return infile + "_extr" + + +orcasong.from_toml.EXTRACTORS["neutrino_mc"] = test_extr + + +class TestSetupProcessorExampleConfig(TestCase): + def setUp(self): + self.processor = from_toml.setup_processor( + infile="test_in", + toml_file=os.path.join(EXAMPLES, "orcasong_example.toml"), + detx_file="test_det", + ) + + def test_time_window(self): + self.assertEqual(self.processor.time_window, [-100, 5000]) + + def test_max_n_hits(self): + self.assertEqual(self.processor.max_n_hits, 2000) + + def test_chunksize(self): + self.assertEqual(self.processor.chunksize, 16) + + def test_extractor_is_dummy_extractor(self): + self.assertEqual(self.processor.extractor, "test_in_extr")