add run ranges and special ledcal

amstrax/auto_processing_new/db_utils.py

@@ -120,16 +120,17 @@ def add_data_entry(
     else:
         log.info(f"Would add data entry to run {run_id} (dry run).")
 
-def query_runs(query):
+
+def get_run_doc(run_id):
     """
-    Query the rundb to retrieve runs based on a specific filter.
+    Get the document for a specific run from the rundb.
     
-    :param query: MongoDB query object.
-    :return: List of runs matching the query.
+    :param run_id: ID of the run to retrieve.
+    :return: Document for the run.
     """
 
     runsdb = amstrax.get_mongo_collection()
 
-    results = list(runsdb.find(query))
-    log.info(f"Found {len(results)} runs matching query.")
-    return results
+    run_doc = runsdb.find_one({'number': int(run_id)})
+    log.info(f"Found run {run_id} in the database.")
+    return run_doc

amstrax/auto_processing_new/offline_processing.py

@@ -23,6 +23,7 @@ def parse_args():
     run_selection.add_argument(
         "--run_file", type=str, help="File with run IDs to process. It should contain one run ID per line."
     )
+    run_selection.add_argument("--run_range", type=str, help="Range of run IDs to process (e.g., 1000-2000).")
 
     # Arguments for processing
     parser.add_argument(
@@ -99,27 +100,30 @@ def check_for_production(args):
             raise ValueError("Output folder not specified.")
 
 
-def main(args):
-    """
-    Main function for offline job submission of selected runs.
-    """
+def get_run_ids_from_args(args):
 
-    # Check the run selection method
+    run_ids = []
     if args.run_id:
-        run_docs = [{"number": int(run_id)} for run_id in args.run_id]
+        run_ids = args.run_id
     elif args.run_file:
-        # It should contain a list of run numbers, one per line
         with open(args.run_file, "r") as f:
-            run_numbers = f.readlines()
-        run_docs = [{"number": int(run_number)} for run_number in run_numbers]
-    else:
-        log.error("Either --run_id or --run_file must be provided.")
-        return
+            run_ids = f.readlines()
+    elif args.run_range:
+        run_range = args.run_range.split("-")
+        run_ids = list(range(int(run_range[0]), int(run_range[1]) + 1))
 
-    # Limit the number of runs to process
-    run_docs = run_docs[args.start_from : args.start_from + args.max_runs]
+    return run_ids
 
 
+def main(args):
+    """
+    Main function for offline job submission of selected runs.
+    """
+
+    # Get the run IDs from the arguments
+    run_docs = get_run_ids_from_args(args)
+    # Limit the number of runs to process
+    run_docs = run_docs[args.start_from : args.start_from + args.max_runs]
     log.info(f"We are about to process {len(run_docs)} runs.")
 
     # Submit jobs for each run

amstrax/auto_processing_new/process.py

@@ -38,6 +38,9 @@ def __init__(self, args):
 
         self.setup_amstrax()
         self.setup_production()
+        self.run_doc = self.db_utils.get_run_doc(self.run_id)
+        self.infer_special_modes()
+
 
     def setup_amstrax(self):
 
@@ -52,6 +55,39 @@ def setup_amstrax(self):
 
         self.db_utils = self.amstrax.db_utils
 
+    def get_run_doc_info(self):
+
+        # let's just print out the main info from the rundoc
+        # like mode, start, end, duration, user, comments, tags
+        # we'll need to format the date and time properly, and calculate the duration
+        # use datetime to convert the date to a readable format
+        log.info(f"Run document for run {self.run_id}:")
+        log.info(f" *** Mode: {self.run_doc.get('mode')}")
+        log.info(f" *** Start: {datetime.datetime.fromtimestamp(self.run_doc.get('start') / 1e9)}")
+        log.info(f" *** End: {datetime.datetime.fromtimestamp(self.run_doc.get('end') / 1e9)}")
+        log.info(f" *** Duration: {(self.run_doc.get('end') - self.run_doc.get('start'))/1e9:.2f} seconds")
+        log.info(f" *** User: {self.run_doc.get('user')}")
+        log.info(f" *** Comments: {self.run_doc.get('comments')}")
+        log.info(f" *** Tags: {self.run_doc.get('tags')}")
+
+    def infer_special_modes(self):
+
+        # Check if there is led in the run_doc
+        if "ledcalibration" in self.run_doc.get('mode'):
+            # Add the LEDCalibration plugin to the context
+            log.info("Detected LED calibration run.")
+            log.info("Adding LEDCalibration plugin to the context.")
+            ax = self.amstrax
+            self.st.register([
+                ax.DAQReader,
+                ax.RecordsLED,
+                ax.LEDCalibration
+            ])
+
+            # override the targets to process only the LEDCalibration
+            self.targets = ["raw_records", "records_led", "led_calibration"]
+
+
     def setup_production(self):
 
         if self.production: