feat: extract A/E model parameters and compute it in hit.py

.github/workflows/main.yml

@@ -51,6 +51,11 @@ jobs:
         uses: actions/setup-python@v6
         with:
           python-version: ${{ matrix.python-version }}
+      - name: Install non-python (homebrew) dependencies
+        if: ${{ matrix.os == 'macOS-latest' }}
+        run: |
+          brew update
+          brew install opencascade cgal gmp mpfr boost
       - name: Get dependencies and install legend-simflow
         run: |
           python -m pip install --upgrade pip wheel setuptools

docs/source/conf.py

@@ -79,6 +79,7 @@
     "dbetto": ("https://dbetto.readthedocs.io/en/stable", None),
     "pylegendmeta": ("https://pylegendmeta.readthedocs.io/en/stable", None),
     "reboost": ("https://reboost.readthedocs.io/en/stable", None),
+    "scipy": ("https://docs.scipy.org/doc/scipy", None),
 }  # add new intersphinx mappings here
 
 # sphinx-autodoc

simflow-config.yaml

@@ -27,6 +27,7 @@ benchmark:
     stp: 10_000
 
 paths:
+  l200data: <...>/public/prodenv/prod-blind/ref-v1.0.0
   benchmarks: $_/generated/benchmarks
   log: $_/generated/log
 

tests/conftest.py

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

tests/test_hpge_pars.py

@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+from lgdo import lh5
+from matplotlib.axes import Axes
+from matplotlib.figure import Figure
+from scipy.stats import norm
+
+from legendsimflow import hpge_pars
+
+
+def test_fit():
+    t = np.linspace(-1000, 2000, 3001)
+    A = norm.pdf(t, loc=0, scale=50)
+
+    res = hpge_pars.fit_currmod(t, A)
+
+    assert isinstance(res[0], np.ndarray)
+    assert isinstance(res[1], np.ndarray)
+    assert isinstance(res[2], np.ndarray)
+
+
+def test_get_index(legend_testdata):
+    ref_path = legend_testdata.get_path("lh5/prod-ref-l200/")
+    path = ref_path / Path("generated/tier/hit/cal/p03/r001")
+    files = [str(p) for p in path.glob("*")]
+
+    # we have to be very generous with the (low stats) test file
+    idx, file_idx = hpge_pars.lookup_currmod_fit_data(
+        files, "ch1084803/hit", ewin_center=100, ewin_width=20
+    )
+
+    assert idx > -1
+    assert file_idx > -1
+
+    # now read back in and check
+
+    energy = lh5.read(
+        "ch1084803/hit/cuspEmax_ctc_cal", files[file_idx], idx=[idx]
+    ).view_as("np")
+    AoE = lh5.read("ch1084803/hit/AoE_Classifier", files[file_idx], idx=[idx]).view_as(
+        "np"
+    )
+
+    assert (energy > 90) & (energy < 110)
+    assert abs(AoE) < 1.5
+
+
+def test_get_waveform(legend_testdata):
+    times, wf = hpge_pars.get_current_pulse(
+        legend_testdata.get_path(
+            "lh5/prod-ref-l200/generated/tier/raw/cal/p03/r001/l200-p03-r001-cal-20230318T012144Z-tier_raw.lh5"
+        ),
+        "ch1084803/raw",
+        idx=0,
+        dsp_config=None,
+        dsp_output="waveform",
+        align=None,
+    )
+
+    assert isinstance(times, np.ndarray)
+    assert isinstance(wf, np.ndarray)
+    assert len(times) == len(wf)
+
+
+def test_plot():
+    fig, ax = hpge_pars.plot_currmod_fit_result(
+        [1, 2, 3], [0, 10, 20], np.linspace(0, 3, 1000), np.linspace(0, 3, 1000)
+    )
+
+    assert isinstance(fig, Figure)
+    assert isinstance(ax, Axes)

tests/test_metadata.py

@@ -55,3 +55,16 @@ def test_run_stuff(config):
     assert metadata.get_runlist(config, "sis1_z8224_src1_Ra224_to_Pb208") == [
         f"l200-p02-r00{r}-phy" for r in range(8)
     ]
+
+    assert (
+        metadata.reference_cal_run(config.metadata, "l200-p16-r006-phy")
+        == "l200-p16-r006-cal"
+    )
+    assert (
+        metadata.reference_cal_run(config.metadata, "l200-p16-r008-ssc")
+        == "l200-p16-r007-cal"
+    )
+    assert (
+        metadata.reference_cal_run(config.metadata, "l200-p16-r009-ssc")
+        == "l200-p16-r007-cal"
+    )

tests/test_utils.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+
+from legendsimflow import utils
+
+
+def test_get_parameter_dict():
+    popt = [100, 10, 60, 0.6, 100, 0.2, 60]
+
+    popt_dict = utils._curve_fit_popt_to_dict(popt)
+    assert len(popt_dict) == 8

workflow/rules/hit.smk

@@ -1,6 +1,8 @@
+import re
+
 from dbetto.utils import load_dict
 
-from legendsimflow import patterns, aggregate
+from legendsimflow import patterns, aggregate, hpge_pars
 from legendsimflow import metadata as mutils
 from legendsimflow import nersc
 
@@ -68,9 +70,10 @@ rule build_tier_hit:
     input:
         geom=patterns.geom_gdml_filename(config, tier="stp"),
         stp_file=patterns.output_simjob_filename(config, tier="stp"),
-        # NOTE: we pass here the full list of maps, but likely not all of them
-        # will be used. room for improvement
+        # NOTE: we pass here the full list of maps/current models, but likely
+        # not all of them will be used. room for improvement
         hpge_dtmaps=lambda wc: aggregate.gen_list_of_merged_dtmaps(config, wc.simid),
+        hpge_currmods=lambda wc: aggregate.gen_list_of_merged_currmods(config, wc.simid),
         # NOTE: technically this rule only depends on one block in the
         # partitioning file, but in practice the full file will always change
         simstat_part_file=rules.make_simstat_partition_file.output[0],
@@ -87,6 +90,7 @@ rule build_tier_hit:
 
 
 def smk_hpge_drift_time_map_inputs(wildcards):
