XY-dependent SE gain and extraction efficiency (#363)

* Add extraction efficiency and garfield gas gap map
* Update load_resource.py
* got rid of the xy confinement
* s2_correction_map fix
* Update load_resource.py
* Update load_resource.py
* Update load_resource.py
* draw se gain from map
* make it work without an SE gain map
* make compatible with previous wfsim
* make compatible with previous wfsim
* fix typo
* fix conditional issue with se_gain_map
* removed print
* fix conditional for garfield stuff
* remove unused config
* change directory
* change directory
* add sr0 to filenames, change fname check
* interpolate between excitation hists
* fix bool
* add print statements for testing
* add print statements to test
* add print statements for test
* fix divide by zero error
* get rid of print statements
* update straxen
* add param nph to docstring
* add comments, copy things to the 1T config
* Test without 1T additions
* add g2_mean to 1T for test
* comment
* add _gg to s2_luminescence when using garfield_gas_gap
* Add _gg to s2_luminescence when using garfield_gas_gap
* allow for constant extraction efficiency
* refactor
* fix tests
* tweak 1
* more debug
* fix it more and more and more
* fix this mess
* one more test
* add load resource
* add gitignor
* fix fail fast
* what is the issue?
* Update tests/test_wfsim.py
* Update wfsim/core/s2.py
* Fixing extraction efficiency calculation when no SE gain map is used
* Correcting extraction efficiency when SE gain is used from fax config
* add print statements to see why tests fail
* add another print
* flatten correction map from mc
* get rid of print statements
* change sampling index

Co-authored-by: Melih Kara <kara@kit.edu>
Co-authored-by: Diego Ramírez García <diego.ramirez@physik.uni-freiburg.de>
Co-authored-by: Joran Angevaare <jorana@nikhef.nl>
Co-authored-by: Andrii Terliuk <andrey.terlyuk@gmail.com>

.github/workflows/pytest.yml

@@ -18,9 +18,6 @@ on:
   release:
     types: [created]
   pull_request:
-    branches:
-      - master
-      - stable
   push:
     branches:
       - master
@@ -30,7 +27,7 @@ jobs:
     name: "${{ matrix.test }}_py${{ matrix.python-version }}"
     runs-on: ubuntu-latest
     strategy:
-      fail-fast: True
+      fail-fast: False
       matrix:
         python-version: [3.8, 3.9, "3.10"]
         test: ['coveralls', 'pytest', 'pytest_no_database']
@@ -84,7 +81,7 @@ jobs:
         env:
           NUMBA_DISABLE_JIT: 1
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-          # We need to check if we had access to the secrets, otherwise coveralls
+          # We need to check if we have access to the secrets, otherwise coveralls
           # will yield a low coverage because of the lack of interfacing with the
           # database.
           HAVE_ACCESS_TO_SECTETS: ${{ secrets.RUNDB_API_URL }}

.gitignore

@@ -21,6 +21,9 @@ custom_data
 test_input_data
 *.zip
 last_bootstrax_exception.txt
+.xenon_config
+test.root
+pytestdebug.log
 
 # cProfile output
 *.prof
@@ -58,4 +61,3 @@ docs/tutorials/*
 
 # coverage
 .coverage
-

tests/test_contexts.py

@@ -1,8 +1,9 @@
 import strax
 import straxen
 import wfsim
+from unittest import skipIf
 
-
+@skipIf(not straxen.utilix_is_configured(), 'utilix is not configured')
 def test_nt_context(register=None, context=None):
     """
     Test a context if it is properly setup. To this end, we perform a
@@ -13,9 +14,6 @@ def test_nt_context(register=None, context=None):
     :param context: Test with some other context than the nT simulation
     context
     """
-    if not straxen.utilix_is_configured():
-        return
-
     if context is None:
         context = straxen.contexts.xenonnt_simulation(cmt_run_id_sim='010000', cmt_version='global_ONLINE')
     assert isinstance(context, strax.Context), f'{context} is not a context'
@@ -25,15 +23,3 @@ def test_nt_context(register=None, context=None):
 
     # Search all plugins for the time field (each should have one)
     context.search_field('time')
-
-
-# def test_mc_chain():
-#     test_nt_context(wfsim.RawRecordsFromMcChain)
-
-
-# def test_fax_nveto():
-#     test_nt_context(wfsim.RawRecordsFromFaxnVeto)
-
-
-# def test_1t_context():
-#     test_nt_context(context=straxen.contexts.xenon1t_simulation())

tests/test_load_resource.py

@@ -1,3 +1,4 @@
+import straxen
 from wfsim.load_resource import load_config
 
 
@@ -11,6 +12,8 @@ def test_load_1t():
         "enable_electron_afterpulses": True,
         "enable_noise": False,
         "field_distortion_on": True,
+        "g2_mean": 32.3,
+        's2_time_model': 's2_time_spread around zero',
     }
     result = load_config(config)
     return result, config
@@ -36,7 +39,8 @@ def test_load_nt():
         "s2_pattern_map": ["constant dummy", 30e-5, [494,]],
         "s2_correction_map": ["constant dummy", 1, []],
         "field_dependencies_map": ["constant dummy", 1, []],
-        "gains": [1 for i in range(494)],
+        "gains": [1 for _ in range(straxen.n_tpc_pmts)],
+        "se_gain_map": ["constant dummy", 1, []],
     }
     result = load_config(config)
     return result, config

tests/test_wfsim.py

@@ -7,7 +7,8 @@
 import straxen
 
 import wfsim
-from .test_load_resource import test_load_nt
+from .test_load_resource import test_load_nt, test_load_1t
+from unittest import skipIf
 
 logging.basicConfig(
     level=logging.DEBUG,
@@ -30,27 +31,31 @@ def test_sim_1T():
     """Test the 1T simulator (should always work with the publicly available files)"""
     with tempfile.TemporaryDirectory() as tempdir:
         log.debug(f'Working in {tempdir}')
-        testing_config_1T = dict(
+        _, conf_1t = test_load_1t()
+        testing_config_1t = dict(
             hev_gain_model=("1T_to_pe_placeholder", True),
             gain_model=("1T_to_pe_placeholder", True),
-            gain_model_mc=("1T_to_pe_placeholder", True),)
+            gain_model_mc=("1T_to_pe_placeholder", True),
+        )
 
         st = strax.Context(
             storage=tempdir,
             config=dict(
                 nchunk=2, event_rate=1, chunk_size=1,
                 detector='XENON1T',
                 fax_config=('https://raw.githubusercontent.com/XENONnT/strax_auxiliary_files'
-                            '/36d352580b328ff057b1588b8af8c9a6ed8ae704/sim_files/fax_config_1t.json'),  # noqa
+                            '/a5b92102505d6d0bfcdb563b6117bd4040a93435/sim_files/fax_config_1t.json'),  # noqa
                 fax_config_override=dict(
                     url_base=("https://raw.githubusercontent.com/XENONnT/strax_auxiliary_files"
-                              "/36d352580b328ff057b1588b8af8c9a6ed8ae704/sim_files/"), ),
+                              "/a5b92102505d6d0bfcdb563b6117bd4040a93435/sim_files/"),
+                    **conf_1t,
+                ),
                 **straxen.contexts.x1t_common_config),
             **straxen.contexts.x1t_context_config,
         )
         st.register(wfsim.RawRecordsFromFax1T)
-        log.debug(f'Setting testing config {testing_config_1T}')
-        st.set_config(testing_config_1T)
+        log.debug(f'Setting testing config {testing_config_1t}')
+        st.set_config(testing_config_1t)
 
         log.debug(f'Getting raw-records')
         rr = st.get_array(run_id, 'raw_records')
@@ -59,18 +64,19 @@ def test_sim_1T():
         log.info(f'All done')
 
 
-def test_sim_nT_basics():
+@skipIf(not straxen.utilix_is_configured(), 'utilix is not configured')
+def test_sim_nt_basics():
     """Test the nT simulator. Uses basic config so complicated steps are skipped. So this will test
        the simple s1 model and the simple s2 model"""
 
     with tempfile.TemporaryDirectory() as tempdir:
         log.debug(f'Working in {tempdir}')
         conf = copy.deepcopy(straxen.contexts.xnt_common_config)
+        resource, conf_override = test_load_nt()
         conf['gain_model'] = ("to_pe_placeholder", True)
         conf['gain_model_mc'] = ("to_pe_placeholder", True)
         conf['hev_gain_model'] = ("to_pe_placeholder", True)
         conf['hit_min_amplitude'] = 'pmt_commissioning_initial'
-        resource, conf_override = test_load_nt()
 
         # The SPE table in this package is for a single channel
         # We generate the full SPE file for testing here
@@ -101,32 +107,27 @@ def test_sim_nT_basics():
         log.info(f'All done')
 
 
-def test_sim_nT_advanced():
-    """Test the nT simulator. Works only if one has access to the XENONnT databases.
-        Clone the repo to dali and type 'pytest' to run. The first run will test simple s1,
-        garfield s2 and noise/afterpulses. The second run will test the s1 spline model"""
-
-    if not straxen.utilix_is_configured():
-        log.warning(f"Utilix is not configured, skipping database-requiring tests!")
-        return
-
+@skipIf(not straxen.utilix_is_configured(), 'utilix is not configured')
+def test_sim_nt_advanced(
+        config = None
+):
+    """
+    Test the nT simulator. Works only if one has access to the XENONnT databases.
+    Clone the repo to dali and type 'pytest' to run.
+    """
     with tempfile.TemporaryDirectory() as tempdir:
         log.debug(f'Working in {tempdir}')
+
         st = straxen.contexts.xenonnt_simulation(cmt_run_id_sim='010000',
                                                  cmt_version='global_ONLINE',
-                                                 _config_overlap={},)
-        st.set_config(dict(gain_model_mc=("to_pe_placeholder", True),
-                           gain_model=("to_pe_placeholder", True),
-                           hit_min_amplitude='pmt_commissioning_initial'
-                          ))
+                                                 _config_overlap={},
+                                                 fax_config='fax_config_nt_sr0_v0.json'
+                                                 )
         st.set_config(dict(nchunk=1, event_rate=1, chunk_size=2,))
 
-        st.set_config({'fax_config_override': dict(s2_luminescence_model='simple',
-                                                   s2_time_model="s2_time_spread around zero",
-                                                   s1_lce_correction_map='XENONnT_s1_xyz_LCE_corrected_qes_MCva43fa9b_wires.json.gz',
-                                                   s1_time_spline='XENONnT_s1_proponly_va43fa9b_wires_20200625.json.gz',
-                                                   s1_model_type='optical_propagation+simple',)})
-
+        if config is not None:
+            log.warning(f'Update config with {config}')
+            st.set_config(config)
         log.debug(f'Getting raw-records')
         rr = st.get_array(run_id, 'raw_records')
         log.debug(f'Getting peaks')
@@ -135,13 +136,52 @@ def test_sim_nT_advanced():
         log.info(f'All done')
 
 
+@skipIf(not straxen.utilix_is_configured(), 'utilix is not configured')
+def test_nt_advanced_alt_s2_model():
+    config = dict(
+        fax_config_override=dict(
+            s2_luminescence_model='simple',
+            s2_time_model="s2_time_spread around zero",
+            s1_lce_correction_map='XENONnT_s1_xyz_LCE_corrected_qes_MCva43fa9b_wires.json.gz',
+            s1_time_spline='XENONnT_s1_proponly_va43fa9b_wires_20200625.json.gz',
+            s1_model_type='optical_propagation+simple',
+        )
+    )
+    test_sim_nt_advanced(config)
+
+
+@skipIf(not straxen.utilix_is_configured(), 'utilix is not configured')
+def test_nt_advanced_garfield():
+    config = dict(
+        fax_config_override=dict(
+            s2_luminescence_model='garfield',
+            s2_correction_map=False,
+            s2_time_model="s2_time_spread around zero",
+            s1_lce_correction_map='XENONnT_s1_xyz_LCE_corrected_qes_MCva43fa9b_wires.json.gz',
+            s1_time_spline='XENONnT_s1_proponly_va43fa9b_wires_20200625.json.gz',
+            s1_model_type='optical_propagation+simple',
+        )
+    )
+    test_sim_nt_advanced(config)
+
+@skipIf(not straxen.utilix_is_configured(), 'utilix is not configured')
+def test_nt_advanced_gas_gap_garfield():
+    config = dict(
+        fax_config_override=dict(
+            s2_luminescence_model='garfield_gas_gap',
+            s2_correction_map="XENONnT_s2_xy_map_v4_210503_mlp_3_in_1_iterated.json",
+            s2_luminescence_gg= "garfield_timing_map_gas_gap_sr0.npy",
+            garfield_gas_gap_map= "garfield_gas_gap_map_sr0.json",
+            se_gain_map="XENONnT_se_xy_map_v1_mlp.json",
+    )
+    )
+    test_sim_nt_advanced(config)
+
+
+@skipIf(not straxen.utilix_is_configured(), 'utilix is not configured')
 def test_sim_mc_chain():
     """Test the nT simulator by Geant4 output file"""
 
-    if not straxen.utilix_is_configured():
-        log.warning(f"Utilix is not configured, skipping database-requiring tests!")
-        return
-
     with tempfile.TemporaryDirectory() as tempdir:
         log.debug(f'Working in {tempdir}')
 
@@ -154,10 +194,7 @@ def test_sim_mc_chain():
         st = straxen.contexts.xenonnt_simulation(cmt_run_id_sim='010000',
                                                  cmt_version='global_ONLINE',
                                                  _config_overlap={},)
-        st.set_config(dict(gain_model_mc=("to_pe_placeholder", True),
-                           gain_model=("to_pe_placeholder", True),
-                           gain_model_nv=("adc_nv", True),
-                           hit_min_amplitude='pmt_commissioning_initial',
+        st.set_config(dict(gain_model_nv=("adc_nv", True),
                           ))
 
         epix_config = {'cut_by_eventid': True, 'debug': True, 'source_rate': 0, 'micro_separation_time': 10.,
@@ -175,8 +212,6 @@ def test_sim_mc_chain():
             fax_config_override=dict(
                 s1_model_type='nest',
                 s2_time_model="s2_time_spread around zero",
-                url_base='https://raw.githubusercontent.com/XENONnT/private_nt_aux_files/master/sim_files',
-                s1_lce_correction_map=["constant dummy", 1, []],
                 enable_electron_afterpulses=False),
             epix_config=epix_config,
             fax_config_nveto='fax_config_nt_nveto.json',

wfsim/core/rawdata.py

@@ -377,6 +377,8 @@ def get_mean_xy_electron(self, peak_type, instruction, tb):
                 _, xy_tmp = self.pulses['s2'].field_distortion_comsol(instruction['x'], instruction['y'], instruction['z'], self.resource)
             elif self.config.get('field_distortion_model', "none") == 'inverse_fdc':
                 _, xy_tmp = self.pulses['s2'].inverse_field_distortion_correction(instruction['x'], instruction['y'], instruction['z'], self.resource)
+            else:
+                raise ValueError(f'{self.config.get("field_distortion_model", "none")} is invalid!')
             tb['x_mean_electron'] = np.mean(xy_tmp.T[0])
             tb['y_mean_electron'] = np.mean(xy_tmp.T[1])
         else:
@@ -455,7 +457,7 @@ def digitizer_saturation(data, channel_mask):
 @export
 class RawDataOptical(RawData):
 
-    def __init__(self, config, channels=[], timings=[]):
+    def __init__(self, config, channels=tuple(), timings=tuple()):
         self.config = config
         self.pulses = dict(
             s1=Pulse(config),