Class structure for COMPASS measurements

libra_toolbox/neutron_detection/activation_foils/compass.py

@@ -1,10 +1,14 @@
 import numpy as np
+from numpy.typing import NDArray
 import os
 from pathlib import Path
 import pandas as pd
-from typing import Tuple, Dict
+from typing import Tuple, Dict, List, Union
 import datetime
 import uproot
+import glob
+
+import warnings
 
 
 def get_channel(filename):
@@ -122,7 +126,9 @@
         with open(info_file, "r") as file:
             lines = file.readlines()
     else:
-        raise FileNotFoundError(f"Could not find run.info file in {directory}")
+        raise FileNotFoundError(
+            f"Could not find run.info file in parent directory {Path(directory).parent}"
+        )
 
     start_time, stop_time = None, None
     for line in lines:
@@ -141,16 +147,147 @@
         return start_time, stop_time
 
 
-def get_live_time_from_root(filename, channel: int):
+def get_live_time_from_root(root_filename: str, channel: int) -> Tuple[float, float]:
     """
     Gets live and real count time from Compass root file.
-    Times are given in seconds.
     Live time is defined as the difference between the actual time that
     a count is occurring and the "dead time," in which the output of detector
-    pulses is saturated such that additional signals cannot be processed.
-    """
+    pulses is saturated such that additional signals cannot be processed."""
 
-    with uproot.open(filename) as root_file:
+    with uproot.open(root_filename) as root_file:
         live_count_time = root_file[f"LiveTime_{channel}"].members["fMilliSec"] / 1000
         real_count_time = root_file[f"RealTime_{channel}"].members["fMilliSec"] / 1000
     return live_count_time, real_count_time
+
+
+class Detector:
+    events: NDArray[Tuple[float, float]]  # type: ignore # Array of (time in ps, energy) pairs
+    channel_nb: int
+    live_count_time: float
+    real_count_time: float
+
+    def __init__(self, channel_nb) -> None:
+        """
+        Initialize a Detector object.
+        Args:
+            channel_nb: channel number of the detector
+        """
+        self.channel_nb = channel_nb
+        self.events = np.empty((0, 2))  # Initialize as empty 2D array with 2 columns
+        self.live_count_time = None
+        self.real_count_time = None
+
+    def get_energy_hist(
+        self, bins: Union[int, NDArray[np.float64], None] = None
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        Get the energy histogram of the detector events.
+        Args:
+            bins: number of bins, can be a numpy array, if None, it will be set to the
+                maximum energy value in the events (one bin per energy value)
+        Returns:
+            Tuple of histogram values and bin edges
+        """
+
+        energy_values = self.events[:, 1].copy()
+        time_values = self.events[:, 0].copy()
+
+        # sort data based on timestamp
+        inds = np.argsort(time_values)
+        time_values = time_values[inds]
+        energy_values = energy_values[inds]
+
+        energy_values = np.nan_to_num(energy_values, nan=0)
+
+        if bins is None:
+            bins = int(np.nanmax(energy_values))
+
+        return np.histogram(energy_values, bins=bins)
+
+
+class Measurement:
+    start_time: datetime.datetime
+    stop_time: datetime.datetime
+    name: str
+    detectors: List[Detector]
+
+    def __init__(self, name: str) -> None:
+        """
+        Initialize a Measurement object.
+        Args:
+            name: name of the measurement
+        """
+        self.start_time = None
+        self.stop_time = None
+        self.name = name
+        self.detectors = []
+
+    @classmethod
+    def from_directory(
+        cls, source_dir: str, name: str, info_file_optional: bool = False
+    ) -> "Measurement":
+        """
+        Create a Measurement object from a directory containing Compass data.
+        Args:
+            source_dir: directory containing Compass data
+            name: name of the measurement
+            info_file_optional: if True, the function will not raise an error
+                if the run.info file is not found
+        Returns:
+            Measurement object
+        """
+        measurement_object = cls(name=name)
+
+        # Get events
+        time_values, energy_values = get_events(source_dir)
+
+        # Get start and stop time
+        try:
+            start_time, stop_time = get_start_stop_time(source_dir)
+            measurement_object.start_time = start_time
+            measurement_object.stop_time = stop_time
+        except FileNotFoundError as e:
+            if info_file_optional:
+                warnings.warn(
+                    "run.info file not found. Assuming start and stop time are not needed."
+                )
+            else:
+                raise FileNotFoundError(e)
+
+        # Create detectors
+        detectors = [Detector(channel_nb=nb) for nb in time_values.keys()]
+
+        # Get live and real count times
+        all_root_filenames = glob.glob(os.path.join(source_dir, "*.root"))
+        if len(all_root_filenames) == 1:
+            root_filename = all_root_filenames[0]
+        else:
+            root_filename = None
+            print("No root file found, assuming all counts are live")
+
+        for detector in detectors:
+            detector.events = np.column_stack(
+                (time_values[detector.channel_nb], energy_values[detector.channel_nb])
+            )
+
+            if root_filename:
+                live_count_time, real_count_time = get_live_time_from_root(
+                    root_filename, detector.channel_nb
+                )
+                detector.live_count_time = live_count_time
+                detector.real_count_time = real_count_time
+            else:
+                real_count_time = (stop_time - start_time).total_seconds()
+                # Assume first and last event correspond to start and stop time of live counts
+                # and convert from picoseconds to seconds
+                ps_to_seconds = 1e-12
+                live_count_time = (
+                    time_values[detector.channel_nb][-1]
+                    - time_values[detector.channel_nb][0]
+                ) * ps_to_seconds
+                detector.live_count_time = live_count_time
+                detector.real_count_time = real_count_time
+
+        measurement_object.detectors = detectors
+
+        return measurement_object

