From dfdf4b0de17b15fbc8464b36aa0ca4df41efbbcd Mon Sep 17 00:00:00 2001
From: Lanqing Yuan <yuanlq@uchicago.edu>
Date: Tue, 23 Apr 2024 13:15:20 -0500
Subject: [PATCH] Replaced print with logging in `contexts.py` (#125)

* log instead of print

* fixed docstr

* better annotation
---
 jobs/job.py                      | 10 ++++++----
 saltax/contexts.py               | 26 +++++++++++++-------------
 saltax/instructions/generator.py |  2 +-
 3 files changed, 20 insertions(+), 18 deletions(-)

diff --git a/jobs/job.py b/jobs/job.py
index b6aab3b..d90e49f 100644
--- a/jobs/job.py
+++ b/jobs/job.py
@@ -175,15 +175,16 @@ def main():
     print_versions()
     _, runid = sys.argv
     runid = int(runid)
-    settings = load_config()
-
+    
+    # Process the saltax desired mode
     logging.info("Loading context...")
+    settings = load_config()
     st = create_context(settings, runid)
     data_types = get_data_types(settings)
     process_data_types(st, str(runid).zfill(6), data_types)
 
     # Process data-only mode if required
-    if settings["process_data"]:
+    if settings["process_data"] and settings["saltax_mode"] == "salt":
         logging.info("====================")
         logging.info("Now starting data-only context for run %d" % runid)
         settings_temp = settings.copy()
@@ -192,7 +193,8 @@ def main():
         process_data_types(st_data, str(runid).zfill(6), data_types)
         logging.info("Finished processing for data-only mode.")
 
-    if settings["process_simu"]:
+    # Process simu-only mode if required
+    if settings["process_simu"] and settings["saltax_mode"] == "salt":
         logging.info("====================")
         logging.info("Now starting simu-only context for run %d" % runid)
         settings_temp = settings.copy()
diff --git a/saltax/contexts.py b/saltax/contexts.py
index f4c613d..94222f9 100644
--- a/saltax/contexts.py
+++ b/saltax/contexts.py
@@ -223,7 +223,7 @@ def xenonnt_salted_fuse(
         if processing_config in st.config:
             st.config[mc_config] = st.config[processing_config]
         else:
-            print(f"Warning! {processing_config} not in context config, skipping...")
+            log.warning(f"Warning! {processing_config} not in context config, skipping...")
 
     # No blinding in simulations
     st.config["event_info_function"] = "disabled"
@@ -252,12 +252,12 @@ def xenonnt_salted_fuse(
         # Try to load instruction from file and generate if not found
         try:
             instr = pd.read_csv(instr_file_name)
-            print("Loaded instructions from file", instr_file_name)
+            log.info("Loaded instructions from file", instr_file_name)
         except:
-            print(f"Instruction file {instr_file_name} not found. Generating instructions...")
+            log.info(f"Instruction file {instr_file_name} not found. Generating instructions...")
             instr = generator_func(runid=runid, **kwargs)
             pd.DataFrame(instr).to_csv(instr_file_name, index=False)
-            print(f"Instructions saved to {instr_file_name}")
+            log.info(f"Instructions saved to {instr_file_name}")
 
         # Load instructions into config
         st.set_config(
@@ -324,12 +324,12 @@ def xenonnt_salted_wfsim(
     if runid is not None:
         try:
             instr = pd.read_csv(instr_file_name)
-            print("Loaded instructions from file", instr_file_name)
+            log.info("Loaded instructions from file", instr_file_name)
         except:
-            print(f"Instruction file {instr_file_name} not found. Generating instructions...")
+            log.info(f"Instruction file {instr_file_name} not found. Generating instructions...")
             instr = generator_func(runid=runid, **kwargs)
             pd.DataFrame(instr).to_csv(instr_file_name, index=False)
-            print(f"Instructions saved to {instr_file_name}")
+            log.info(f"Instructions saved to {instr_file_name}")
 
     # Based on cutax.xenonnt_sim_base()
     fax_conf = "fax_config_nt_{:s}.json".format(simu_config_version)
@@ -470,12 +470,12 @@ def fxenonnt(
     """
     assert saltax_mode in SALTAX_MODES, "saltax_mode must be one of %s" % (SALTAX_MODES)
     if runid is None:
-        print(
+        log.warning(
             "Since you specified runid=None, this context will not be able to compute raw_records_simu."
         )
-        print("Welcome to data-loading only mode!")
+        log.warning("Welcome to data-loading only mode!")
     else:
-        print("Welcome to computation mode which only works for run %s!" % (runid))
+        log.warning("Welcome to computation mode which only works for run %s!" % (runid))
 
     return xenonnt_salted_fuse(
         runid=runid,
@@ -539,12 +539,12 @@ def sxenonnt(
     """
     assert saltax_mode in SALTAX_MODES, "saltax_mode must be one of %s" % (SALTAX_MODES)
     if runid is None:
-        print(
+        log.warning(
             "Since you specified runid=None, this context will not be able to compute raw_records_simu."
         )
-        print("Welcome to data-loading only mode!")
+        log.warning("Welcome to data-loading only mode!")
     else:
-        print("Welcome to computation mode which only works for run %s!" % (runid))
+        log.warning("Welcome to computation mode which only works for run %s!" % (runid))
 
     return xenonnt_salted_wfsim(
         runid=runid,
diff --git a/saltax/instructions/generator.py b/saltax/instructions/generator.py
index e34a57d..f3183e5 100644
--- a/saltax/instructions/generator.py
+++ b/saltax/instructions/generator.py
@@ -385,7 +385,7 @@ def generator_flat(
     """Generate instructions for a run with flat energy spectrum.
 
     :param runid: run number in integer
-    :param en_range: (en_min, en_max) in keV, default: (0, 30.0)
+    :param en_range: (en_min, en_max) in keV, default: (0.2, 15.0)
     :param recoil: NEST recoil type, default: 8 (beta ER)
     :param n_tot: total number of events to generate, default: None i.e.
         generate events until end_time