Modifications of nT simulation context

straxen/contexts.py

@@ -220,37 +220,61 @@ def xenonnt_led(**kwargs):
     return st
 
 
-def xenonnt_simulation(output_folder='./strax_data',
-                       cmt_run_id='020280',
-                       cmt_version='v3',
-                       fax_config='fax_config_nt_design.json',
-                       cmt_option_overwrite=immutabledict(),
-                       overwrite_from_fax_file=False,
-                       _cmt_run_id_proc_only=None,
-                       _forbid_creation_of=None,
-                       _config_overlap=immutabledict(drift_time_gate='electron_drift_time_gate',
-                                                     drift_velocity_liquid='electron_drift_velocity',
-                                                     electron_lifetime_liquid='elife_conf'),
-                       **kwargs):
+
+def xenonnt_simulation(
+                output_folder                = './strax_data',
+                cmt_run_id_sim               = None,
+                cmt_run_id_proc              = None,
+                cmt_version                  = 'v3',
+                fax_config                   = 'fax_config_nt_design.json',
+                overwrite_from_fax_file_sim  = False,
+                overwrite_from_fax_file_proc = False,
+                cmt_option_overwrite_sim     = immutabledict(),
+                cmt_option_overwrite_proc    = immutabledict(),
+                _forbid_creation_of          = None,
+                _config_overlap              = immutabledict(
+                                drift_time_gate='electron_drift_time_gate',
+                                drift_velocity_liquid='electron_drift_velocity',
+                                electron_lifetime_liquid='elife_conf'),
+               **kwargs):
     """
-    Context with CMT options updated with cmt_run_id.
+
+    The most generic context that allows for setting full divergent 
+    settings for simulation purposes
+
+    It makes full divergent setup, allowing to set detector simulation 
+    part (i.e. for wfsim up to truth and  raw_records). Parameters _sim 
+    refer to detector simulation parameters. 
+
+    Arguments having _proc in their name refer to detector parameters that 
+    are used for processing of simulations as done to the real datector 
+    data. This means starting from already existing raw_records and finishing 
+    with higher level data, such as peaks, events etc.
+
     CMT options can also be overwritten via fax config file.
     :param output_folder: Output folder for strax data.
-    :param cmt_run_id: Run id which should be used to load corrections.
+    :param cmt_run_id_sim: Run id for detector parameters from CMT to be used in 
+        for creation of raw_records. 
+    :param cmt_run_id_proc: Run id for detector parameters from CMT to be used in 
+        for processing from raw_records to higher level data.
     :param cmt_version: Global version for corrections to be loaded.
     :param fax_config: Fax config file to use.
-    :param cmt_option_overwrite: Dictionary to overwrite CMT settings. Keys
-        must be valid CMT option keys.
-    :param overwrite_from_fax_file: If true overwrites CMT options with
-        constants from fax file.
-    :param _cmt_run_id_proc_only: Run id just for > raw_records processing if diverge from
-    cmt_run_id.
+    :param overwrite_from_fax_file_sim: If true sets detector simulation 
+        parameters for truth/raw_records from from fax_config file istead of CMT
+    :param overwrite_from_fax_file_proc:  If true sets detector processing 
+        parameters after raw_records(peaklets/events/etc) from from fax_config 
+        file istead of CMT
+    :param cmt_option_overwrite_sim: Dictionary to overwrite CMT settings for 
+        the detector simulation part. 
+    :param cmt_option_overwrite_proc: Dictionary to overwrite CMT settings for 
+        the data processing part. 
     :param _forbid_creation_of: str/tuple, of datatypes to prevent form
         being written (e.g. 'raw_records' for read only simulation context).
     :param _config_overlap: Dictionary of options to overwrite. Keys
         must be simulation config keys, values must be valid CMT option keys.
     :param kwargs: Additional kwargs taken by strax.Context.
     :return: strax.Context instance
+
     """
     import wfsim
     st = strax.Context(
@@ -261,56 +285,87 @@ def xenonnt_simulation(output_folder='./strax_data',
                     **straxen.contexts.xnt_common_config,),
         **straxen.contexts.xnt_common_opts, **kwargs)
     st.register(wfsim.RawRecordsFromFaxNT)
-
+    st.apply_cmt_version(f'global_{cmt_version}')
+
     if _forbid_creation_of is not None:
         st.context_config['forbid_creation_of'] += strax.to_str_tuple(_forbid_creation_of)
-
-    st.apply_cmt_version(f'global_{cmt_version}')
-    if not cmt_run_id:
-        raise ValueError('You have to specify a run_id which should be used to initialize '
-                         'the corrections.')
-    if _cmt_run_id_proc_only is None:
-        _cmt_run_id_proc_only = cmt_run_id
-
-    # Setup processing
+
+    # doing sanity checks for cmt run ids for simulation and processing
+    if not (cmt_run_id_sim is None) and not (cmt_run_id_proc is None):
+        if not (cmt_run_id_sim==cmt_run_id_proc):
+            print("INFO : divergent CMT runs for simulation and processing")
+            print("    cmt_run_id_sim".ljust(25), cmt_run_id_sim)
+            print("    cmt_run_id_proc".ljust(25), cmt_run_id_proc)
+    elif (cmt_run_id_sim is None) and not (cmt_run_id_proc is None):
+        cmt_run_id_sim=cmt_run_id_proc
+    elif not (cmt_run_id_sim is None) and (cmt_run_id_proc is None):
+        cmt_run_id_proc=cmt_run_id_sim
+    else:
+        raise RuntimeError("Trying to set both cmt_run_id_sim and cmt_run_id_proc to None")
+
     # Replace default cmt options with cmt_run_id tag + cmt run id
-    cmt_options = straxen.get_corrections.get_cmt_options(st)
+    cmt_options = straxen.get_corrections.get_cmt_options(st)    
+
+    # First, fix gain model for simulation
+    st.set_config({'gain_model_mc': tuple(['cmt_run_id', 
+                                           cmt_run_id_sim, 
+                                           *cmt_options['gain_model']])})
+    fax_config_override_from_cmt = dict()
+    for fax_field, cmt_field in _config_overlap.items():
+        fax_config_override_from_cmt[fax_field] = tuple(
+                                                    ['cmt_run_id', 
+                                                     cmt_run_id_sim, 
+                                                     *cmt_options[cmt_field]]
+                                                    )
+    st.set_config({'fax_config_override_from_cmt': fax_config_override_from_cmt})
+
+    # and all other parameters for processing
     for option in cmt_options:
-        st.config[option] = tuple(['cmt_run_id', _cmt_run_id_proc_only, *cmt_options[option]])
-
+        st.config[option] = tuple(['cmt_run_id', cmt_run_id_proc, *cmt_options[option]])
+
+    # Done with "default" usage, now to overwrites from file
+    #
     # Take fax config and put into context option
-    if overwrite_from_fax_file:
+    if (overwrite_from_fax_file_proc or 
+                overwrite_from_fax_file_sim):
         fax_config = straxen.get_resource(fax_config, fmt='json')
-
-        for fax_field, cmt_field in _config_overlap.items():
-            st.config[cmt_field] = tuple([cmt_options[cmt_field][0] + '_constant',
-                                          fax_config[fax_field]])
-
-    # Setup simulation
-    # Pass the CMT options to simulation (override all other config input methods)
-    else:
-        fax_config_override_from_cmt = dict()
-
         for fax_field, cmt_field in _config_overlap.items():
-            fax_config_override_from_cmt[fax_field] = tuple(['cmt_run_id', cmt_run_id, *cmt_options[cmt_field]])
-        st.set_config({'fax_config_override_from_cmt': fax_config_override_from_cmt})
-
-    # Fix gain model
-    st.set_config({'gain_model_mc': tuple(['cmt_run_id', cmt_run_id, *cmt_options['gain_model']])})
-
-    # User customized overwrite
-    cmt_options = straxen.get_corrections.get_cmt_options(st)  # This includes gain_model_mc
-    for option in cmt_option_overwrite:
+            if overwrite_from_fax_file_proc:
+                st.config[cmt_field] = tuple([
+                    cmt_options[cmt_field][0] + '_constant',fax_config[fax_field]])
+            if overwrite_from_fax_file_sim:
+                st.config['fax_config_override_from_cmt'][fax_field] = tuple(
+                    [cmt_options[cmt_field][0] + '_constant',fax_config[fax_field]])   
+
+    # And as the last step - manual overrrides, since they have the highest priority 
+    # User customized for simulation
+    for option in cmt_option_overwrite_sim:
         if option not in cmt_options:
             raise ValueError(f'Overwrite option {option} is not using CMT by default '
                              'you should just use set config')
+        if not option in _config_overlap.values():
+            raise ValueError(f'Overwrite option {option} does not have mapping from '
+                              'CMT to fax config! ')      
+        for fax_key,cmt_key in _config_overlap.items():
+            if cmt_key==option: continue
         _name_index = 2 if 'cmt_run_id' in cmt_options[option] else 0
-        st.config[option] = (cmt_options[option][_name_index] + '_constant', cmt_option_overwrite[option])
-
+        st.config['fax_config_override_from_cmt'][fax_key] = (
+                                                    cmt_options[option][_name_index] + '_constant',
+                                                    cmt_option_overwrite_sim[option]
+                                                             )
+        del(fax_key,cmt_key,_name_index)
+    # User customized for simulation
+    for option in cmt_option_overwrite_proc:
+        if option not in cmt_options:
+            raise ValueError(f'Overwrite option {option} is not using CMT by default '
+                             'you should just use set config')
+        _name_index = 2 if 'cmt_run_id' in cmt_options[option] else 0
+        st.config[option] = (cmt_options[option][_name_index] + '_constant', cmt_option_overwrite_proc[option])
+        del(_name_index)
     # Only for simulations
     st.set_config({"event_info_function": "disabled"})
-
-    return st
+    
+    return(st)
 
 
 ##