added an attribute to LSCReader

libra_toolbox/tritium/lsc_measurements.py

@@ -11,6 +11,8 @@
 
 
 class LSCFileReader:
+    quench_set: str
+
     def __init__(
         self,
         file_path: str,
@@ -29,6 +31,7 @@ def __init__(
         self.labels_column = labels_column
         self.data = None
         self.header_content = None
+        self.quench_set = None
 
     def read_file(self):
         """
@@ -50,12 +53,14 @@ def read_file(self):
         with open(self.file_path, "r") as file:
             lines = file.readlines()
             for i, line in enumerate(lines):
+                if line.startswith("Quench Set:"):
+                    quench_set_line_idx = i + 1
                 if line.startswith("S#"):
                     start = i
                     break
                 header_lines.append(line)
         self.header_content = "".join(header_lines)
-
+        self.quench_set = lines[quench_set_line_idx].strip()
         # read the file with dataframe starting from the line with S#
         self.data = pd.read_csv(self.file_path, skiprows=start)
 

test/tritium/test_lsc_reader.py

@@ -47,3 +47,13 @@ def test_lsc_reader_label_automated():
 
     assert len(bq1_values) == 10
     assert len(bq1_values_with_labels) == 10
+
+
+def test_read_quench_set():
+    """Tests reading the quench set from the LSC file"""
+    filename = Path(__file__).parent / "test_lsc_file_with_labels.csv"
+
+    csv_reader = LSCFileReader(filename, labels_column="SMPL_ID")
+    csv_reader.read_file()
+
+    assert csv_reader.quench_set == "Low Energy: 3H-UG"