diff --git a/.github/workflows/Pipi.yml b/.github/workflows/Pipi.yml index d1b93bac..8b3f02e3 100644 --- a/.github/workflows/Pipi.yml +++ b/.github/workflows/Pipi.yml @@ -13,7 +13,7 @@ jobs: - name: Setup python uses: actions/setup-python@v4 with: - python-version: '3.8' + python-version: '3.9' - name: Checkout repo uses: actions/checkout@v3 - name: Install dependencies diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml index 0b77ef78..ad795dea 100644 --- a/.github/workflows/pytest.yml +++ b/.github/workflows/pytest.yml @@ -29,15 +29,15 @@ jobs: strategy: fail-fast: False matrix: - python-version: [3.8, 3.9, "3.10"] + python-version: [ "3.9", "3.10" ] test: ['coveralls', 'pytest', 'pytest_no_database'] # Only run coverage / no_database on py3.8 exclude: - - python-version: 3.9 - test: pytest_no_database - python-version: "3.10" + test: pytest_no_database + - python-version: "3.11" test: coveralls - - python-version: "3.10" + - python-version: "3.11" test: pytest_no_database steps: # Setup and installation @@ -84,8 +84,8 @@ jobs: # 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 }} - if: matrix.test == 'coveralls' && env.HAVE_ACCESS_TO_SECTETS != null + HAVE_ACCESS_TO_SECRETS: ${{ secrets.RUNDB_API_URL }} + if: matrix.test == 'coveralls' && env.HAVE_ACCESS_TO_SECRETS != null run: | coverage run --source=wfsim setup.py test -v coveralls --service=github diff --git a/.readthedocs.yml b/.readthedocs.yml index 5bf2645f..992eba7e 100644 --- a/.readthedocs.yml +++ b/.readthedocs.yml @@ -9,7 +9,7 @@ sphinx: build: os: ubuntu-22.04 tools: - python: "3.8" + python: "3.9" python: install: diff --git a/extra_requirements/requirements-tests.txt b/extra_requirements/requirements-tests.txt index fd7eb036..00af8022 100644 --- a/extra_requirements/requirements-tests.txt +++ b/extra_requirements/requirements-tests.txt @@ -1,4 +1,4 @@ -strax==1.5.5 -straxen==2.1.6 -epix==0.3.5 +strax>=1.6.0 +straxen>=2.2.0 +epix>=0.3.6 git+https://github.com/XENONnT/ax_env diff --git a/setup.py b/setup.py index d91494de..4374fb6e 100644 --- a/setup.py +++ b/setup.py @@ -26,7 +26,7 @@ 'flake8', 'pytest-cov', ], - python_requires=">=3.8", + python_requires=">=3.9", extras_require={ 'docs': [ 'sphinx', @@ -44,7 +44,6 @@ 'Development Status :: 5 - Production/Stable', 'License :: OSI Approved :: BSD License', 'Natural Language :: English', - 'Programming Language :: Python :: 3.8', 'Programming Language :: Python :: 3.9', 'Programming Language :: Python :: 3.10', 'Intended Audience :: Science/Research', diff --git a/tests/test_contexts.py b/tests/test_contexts.py index 3a77e037..43b30945 100644 --- a/tests/test_contexts.py +++ b/tests/test_contexts.py @@ -1,7 +1,10 @@ +import os import strax import straxen +from unittest import TestCase, skipIf +from straxen.contexts import xenonnt import wfsim -from unittest import skipIf + @skipIf(not straxen.utilix_is_configured(), 'utilix is not configured') def test_nt_context(register=None, context=None): @@ -15,7 +18,7 @@ def test_nt_context(register=None, context=None): context """ if context is None: - context = straxen.contexts.xenonnt_simulation(cmt_run_id_sim='010000', cmt_version='global_ONLINE') + context = wfsim.contexts.xenonnt_simulation(cmt_run_id_sim='010000', cmt_version='global_ONLINE') assert isinstance(context, strax.Context), f'{context} is not a context' if register is not None: @@ -23,3 +26,72 @@ def test_nt_context(register=None, context=None): # Search all plugins for the time field (each should have one) context.search_field('time') + + +# Simulation contexts are only tested when special flags are set + + +class TestSimContextNT(TestCase): + @staticmethod + def context(*args, **kwargs): + kwargs.setdefault("cmt_version", "global_ONLINE") + return wfsim.contexts.xenonnt_simulation(*args, **kwargs) + + @skipIf(not straxen.utilix_is_configured(), "No db access, cannot test!") + def test_nt_sim_context_main(self): + st = self.context(cmt_run_id_sim="008000") + st.search_field("time") + + @skipIf(not straxen.utilix_is_configured(), "No db access, cannot test!") + def test_nt_sim_context_alt(self): + """Some examples of how to run with a custom WFSim context.""" + self.context(cmt_run_id_sim="008000", cmt_run_id_proc="008001") + self.context(cmt_run_id_sim="008000", cmt_option_overwrite_sim={"elife": 1e6}) + + self.context(cmt_run_id_sim="008000", overwrite_fax_file_sim={"elife": 1e6}) + + @skipIf(not straxen.utilix_is_configured(), "No db access, cannot test!") + def test_nt_diverging_context_options(self): + """Test diverging options. Idea is that you can use different settings for processing and + generating data, should have been handled by RawRecordsFromWFsim but is now hacked into the + xenonnt_simulation context. + + Just to show how convoluted this syntax for the + xenonnt_simulation context / CMT is... + + """ + self.context( + cmt_run_id_sim="008000", + cmt_option_overwrite_sim={"elife": ("elife_constant", 1e6, True)}, + cmt_option_overwrite_proc={"elife": ("elife_constant", 1e5, True)}, + overwrite_from_fax_file_proc=True, + overwrite_from_fax_file_sim=True, + _config_overlap={"electron_lifetime_liquid": "elife"}, + ) + + def test_nt_sim_context_bad_inits(self): + with self.assertRaises(RuntimeError): + self.context( + cmt_run_id_sim=None, + cmt_run_id_proc=None, + ) + + +def test_sim_context(): + wfsim.contexts.xenon1t_simulation() + + +@skipIf(not straxen.utilix_is_configured(), "No db access, cannot test!") +def test_sim_offline_context(): + wfsim.contexts.xenonnt_simulation_offline( + run_id="026000", + global_version="global_v11", + fax_config="fax_config_nt_sr0_v4.json", + ) + + +@skipIf(not straxen.utilix_is_configured(), "No db access, cannot test!") +def test_offline(): + st = xenonnt("latest") + st.provided_dtypes() + diff --git a/tests/test_wfsim.py b/tests/test_wfsim.py index bbd4948c..623ac5e5 100644 --- a/tests/test_wfsim.py +++ b/tests/test_wfsim.py @@ -51,7 +51,7 @@ def test_sim_1T(): **conf_1t, ), **straxen.legacy.x1t_common_config), - **straxen.legacy.contexts_1t.get_x1t_context_config(), + **straxen.legacy.get_x1t_context_config(), ) st.register(wfsim.RawRecordsFromFax1T) log.debug(f'Setting testing config {testing_config_1t}') @@ -118,7 +118,7 @@ def test_sim_nt_advanced( with tempfile.TemporaryDirectory() as tempdir: log.debug(f'Working in {tempdir}') - st = straxen.contexts.xenonnt_simulation(cmt_run_id_sim='010000', + st = wfsim.contexts.xenonnt_simulation(cmt_run_id_sim='010000', cmt_version='global_ONLINE', fax_config='fax_config_nt_sr0_v0.json', _config_overlap={},) @@ -199,7 +199,7 @@ def test_sim_mc_chain(): url_data = requests.get(test_g4).content with open('test.root', mode='wb') as f: f.write(url_data) - st = straxen.contexts.xenonnt_simulation(cmt_run_id_sim='010000', + st = wfsim.contexts.xenonnt_simulation(cmt_run_id_sim='010000', cmt_version='global_ONLINE', _config_overlap={},) st.set_config(dict(gain_model_nv="legacy-to-pe://adc_nv", diff --git a/wfsim/__init__.py b/wfsim/__init__.py index c2197ebb..9c3fa797 100644 --- a/wfsim/__init__.py +++ b/wfsim/__init__.py @@ -11,3 +11,5 @@ from .load_resource import * from .utils import * + +from .contexts import * diff --git a/wfsim/contexts.py b/wfsim/contexts.py new file mode 100644 index 00000000..c3a48586 --- /dev/null +++ b/wfsim/contexts.py @@ -0,0 +1,292 @@ +from typing import Optional +from immutabledict import immutabledict +import strax +import straxen +import wfsim +from .strax_interface import RawRecordsFromFax1T + + +def xenonnt_simulation_offline( + output_folder: str = "./strax_data", + wfsim_registry: str = "RawRecordsFromFaxNT", + run_id: Optional[str] = None, + global_version: Optional[str] = None, + fax_config: Optional[str] = None, + **kwargs, +): + """ + :param output_folder: strax_data folder + :param wfsim_registry: Raw_records generation mechanism, + 'RawRecordsFromFaxNT', 'RawRecordsFromMcChain', etc, + https://github.com/XENONnT/WFSim/blob/master/wfsim/strax_interface.py + :param run_id: Real run_id to use to fetch the corrections + :param global_version: Global versions + https://github.com/XENONnT/corrections/tree/master/XENONnT/global_versions + :param fax_config: WFSim configuration files + https://github.com/XENONnT/private_nt_aux_files/blob/master/sim_files/fax_config_nt_sr0_v4.json + :return: strax context for simulation + """ + if run_id is None: + raise ValueError("Specify a run_id to load the corrections") + if global_version is None: + raise ValueError("Specify a correction global version") + if fax_config is None: + raise ValueError("Specify a simulation configuration file") + + # General strax context, register common plugins + st = strax.Context( + storage=strax.DataDirectory(output_folder), + **straxen.contexts.xnt_common_opts, + **kwargs, + ) + # Register simulation configs required by WFSim plugins + st.config.update( + dict( + detector="XENONnT", + fax_config=fax_config, + check_raw_record_overlaps=True, + **straxen.contexts.xnt_common_config, + ) + ) + # Register WFSim raw_records plugin to overwrite real data raw_records + wfsim_plugin = getattr(wfsim, wfsim_registry) + st.register(wfsim_plugin) + for plugin_name in wfsim_plugin.provides: + assert "wfsim" in str(st._plugin_class_registry[plugin_name]) + # Register offline global corrections same as real data + st.apply_xedocs_configs(version=global_version) + # Real data correction is run_id dependent, + # but in sim we often use run_id not in the runDB + # So we switch the run_id dependence to a specific run -> run_id + local_versions = st.config + for config_name, url_config in local_versions.items(): + if isinstance(url_config, str): + if "run_id" in url_config: + local_versions[config_name] = straxen.URLConfig.format_url_kwargs( + url_config, run_id=run_id + ) + st.config = local_versions + # In simulation, the raw_records generation depends on gain measurement + st.config["gain_model_mc"] = st.config["gain_model"] + # No blinding in simulations + st.config["event_info_function"] = "disabled" + return st + + +def xenonnt_simulation( + output_folder="./strax_data", + wfsim_registry="RawRecordsFromFaxNT", + cmt_run_id_sim=None, + cmt_run_id_proc=None, + cmt_version="global_ONLINE", + 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", + ), + **kwargs, +): + """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 detector + data. This means starting from already existing raw_records and finishing + with higher level data, such as peaks, events etc. + + If only one cmt_run_id is given, the second one will be set automatically, + resulting in CMT match between simulation and processing. However, detector + parameters can be still overwritten from fax file or manually using cmt + config overwrite options. + + CMT options can also be overwritten via fax config file. + :param output_folder: Output folder for strax data. + :param wfsim_registry: Name of WFSim plugin used to generate data. + :param cmt_run_id_sim: Run id for detector parameters from CMT to be used + for creation of raw_records. + :param cmt_run_id_proc: Run id for detector parameters from CMT to be used + 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 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 instead 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 + + """ + + st = strax.Context( + storage=strax.DataDirectory(output_folder), + config=dict( + detector="XENONnT", + fax_config=fax_config, + check_raw_record_overlaps=True, + **straxen.contexts.xnt_common_config, + ), + **straxen.contexts.xnt_common_opts, + **kwargs, + ) + st.register(getattr(wfsim, wfsim_registry)) + + # Make sure that the non-simulated raw-record types are not requested + st.deregister_plugins_with_missing_dependencies() + + if straxen.utilix_is_configured( + warning_message="Bad context as we cannot set CMT since we have no database access" + ): + st.apply_cmt_version(cmt_version) + + if _forbid_creation_of is not None: + st.context_config["forbid_creation_of"] += strax.to_str_tuple(_forbid_creation_of) + + # 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( + "cmt_run_id_sim and cmt_run_id_proc are None. " + "You have to specify at least one CMT run id. " + ) + 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 + + # Replace default cmt options with cmt_run_id tag + cmt run id + cmt_options_full = straxen.get_corrections.get_cmt_options(st) + + # prune to just get the strax options + cmt_options = {key: val["strax_option"] for key, val in cmt_options_full.items()} + + # 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(): + value = cmt_options[cmt_field] + + # URL configs need to be converted to the expected format + if isinstance(value, str): + opt_cfg = cmt_options_full[cmt_field] + version = straxen.URLConfig.kwarg_from_url(value, "version") + # We now allow the cmt name to be different from the config name + # WFSim expects the cmt name + value = (opt_cfg["correction"], version, True) + + fax_config_override_from_cmt[fax_field] = ("cmt_run_id", cmt_run_id_sim, *value) + st.set_config({"fax_config_override_from_cmt": fax_config_override_from_cmt}) + + # and all other parameters for processing + for option in cmt_options: + value = cmt_options[option] + if isinstance(value, str): + # for URL configs we can just replace the run_id keyword argument + # This will become the proper way to override the run_id for cmt configs + st.config[option] = straxen.URLConfig.format_url_kwargs(value, run_id=cmt_run_id_proc) + else: + # FIXME: Remove once all cmt configs are URLConfigs + st.config[option] = ("cmt_run_id", cmt_run_id_proc, *value) + + # Done with "default" usage, now to overwrites from file + # + # Take fax config and put into context option + 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(): + if overwrite_from_fax_file_proc: + if isinstance(cmt_options[cmt_field], str): + # URLConfigs can just be set to a constant + st.config[cmt_field] = fax_config[fax_field] + else: + # FIXME: Remove once all cmt configs are URLConfigs + st.config[cmt_field] = ( + cmt_options[cmt_field][0] + "_constant", + fax_config[fax_field], + ) + if overwrite_from_fax_file_sim: + # CMT name allowed to be different from the config name + # WFSim needs the cmt name + cmt_name = cmt_options_full[cmt_field]["correction"] + + st.config["fax_config_override_from_cmt"][fax_field] = ( + cmt_name + "_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 option not 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: + cmt_name = cmt_options_full[option]["correction"] + st.config["fax_config_override_from_cmt"][fax_key] = ( + cmt_name + "_constant", + cmt_option_overwrite_sim[option], + ) + del (fax_key, cmt_key) + # 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" + ) + + if isinstance(cmt_options[option], str): + # URLConfig options can just be set to constants, no hacks needed + # But for now lets keep things consistent for people + st.config[option] = cmt_option_overwrite_proc[option] + else: + # CMT name allowed to be different from the config name + # WFSim needs the cmt name + cmt_name = cmt_options_full[option]["correction"] + st.config[option] = (cmt_name + "_constant", cmt_option_overwrite_proc[option]) + # Only for simulations + st.set_config({"event_info_function": "disabled"}) + + return st + + +def xenon1t_simulation(output_folder="./strax_data"): + st = strax.Context( + storage=strax.DataDirectory(output_folder), + config=dict( + fax_config="fax_config_1t.json", detector="XENON1T", + **straxen.legacy.x1t_common_config, + ), + **straxen.legacy.get_x1t_context_config(), + ) + st.register(RawRecordsFromFax1T) + st.deregister_plugins_with_missing_dependencies() + return st