Optmap rebin

docs/source/conf.py

@@ -61,6 +61,7 @@
     "pyg4ometry": ("https://pyg4ometry.readthedocs.io/en/stable", None),
     "legendmeta": ("https://pylegendmeta.readthedocs.io/en/stable/", None),
     "lgdo": ("https://legend-pydataobj.readthedocs.io/en/stable/", None),
+    "dbetto": ("https://dbetto.readthedocs.io/en/stable/", None),
 }  # add new intersphinx mappings here
 
 # sphinx-autodoc

src/reboost/__init__.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from reboost import build_hit, core, iterator, math, optmap, shape
+from reboost import build_hit, core, iterator, math, shape
 from reboost._version import version as __version__
 
-__all__ = ["__version__", "build_hit", "core", "iterator", "math", "optmap", "shape"]
+__all__ = ["__version__", "build_hit", "core", "iterator", "math", "shape"]

src/reboost/build_hit.py

@@ -195,7 +195,7 @@ def build_hit(
     config
         dictionary or path to YAML file containing the processing chain.
     args
-        dictionary or :class:`legendmeta.AttrsDict` of the global arguments.
+        dictionary or :class:`dbetto.AttrsDict` of the global arguments.
     stp_files
         list of strings or string of the stp file path.
     glm_files

src/reboost/optmap/__init__.py

@@ -0,0 +1,5 @@
+from __future__ import annotations
+
+from .optmap import OpticalMap
+
+__all__ = ["OpticalMap"]

src/reboost/optmap/cli.py

@@ -168,6 +168,12 @@ def optical_cli() -> None:
     )
     convolve_parser.add_argument("--output", help="output hit LH5 file", metavar="OUTPUT_HIT")
 
+    # STEP X: rebin maps
+    rebin_parser = subparsers.add_parser("rebin", help="rebin optical maps")
+    rebin_parser.add_argument("input", help="input map LH5 files", metavar="INPUT_MAP")
+    rebin_parser.add_argument("output", help="output map LH5 file", metavar="OUTPUT_MAP")
+    rebin_parser.add_argument("--factor", type=int, help="integer scale-down factor")
+
     args = parser.parse_args()
 
     log_level = (None, logging.INFO, logging.DEBUG)[min(args.verbose, 2)]
@@ -261,3 +267,11 @@ def optical_cli() -> None:
         _check_input_file(parser, [args.map, args.edep])
         _check_output_file(parser, args.output)
         convolve(args.map, args.edep, args.edep_lgdo, args.material, args.output, args.bufsize)
+
+    # STEP X: rebin maps
+    if args.command == "rebin":
+        from reboost.optmap.create import rebin_optical_maps
+
+        _check_input_file(parser, args.input)
+        _check_output_file(parser, args.output)
+        rebin_optical_maps(args.input, args.output, args.factor)

src/reboost/optmap/create.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import copy
 import gc
 import logging
 import multiprocessing as mp
@@ -430,7 +431,7 @@ def check_optical_map(map_l5_file: str):
     for submap in list_optical_maps(map_l5_file):
         try:
             om = OpticalMap.load_from_file(map_l5_file, submap)
-        except BaseException:
+        except Exception:
             log.exception("error while loading optical map %s", submap)
             continue
         om.check_histograms(include_prefix=True)
@@ -439,3 +440,42 @@ def check_optical_map(map_l5_file: str):
             log.error("edges of optical map %s differ", submap)
         else:
             all_binning = om.binning
