Fix get_events() for CoMPASS csv files with waveforms

libra_toolbox/neutron_detection/activation_foils/compass.py

@@ -938,26 +938,50 @@ def get_events(directory: str) -> Tuple[Dict[int, np.ndarray], Dict[int, np.ndar
         time_values[ch] = np.empty(0)
         energy_values[ch] = np.empty(0)
         for i, filename in enumerate(data_filenames[ch]):
+            print(f'Processing File {i}')
+
+            csv_file_path = os.path.join(directory, filename)
 
             # only the first file has a header
             if i == 0:
-                header = 0
+                # determine the column names
+                # 
+                # Typically, setting the header argument to 1
+                # would normally work, but on some CoMPASS csv
+                # files, specifically those with waveform data,
+                # the column header has far fewer entries
+                # than the number of columns in the csv file.
+                # This is due to the "SAMPLES" column, which 
+                # contains the waveform data actually being made
+                # up of the 7th-nth column of an n column csv file.
+                #
+                # So to mitigate this, we will read in the header
+                # manually and determine which column of 
+                # the dataset to read in. 
+                first_row_df = pd.read_csv(csv_file_path,
+                                           delimiter=";",
+                                           header=None,
+                                           nrows=1)
+                column_names = first_row_df.to_numpy()[0]
+                # Determine which column applies to time and energy
+                time_col = np.where(column_names=="TIMETAG")[0][0]
+                energy_col = np.where(column_names=="ENERGY")[0][0]
+                # First csv file has header, so skip it
+                # because we already read it in
+                skiprows=1
             else:
-                header = None
-
-            csv_file_path = os.path.join(directory, filename)
+                # For subsequent csv files, don't skip any rows
+                # as there won't be any header
+                skiprows=0
 
-            df = pd.read_csv(csv_file_path, delimiter=";", header=header)
 
-            # read the header and store in names
-            if i == 0:
-                names = df.columns.values
-            else:
-                # apply the column names if not the first file
-                df.columns = names
+            df = pd.read_csv(csv_file_path, 
+                             delimiter=";", 
+                             header=None,
+                             skiprows=skiprows)
 
-            time_data = df["TIMETAG"].to_numpy()
-            energy_data = df["ENERGY"].to_numpy()
+            time_data = df[time_col].to_numpy()
+            energy_data = df[energy_col].to_numpy()
 
             # Extract and append the energy data to the list
             time_values[ch] = np.concatenate([time_values[ch], time_data])

test/neutron_detection/compass_test_data/events/waveforms/Data_CH4@DT5725_1360_Co60_0_872uCi_19Mar2014_20241213_4in.CSV

@@ -0,0 +1,7 @@
+BOARD;CHANNEL;TIMETAG;ENERGY;ENERGYSHORT;FLAGS;PROBE_CODE;SAMPLES
+0;4;80413091;1727;1407;0x4000;1;13153;13152;13149;13150;13156;13146;13154;13150;13153;13152;13152;13152;13148;13153;13153;13156;13157;13145;13154;13143;13151;13134;13119;13074;12992;12926;12834;12773;12707;12646;12603;12544;12520;12461;12438;12397;12383;12374;12388;12413;12404;12381;12386;12390;12361;12303;12308;12361;12379;12373;12394;12380;12420;12470;12494;12555;12589;12586;12614;12605;12597;12586;12602;12599;12594;12604;12637;12683;12688;12661;12646;12664;12692;12688;12669;12673;12697;12709;12720;12700;12706;12719;12704;12699;12709;12722;12724;12744;12773;12786;12798;12776;12768;12796;12804;12803;12801;12810;12818;12812;12800;12815;12803;12818;12842;12821;12818;12823;12856;12874;12859;12880;12905;12897;12899;12900;12909;12917;12916;12922;12924;12930;12956;12954;12963;12967;12964;12970;12971;12956;12953;12970;12987;12982;12991;12988;12981;12989;13004;12999;13013;12998;13004;13020;13018;13020;13018;13020;13030;13029;13046;13038;13035;13045;13037;13048;13041;13055;13051;13035;13045;13039;13050;13066;13066;13065;13069;13059;13062;13052;13058;13069;13069;13081;13073;13084;13083;13064;13084;13090;13087;13080;13078;13084;13081;13085;13099;13079;13082;13087;13088;13093;13087;13080;13067;13071;13072;13080;13091;13081
+0;4;849882747;613;499;0x4000;1;13147;13157;13147;13158;13155;13152;13153;13155;13153;13150;13157;13155;13152;13162;13144;13149;13159;13146;13153;13145;13143;13137;13092;13070;13043;13012;12989;12954;12935;12928;12933;12932;12921;12923;12915;12930;12932;12924;12917;12916;12915;12922;12915;12911;12907;12905;12885;12887;12910;12906;12910;12915;12895;12879;12873;12894;12904;12920;12940;12934;12943;12957;12956;12970;12965;12966;12963;12969;12971;12982;12978;12963;12988;13011;13003;12994;12993;12991;13009;13011;13012;13021;13020;13015;13010;12997;13006;13008;13016;13005;12999;13011;13007;13012;13018;13027;13029;13043;13058;13044;13044;13030;13035;13050;13041;13033;13028;13026;13036;13050;13043;13058;13068;13062;13063;13070;13069;13052;13064;13072;13073;13068;13070;13089;13087;13091;13086;13096;13094;13084;13093;13084;13084;13101;13112;13100;13088;13093;13090;13104;13099;13094;13100;13103;13095;13095;13099;13110;13103;13120;13109;13113;13112;13119;13113;13119;13125;13121;13130;13114;13126;13122;13117;13124;13108;13123;13120;13114;13112;13120;13125;13107;13122;13121;13122;13118;13117;13124;13126;13131;13129;13127;13117;13130;13127;13122;13135;13132;13137;13129;13123;13124;13127;13138;13122;13133;13134;13138;13147;13139
+0;4;2850906749;1539;1239;0x4000;1;13155;13154;13152;13155;13156;13148;13155;13155;13146;13160;13150;13148;13158;13159;13153;13148;13158;13150;13154;13148;13146;13138;13120;13082;13038;12986;12900;12826;12739;12676;12627;12574;12562;12541;12517;12502;12475;12464;12472;12489;12501;12512;12520;12530;12525;12510;12462;12460;12463;12482;12502;12505;12525;12550;12555;12566;12561;12577;12561;12574;12623;12657;12659;12661;12658;12673;12675;12659;12659;12674;12695;12701;12691;12684;12698;12727;12736;12747;12753;12769;12765;12749;12738;12774;12811;12814;12832;12847;12850;12864;12845;12841;12833;12832;12838;12838;12846;12861;12860;12865;12862;12851;12861;12894;12886;12900;12910;12913;12908;12915;12924;12920;12918;12912;12924;12923;12896;12911;12921;12937;12926;12922;12935;12961;12964;12972;12983;12999;12988;12975;12991;13012;13022;12998;13007;13004;12990;12997;13008;13002;13010;13029;13014;13018;13025;13039;13027;13029;13023;13019;13026;13040;13061;13059;13056;13051;13048;13057;13055;13059;13063;13062;13063;13060;13053;13054;13062;13080;13084;13087;13084;13073;13070;13079;13072;13087;13089;13094;13095;13086;13090;13081;13074;13077;13078;13088;13087;13105;13099;13106;13098;13101;13100;13086;13100;13101;13087;13090;13081;13090
+0;4;5758064121;1563;1258;0x4000;1;13155;13146;13148;13158;13143;13152;13146;13153;13157;13148;13159;13145;13154;13154;13148;13155;13148;13153;13154;13148;13145;13136;13105;13054;12994;12928;12847;12781;12722;12681;12621;12557;12514;12459;12437;12464;12459;12437;12427;12437;12421;12399;12408;12428;12442;12460;12476;12483;12496;12516;12546;12540;12528;12561;12611;12620;12616;12624;12615;12617;12637;12659;12665;12651;12659;12666;12649;12680;12661;12656;12686;12676;12682;12701;12745;12774;12779;12786;12771;12774;12766;12744;12749;12776;12791;12795;12801;12792;12806;12815;12817;12804;12809;12853;12848;12869;12874;12877;12883;12882;12911;12920;12918;12897;12892;12908;12894;12895;12877;12864;12906;12909;12927;12925;12931;12937;12930;12935;12934;12952;12957;12945;12940;12935;12924;12949;12949;12937;12952;12957;12979;12995;13003;13016;13015;13020;13013;13034;13014;13030;13033;13024;13040;13032;13016;13020;13013;13026;13020;13027;13038;13044;13037;13034;13040;13036;13040;13039;13034;13049;13051;13057;13064;13067;13065;13075;13073;13083;13072;13071;13083;13073;13079;13080;13074;13078;13072;13088;13078;13075;13081;13069;13079;13078;13062;13084;13071;13068;13080;13073;13094;13089;13085;13093;13093;13095;13097;13091;13095;13102
+0;4;6286463248;246;204;0x4000;1;13153;13153;13152;13158;13152;13153;13153;13155;13155;13151;13155;13155;13150;13147;13137;13111;13108;13101;13091;13088;13075;13083;13078;13064;13050;13036;13032;13030;13047;13046;13054;13061;13052;13066;13056;13062;13058;13045;13058;13064;13074;13070;13055;13064;13083;13066;13077;13068;13069;13077;13068;13089;13083;13074;13084;13080;13094;13085;13081;13072;13067;13072;13078;13086;13097;13098;13088;13075;13068;13065;13088;13080;13088;13099;13090;13098;13092;13084;13086;13073;13066;13073;13073;13092;13100;13102;13113;13098;13102;13103;13105;13108;13108;13117;13115;13120;13118;13125;13133;13116;13125;13120;13122;13132;13132;13128;13116;13113;13114;13111;13109;13123;13121;13118;13127;13122;13133;13123;13127;13130;13130;13140;13124;13125;13108;13131;13137;13138;13121;13108;13106;13111;13133;13139;13137;13151;13123;13125;13125;13124;13148;13144;13151;13131;13126;13134;13129;13137;13135;13146;13135;13139;13134;13135;13142;13141;13151;13143;13144;13139;13138;13141;13142;13146;13136;13150;13146;13147;13143;13134;13146;13142;13146;13146;13135;13150;13136;13141;13144;13136;13147;13142;13134;13147;13142;13137;13150;13140;13143;13157;13144;13157;13140;13141;13144;13136;13149;13145;13142;13142
+0;4;6518702279;1724;1404;0x4000;1;13155;13156;13159;13152;13161;13143;13161;13153;13150;13160;13150;13160;13154;13154;13145;13156;13158;13152;13158;13144;13149;13138;13118;13097;13020;12943;12839;12764;12694;12617;12574;12514;12456;12410;12422;12445;12424;12406;12396;12408;12408;12410;12414;12412;12434;12442;12465;12452;12489;12528;12550;12563;12541;12533;12504;12489;12493;12486;12493;12486;12505;12518;12541;12544;12550;12575;12596;12622;12639;12617;12604;12638;12666;12663;12674;12650;12629;12641;12678;12700;12692;12696;12705;12689;12696;12697;12692;12701;12730;12747;12773;12770;12785;12803;12802;12794;12801;12791;12793;12814;12851;12850;12840;12841;12856;12868;12887;12893;12900;12909;12909;12908;12889;12896;12898;12921;12933;12938;12940;12904;12905;12921;12946;12963;12959;12965;12969;12960;12975;12979;12996;13004;13008;13006;13002;13015;13024;13013;13013;13001;12996;13011;13000;13005;13013;13006;13018;13003;13011;13004;13021;13043;13037;13034;13038;13046;13048;13045;13049;13052;13059;13052;13050;13053;13057;13055;13051;13039;13038;13030;13055;13062;13051;13067;13084;13077;13083;13078;13078;13083;13081;13094;13085;13094;13100;13086;13080;13078;13078;13080;13094;13092;13096;13096;13104;13102;13105;13098;13102;13098

test/neutron_detection/compass_test_data/events/waveforms/Data_CH4@DT5725_1360_Co60_0_872uCi_19Mar2014_20241213_4in_1.CSV

@@ -0,0 +1,2 @@
+0;4;14300873206559;1700;1364;0x4000;1;13161;13144;13160;13152;13155;13156;13155;13155;13148;13152;13154;13152;13150;13150;13158;13150;13151;13156;13154;13149;13149;13135;13142;13093;13029;12969;12897;12819;12739;12674;12615;12580;12585;12553;12495;12494;12484;12475;12453;12427;12431;12419;12411;12396;12413;12441;12430;12422;12429;12442;12444;12449;12470;12497;12496;12483;12490;12497;12501;12535;12547;12530;12539;12540;12562;12574;12608;12622;12629;12650;12640;12631;12645;12661;12657;12661;12667;12668;12678;12673;12669;12675;12700;12709;12728;12746;12754;12768;12777;12801;12823;12831;12844;12810;12805;12817;12830;12851;12852;12858;12850;12843;12845;12851;12854;12873;12893;12889;12878;12869;12887;12876;12876;12896;12899;12923;12919;12906;12891;12881;12885;12901;12924;12924;12936;12920;12931;12946;12939;12942;12949;12946;12978;12950;12959;12970;12968;12980;12986;12994;13013;13029;13029;13034;13037;13040;13020;13004;13007;13007;13021;13018;13013;13019;13040;13029;13037;13047;13049;13061;13048;13044;13052;13060;13068;13074;13065;13067;13077;13076;13086;13065;13079;13081;13078;13073;13067;13082;13080;13078;13082;13082;13080;13084;13087;13087;13094;13087;13078;13078;13091;13085;13096;13096;13104;13105;13102;13102;13107;13092
+0;4;14301010164183;1498;1202;0x4000;1;13153;13154;13153;13156;13151;13151;13155;13154;13154;13148;13162;13145;13159;13150;13152;13155;13155;13151;13145;13145;13145;13142;13109;13054;13012;12935;12863;12794;12736;12692;12622;12582;12553;12532;12539;12529;12514;12501;12492;12485;12474;12446;12463;12478;12496;12530;12561;12579;12548;12542;12559;12574;12592;12585;12584;12605;12628;12636;12628;12634;12648;12665;12674;12678;12696;12675;12671;12697;12689;12693;12722;12750;12750;12741;12743;12735;12756;12783;12803;12782;12775;12772;12775;12772;12772;12784;12786;12795;12810;12803;12814;12838;12852;12849;12857;12848;12845;12855;12879;12883;12891;12891;12890;12887;12887;12905;12908;12915;12912;12917;12925;12923;12947;12934;12942;12946;12939;12941;12941;12939;12931;12940;12954;12973;12970;12966;12976;12957;12966;12985;12993;12997;12997;13000;13008;13008;13005;13008;13004;13001;13009;13008;13018;13023;13030;13049;13030;13045;13041;13037;13033;13034;13035;13028;13026;13044;13064;13061;13056;13055;13068;13076;13077;13077;13064;13059;13069;13070;13070;13081;13088;13080;13073;13063;13069;13068;13075;13071;13079;13087;13085;13090;13097;13096;13088;13080;13084;13087;13093;13094;13097;13100;13108;13107;13100;13106;13099;13116;13117;13120

test/neutron_detection/test_compass.py

@@ -97,32 +97,35 @@ def test_sort_compass_files(tmpdir, base_name: str, expected_filenames: dict):
 
 
 @pytest.mark.parametrize(
-    "expected_time, expected_energy, expected_idx",
+    "waveform_directory, expected_time, expected_energy, expected_idx, expected_keys, test_ch",
     [
-        (6685836624, 515, 5),
-        (11116032249, 568, 6),
-        (1623550122, 589, -1),
-        (535148093, 1237, -2),
+        ("no_waveforms", 6685836624, 515, 5, [5, 15], 5),
+        ("no_waveforms", 11116032249, 568, 6, [5, 15], 5),
+        ("no_waveforms", 1623550122, 589, -1, [5, 15], 5),
+        ("no_waveforms", 535148093, 1237, -2, [5, 15], 5),
+        ("waveforms", 80413091, 1727, 0, [4], 4),
+        ("waveforms", 2850906749, 1539, 2, [4], 4),
+        ("waveforms", 14300873206559, 1700, 6, [4], 4)
     ],
 )
-def test_get_events(expected_time, expected_energy, expected_idx):
+def test_get_events(waveform_directory, expected_time, 
+                    expected_energy, expected_idx,
+                    expected_keys, test_ch):
     """
     Test the get_events function from the compass module.
     Checks that specific time and energy values are returned for a given channel
     """
-    test_directory = Path(__file__).parent / "compass_test_data/events"
+    test_directory = Path(__file__).parent / "compass_test_data/events" / waveform_directory
     times, energies = compass.get_events(test_directory)
     assert isinstance(times, dict)
     assert isinstance(energies, dict)
 
-    expected_keys = [5, 15]
     for key in expected_keys:
         assert key in times
         assert key in energies
 
-    ch = 5
-    assert times[ch][expected_idx] == expected_time
-    assert energies[ch][expected_idx] == expected_energy
+    assert times[test_ch][expected_idx] == expected_time
+    assert energies[test_ch][expected_idx] == expected_energy
 
 
 utc_minus5 = datetime.timezone(datetime.timedelta(hours=-5))