diff --git a/docs/source/conf.py b/docs/source/conf.py index a9f1c10..15a6a9c 100644 --- a/docs/source/conf.py +++ b/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 diff --git a/src/reboost/__init__.py b/src/reboost/__init__.py index a5b4ecf..1d3a8d1 100644 --- a/src/reboost/__init__.py +++ b/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"] diff --git a/src/reboost/build_hit.py b/src/reboost/build_hit.py index cfbaa0b..ddee361 100644 --- a/src/reboost/build_hit.py +++ b/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 diff --git a/src/reboost/optmap/__init__.py b/src/reboost/optmap/__init__.py index e69de29..818cb66 100644 --- a/src/reboost/optmap/__init__.py +++ b/src/reboost/optmap/__init__.py @@ -0,0 +1,5 @@ +from __future__ import annotations + +from .optmap import OpticalMap + +__all__ = ["OpticalMap"] diff --git a/src/reboost/optmap/cli.py b/src/reboost/optmap/cli.py index d5db27f..39b4f02 100644 --- a/src/reboost/optmap/cli.py +++ b/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) diff --git a/src/reboost/optmap/create.py b/src/reboost/optmap/create.py index edcdc72..dff1529 100644 --- a/src/reboost/optmap/create.py +++ b/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) diff --git a/src/reboost/optmap/optmap.py b/src/reboost/optmap/optmap.py index f6bbefc..88ebdaa 100644 --- a/src/reboost/optmap/optmap.py +++ b/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() @@ -115,7 +116,7 @@ def _fill_histogram( assert ax.is_range assert ax.closedleft oor_mask &= (ax.first <= col) & (col < ax.last) - idx_s = np.floor((col - ax.first).astype(np.float64) / ax.step).astype(np.int64) + idx_s = np.floor((col.astype(np.float64) - ax.first) / ax.step).astype(np.int64) assert np.all(idx_s[oor_mask] < self._single_shape[dim]) idx += s * idx_s @@ -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 + " - " diff --git a/tests/test_optmap.py b/tests/test_optmap.py index ba79212..b050f4a 100644 --- a/tests/test_optmap.py +++ b/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