+
+
+def rebin_optical_maps(map_l5_file: str, output_lh5_file: str, factor: int):
+    """Rebin the optical map by an integral factor.
+
+    .. note ::
+
+        the factor has to divide the bincounts on all axes.
+    """
+    if not isinstance(factor, int) or factor <= 1:
+        msg = f"invalid rebin factor {factor}"
+        raise ValueError(msg)
+
+    def _rebin_map(large: NDArray, factor: int) -> NDArray:
+        factor = np.full(3, factor, dtype=int)
+        sh = np.column_stack([np.array(large.shape) // factor, factor]).ravel()
+        return large.reshape(sh).sum(axis=(1, 3, 5))
+
+    for submap in list_optical_maps(map_l5_file):
+        log.info("rebinning optical map group: %s", submap)
+
+        om = OpticalMap.load_from_file(map_l5_file, submap)
+
+        settings = om.get_settings()
+        if not all(b % factor == 0 for b in settings["bins"]):
+            msg = f"invalid factor {factor}, not a divisor"
+            raise ValueError(msg)
+        settings = copy.copy(settings)
+        settings["bins"] = [b // factor for b in settings["bins"]]
+
+        om_new = OpticalMap.create_empty(om.name, settings)
+        om_new.h_vertex = _rebin_map(om.h_vertex, factor)
+        om_new.h_hits = _rebin_map(om.h_hits, factor)
+        om_new.create_probability()
+        om_new.write_lh5(lh5_file=output_lh5_file, group=submap, wo_mode="write_safe")
+
+    # just copy hitcounts exponent.
+    for dset in ("_hitcounts_exp", "_hitcounts"):
+        lh5.write(lh5.read(dset, lh5_file=map_l5_file), dset, lh5_file=output_lh5_file)

src/reboost/optmap/optmap.py

@@ -48,6 +48,7 @@ def __init__(self, name: str, settings: Mapping[str, str], use_shmem: bool = Fal
                 for i in range(3)
             ]
 
+    @staticmethod
     def create_empty(name: str, settings: Mapping[str, str]) -> OpticalMap:
         om = OpticalMap(name, settings)
         om.h_vertex = om._prepare_hist()
@@ -215,9 +216,11 @@ def _divide_hist(self, h1: NDArray, h2: NDArray) -> tuple[NDArray, NDArray]:
         return ratio_0, ratio_err_0
 
     def create_probability(self) -> None:
+        """Compute probability map (and map uncertainty) from vertex and hit map."""
         self.h_prob, self.h_prob_uncert = self._divide_hist(self.h_hits, self.h_vertex)
 
     def write_lh5(self, lh5_file: str, group: str = "all", wo_mode: str = "write_safe") -> None:
+        """Write this map to a LH5 file."""
         if wo_mode not in ("write_safe", "overwrite_file"):
             msg = f"invalid wo_mode {wo_mode} for optical map"
             raise ValueError(msg)
@@ -237,6 +240,25 @@ def write_hist(h: NDArray, name: str, fn: str, group: str, wo_mode: str):
         write_hist(self.h_prob, "p_det", lh5_file, group, "write_safe")
         write_hist(self.h_prob_uncert, "p_det_err", lh5_file, group, "write_safe")
 
+    def get_settings(self) -> dict:
+        """Get the binning settings that were used to create this optical map instance."""
+        if self.settings is not None:
+            return self.settings
+
+        range_in_m = []
+        bins = []
+        for b in self.binning:
+            if not b.is_range:
+                msg = "cannot get binning settings for variable binning map"
+                raise RuntimeError(msg)
+            if b.get_binedgeattrs().get("units") != "m":
+                msg = "invalid units. can only work with optical maps in meter"
+                raise RuntimeError(msg)
+            range_in_m.append([b.first, b.last])
+            bins.append(b.nbins)
+
+        return {"range_in_m": np.array(range_in_m), "bins": np.array(bins)}
+
     def check_histograms(self, include_prefix: bool = False) -> None:
         log_prefix = "" if not include_prefix else self.name + " - "
 

tests/test_optmap.py

@@ -10,6 +10,7 @@
     create_optical_maps,
     list_optical_maps,
     merge_optical_maps,
+    rebin_optical_maps,
 )
 from reboost.optmap.evt import build_optmap_evt
 from reboost.optmap.optmap import OpticalMap
@@ -156,6 +157,25 @@ def test_optmap_merge(tbl_evt_fns, tmptestdir):
     merge_optical_maps([map1_fn, map2_fn], map_merged_fn, settings, n_procs=2)
 
 
+@pytest.mark.filterwarnings("ignore::scipy.optimize._optimize.OptimizeWarning")
+def test_optmap_rebin(tbl_evt_fns, tmptestdir):
+    settings = {
+        "range_in_m": [[0, 1], [0, 1], [0, 1]],
+        "bins": [10, 10, 10],
+    }
+
+    map1_fn = str(tmptestdir / "map-to-rebin.lh5")
+    create_optical_maps(
+        tbl_evt_fns,
+        settings,
+        chfilter=("001", "002", "003"),
+        output_lh5_fn=map1_fn,
+    )
+
+    map_rebinned_fn = str(tmptestdir / "map-rebinned.lh5")
+    rebin_optical_maps(map1_fn, map_rebinned_fn, factor=2)
+
+
 @pytest.fixture
 def tbl_edep(tmptestdir):
     evt_count = 100