test/neutron_detection/compass_test_data/complete_measurement/data/Data_CH1@V1725_292_Co60_0_872uCi_19Mar2014_250318_run2.csv

@@ -0,0 +1,11 @@
+BOARD;CHANNEL;TIMETAG;ENERGY;ENERGYSHORT;FLAGS
+0;5;234859459;2;2;0x4000
+0;5;421999310;0;1;0x4000
+0;5;535148093;1237;810;0x4000
+0;5;1623550122;589;396;0x4000
+0;5;5997211248;375;251;0x4000
+0;5;6685836624;515;340;0x4000
+0;5;11116032249;568;380;0x4000
+0;5;11281099382;1;0;0x4000
+0;5;12783039350;5;0;0x4000
+0;5;18306299412;2;0;0x4000

test/neutron_detection/compass_test_data/complete_measurement/data/Data_CH1@V1725_292_Co60_0_872uCi_19Mar2014_250318_run2_1.csv

@@ -0,0 +1,4 @@
+0;5;234859459;2;2;0x4000
+0;5;421999310;0;1;0x4000
+0;5;535148093;1237;810;0x4000
+0;5;1623550122;589;396;0x4000

test/neutron_detection/compass_test_data/complete_measurement/run.info

@@ -0,0 +1,31 @@
+id=Co60_0_872uCi_19Marc2014_250319_run3
+time.start=2025/03/19 09:47:46.724-0400
+time.stop=2025/03/19 09:53:05.651-0400
+time.real=00:05:18
+board.0-14-292.readout.rate=51.149 kb/s
+board.0-14-292.4.rejections.singles=0.0
+board.0-14-292.4.rejections.pileup=0.0
+board.0-14-292.4.rejections.saturation=1301.87
+board.0-14-292.4.rejections.energy=0.0
+board.0-14-292.4.rejections.psd=0.0
+board.0-14-292.4.rejections.timedistribution=0.0
+board.0-14-292.4.throughput=2627.31
+board.0-14-292.4.icr=3059.24
+board.0-14-292.4.ocr=1293.91
+board.0-14-292.4.calibration.energy.c0=0.0
+board.0-14-292.4.calibration.energy.c1=1.0
+board.0-14-292.4.calibration.energy.c2=0.0
+board.0-14-292.4.calibration.energy.uom=keV
+board.0-14-292.5.rejections.singles=0.0
+board.0-14-292.5.rejections.pileup=0.0
+board.0-14-292.5.rejections.saturation=717.247
+board.0-14-292.5.rejections.energy=0.0
+board.0-14-292.5.rejections.psd=0.0
+board.0-14-292.5.rejections.timedistribution=0.0
+board.0-14-292.5.throughput=1694.65
+board.0-14-292.5.icr=1703.71
+board.0-14-292.5.ocr=984.476
+board.0-14-292.5.calibration.energy.c0=0.0
+board.0-14-292.5.calibration.energy.c1=1.0
+board.0-14-292.5.calibration.energy.c2=0.0
+board.0-14-292.5.calibration.energy.uom=keV

test/neutron_detection/compass_test_data/complete_measurement_no_root/data/Data_CH1@V1725_292_Co60_0_872uCi_19Mar2014_250318_run2.csv

@@ -0,0 +1,11 @@
+BOARD;CHANNEL;TIMETAG;ENERGY;ENERGYSHORT;FLAGS
+0;5;234859459;2;2;0x4000
+0;5;421999310;0;1;0x4000
+0;5;535148093;1237;810;0x4000
+0;5;1623550122;589;396;0x4000
+0;5;5997211248;375;251;0x4000
+0;5;6685836624;515;340;0x4000
+0;5;11116032249;568;380;0x4000
+0;5;11281099382;1;0;0x4000
+0;5;12783039350;5;0;0x4000
+0;5;18306299412;2;0;0x4000

test/neutron_detection/compass_test_data/complete_measurement_no_root/data/Data_CH1@V1725_292_Co60_0_872uCi_19Mar2014_250318_run2_1.csv

@@ -0,0 +1,4 @@
+0;5;234859459;2;2;0x4000
+0;5;421999310;0;1;0x4000
+0;5;535148093;1237;810;0x4000
+0;5;1623550122;589;396;0x4000

test/neutron_detection/compass_test_data/complete_measurement_no_root/run.info

@@ -0,0 +1,31 @@
+id=Co60_0_872uCi_19Marc2014_250319_run3
+time.start=2025/03/19 09:47:46.724-0400
+time.stop=2025/03/19 09:53:05.651-0400
+time.real=00:05:18
+board.0-14-292.readout.rate=51.149 kb/s
+board.0-14-292.4.rejections.singles=0.0
+board.0-14-292.4.rejections.pileup=0.0
+board.0-14-292.4.rejections.saturation=1301.87
+board.0-14-292.4.rejections.energy=0.0
+board.0-14-292.4.rejections.psd=0.0
+board.0-14-292.4.rejections.timedistribution=0.0
+board.0-14-292.4.throughput=2627.31
+board.0-14-292.4.icr=3059.24
+board.0-14-292.4.ocr=1293.91
+board.0-14-292.4.calibration.energy.c0=0.0
+board.0-14-292.4.calibration.energy.c1=1.0
+board.0-14-292.4.calibration.energy.c2=0.0
+board.0-14-292.4.calibration.energy.uom=keV
+board.0-14-292.5.rejections.singles=0.0
+board.0-14-292.5.rejections.pileup=0.0
+board.0-14-292.5.rejections.saturation=717.247
+board.0-14-292.5.rejections.energy=0.0
+board.0-14-292.5.rejections.psd=0.0
+board.0-14-292.5.rejections.timedistribution=0.0
+board.0-14-292.5.throughput=1694.65
+board.0-14-292.5.icr=1703.71
+board.0-14-292.5.ocr=984.476
+board.0-14-292.5.calibration.energy.c0=0.0
+board.0-14-292.5.calibration.energy.c1=1.0
+board.0-14-292.5.calibration.energy.c2=0.0
+board.0-14-292.5.calibration.energy.uom=keV

test/neutron_detection/test_compass.py

@@ -275,3 +275,59 @@ def test_get_live_time_from_root(root_filename, channel, live_time, real_time):
     )
     assert live_time_out == live_time
     assert real_time_out == real_time
+
+
+@pytest.mark.parametrize("no_root", [True, False])
+def test_measurement_object_from_directory(no_root):
+    """
+    Test the Measurement object creation from a directory.
+    """
+    if no_root:
+        test_directory = (
+            Path(__file__).parent
+            / "compass_test_data/complete_measurement_no_root/data"
+        )
+    else:
+        test_directory = (
+            Path(__file__).parent / "compass_test_data/complete_measurement/data"
+        )
+
+    measurement = compass.Measurement.from_directory(test_directory, name="test")
+
+    assert len(measurement.detectors) == 1
+    assert isinstance(measurement.detectors[0], compass.Detector)
+    assert measurement.detectors[0].channel_nb == 1
+
+    assert measurement.detectors[0].events.shape[1] == 2
+
+    measurement.detectors[0].get_energy_hist(bins=None)
+
+
+@pytest.mark.parametrize(
+    "bins",
+    [
+        10,
+        20,
+        50,
+        100,
+        None,
+        np.arange(0, 10, 1),
+        np.linspace(0, 10, num=100),
+    ],
+)
+def test_detector_get_energy_hist(bins):
+    """
+    Test the get_energy_hist method of the Detector class.
+    """
+    my_detector = compass.Detector(channel_nb=1)
+    my_detector.events = np.array(
+        [
+            [1, 2],
+            [3, 4],
+            [5, 6],
+            [7, 8],
+            [9, 10],
+        ]
+    )
+
+    my_detector.get_energy_hist(bins=bins)