reboost 0.2.6__py3-none-any.whl → 0.3.0__py3-none-any.whl

reboost/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.6'
-__version_tuple__ = version_tuple = (0, 2, 6)
+__version__ = version = '0.3.0'
+__version_tuple__ = version_tuple = (0, 3, 0)

reboost/build_glm.py

@@ -238,7 +238,7 @@ def build_glm(
         # start row for each table
         start_row = dict.fromkeys(lh5_table_list, 0)
 
-        vfield = f"stp/vertices/{id_name}"
+        vfield = f"vtx/{id_name}"
 
         # iterate over the vertex table
         for vert_obj, vidx, n_evtid in LH5Iterator(stp_file, vfield, buffer_len=evtid_buffer):

reboost/build_hit.py

@@ -290,7 +290,7 @@ def build_hit(
                     start_row=start_evtid,
                     stp_field=in_field,
                     n_rows=n_evtid,
-                    read_vertices=True,
+                    read_vertices=False,
                     buffer=buffer,
                     time_dict=time_dict[proc_name],
                 )
@@ -299,8 +299,15 @@ def build_hit(
                     if stps is None:
                         continue
 
+                    # convert to awkward
+                    if time_dict is not None:
+                        start_time = time.time()
+
                     ak_obj = stps.view_as("ak")
 
+                    if time_dict is not None:
+                        time_dict[proc_name].update_field("conv", start_time)
+
                     # produce the hit table
                     for out_det_idx, out_detector in enumerate(out_detectors):
                         # loop over the rows

reboost/iterator.py

@@ -4,6 +4,7 @@ import logging
 import time
 import typing
 
+import awkward as ak
 from lgdo.lh5 import LH5Store
 from lgdo.types import LGDO, Table
 
@@ -94,7 +95,7 @@ class GLMIterator:
         # read the glm rows]
         if self.glm_file is not None:
             glm_rows, n_rows_read = self.sto.read(
-                f"glm/{self.lh5_group}", self.glm_file, start_row=self.start_row_tmp, n_rows=n_rows
+                f"/glm/{self.lh5_group}", self.glm_file, start_row=self.start_row_tmp, n_rows=n_rows
             )
         else:
             # get the maximum row to read
@@ -124,12 +125,13 @@ class GLMIterator:
         if len(glm_ak) > 0:
             # extract range of stp rows to read
             start = glm_ak.start_row[0]
-            n = sum(glm_ak.n_rows)
+            n = ak.sum(glm_ak.n_rows)
+
             if self.time_dict is not None:
                 time_start = time.time()
 
             stp_rows, n_steps = self.sto.read(
-                f"{self.stp_field}/{self.lh5_group}",
+                f"/{self.stp_field}/{self.lh5_group}",
                 self.stp_file,
                 start_row=int(start),
                 n_rows=int(n),
@@ -143,7 +145,7 @@ class GLMIterator:
 
             if self.read_vertices:
                 vert_rows, _ = self.sto.read(
-                    f"{self.stp_field}/vertices",
+                    "/vtx",
                     self.stp_file,
                     start_row=self.start_row,
                     n_rows=n_rows,

reboost/optmap/cli.py

@@ -166,7 +166,12 @@ def optical_cli() -> None:
         action="store",
         required=True,
         metavar="LGDO_PATH",
-        help="path to LGDO inside non-optical LH5 hit file (e.g. /hit/detXX)",
+        help="path to LGDO inside non-optical LH5 hit file (e.g. /stp/detXX)",
+    )
+    convolve_parser.add_argument(
+        "--dist-mode",
+        action="store",
+        default="multinomial+no-fano",
     )
     convolve_parser.add_argument("--output", help="output hit LH5 file", metavar="OUTPUT_HIT")
 
@@ -267,8 +272,16 @@ def optical_cli() -> None:
         from reboost.optmap.convolve import convolve
 
         _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)
+        _check_output_file(parser, args.output, optional=True)
+        convolve(
+            args.map,
+            args.edep,
+            args.edep_lgdo,
+            args.material,
+            args.output,
+            args.bufsize,
+            dist_mode=args.dist_mode,
+        )
 
     # STEP X: rebin maps
     if args.command == "rebin":

reboost/optmap/convolve.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import logging
 
 import legendoptics.scintillate as sc
+import numba
 import numpy as np
 import pint
 from legendoptics import lar
@@ -23,7 +24,8 @@ OPTMAP_SUM_CH = -2
 
 def open_optmap(optmap_fn: str):
     maps = lh5.ls(optmap_fn)
-    det_ntuples = [m for m in maps if m not in ("all", "_hitcounts", "_hitcounts_exp")]
+    # TODO: rewrite logic to only accept _<number> instead of a blacklist
+    det_ntuples = [m for m in maps if m not in ("all", "_hitcounts", "_hitcounts_exp", "all_orig")]
     detids = np.array([int(m.lstrip("_")) for m in det_ntuples])
     detidx = np.arange(0, detids.shape[0])
 
@@ -69,15 +71,16 @@ def iterate_stepwise_depositions(
     optmap_for_convolve,
     scint_mat_params: sc.ComputedScintParams,
     rng: np.random.Generator = None,
+    dist: str = "multinomial",
+    mode: str = "no-fano",
 ):
     # those np functions are not supported by numba, but needed for efficient array access below.
     x0 = structured_to_unstructured(edep_df[["xloc_pre", "yloc_pre", "zloc_pre"]], np.float64)
     x1 = structured_to_unstructured(edep_df[["xloc_post", "yloc_post", "zloc_post"]], np.float64)
 
     rng = np.random.default_rng() if rng is None else rng
-
     output_map, res = _iterate_stepwise_depositions(
-        edep_df, x0, x1, rng, *optmap_for_convolve, scint_mat_params
+        edep_df, x0, x1, rng, *optmap_for_convolve, scint_mat_params, dist, mode
     )
     if res["any_no_stats"] > 0 or res["det_no_stats"] > 0:
         log.warning(
@@ -119,6 +122,9 @@ def _pdgid_to_particle(pdgid: int) -> sc.ParticleIndex:
     return sc.PARTICLE_INDEX_ELECTRON
 
 
+__counts_per_bin_key_type = numba.types.UniTuple(numba.types.int64, 3)
+
+
 # - run with NUMBA_FULL_TRACEBACKS=1 NUMBA_BOUNDSCHECK=1 for testing/checking
 # - cache=True does not work with outer prange, i.e. loading the cached file fails (numba bug?)
 # - the output dictionary is not threadsafe, so parallel=True is not working with it.
@@ -134,6 +140,8 @@ def _iterate_stepwise_depositions(
     optmap_weights,
     optmap_multi_det_exp,
     scint_mat_params: sc.ComputedScintParams,
+    dist: str,
+    mode: str,
 ):
     pdgid_map = {}
     output_map = {}
@@ -160,6 +168,7 @@ def _iterate_stepwise_depositions(
             charge,
             t.edep,
             rng,
+            emission_term_model=("poisson" if mode == "no-fano" else "normal_fano"),
         )
         if scint_times.shape[0] == 0:  # short-circuit if we have no photons at all.
             continue
@@ -175,6 +184,11 @@ def _iterate_stepwise_depositions(
         # there are _much_ less unique bins, unfortunately np.unique(..., axis=n) does not work
         # with numba; also np.sort(..., axis=n) also does not work.
 
+        counts_per_bin = numba.typed.Dict.empty(
+            key_type=__counts_per_bin_key_type,
+            value_type=np.int64,
+        )
+
         # get probabilities from map.
         hitcount = np.zeros((detidx.shape[0], bins.shape[0]), dtype=np.int64)
         for j in prange(bins.shape[0]):
@@ -191,35 +205,65 @@ def _iterate_stepwise_depositions(
                 continue
             if px_any == 0.0:
                 continue
-            if rng.uniform() >= px_any:
-                continue
-            ph_det += 1
-            # we detect this energy deposition; we should at least get one photon out here!
-
-            detsel_size = 1
-            if np.isfinite(optmap_multi_det_exp):
-                detsel_size = rng.geometric(1 - np.exp(-optmap_multi_det_exp))
-
-            px_sum = optmap_weights[OPTMAP_SUM_CH, cur_bins[0], cur_bins[1], cur_bins[2]]
-            assert px_sum >= 0.0  # should not be negative.
-            detp = np.empty(detidx.shape, dtype=np.float64)
-            had_det_no_stats = 0
-            for d in detidx:
-                # normalize so that sum(detp) = 1
-                detp[d] = optmap_weights[d, cur_bins[0], cur_bins[1], cur_bins[2]] / px_sum
-                if detp[d] < 0.0:
-                    had_det_no_stats = 1
-                    detp[d] = 0.0
-            det_no_stats += had_det_no_stats
-
-            # should be equivalent to rng.choice(detidx, size=(detsel_size, p=detp)
-            detsel = detidx[
-                np.searchsorted(np.cumsum(detp), rng.random(size=(detsel_size,)), side="right")
-            ]
-            for d in detsel:
-                hitcount[d, j] += 1
-            ph_det2 += detsel.shape[0]
 
+            if dist == "multinomial":
+                if rng.uniform() >= px_any:
+                    continue
+                ph_det += 1
+                # we detect this energy deposition; we should at least get one photon out here!
+
+                detsel_size = 1
+                if np.isfinite(optmap_multi_det_exp):
+                    detsel_size = rng.geometric(1 - np.exp(-optmap_multi_det_exp))
+
+                px_sum = optmap_weights[OPTMAP_SUM_CH, cur_bins[0], cur_bins[1], cur_bins[2]]
+                assert px_sum >= 0.0  # should not be negative.
+                detp = np.empty(detidx.shape, dtype=np.float64)
+                had_det_no_stats = 0
+                for d in detidx:
+                    # normalize so that sum(detp) = 1
+                    detp[d] = optmap_weights[d, cur_bins[0], cur_bins[1], cur_bins[2]] / px_sum
+                    if detp[d] < 0.0:
+                        had_det_no_stats = 1
+                        detp[d] = 0.0
+                det_no_stats += had_det_no_stats
+
+                # should be equivalent to rng.choice(detidx, size=(detsel_size, p=detp)
+                detsel = detidx[
+                    np.searchsorted(np.cumsum(detp), rng.random(size=(detsel_size,)), side="right")
+                ]
+                for d in detsel:
+                    hitcount[d, j] += 1
+                ph_det2 += detsel.shape[0]
+
+            elif dist == "poisson":
+                # store the photon count in each bin, to sample them all at once below.
+                if cur_bins not in counts_per_bin:
+                    counts_per_bin[cur_bins] = 1
+                else:
+                    counts_per_bin[cur_bins] += 1
+
+            else:
+                msg = "unknown distribution"
+                raise RuntimeError(msg)
+
+        if dist == "poisson":
+            for j, (cur_bins, ph_counts_to_poisson) in enumerate(counts_per_bin.items()):
+                had_det_no_stats = 0
+                had_any = 0
+                for d in detidx:
+                    detp = optmap_weights[d, cur_bins[0], cur_bins[1], cur_bins[2]]
+                    if detp < 0.0:
+                        had_det_no_stats = 1
+                        continue
+                    pois_cnt = rng.poisson(lam=ph_counts_to_poisson * detp)
+                    hitcount[d, j] += pois_cnt
+                    ph_det2 += pois_cnt
+                    had_any = 1
+                ph_det += had_any
+                det_no_stats += had_det_no_stats
+
+        assert scint_times.shape[0] >= hitcount.shape[1]  # TODO: use the right assertion here.
         out_hits_len = np.sum(hitcount)
         if out_hits_len > 0:
             out_times = np.empty(out_hits_len, dtype=np.float64)
@@ -285,6 +329,7 @@ def convolve(
     material: str,
     output_file: str | None = None,
     buffer_len: int = int(1e6),
+    dist_mode: str = "multinomial+no-fano",
 ):
     if material not in ["lar", "pen"]:
         msg = f"unknown material {material} for scintillation"
@@ -304,6 +349,11 @@ def convolve(
     log.info("opening map %s", map_file)
     optmap_for_convolve = open_optmap(map_file)
 
+    # special handling of distributions and flags.
+    dist, mode = dist_mode.split("+")
+    assert dist in ("multinomial", "poisson")
+    assert mode in ("", "no-fano")
+
     log.info("opening energy deposition hit output %s", edep_file)
     it = LH5Iterator(edep_file, edep_path, buffer_len=buffer_len)
     for it_count, (edep_lgdo, edep_events, edep_n_rows) in enumerate(it):
@@ -311,7 +361,9 @@ def convolve(
         edep_df = edep_lgdo.view_as("pd").iloc[0:edep_n_rows].to_records()
 
         log.info("start event processing (%d)", it_count)
-        output_map = iterate_stepwise_depositions(edep_df, optmap_for_convolve, scint_mat_params)
+        output_map = iterate_stepwise_depositions(
+            edep_df, optmap_for_convolve, scint_mat_params, dist=dist, mode=mode
+        )
 
         log.info("store output photon hits (%d)", it_count)
         ph_count_o, tbl = get_output_table(output_map)
@@ -319,4 +371,4 @@ def convolve(
             "output photons: %d energy depositions -> %d photons", len(output_map), ph_count_o
         )
         if output_file is not None:
-            lh5.write(tbl, "optical", lh5_file=output_file, group="hit", wo_mode="append")
+            lh5.write(tbl, "optical", lh5_file=output_file, group="stp", wo_mode="append")

reboost/optmap/evt.py

@@ -26,7 +26,7 @@ def build_optmap_evt(
         msg = f"temporary output file {lh5_out_file_tmp} already exists"
         raise RuntimeError(msg)
 
-    vert_it = LH5Iterator(lh5_in_file, "stp/vertices", buffer_len=buffer_len)
+    vert_it = LH5Iterator(lh5_in_file, "vtx", buffer_len=buffer_len)
     opti_it = LH5Iterator(lh5_in_file, "stp/optical", buffer_len=buffer_len)
 
     detectors = [str(d) for d in detectors]

reboost/utils.py

@@ -296,7 +296,10 @@ def _check_input_file(parser, file: str | Iterable[str], descr: str = "input") -
         parser.error(f"{descr} file(s) {''.join(not_existing)} missing")
 
 
-def _check_output_file(parser, file: str | Iterable[str]) -> None:
+def _check_output_file(parser, file: str | Iterable[str] | None, optional: bool = False) -> None:
+    if file is None and optional:
+        return
+
     file = (file,) if isinstance(file, str) else file
     for f in file:
         if Path(f).exists():

{reboost-0.2.6.dist-info → reboost-0.3.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: reboost
-Version: 0.2.6
+Version: 0.3.0
 Summary: New LEGEND Monte-Carlo simulation post-processing
 Author-email: Manuel Huber <info@manuelhu.de>, Toby Dixon <toby.dixon.23@ucl.ac.uk>, Luigi Pertoldi <gipert@pm.me>
 Maintainer: The LEGEND Collaboration

{reboost-0.2.6.dist-info → reboost-0.3.0.dist-info}/RECORD

@@ -1,16 +1,16 @@
 reboost/__init__.py,sha256=RVNl3Qgx_hTUeBGXaWYmiTcmXUDhTfvlAGGC8bo_jP8,316
-reboost/_version.py,sha256=nObnONsicQ3YX6SG5MVBxmIp5dmRacXDauSqZijWQbY,511
+reboost/_version.py,sha256=AGmG_Lx0-9ztFw_7d9mYbaYuC-2abxE1oXOUNAY29YY,511
 reboost/build_evt.py,sha256=zj3wG_kaV3EoRMQ33AkCNa_2Fv8cLtRuhyRyRmSrOYQ,4797
-reboost/build_glm.py,sha256=kSY9hQjEsOE-0PiblhdBy_SvFIlgXLX6CUlgpxW-_OI,9389
-reboost/build_hit.py,sha256=OyXkYdLIpSBmq3MLNXD_kLjVqxeSQcs7RGAOZSy1Bns,14695
+reboost/build_glm.py,sha256=VIzRyCc53FQvcXAVbARszNH5wH5Pr0a9WWlSCZuUY5w,9380
+reboost/build_hit.py,sha256=Zu8WoeFYtFU_xUO3I8OZM5psNaiv1boMJPnBWC58nfQ,14958
 reboost/build_tcm.py,sha256=-PawBHoHj0zsm4XsZu5bco9d9a09STicZchduefSNfI,2951
 reboost/cli.py,sha256=swPJcYzvg18rSOMN-mpe0PCMf1-a9V7osIssX7JP7k0,6459
 reboost/core.py,sha256=7Nclc6RUCOSJ1CWVAX0rFNJGM1LEgqvc4tD04CxEAtg,10766
-reboost/iterator.py,sha256=0KmrekpZwOYZJaP0nmp-SNrr2WmGUKeNUVcqO-OChhY,4757
+reboost/iterator.py,sha256=D0aDCwtUKeo3_JQQKrmYmFc7jbqcKb5jrpMkUuaJ76s,4762
 reboost/log_utils.py,sha256=VqS_9OC5NeNU3jcowVOBB0NJ6ssYvNWnirEY-JVduEA,766
 reboost/profile.py,sha256=EOTmjmS8Rm_nYgBWNh6Rntl2XDsxdyed7yEdWtsZEeg,2598
 reboost/units.py,sha256=3EH8XlpbsObdu5vLgxhm1600L6UNYD5jng4SjJT_1QE,2202
-reboost/utils.py,sha256=eBw0ZzwhlniTLbjz9tnstCXSYrjSeH4FJ0fkJ9-uqps,8450
+reboost/utils.py,sha256=VheyCMcaXT5VPXAotiI5YqLKTf45KhTMpwULg2xyB3o,8531
 reboost/hpge/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 reboost/hpge/psd.py,sha256=868OUJzO9TNja0YSrZ3NDGeEAbUtpDZnmvBDm0jCC9E,6856
 reboost/hpge/surface.py,sha256=SZyTmOCTipf27jYaJhtdInzGF1RZ2wKpbtf6HlOQYwM,3662
@@ -19,10 +19,10 @@ reboost/math/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 reboost/math/functions.py,sha256=OymiYTcA0NXxxm-MBDw5kqyNwHoLCmuv4J48AwnSrbU,5633
 reboost/math/stats.py,sha256=iiOEi87x93kqPWeSmlRiA5Oe-R8XR-plm6Z532PhC9M,1401
 reboost/optmap/__init__.py,sha256=imvuyld-GLw8qdwqW-lXCg2feptcTyQo3wIzPvDHwmY,93
-reboost/optmap/cli.py,sha256=wBexh-zrr5ABherEyk9xigxdArvOAKiiRQwAYon9Sro,9408
-reboost/optmap/convolve.py,sha256=5FksUrVIG8ysn42QbWBrAx8M1HfAVJtaJJyE8oJ1NGM,12043
+reboost/optmap/cli.py,sha256=TszAYZIwHTXwJYQOMKxreCZ6pXlGUb6q6YZ3iW3Bxf4,9670
+reboost/optmap/convolve.py,sha256=jCH_d04yioB8hsJEPunm0zynA0ne4lx0ldSC1GJG_eY,14129
 reboost/optmap/create.py,sha256=Nm5-xEe8M9q2GFQnUv8oN8qpAz9nZArIrQcPboqRmCQ,17153
-reboost/optmap/evt.py,sha256=m3NWuLEk4zDQJO5vXq_XLLnqmkQwmtdKI3fqmZQBBvc,4707
+reboost/optmap/evt.py,sha256=9rfAdN9MqL6UuUxUcMDGVwpcuqRVc2RwmEmd87jgYww,4698
 reboost/optmap/mapview.py,sha256=73kpe0_SKDj9bIhEx1ybX1sBP8TyvufiLfps84A_ijA,6798
 reboost/optmap/numba_pdg.py,sha256=y8cXR5PWE2Liprp4ou7vl9do76dl84vXU52ZJD9_I7A,731
 reboost/optmap/optmap.py,sha256=j4rfbQ84PYSpE-BvP4Rdt96ZjPdwy8P4e4eZz1mATys,12817
@@ -30,9 +30,9 @@ reboost/shape/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 reboost/shape/cluster.py,sha256=RIvBlhHzp88aaUZGofp5SD9bimnoiqIOddhQ84jiwoM,8135
 reboost/shape/group.py,sha256=Q3DhEPxbhw3p4bwvpswSd0A-p224l5vRZnfQIEkOVJE,4475
 reboost/shape/reduction.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-reboost-0.2.6.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
-reboost-0.2.6.dist-info/METADATA,sha256=Irdm3IrpsoXDPWyqeYSa4QHI6MdACh0ZYAnGMQwofbg,44251
-reboost-0.2.6.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
-reboost-0.2.6.dist-info/entry_points.txt,sha256=DxhD6BidSWNot9BrejHJjQ7RRLmrMaBIl52T75oWTwM,93
-reboost-0.2.6.dist-info/top_level.txt,sha256=q-IBsDepaY_AbzbRmQoW8EZrITXRVawVnNrB-_zyXZs,8
-reboost-0.2.6.dist-info/RECORD,,
+reboost-0.3.0.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
+reboost-0.3.0.dist-info/METADATA,sha256=GCn2W1jn_cplxbS1-bJCRjghoEaPgscoTuAhVkizXVY,44251
+reboost-0.3.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+reboost-0.3.0.dist-info/entry_points.txt,sha256=DxhD6BidSWNot9BrejHJjQ7RRLmrMaBIl52T75oWTwM,93
+reboost-0.3.0.dist-info/top_level.txt,sha256=q-IBsDepaY_AbzbRmQoW8EZrITXRVawVnNrB-_zyXZs,8
+reboost-0.3.0.dist-info/RECORD,,

{reboost-0.2.6.dist-info → reboost-0.3.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.7.1)
+Generator: setuptools (80.9.0)
 Root-Is-Purelib: true
 Tag: py3-none-any