Modifications of nT simulation context

straxen/contexts.py

@@ -219,32 +219,52 @@ def xenonnt_led(**kwargs):
     st.register([straxen.DAQReader, straxen.LEDCalibration])
     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
@@ -261,52 +281,78 @@ 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
+    # doing sanity checks for cmt run ids for simulation and processing
+    if (not cmt_run_id_sim ) and (not cmt_run_id_proc ):
+        raise RuntimeError("Setting both cmt_run_id_sim and cmt_run_id_proc to None")
+    if (cmt_run_id_sim and cmt_run_id_proc ) and (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)
+    else:
+        cmt_id = cmt_run_id_sim  or cmt_run_id_proc
+        cmt_run_id_sim  = cmt_id
+        cmt_run_id_proc = cmt_id
 
-    # Setup processing
     # Replace default cmt options with cmt_run_id tag + cmt run id
     cmt_options = straxen.get_corrections.get_cmt_options(st)
+
+    # First, fix gain model for simulation
+    st.set_config({'gain_model_mc': 
+                        ('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] = ('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] = ('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] = ( 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] = (
+                         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"})