Clean and refactor amstrax for XAMS v2.0 (#234)

* now I see peaks * remove straxen dependencies * fix import numba and remove print * fix peaks function * remove interpolation maps * daqreader fixed * refactor plugins * it was carlo.. * fixed records * fix typo * cleaned peak processing * added led calibration * led calibration * line endings and remove unused files * remove deleted plugins from context * worked on autoprocessing * autoprocessing * autoprocess chaos * clean up processing * clean carlo dep * add print statement * small changes led_cal * clean more * clean more * led cal fixed * remove paths * new plugin records led * contexts register * avoid mix config * typo parenthesis * remove path autoprocessing * add targets autoprocessing * add bin * refactor peaks and bin * refactor events * contexts * autoprocess single * bug * handle single run processing * handle single run processing * handle single run processing * handle single run processing * handle single run processing * handle single run processing * handle single run processing * small fix with maricke * remove straxen * remove straxen * remove straxen * fix requirements * fix requirements * fix requirements * init rundb context * bad path to pass tests.. * events basics are todo * tests dont work on github * docs xams little go --------- Authored-by: Carlo Fuselli <[email protected]>
XAMS-nikhef · Aug 29, 2023 · d3669d5 · d3669d5
1 parent 383a0fb
commit d3669d5
Show file tree

Hide file tree

Showing 48 changed files with 1,562 additions and 3,745 deletions.
diff --git a/README.md b/README.md
@@ -33,3 +33,5 @@ First you'll need to get some data. For this you'll need to either run somewhere
 /data/xenon/xams is available or have data downloaded. Either strax processed data or pax raw data
 is fine. Once you have data you can try running the Tutorial notebook (although it's in Dutch), just
 make sure to change the output_folder to where you have data
+
+Test
diff --git a/amstrax/__init__.py b/amstrax/__init__.py
@@ -1,12 +1,15 @@
 __version__ = '1.2.0'
 
-from . import hitfinder_thresholds
 from . import plugins
-from .common import *
-from .hitfinder_thresholds import *
-from .itp_map import *
 from .plugins import *
+
+from .common import *
 from .rundb import *
-from .contexts import *
+
+from . import contexts
+
+from . import hitfinder_thresholds
+from .hitfinder_thresholds import *
 
 from . import auto_processing
+
diff --git a/amstrax/auto_processing/__init__.py b/amstrax/auto_processing/__init__.py
@@ -1,4 +1,4 @@
 from . import amstraxer
+from . import amstraxer_easy
 from . import submit_stbc
 from . import auto_processing
-from . import process_run
diff --git a/amstrax/auto_processing/amstraxer.py b/amstrax/auto_processing/amstraxer.py
@@ -26,15 +26,20 @@ def parse_args():
         help="ID of the run to process; usually the run name.")
     parser.add_argument(
         '--context',
-        default='xams_little',
+        default='xams',
         help="Name of context to use")
     parser.add_argument(
         '--target',
-        default='raw_records_v1724',
+        default=['raw_records'],
+        nargs="*",
         help='Target final data type to produce')
+    parser.add_argument(
+        '--output_folder',
+        default='./strax_data',
+        help='Output folder for context')
     parser.add_argument(
         '--detector',
-        default='xamsl',
+        default='xams',
         help="xamsl or xams")
     parser.add_argument(
         '--from_scratch',
@@ -94,14 +99,16 @@ def main(args):
     # to do them after argparsing (so --help is fast)
     import strax
     print(f"\tstrax {strax.__version__} at {osp.dirname(strax.__file__)}")
+
     import amstrax
     print(f"\tamstrax {amstrax.__version__} at {osp.dirname(amstrax.__file__)}")
 
+
     if args.context_kwargs:
         logging.info(f'set context kwargs {args.context_kwargs}')
-        st = getattr(amstrax.contexts, args.context)(**args.context_kwargs)
+        st = getattr(amstrax.contexts, args.context)(output_folder=args.output_folder, **args.context_kwargs)
     else:
-        st = getattr(amstrax.contexts, args.context)()
+        st = getattr(amstrax.contexts, args.context)(output_folder=args.output_folder)
 
     if args.config_kwargs:
         logging.info(f'set context options to {args.config_kwargs}')
@@ -118,10 +125,12 @@ def main(args):
         testing_rd = args.testing_rundoc
         if testing_rd is not None:
             testing_rd['start'] = datetime.datetime.now()
+        print("I'll try to create raw_records")
         st = amstrax.contexts.context_for_daq_reader(st,
                                                      args.run_id,
                                                      run_doc=testing_rd,
-                                                     detector=args.detector)
+                                                     detector=args.detector,
+                                                     check_exists=False)
 
     if args.from_scratch:
         for q in st.storage:
@@ -130,6 +139,9 @@ def main(args):
             strax.DataDirectory('./strax_data',
                                 overwrite='always',
                                 provide_run_metadata=False))