+    """Prepare inputs for the HPGe drift time map rule."""
     meta = config.metadata.hardware.detectors.germanium.diodes[wildcards.hpge_detector]
     ids = {"bege": "B", "coax": "C", "ppc": "P", "icpc": "V"}
     crystal_name = (
@@ -197,3 +201,85 @@ rule merge_hpge_drift_time_maps:
           done
         done
         """
+
+
+def smk_extract_current_pulse_model_inputs(wildcards):
+    """Prepare inputs for the HPGe current model extraction rule."""
+    raw_file, wf_idx, dsp_cfg_file = hpge_pars.find_current_pulse_model_inputs(
+        config,
+        wildcards.runid,
+        wildcards.hpge_detector,
+        hit_tier_name="hit",
+        use_hpge_name="true",
+    )
+
+    evt_idx_file = patterns.input_currmod_evt_idx_file(
+        config, runid=wildcards.runid, hpge_detector=wildcards.hpge_detector
+    )
+
+    with evt_idx_file.open() as f:
+        f.write(wf_idx)
+
+    return {
+        "raw_file": raw_file,
+        "raw_wf_idx_file": evt_idx_file,
+        "dsp_cfg_file": dsp_cfg_file,
+    }
+
+
+rule extract_current_pulse_model:
+    """Extract the HPGe current signal model.
+
+    Perform a fit of the current waveform and stores the best-fit model
+    parameters in a YAML file.
+
+    :::{warning}
+    This rule does not have the relevant LEGEND-200 data files as input, since
+    they are dynamically discovered and this would therefore slow down the DAG
+    generation. Therefore, remember to force-rerun if the input data is
+    updated!
+    :::
+
+    Uses wildcards `runid` and `hpge_detector`.
+    """
+    message:
+        "Extracting current model for detector {wildcards.hpge_detector} in {wildcards.runid}"
+    # NOTE: we don't list the file dependencies here because they are
+    # dynamically generated, and that would slow down the DAG generation
+    # input:
+    #     unpack(smk_extract_current_pulse_model_inputs),
+    params:
+        # evt_idx=parse_input(input.raw_wf_idx_file, parser=mutils.extract_integer),
+        hit_tier_name="hit",
+    output:
+        pars_file=temp(patterns.output_currmod_filename(config)),
+        plot_file=patterns.plot_currmod_filename(config),
+    log:
+        patterns.log_currmod_filename(config, SIMFLOW_CONTEXT.proctime),
+    script:
+        "../src/legendsimflow/scripts/extract_hpge_current_pulse_model.py"
+
+
+rule merge_current_pulse_model_pars:
+    """Merge the HPGe current signal model parameters in a single file per `runid`.
+
+    Uses wildcard `runid`.
+    """
+    message:
+        "Merging current model parameters in {wildcards.runid}"
+    input:
+        lambda wc: aggregate.gen_list_of_currmods(config, wc.runid),
+    output:
+        patterns.output_currmod_merged_filename(config),
+    run:
+        import dbetto
+        from legendsimflow.aggregate import gen_list_of_hpges_valid_for_dtmap
+
+        # NOTE: this is guaranteed to be sorted as in the input file list
+        hpges = gen_list_of_hpges_valid_for_dtmap(config, wildcards.runid)
+
+        out_dict = {}
+        for i, f in enumerate(input):
+            out_dict[hpges[i]] = dbetto.utils.load_dict(f)
+
+        dbetto.utils.write_dict(out_dict, output[0])

workflow/src/legendsimflow/aggregate.py

@@ -127,7 +127,7 @@ def start_key(config: SimflowConfig, runid: str) -> str:
 
 
 def gen_list_of_hpges_valid_for_dtmap(config: SimflowConfig, runid: str) -> list[str]:
-    """Make a list of HPGe detector for which we want to generate a drift time map.
+    """Make a sorted list of HPGe detector for which we want to generate a drift time map.
 
     It generates the list of deployed detectors in `runid` via the LEGEND
     channelmap, then checks if in the crystal metadata there's all the
@@ -139,6 +139,10 @@ def gen_list_of_hpges_valid_for_dtmap(config: SimflowConfig, runid: str) -> list
 
     hpges = []
     for _, hpge in chmap.group("system").geds.items():
+        # TEMPORARY HACK
+        if hpge.name in ("V00050A", "V13046A", "V00048B", "V14654A"):
+            continue
+
         m = crystal_meta(
             config, config.metadata.hardware.detectors.germanium.diodes[hpge.name]
         )
@@ -157,7 +161,7 @@ def gen_list_of_hpges_valid_for_dtmap(config: SimflowConfig, runid: str) -> list
         msg = f"the list of HPGes valid for drift time map generation in {runid} is empty!"
         log.warning(msg)
 
-    return hpges
+    return sorted(hpges)
 
 
 def gen_list_of_all_runids(config):
@@ -207,6 +211,23 @@ def gen_list_of_all_dtmap_plots_outputs(config: SimflowConfig) -> set[str]:
     return files
 
 
+def gen_list_of_currmods(config: SimflowConfig, runid: str) -> list[str]:
+    """Generate the list of HPGe current model parameter files for a `runid`."""
+    hpges = gen_list_of_hpges_valid_for_dtmap(config, runid)
+    return [
+        patterns.output_currmod_filename(config, hpge_detector=hpge, runid=runid)
+        for hpge in hpges
+    ]
+
+
+def gen_list_of_merged_currmods(config: SimflowConfig, simid: str) -> list[str]:
+    r"""Generate the list of (merged) HPGe current model parameter files for all requested `runid`\ s."""
+    return [
+        patterns.output_currmod_merged_filename(config, runid=runid)
+        for runid in get_runlist(config, simid)
+    ]
+
+
 def process_simlist(
     config: SimflowConfig, simlist: Iterable[str] | None = None
 ) -> list[Path]: