Merge pull request #92 from cdunn314/update_channel_reading

Fix bug in custom peak_kwargs input into find_peaks

Author: RemDelaporteMathurin

libra_toolbox/neutron_detection/activation_foils/compass.py

@@ -536,19 +536,18 @@ def compute_detection_efficiency(
         detection_efficiency = nb_counts_measured / expected_nb_counts
 
         return detection_efficiency
-
-    def get_peaks(self, hist: np.ndarray, **kwargs) -> np.ndarray:
-        """Returns the peak indices of the histogram
+    
+    def get_peak_fitting_parameters(self, hist: np.ndarray, **kwargs) -> Dict[str, Union[float, List[float]]]:
+        """Returns the peak fitting parameters for the given histogram and check source nuclide.
 
         Args:
             hist: a histogram
             kwargs: optional parameters for the peak finding algorithm
                 see scipy.signal.find_peaks for more information
 
         Returns:
-            the peak indices in ``hist``
+            the peak fitting parameters
         """
-
         if self.detector_type.lower() == 'nai':
             # peak finding parameters
             start_index = 100
@@ -604,22 +603,46 @@ def get_peaks(self, hist: np.ndarray, **kwargs) -> np.ndarray:
 
         # update the parameters if kwargs are provided
         if kwargs:
+            print("Using custom peak finding parameters from kwargs: ", kwargs)
             start_index = kwargs.get("start_index", start_index)
             prominence = kwargs.get("prominence", prominence)
             height = kwargs.get("height", height)
             width = kwargs.get("width", width)
             distance = kwargs.get("distance", distance)
 
+        return {
+            "start_index": start_index,
+            "prominence": prominence,
+            "height": height,
+            "width": width,
+            "distance": distance
+        }
+
+    def get_peaks(self, hist: np.ndarray, **kwargs) -> np.ndarray:
+        """Returns the peak indices of the histogram
+
+        Args:
+            hist: a histogram
+            kwargs: optional parameters for the peak finding algorithm
+                see scipy.signal.find_peaks for more information
+
+        Returns:
+            the peak indices in ``hist``
+        """
+
+        # get the peak fitting parameters based on the check source nuclide and the detector type
+        peak_fitting_parameters = self.get_peak_fitting_parameters(hist, **kwargs)
+
         # run the peak finding algorithm
         # NOTE: the start_index is used to ignore the low energy region
         peaks, peak_data = find_peaks(
-            hist[start_index:],
-            prominence=prominence,
-            height=height,
-            width=width,
-            distance=distance,
+            hist[peak_fitting_parameters["start_index"]:],
+            prominence=peak_fitting_parameters["prominence"],
+            height=peak_fitting_parameters["height"],
+            width=peak_fitting_parameters["width"],
+            distance=peak_fitting_parameters["distance"],
         )
-        peaks = np.array(peaks) + start_index
+        peaks = np.array(peaks) + peak_fitting_parameters["start_index"]
 
         # special case for Mn-54, only keep the first high count energy peak
         if self.check_source.nuclide == mn54 and len(peaks) > 1:
@@ -798,6 +821,13 @@ def get_neutron_rate(
 
         return flux
 
+def _get_nuclide_peak_kwargs(measurement: CheckSourceMeasurement, peak_kwargs: dict = None) -> dict:
+    kwargs = {}
+    if peak_kwargs is not None:
+        if measurement.check_source.nuclide.name in peak_kwargs.keys():
+            kwargs = peak_kwargs[measurement.check_source.nuclide.name]
+    return kwargs
+
 
 def get_calibration_data(
     check_source_measurements: List[CheckSourceMeasurement],
@@ -838,11 +868,9 @@ def get_calibration_data(
             hist, bin_edges = detector.get_energy_hist_background_substract(
                 background_detector, bins=None
             )
-            kwargs = {}
-            if peak_kwargs is not None:
-                if measurement.check_source.nuclide in peak_kwargs.keys():
-                    kwargs = peak_kwargs[measurement.check_source.nuclide]
-                    found_a_nuclide = True
+            kwargs = _get_nuclide_peak_kwargs(measurement, peak_kwargs)
+            if kwargs:
+                found_a_nuclide = True
 
             peaks_ind = measurement.get_peaks(hist, **kwargs)
             peaks = bin_edges[peaks_ind]

test/neutron_detection/test_compass.py

@@ -33,6 +33,7 @@ def test_get_channel(filename, expected_channel):
     assert ch == expected_channel
 
 
+
 def create_empty_csv_files(directory, base_name, count, channel):
     """
     Creates empty CSV files in a specified directory with a specific pattern.
@@ -1724,3 +1725,102 @@ def test_get_peaks(detector_type, nuclide, signal_to_background_ratio):
     for test_energy, expected_energy in zip(test_peaks, nuclide.energy):
         print(f"Detected peak at {test_energy:.2f} keV, expected at {expected_energy} keV")
         assert np.isclose(test_energy, expected_energy, rtol=0.05)
+
+
+@pytest.mark.parametrize(
+    "peak_kwargs, expected",
+    [
+        (None, {}),
+        ({na22.name: {"start_index": 123, "height": 4.5},
+          co60.name: {"start_index": 456, "height": 7.8}}, 
+          {"start_index": 123, "height": 4.5}),
+        ({co60.name: {"distance": 99}}, {}),
+    ],
+)
+def test_get_nuclide_peak_kwargs(peak_kwargs, expected):
+    measurement = compass.CheckSourceMeasurement("test")
+    measurement.check_source = CheckSource(
+        nuclide=na22,
+        activity_date=datetime.datetime(2024, 1, 1),
+        activity=1.0,
+    )
+
+    result = compass._get_nuclide_peak_kwargs(measurement, peak_kwargs)
+
+    assert result == expected
+
+    
+@pytest.mark.parametrize(
+    "detector_type, nuclide",
+    [
+        ("NaI", na22),
+        ("NaI", co60),
+        ("NaI", ba133),
+        ("NaI", mn54),
+        ("HPGe", na22),
+        ("HPGe", co60),
+        ("HPGe", ba133),
+        ("HPGe", mn54),
+    ],
+)
+def test_get_peak_fitting_parameters_with_kwargs(detector_type, nuclide):
+    """Test that kwargs properly override default parameters in get_peak_fitting_parameters."""
+    
+    # Create a test histogram (random values)
+    hist = np.random.rand(1000) * 100
+    hist[500:600] = np.random.rand(100) * 1000  # Add a peak
+    
+    # Create a CheckSourceMeasurement
+    check_source = CheckSource(
+        nuclide=nuclide,
+        activity_date=datetime.datetime(2024, 1, 1),
+        activity=1.0
+    )
+    
+    measurement = compass.CheckSourceMeasurement(name="test_measurement")
+    measurement.check_source = check_source
+    measurement.detector_type = detector_type
+    
+    # Test 1: Get default parameters (no kwargs)
+    default_params = measurement.get_peak_fitting_parameters(hist)
+    
+    # Check that all expected keys are present
+    assert "start_index" in default_params
+    assert "prominence" in default_params
+    assert "height" in default_params
+    assert "width" in default_params
+    assert "distance" in default_params
+    
+    # Test 2: Override with custom kwargs
+    custom_kwargs = {
+        "start_index": 200,
+        "prominence": 50.0,
+        "height": 75.0,
+        "width": [5, 100],
+        "distance": 50,
+    }
+    
+    custom_params = measurement.get_peak_fitting_parameters(hist, **custom_kwargs)
+    
+    # Verify all kwargs were applied
+    assert custom_params["start_index"] == 200, "start_index kwarg not applied"
+    assert custom_params["prominence"] == 50.0, "prominence kwarg not applied"
+    assert custom_params["height"] == 75.0, "height kwarg not applied"
+    assert custom_params["width"] == [5, 100], "width kwarg not applied"
+    assert custom_params["distance"] == 50, "distance kwarg not applied"
+    
+    # Test 3: Partial kwargs override (some default, some custom)
+    partial_kwargs = {
+        "start_index": 300,
+        "height": 100.0,
+    }
+    
+    partial_params = measurement.get_peak_fitting_parameters(hist, **partial_kwargs)
+    
+    # Verify partial overrides work
+    assert partial_params["start_index"] == 300, "partial start_index kwarg not applied"
+    assert partial_params["height"] == 100.0, "partial height kwarg not applied"
+    # Other parameters should remain as defaults
+    assert partial_params["prominence"] == default_params["prominence"], "Other params should not change"
+    assert partial_params["width"] == default_params["width"], "Other params should not change"
+    assert partial_params["distance"] == default_params["distance"], "Other params should not change"