+
+    print(st._plugin_class_registry)
+
     if st.is_stored(args.run_id, args.target):
         print("This data is already available.")
         return 1

diff --git a/amstrax/auto_processing/amstraxer_easy.py b/amstrax/auto_processing/amstraxer_easy.py
@@ -0,0 +1,108 @@
+#!/usr/bin/env python
+import sys, os
+import amstrax
+import strax
+import numpy as np
+import pandas as pd
+import datetime
+import argparse
+import json
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description='Process a single run with amstrax',
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument(
+        'run_id',
+        metavar='RUN_ID',
+        type=str,
+        help="ID of the run to process; usually the run name.")
+    parser.add_argument(
+        '--context',
+        default='xams',
+        help="Name of context to use")
+    parser.add_argument(
+        '--target',
+        default=['raw_records'],
+        nargs="*",
+        help='Target final data type to produce')
+    parser.add_argument(
+        '--output_folder',
+        default='./strax_data',
+        help='Output folder for context')
+    parser.add_argument(
+        '--detector',
+        default='xams',
+        help="xamsl or xams")
+    parser.add_argument(
+        '--from_scratch',
+        action='store_true',
+        help='Start processing at raw_records, regardless of what data is available. '
+             'Saving will ONLY occur to ./strax_data! If you already have the target'
+             'data in ./strax_data, you need to delete it there first.')
+    parser.add_argument(
+        '--config_kwargs',
+        type=json.loads,
+        help='Use a json-dict to set the context to. For example:'
+             '--config_kwargs '
+             '\'{'
+             '"allow_multiprocess": true, '
+             '"max_messages":4, '
+             '"allow_shm": true, '
+             '"allow_lazy": true}\''
+    )
+    parser.add_argument(
+        '--testing_rundoc',
+        type=json.loads,
+        help='This is only used for testing, do not use!'
+    )
+    parser.add_argument(
+        '--context_kwargs',
+        type=json.loads,
+        help='Use a json-file to load the context with (see config_kwargs for JSON-example)')
+    parser.add_argument(
+        '--timeout',
+        default=None, type=int,
+        help="Strax' internal mailbox timeout in seconds")
+    parser.add_argument(
+        '--workers',
+        default=1, type=int,
+        help=("Number of worker threads/processes. "
+              "Strax will multithread (1/plugin) even if you set this to 1."))
+    parser.add_argument(
+        '--debug',
+        action='store_true',
+        help="Enable debug logging to stdout")
+    parser.add_argument(
+        '--build_lowlevel',
+        action='store_true',
+        help='Build low-level data even if the context forbids it.')
+    return parser.parse_args()
+
+
+
+def main(args):
+
+    run_id = args.run_id
+    client = amstrax.get_mongo_client()
+    run_col = amstrax.get_mongo_collection('xams')
+    run_doc = run_col.find_one({'number': int(run_id)})
+    live_data = run_doc['daq_config']['strax_output_path']
+    output_folder = args.output_folder
+
+    st = amstrax.contexts.xams(output_folder=output_folder, init_rundb=False)
+    st.storage += [strax.DataDirectory(live_data,
+                                 provide_run_metadata=False,
+                                 deep_scan=False,
+                                 readonly=True)]
+
+    st.set_config({'live_data_dir': f'{live_data}'})
+    daqst = amstrax.contexts.context_for_daq_reader(st, run_id, 'xams', run_doc=run_doc, check_exists=False)
+
+    for t in args.target:
+        print(f'Processing {t} in amstraxer_easy')
+        daqst.make(run_id, t, progress_bar=True)
+
+if __name__ == '__main__':
+    args = parse_args()
+    sys.exit(main(args))
diff --git a/amstrax/auto_processing/auto_processing.py b/amstrax/auto_processing/auto_processing.py
@@ -1,23 +1,34 @@
 import argparse
 import time
