hit: fix PSD observables

tests/conftest.py

@@ -17,7 +17,7 @@
 @pytest.fixture(scope="session")
 def legend_testdata():
     ldata = LegendTestData()
-    ldata.checkout("dca4e16")
+    ldata.checkout("0bb2dfc")
     return ldata
 
 

tests/test_reboost.py

@@ -1,5 +1,8 @@
 from __future__ import annotations
 
+import awkward as ak
+import pyg4ometry
+import reboost
 from lgdo import lh5
 
 from legendsimflow import reboost as rutils
@@ -35,3 +38,98 @@ def test_remage_hit_range(legend_testdata):
                 n += nen
 
         assert n == n_rows
+
+
+def test_psd_stuff(legend_testdata):
+    dt_map = {}
+    for angle in ("000", "045"):
+        dt_map[angle] = reboost.hpge.utils.get_hpge_rz_field(
+            legend_testdata["lh5/V00048A-drift-time-maps-xtal-axes.lh5"],
+            "V00048A",
+            f"drift_time_{angle}_deg",
+            bounds_error=False,
+        )
+
+    xloc = [
+        [0.162, 0.162, 0.162, 0.162, 0.162, 0.162, 0.162, 0.162, 0.162],
+        [0.183, 0.183],
+        [0.197, 0.197, 0.197, 0.197, 0.197, 0.197, 0.197, 0.197],
+        [0.201, 0.201, 0.201, 0.201, 0.201, 0.201, 0.201, 0.201],
+        [0.213, 0.213, 0.213, 0.213, 0.213],
+        [0.176, 0.176, 0.176, 0.176, 0.176, 0.176, 0.177],
+        [0.193, 0.193, 0.193, 0.193, 0.193, 0.193],
+        [0.201, 0.201, 0.201],
+        [0.157, 0.157, 0.157, 0.157, 0.157],
+        [0.164, 0.164, 0.164],
+    ]
+
+    yloc = [
+        [0.107, 0.107, 0.107, 0.107, 0.107, 0.107, 0.107, 0.107, 0.107],
+        [0.0906, 0.0906],
+        [0.122, 0.122, 0.122, 0.122, 0.122, 0.122, 0.122, 0.122],
+        [0.11, 0.11, 0.11, 0.11, 0.11, 0.11, 0.11, 0.11],
+        [0.131, 0.131, 0.131, 0.131, 0.131],
+        [0.0946, 0.0946, 0.0946, 0.0947, 0.0947, 0.0947, 0.0947],
+        [0.157, 0.157, 0.157, 0.157, 0.157, 0.157],
+        [0.145, 0.145, 0.145],
+        [0.124, 0.124, 0.124, 0.124, 0.124],
+        [0.144, 0.144, 0.144],
+    ]
+
+    zloc = [
+        [0.535, 0.535, 0.536, 0.536, 0.536, 0.536, 0.536, 0.536, 0.536],
+        [0.509, 0.509],
+        [0.538, 0.538, 0.538, 0.538, 0.538, 0.538, 0.538, 0.538],
+        [0.533, 0.533, 0.533, 0.533, 0.533, 0.533, 0.533, 0.533],
+        [0.52, 0.52, 0.52, 0.52, 0.52],
+        [0.506, 0.506, 0.506, 0.507, 0.507, 0.507, 0.507],
+        [0.557, 0.557, 0.557, 0.557, 0.557, 0.557],
+        [0.516, 0.516, 0.516],
+        [0.519, 0.519, 0.519, 0.519, 0.519],
+        [0.54, 0.54, 0.54],
+    ]
+
+    edep = ak.Array(
+        [
+            [130, 97.2, 233, 179, 133, 85.9, 173, 129, 331],
+            [342, 125],
+            [42.7, 75.9, 92.6, 261, 193, 49.6, 118, 249],
+            [179, 35.6, 72.4, 174, 167, 241, 134, 126],
+            [173, 37.3, 184, 156, 259],
+            [13.9, 29.9, 91.1, 107, 106, 193, 70.8],
+            [218, 13.2, 111, 111, 143, 118],
+            [84.3, 56.9, 314],
+            [63.7, 17.8, 92.4, 213, 186],
+            [51, 95.3, 312],
+        ]
+    )
+
+    chunk = ak.Array({"xloc": xloc, "yloc": yloc, "zloc": zloc})
+    det_loc = pyg4ometry.gdml.Defines.Position(
+        "Position",
+        183.415,
+        125.070,
+        490.044,
+        unit="mm",
+    )
+
+    dt = rutils.hpge_corrected_drift_time(
+        chunk,
+        dt_map,
+        det_loc,
+    )
+
+    assert ak.all((dt > 0) & (dt < 3000))
+
+    pars = {
+        "amax": 852.0,
+        "mu": 0.0999,
+        "sigma": 54.5,
+        "tail_fraction": 0.223,
+        "tau": 507.0,
+        "high_tail_fraction": 0.00611,
+        "high_tau": 593.0,
+    }
+    amax = rutils.hpge_max_current(edep, dt, pars)
+
+    assert ak.all((amax > 0) & (amax < 3000))

tests/test_utils.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import numpy as np
+import pytest
 from dbetto import AttrsDict
 from lgdo import Array, Table
 
@@ -35,3 +36,12 @@ def test_get_dataflow_config(test_l200data):
     assert isinstance(config, AttrsDict)
     assert "paths" in config
     assert "$_" not in config.paths.tier
+
+
+def test_check_nans_leq():
+    utils.check_nans_leq([1, 2, 3, 4], "boh", 0.1)
+    utils.check_nans_leq([[1], [2, 3], 4], "boh", 0.1)
+    utils.check_nans_leq([[1], [2, 3], np.nan], "boh", 0.5)
+
+    with pytest.raises(RuntimeError):
+        utils.check_nans_leq([[1], [2, 3], np.nan], "boh", 0.1)

workflow/src/legendsimflow/reboost.py

@@ -178,10 +178,11 @@ def get_senstables(
 
 def load_hpge_dtmaps(
     config: SimflowConfig, det_name: str, runid: str
-) -> dict[str, reboost.hpge.utils.HPGeRZField]:
+) -> dict[str, reboost.hpge.utils.HPGeRZField] | None:
     """Load HPGe drift time maps from disk.
 
     Automatically finds and loads drift time maps for crystal axes <100> <110>.
+    If no map is found, None is returned.
 
     Note
     ----
@@ -197,7 +198,10 @@ def load_hpge_dtmaps(
         dt_map = {}
         for angle in ("000", "045"):
             dt_map[angle] = reboost.hpge.utils.get_hpge_rz_field(
-                hpge_dtmap_file, det_name, f"drift_time_{angle}_deg"
+                hpge_dtmap_file,
+                det_name,
+                f"drift_time_{angle}_deg",
+                bounds_error=False,
             )
     else:
         msg = (
@@ -207,6 +211,8 @@ def load_hpge_dtmaps(
         log.warning(msg)
         dt_map = None
 
+    return dt_map
+
 
 def get_remage_hit_range(
     tcm: ak.Array, det_name: str, uid: int, evt_idx_range: list[int]
@@ -269,7 +275,7 @@ def get_remage_hit_range(
 
 def hpge_corrected_drift_time(
     chunk: ak.Array,
-    dt_map: reboost.hpge.utils.HPGeRZField,
+    dt_map: dict[str, reboost.hpge.utils.HPGeRZField],
     det_loc: pyg4ometry.gdml.Defines.Position,
 ) -> ak.Array:
     """HPGe drift time heuristic corrected for crystal axis effects.
@@ -320,7 +326,7 @@ def hpge_max_current(
     t_domain = {"low": -1000, "high": 4000, "step": 1}
 
     # instantiate the template
-    a_tmpl = reboost.hpge.psd.get_current_template(
+    a_tmpl, times = reboost.hpge.psd.get_current_template(
         **t_domain,
         mean_aoe=1,  # set the maximum of the template to unity, so the A/E will be calibrated
         **currmod_pars,
@@ -330,8 +336,8 @@ def hpge_max_current(
         edep,
         drift_time,
         template=a_tmpl,
-        times=np.arange(t_domain["low"], t_domain["high"]),
-    )
+        times=times,
+    ).view_as("ak")
 
 
 def build_tcm(

workflow/src/legendsimflow/scripts/tier/evt.py

@@ -153,7 +153,7 @@ def _read_hits(tcm_ak, tier, field):
     )
 
     # simply forward some fields
-    for field in ["aoe", "drift_time_heuristic"]:
+    for field in ["aoe"]:
         field_data = _read_hits(tcm, "hit", field)
         out_table.add_field(f"geds/{field}", VectorOfVectors(field_data[hitsel]))
 

workflow/src/legendsimflow/scripts/tier/hit.py

@@ -173,7 +173,7 @@ def DEFAULT_ENERGY_RES_FUNC(energy):
         # load parameters of the current model
         pars = currmod_pars_all.get(det_name, None)
         currmod_pars = (
-            pars.get("current_pulse_shape", None) if pars is not None else None
+            pars.get("current_pulse_pars", None) if pars is not None else None
         )
 
         n_tot = 0
@@ -218,32 +218,44 @@ def DEFAULT_ENERGY_RES_FUNC(energy):
             energy = reboost.math.stats.gaussian_sample(
                 energy_true,
                 energy_res / 2.35482,
+            ).view_as("ak")
+
+            # energy can't be negative as a result of smearing and later we
+            # divide for it
+            energy = ak.where(
+                (energy <= 0) & (energy_true >= 0), np.finfo(float).tiny, energy
             )
 
             # PSD: if the drift time map is none, it means that we don't
             # have the detector model to simulate PSD in a more advanced
             # way
 
             # default to NaN
-            dt_heuristic = np.full(len(chunk), np.nan)
+            drift_time = ak.full_like(chunk.xloc, fill_value=np.nan)
             aoe = np.full(len(chunk), np.nan)
 
             if dt_map is not None:
-                _drift_time = reboost_utils.hpge_corrected_drift_time(
+                msg = "computing PSD observables"
+                log.info(msg)
+
+                drift_time = reboost_utils.hpge_corrected_drift_time(
                     chunk, dt_map, det_loc
                 )
-                dt_heuristic = reboost.hpge.psd.drift_time_heuristic(
-                    _drift_time, chunk.edep
-                )
+                utils.check_nans_leq(drift_time, "drift_time", 0.1)
+
                 _a_max = reboost_utils.hpge_max_current(
-                    edep_active, _drift_time, currmod_pars
+                    edep_active, drift_time, currmod_pars
                 )
+                utils.check_nans_leq(_a_max, "max_current", 0.1)
+
                 aoe = _a_max / energy
 
             out_table = reboost_utils.make_output_chunk(lgdo_chunk)
 
             out_table.add_field("energy", lgdo.Array(energy, attrs={"units": "keV"}))
-            out_table.add_field("drift_time_heuristic", lgdo.Array(dt_heuristic))
+            out_table.add_field(
+                "drift_time", lgdo.VectorOfVectors(drift_time, attrs={"units": "ns"})
+            )
             out_table.add_field("aoe", lgdo.Array(aoe))
 
             _, period, run, _ = mutils.parse_runid(runid)

workflow/src/legendsimflow/utils.py

@@ -23,6 +23,7 @@
 from datetime import datetime
 from pathlib import Path
 
+import awkward as ak
 import dbetto
 import h5py
 import legenddataflowscripts as ldfs
@@ -290,3 +291,23 @@ def add_field_string(name: str, chunk: lgdo.Table, data: str) -> None:
     dtype = h5py.string_dtype(encoding="utf-8", length=len(data))
     data_array = np.full(len(chunk), fill_value=data, dtype=dtype)
     chunk.add_field(name, lgdo.Array(data_array))
+
+
+def check_nans_leq(array: ArrayLike, name: str, less_than_frac: float = 0.1) -> None:
+    """Raise an exception if the fraction of NaN values in `array` is above threshold.
+
+    Parameters
+    ----------
+    array
+        the array to analyze.
+    name
+        array name for exception message.
+    less_than_frac
+        raise exception if fraction of NaNs is above this threshold.
+    """
+    flat = ak.ravel(array)
+    n_el = len(flat)
+    n_nans = ak.sum(ak.is_none(ak.nan_to_none(flat)))
+    if (n_nans / n_el) > less_than_frac:
+        msg = f"more than {100 * less_than_frac}% of NaNs detected in array {name}!"
+        raise RuntimeError(msg)