-
+from datetime import datetime
+import subprocess 
 
 def parse_args():
     parser = argparse.ArgumentParser(
-        description='Autoprocess xams(l) data',
+        description='Autoprocess xams data',
         formatter_class=argparse.ArgumentDefaultsHelpFormatter)
     parser.add_argument(
         '--target',
-        default='raw_records_v1724',
+        nargs="*",
+        default=['peak_basics'],
         help="Target final data type to produce.")
+    parser.add_argument(
+        '--output_folder',
+        default='./strax_data',
+        help="Path where to save processed data")   
     parser.add_argument(
         '--timeout',
-        default=60,
+        default=20,
         type=int,
         help="Sleep this many seconds in between")
+    parser.add_argument(
+        '--process_stomboot',
+        default=False,
+        type=bool,
+        help="False: process locally. True: submit and process on stomboot. Default: False.")
     parser.add_argument(
         '--max_jobs',
-        default=None,
+        default=150,
         type=int,
         help="Max number of jobs to submit, if you exceed this number, break submitting new jobs")
     parser.add_argument(
@@ -26,45 +37,81 @@ def parse_args():
         help="xams_little or xams")
     parser.add_argument(
         '--detector',
-        default='xamsl',
+        default='xams',
         help="xamsl or xams")
+    parser.add_argument(
+        '--run_id', 
+        default=None,
+        help="Single run ID to process")
     return parser.parse_args()
 
 
 if __name__ == '__main__':
+
     args = parse_args()
     version = '2.1.0'
     print('Starting autoprocess version %s...' % version)
 
     # Later import to prevent slow --help
+
+    import sys, os
+
     import amstrax
-    from amstrax.auto_processing import submit_stbc
+    amstrax_dir = amstrax.amstrax_dir
 
     # settings
     nap_time = int(args.timeout)
     max_jobs = int(args.max_jobs) if args.max_jobs is not None else None
     context = args.context
+    output_folder = args.output_folder
+    process_stomboot = args.process_stomboot
     detector = args.detector
     target = args.target
     runs_col = amstrax.get_mongo_collection(detector)
 
-    while 1:
+    print('Correctly connected, starting loop')
+
+    infinite = True
+
+    while infinite:
         # Update task list
-        run_docs_to_do = list(runs_col.find({'processing_status': 'pending'}))
+        # Probably want to add here some max retry if fail
+        run_docs_to_do = list(runs_col.find({
+            'processing_status':{'$ne': 'done'},
+            'end':{"$ne":None},
+            'start':{'$gt': datetime(2023,1,25)},
+            'processing_failed':{'$not': {'$gt': 3}},
+            }).sort('start', -1))
+
+        if args.run_id is not None:
+            infinite = False
+            run_docs_to_do = [runs_col.find_one({'number': int(args.run_id)})]
+
         if len(run_docs_to_do) > 0:
             print('I found %d runs to process, time to get to work!' % len(run_docs_to_do))
             print('These runs I will do:')
             print([run_doc['number'] for run_doc in run_docs_to_do])
 
         for run_doc in run_docs_to_do[:max_jobs]:
             run_name = f'{int(run_doc["number"]):06}'
-            submit_stbc.submit_job(run_name, target=target, context=context, detector=detector,script='process_run')
-            runs_col.find_one_and_update({'number': run_name},
-                                         {'$set': {'processing_status': 'submitted_job'
-                                                  }})
+
+            if process_stomboot:
+                submit_stbc.submit_job(run_name, target=target, context=context, detector=detector,script='process_run')
+                runs_col.find_one_and_update({'number': run_name},
+                                            {'$set': {'processing_status': 'submitted_job' }})
+
+            else: #process locally
+                runs_col.find_one_and_update({'number': run_name},
+                                            {'$set': {'processing_status': 'processing'}})
+                target = " ".join(target)
+                subprocess.run(f"process_run {run_name} --target {target} --output_folder {output_folder}", shell=True)
             time.sleep(2)
+
         if max_jobs is not None and len(run_docs_to_do) > max_jobs:
             print(f'Got {len(run_docs_to_do)} which is larger than {max_jobs}, I quit!')
             break
         print("Waiting %d seconds before rechecking, press Ctrl+C to quit..." % nap_time)
         time.sleep(nap_time)
+
+
+print('Done!')
diff --git a/amstrax/auto_processing/submit_stbc.py b/amstrax/auto_processing/submit_stbc.py
@@ -64,7 +64,7 @@ def parse_args():
         help="ID of the run to process; usually the run name.")
     parser.add_argument(
         '--target',
-        default='raw_records_v1724',
+        default='raw_records',
         help="Target final data type to produce.")
     parser.add_argument(
         '--context',