diff --git a/.bumpversion.cfg b/.bumpversion.cfg index 1d6628fa8..091be5ef8 100644 --- a/.bumpversion.cfg +++ b/.bumpversion.cfg @@ -1,5 +1,5 @@ [bumpversion] -current_version = 0.19.3 +current_version = 1.1.2 files = setup.py straxen/__init__.py commit = True tag = True diff --git a/.github/scripts/filter_strax_from_requirements.sh b/.github/scripts/filter_strax_from_requirements.sh new file mode 100644 index 000000000..2253577b9 --- /dev/null +++ b/.github/scripts/filter_strax_from_requirements.sh @@ -0,0 +1,3 @@ +echo `pwd` +cat requirements.txt | grep -v 'strax' &> requirements.txt +cat requirements.txt diff --git a/.github/scripts/preinstall_requirements.sh b/.github/scripts/preinstall_requirements.sh deleted file mode 100644 index 4361ea4bb..000000000 --- a/.github/scripts/preinstall_requirements.sh +++ /dev/null @@ -1,11 +0,0 @@ -#!/bin/bash - -echo "download requirements from base_environment" -wget -O pre_requirements.txt https://raw.githubusercontent.com/XENONnT/base_environment/master/requirements.txt &> /dev/null -echo "select important dependencies for strax(en)" -grep 'numpy\|numba\|blosc\|scikit-learn\|holoviews\|zstd' pre_requirements.txt &> sel_pre_requirements.txt -echo "Will pre-install:" -cat sel_pre_requirements.txt -echo "Start preinstall and rm pre-requirements:" -pip install -r sel_pre_requirements.txt -rm pre_requirements.txt sel_pre_requirements.txt diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml index 42b11ce39..a8c7e6bfa 100644 --- a/.github/workflows/pytest.yml +++ b/.github/workflows/pytest.yml @@ -18,21 +18,17 @@ on: release: types: [created] pull_request: - branches: - - master - - stable - - joran_master push: branches: - master - - joran_master + - stable jobs: update: name: "${{ matrix.test }}_py${{ matrix.python-version }}" runs-on: ubuntu-latest strategy: - fail-fast: True + fail-fast: False matrix: python-version: [3.7, 3.8, 3.9] test: ['coveralls', 'pytest', 'pytest_no_database'] @@ -54,11 +50,21 @@ jobs: python-version: ${{ matrix.python-version }} - name: Checkout repo uses: actions/checkout@v2 + - name: Remove strax from reqs. + if: matrix.test == 'coveralls' && env.HAVE_ACCESS_TO_SECTETS != null + env: + HAVE_ACCESS_TO_SECTETS: ${{ secrets.RUNDB_API_URL }} + run: | + bash .github/scripts/filter_strax_from_requirements.sh - name: Install requirements for tests and latest strax run: | pip install -r extra_requirements/requirements-tests.txt - pip install git+https://github.com/AxFoundation/strax.git - + git clone https://github.com/AxFoundation/strax ../strax + pip install -e ../strax + - name: Start MongoDB + uses: supercharge/mongodb-github-action@1.6.0 + with: + mongodb-version: 4.2 # Secrets and required files - name: patch utilix file # Patch this file if we want to have access to the database @@ -85,18 +91,20 @@ jobs: env: HAVE_ACCESS_TO_SECTETS: ${{ secrets.RUNDB_API_URL }} if: env.HAVE_ACCESS_TO_SECTETS == null || matrix.test == 'pytest_no_database' - run: | - bash .github/scripts/create_pre_apply_function.sh $HOME - - # Run tests + run: bash .github/scripts/create_pre_apply_function.sh $HOME - name: Test package # This is running a normal test if: matrix.test == 'pytest_no_database' || matrix.test == 'pytest' + env: + ALLOW_WFSIM_TEST: 1 + TEST_MONGO_URI: 'mongodb://localhost:27017/' run: | - python setup.py test -v + pytest -rsxv --durations 0 - name: Coveralls # Make the coverage report and upload env: + TEST_MONGO_URI: 'mongodb://localhost:27017/' + ALLOW_WFSIM_TEST: 1 NUMBA_DISABLE_JIT: 1 GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} # We need to check if we had access to the secrets, otherwise coveralls @@ -107,7 +115,6 @@ jobs: run: | coverage run --source=straxen setup.py test -v coveralls --service=github - # Done - name: goodbye run: echo "tests done, bye bye" diff --git a/.github/workflows/test_install.yml b/.github/workflows/test_install.yml new file mode 100644 index 000000000..8316a4482 --- /dev/null +++ b/.github/workflows/test_install.yml @@ -0,0 +1,36 @@ +# Test if we can actually install strax by installing +name: Installation test + +on: + workflow_dispatch: + release: + types: [created] + pull_request: + branches: + - master + - stable + push: + branches: + - master + +jobs: + update: + name: "py${{ matrix.python-version }}" + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: [3.8] + steps: + - name: Setup python + uses: actions/setup-python@v2 + with: + python-version: ${{ matrix.python-version }} + - name: Checkout repo + uses: actions/checkout@v2 + - name: Install straxen + run: python setup.py install + - name: Test import + run: python -c "import straxen; straxen.print_versions()" + - name: goodbye + run: echo goodbye diff --git a/HISTORY.md b/HISTORY.md index bd5794170..bdf69d9f2 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -1,3 +1,94 @@ +1.1.2 / 2021-10-27 +------------------- +minor / patches: +- Plugin for afterpulse processing (#549) +- Veto online monitor (#707) +- Refactor straxen tests (#703) +- WFSim registry as argument for simulations context (#713) +- Update S1 AFT map in event pattern fit (#697) + +fixes/tests: +- Set default drift time as nan (#700) +- Revert auto inclusion of rucio remote #688 (#701) +- fix bug in CMT (#710) +- Fix one year querries (#711 +- Test new numba (#702) +- Unify CMT call in contexts (#717) +- Small codefactor patch (#714) +- test nv with nv data (#709) +- Add small test for wfsim (#716) + + +1.1.1 / 2021-10-19 +------------------- + - Fix to test for RunDB frontend when no test DB is sourced (6da2233) + + +1.1.0 / 2021-10-18 +------------------- +major / minor: + +- Previous S2 Shadow Plugin draft (#664) +- Use admix in straxen (#688) +- Add posdiff plugin (#669) +- updated S2 corrected area (#686) +- Version bump of hitlets (#690) +- Add n saturated channels (#691) +- add small tool to extract run comments from database (#692) +- Update online_monitor_nv to v0.0.3 (#696) + + +patches and fixes: + +- Use read by index and check for NaNs (#661) +- Add small feature for printing versions of git (#665) +- Fix minianalyses from apply_selection (#666) +- fix some warnings from testing (#667) +- Add source to runs table (#673) +- Pbar patch for rundb query (#685) +- Implement SDSC as a local RSE for Expanse (#687) +- Skips superruns in rucio frontend (#689) +- Warn about non-loadable loggers (#693) +- Add RunDb read/write-test (#695) +- Fix bug in rucio frontend (#699) + + + +1.0.0 / 2021-09-01 +------------------- +major / minor: + +- merge s2 without s1 (#645) +- First nVeto monitor plugin (#634) +- Peak event veto tagging (#618) +- Fix peaklet area bias (#601) +- Add lone hit information to merged S2s. (#623) + + +patches and fixes: + +- Fix n_hits of peaks (#646) +- Update requirements for strax (#644) +- Modifications of nT simulation context (#602) +- Straxer for other packages (#595) +- [Bug fix] alt_s{i}_delay computation (#598) +- Bump version refactor code for cleanliness. (#597) +- Increase buffer size (#604) +- Stop testing py3.6 (#621) +- Remove online event monitor (#620) +- Add matplotlib to test requirements (#626) +- Fix rundb select runs with superruns (#627) +- Change EventInfo to save when explicit (#628) +- Update test data (#631) +- Allow database to not be initialized (#636) +- new plot_pmts (#637) +- Speed up event pattern fit (#625) +- kwargs for saver (#639) +- Add a plugin for external trigger run on nVeto calibration (#630) +- Fix veto event positions (#641) +- Use rucio from straxen & nest RucioRemote imports (#592) + + 0.19.3 / 2021-07-16 ------------------- - Rewrite EventBasics, set event level S1 tight coincidence (#569) diff --git a/README.md b/README.md index 649eb753e..e93da1196 100644 --- a/README.md +++ b/README.md @@ -18,4 +18,3 @@ For installation instructions and usage information, please see the [straxen doc [![Update context collection](https://github.com/XENONnT/straxen/workflows/Update%20context%20collection/badge.svg)](https://github.com/XENONnT/straxen/actions/workflows/contexts.yml) [![Python style](https://github.com/XENONnT/straxen/actions/workflows/code_style.yml/badge.svg)](https://github.com/XENONnT/straxen/actions/workflows/code_style.yml) -[![Coveralls](https://github.com/XENONnT/straxen/actions/workflows/coveralls.yml/badge.svg?branch=master)](https://github.com/XENONnT/straxen/actions/workflows/coveralls.yml) diff --git a/bin/bootstrax b/bin/bootstrax index 2f2d8da94..fe496b513 100755 --- a/bin/bootstrax +++ b/bin/bootstrax @@ -15,7 +15,7 @@ How to use For more info, see the documentation: https://straxen.readthedocs.io/en/latest/bootstrax.html """ -__version__ = '1.1.8' +__version__ = '1.1.9' import argparse from datetime import datetime, timedelta, timezone @@ -38,6 +38,7 @@ import straxen import threading import pandas as pd import typing as ty +import daqnt parser = argparse.ArgumentParser( description="XENONnT online processing manager") @@ -203,12 +204,12 @@ max_queue_new_runs = 2 # this many times. If high level data is hitting some edge case, we # might want to be able to keep the intermediate level data. # NB: string match so events applies e.g. to event_basics, peak to e.g. peaklets -remove_target_after_fails = dict(events=2, +remove_target_after_fails = dict(event=2, hitlets=2, peaks=4, + online=5, peak_basics=5, peaklets=6, - online=5, ALL=7) ## @@ -216,18 +217,17 @@ remove_target_after_fails = dict(events=2, ## hostname = socket.getfqdn() -try: - import daqnt - log_name = 'bootstrax_' + hostname + ('' if args.production else '_TESTING') - log = daqnt.get_daq_logger(log_name, log_name, level=logging.DEBUG) -except ModuleNotFoundError: - logging.basicConfig( - level=logging.INFO, - format='%(asctime)s %(name)s %(levelname)-8s %(message)s', - datefmt='%m-%d %H:%M') - log = logging.getLogger() +log_name = 'bootstrax_' + hostname + ('' if args.production else '_TESTING') +log = daqnt.get_daq_logger(log_name, log_name, level=logging.DEBUG) log.info(f'---\n bootstrax version {__version__}\n---') +log.info( + straxen.print_versions( + modules='strax straxen utilix daqnt numpy tensorflow numba'.split(), + include_git=True, + return_string=True, + )) + # Set the output folder output_folder = '/data/xenonnt_processed/' if args.production else test_data_folder @@ -273,9 +273,10 @@ def new_context(cores=args.cores, allow_multiprocess=cores > 1, allow_shm=cores > 1, allow_lazy=False, - use_rucio=False, max_messages=max_messages, - timeout=timeout) + timeout=timeout, + _rucio_path=None, + ) if not args.production: # Keep the rundb but set it to readonly and local only, delete # all other storage frontends except fo the test folder. @@ -522,7 +523,8 @@ def infer_target(rd: dict) -> dict: # Special modes override target for these led_modes = ['pmtgain'] - diagnostic_modes = ['exttrig', 'noise', 'pmtap'] + diagnostic_modes = ['exttrig', 'noise'] + ap_modes = ['pmtap'] mode = str(rd.get('mode')) detectors = list(rd.get('detectors')) @@ -531,6 +533,10 @@ def infer_target(rd: dict) -> dict: log.debug('led-mode') targets = 'led_calibration' post_process = 'raw_records' + elif np.any([m in mode for m in ap_modes]): + log.debug('afterpulse mode') + targets = 'afterpulses' + post_process = 'raw_records' elif np.any([m in mode for m in diagnostic_modes]): log.debug('diagnostic-mode') targets = 'raw_records' diff --git a/extra_requirements/requirements-tests.txt b/extra_requirements/requirements-tests.txt index 9e49cd6d5..c548396ac 100644 --- a/extra_requirements/requirements-tests.txt +++ b/extra_requirements/requirements-tests.txt @@ -1,30 +1,33 @@ # File for the requirements of straxen with the automated tests -blosc==1.10.4 # Strax dependency +blosc==1.10.6 # Strax dependency boltons==21.0.0 datashader==0.13.0 -dask==2021.7.2 +dask==2021.10.0 dill==0.3.4 # Strax dependency coveralls==3.2.0 commentjson==0.9.0 coverage==5.5 -flake8==3.9.2 +flake8==4.0.1 +gitpython==3.1.18 holoviews==1.14.5 -ipywidgets==7.6.3 -hypothesis==6.14.6 -jupyter-client==6.1.12 # for ipywidgets -matplotlib==3.4.2 +ipywidgets==7.6.4 +hypothesis==6.24.1 +jupyter-client==7.0.6 # for ipywidgets +matplotlib==3.4.3 multihist==0.6.4 npshmex==0.2.1 # Strax dependency -numba==0.53.1 # Strax dependency +numba==0.54.1 # Strax dependency numpy==1.19.5 -pandas==1.3.1 # Strax dependency +pandas==1.3.2 # Strax dependency psutil==5.8.0 # Strax dependency -pytest==6.2.4 -pytest-cov==2.12.1 -scikit-learn==0.24.2 +pytest==6.2.5 +pytest-cov==3.0.0 +scikit-learn==1.0.1 scipy==1.7.1 # Strax dependency -tensorflow==2.5.1 -tqdm==4.62.0 +tensorflow==2.6.0 +typing_extensions==3.7.4.3 # Tensorflow depencency +tqdm==4.62.2 +wfsim==0.5.9 xarray==0.19.0 -utilix==0.6.1 +utilix==0.6.5 zstd==1.5.0.2 # Strax dependency diff --git a/requirements.txt b/requirements.txt index 30613a3bc..239e91901 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,13 +1,15 @@ -strax>=0.16.0 +numpy~=1.19.5 +packaging~=20.8 +strax>=1.1.1 utilix>=0.5.3 # tensorflow>=2.3.0 # Optional, to (re)do posrec # holoviews # Optional, to enable wf display # datashader # Optional, to enable wf display bokeh>=2.2.3 multihist>=0.6.3 -packaging immutabledict numba>=0.50.0 requests commentjson matplotlib +gitpython diff --git a/setup.cfg b/setup.cfg index 0b79c7489..89eb170fe 100644 --- a/setup.cfg +++ b/setup.cfg @@ -61,6 +61,7 @@ exclude = # WPS319 Found bracket in wrong position # WPS326 Found implicit string concatenation # WPS331 Found variables that are only used for return: _time_zone_widget +# WPS335 Found incorrect for loop iter type # WPS420 Found wrong keyword: del # WPS421 Found wrong function call: print # WPS432 Found magic number: 2e-05 @@ -73,7 +74,7 @@ exclude = # WPS519 Found implicit sum() call # WPS602 Found using @staticmethod # Q000 Remove bad quotes -ignore = C409, C819, D100, D205, D400, D401, D403, DAR101, DAR201, DAR401, E128, E226, E251, F403, F541, I001, I003, I004, I005, P101, WPS110, WPS111, WPS211, WPS122, WPS210, WPS218, WPS229, WPS231, WPS237, WPS237, WPS300, WPS303, WPS304, WPS305, WPS306, WPS317, WPS318, WPS319, WPS326, WPS331, WPS336, WPS420, WPS421, WPS432, WPS433, WPS440, WPS435, WPS437, WPS503, WPS504, WPS519, WPS602, Q000 +ignore = C409, C819, D100, D205, D400, D401, D403, DAR101, DAR201, DAR401, E128, E226, E251, F403, F541, I001, I003, I004, I005, P101, WPS110, WPS111, WPS211, WPS122, WPS210, WPS218, WPS229, WPS231, WPS237, WPS237, WPS300, WPS303, WPS304, WPS305, WPS306, WPS317, WPS318, WPS319, WPS326, WPS331, WPS335, WPS336, WPS420, WPS421, WPS432, WPS433, WPS440, WPS435, WPS437, WPS503, WPS504, WPS519, WPS602, Q000 per-file-ignores = diff --git a/setup.py b/setup.py index c564e5642..2677b7840 100644 --- a/setup.py +++ b/setup.py @@ -21,7 +21,7 @@ def open_requirements(path): history = file.read() setuptools.setup(name='straxen', - version='0.19.3', + version='1.1.2', description='Streaming analysis for XENON', author='Straxen contributors, the XENON collaboration', url='https://github.com/XENONnT/straxen', diff --git a/straxen/__init__.py b/straxen/__init__.py index 17d1bcfc7..f00044077 100644 --- a/straxen/__init__.py +++ b/straxen/__init__.py @@ -1,4 +1,4 @@ -__version__ = '0.19.3' +__version__ = '1.1.2' from utilix import uconfig from .common import * @@ -26,6 +26,9 @@ # Otherwise we have straxen.demo() etc. from . import contexts +from . import test_utils +from .test_utils import * + try: from . import holoviews_utils from .holoviews_utils import * diff --git a/straxen/analyses/daq_waveforms.py b/straxen/analyses/daq_waveforms.py index fae7656ca..d112a909f 100644 --- a/straxen/analyses/daq_waveforms.py +++ b/straxen/analyses/daq_waveforms.py @@ -1,8 +1,7 @@ -import numba import pandas -import straxen import numpy as np -import strax +import straxen +import utilix import pymongo import typing import matplotlib.pyplot as plt @@ -20,7 +19,6 @@ def daq_plot(context, Plot with peak, records and records sorted by "link" or "ADC ID" (other items are also possible as long as it is in the channel map). """ - f, axes = plt.subplots(3, 1, figsize=figsize, gridspec_kw={'height_ratios': [1, 1, lower_panel_height]}) @@ -58,7 +56,6 @@ def daq_plot(context, def _get_daq_config( - context: strax.Context, run_id: str, config_name: str = 'daq_config', run_collection: typing.Optional[pymongo.collection.Collection] = None) -> dict: @@ -66,10 +63,10 @@ def _get_daq_config( Query the runs database for the config of the daq during this run. Either use the context of the runs collection. """ - if not context.storage[0].__class__.__name__ == 'RunDB' and run_collection is None: - raise NotImplementedError('Only works with the runs-database') if run_collection is None: - run_collection = context.storage[0].collection + if not straxen.utilix_is_configured(): + raise NotImplementedError('Only works with the runs-database') + run_collection = utilix.rundb.xent_collection() daq_doc = run_collection.find_one({"number": int(run_id)}, projection={config_name: 1}) if daq_doc is None or config_name not in daq_doc: @@ -123,12 +120,12 @@ def _group_channels_by_index(cable_map: pandas.DataFrame, return np.array(labels), idx -def group_by_daq(context, run_id, group_by: str): +def group_by_daq(run_id, group_by: str): """From the channel map, get the mapping of channel number -> group by""" cable_map = _get_cable_map() if group_by == 'link': labels, idx = _group_channels_by_index(cable_map, group_by='ADC ID') - daq_config = _get_daq_config(context, run_id) + daq_config = _get_daq_config(run_id) labels = [_board_to_host_link(daq_config, l) for l in labels] labels = np.array(labels) order = np.argsort(labels) diff --git a/straxen/analyses/pulse_plots.py b/straxen/analyses/pulse_plots.py index 46a1ea0fb..48f3de1c3 100644 --- a/straxen/analyses/pulse_plots.py +++ b/straxen/analyses/pulse_plots.py @@ -12,7 +12,8 @@ def plot_pulses_tpc(context, raw_records, run_id, time_range, plot_pulses(context, raw_records, run_id, time_range, plot_hits, plot_median, max_plots, store_pdf, path) - + + @straxen.mini_analysis(requires=('raw_records_mv',), warn_beyond_sec=5) def plot_pulses_mv(context, raw_records_mv, run_id, time_range, plot_hits=False, plot_median=False, @@ -79,6 +80,8 @@ def plot_pulses(context, raw_records, run_id, time_range, fname = os.path.join(path, fname) pdf = PdfPages(fname) + hits = None # needed for delete if false + for inds in _yield_pulse_indices(raw_records): # Grouped our pulse so now plot: rr_pulse = raw_records[inds] @@ -121,7 +124,6 @@ def plot_pulses(context, raw_records, run_id, time_range, ls='dotted', color='orange' ) - hits = None # needed for delet if false if plot_hits: min_amplitude = thr diff --git a/straxen/analyses/records_matrix.py b/straxen/analyses/records_matrix.py index efcc4d18b..fca0fd374 100644 --- a/straxen/analyses/records_matrix.py +++ b/straxen/analyses/records_matrix.py @@ -97,7 +97,7 @@ def raw_records_matrix(context, run_id, raw_records, time_range, @numba.njit def _records_to_matrix(records, t0, window, n_channels, dt=10): - n_samples = window // dt + n_samples = (window // dt) + 1 # Use 32-bit integers, so downsampling saturated samples doesn't # cause wraparounds # TODO: amplitude bit shift! @@ -114,7 +114,12 @@ def _records_to_matrix(records, t0, window, n_channels, dt=10): if dt >= samples_per_record * r['dt']: # Downsample to single sample -> store area - y[(r['time'] - t0) // dt, r['channel']] += r['area'] + idx = (r['time'] - t0) // dt + if idx >= len(y): + print(len(y), idx) + raise IndexError('Despite n_samples = window // dt + 1, our ' + 'idx is too high?!') + y[idx, r['channel']] += r['area'] continue # Assume out-of-bounds data has been zeroed, so we do not diff --git a/straxen/analyses/waveform_plot.py b/straxen/analyses/waveform_plot.py index cd7a2058b..f0776d9e2 100644 --- a/straxen/analyses/waveform_plot.py +++ b/straxen/analyses/waveform_plot.py @@ -1,11 +1,9 @@ import numpy as np import matplotlib import matplotlib.pyplot as plt -import warnings import strax import straxen from mpl_toolkits.axes_grid1 import inset_locator -from datetime import datetime from .records_matrix import DEFAULT_MAX_SAMPLES from .daq_waveforms import group_by_daq @@ -150,7 +148,7 @@ def plot_records_matrix(context, run_id, ignore_max_sample_warning=ignore_max_sample_warning, **kwargs) if group_by is not None: - ylabs, wvm_mask = group_by_daq(context, run_id, group_by) + ylabs, wvm_mask = group_by_daq(run_id, group_by) wvm = wvm[:, wvm_mask] plt.ylabel(group_by) else: diff --git a/straxen/common.py b/straxen/common.py index 957e9476a..375fc523a 100644 --- a/straxen/common.py +++ b/straxen/common.py @@ -2,7 +2,6 @@ import configparser import gzip import inspect -import io import commentjson import json import os @@ -11,7 +10,6 @@ import dill import socket import sys -import tarfile import urllib.request import tqdm import numpy as np @@ -25,7 +23,8 @@ export, __all__ = strax.exporter() __all__ += ['straxen_dir', 'first_sr1_run', 'tpc_r', 'tpc_z', 'aux_repo', 'n_tpc_pmts', 'n_top_pmts', 'n_hard_aqmon_start', 'ADC_TO_E', - 'n_nveto_pmts', 'n_mveto_pmts', 'tpc_pmt_radius', 'cryostat_outer_radius'] + 'n_nveto_pmts', 'n_mveto_pmts', 'tpc_pmt_radius', 'cryostat_outer_radius', + 'INFINITY_64BIT_SIGNED'] straxen_dir = os.path.dirname(os.path.abspath( inspect.getfile(inspect.currentframe()))) @@ -53,6 +52,8 @@ LAST_MISCABLED_RUN = 8796 TSTART_FIRST_CORRECTLY_CABLED_RUN = 1596036001000000000 +INFINITY_64BIT_SIGNED = 9223372036854775807 + @export def pmt_positions(xenon1t=False): @@ -241,70 +242,6 @@ def resource_from_url(html: str, fmt='text'): return result -@export -def get_secret(x): - """Return secret key x. In order of priority, we search: - * Environment variable: uppercase version of x - * xenon_secrets.py (if included with your straxen installation) - * A standard xenon_secrets.py located on the midway analysis hub - (if you are running on midway) - """ - warn("xenon_secrets is deprecated, and will be replaced with utilix" - "configuration file instead. See https://github.com/XENONnT/utilix") - env_name = x.upper() - if env_name in os.environ: - return os.environ[env_name] - - message = (f"Secret {x} requested, but there is no environment " - f"variable {env_name}, ") - - # now try using utilix. We need to check that it is not None first! - # this will be main method in a future release - if straxen.uconfig is not None and straxen.uconfig.has_option('straxen', x): - try: - return straxen.uconfig.get('straxen', x) - except configparser.NoOptionError: - warn(f'straxen.uconfig does not have {x}') - - # if that doesn't work, revert to xenon_secrets - try: - from . import xenon_secrets - except ImportError: - message += ("nor was there a valid xenon_secrets.py " - "included with your straxen installation, ") - - # If on midway, try loading a standard secrets file instead - if 'rcc' in socket.getfqdn(): - path_to_secrets = '/project2/lgrandi/xenonnt/xenon_secrets.py' - if os.path.exists(path_to_secrets): - sys.path.append(osp.dirname(path_to_secrets)) - import xenon_secrets - sys.path.pop() - else: - raise ValueError( - message + ' nor could we load the secrets module from ' - f'{path_to_secrets}, even though you seem ' - 'to be on the midway analysis hub.') - - else: - raise ValueError( - message + 'nor are you on the midway analysis hub.') - - if hasattr(xenon_secrets, x): - return getattr(xenon_secrets, x) - raise ValueError(message + " and the secret is not in xenon_secrets.py") - - -@export -def download_test_data(): - """Downloads strax test data to strax_test_data in the current directory""" - blob = get_resource('https://raw.githubusercontent.com/XENONnT/strax_auxiliary_files/609b492e1389369734c7d2cbabb38059f14fc05e/strax_files/strax_test_data_straxv0.9.tar', # noqa - fmt='binary') - f = io.BytesIO(blob) - tf = tarfile.open(fileobj=f) - tf.extractall() - - @export def get_livetime_sec(context, run_id, things=None): """Get the livetime of a run in seconds. If it is not in the run metadata, @@ -340,7 +277,7 @@ def pre_apply_function(data, run_id, target, function_name='pre_apply_function') if function_name not in _resource_cache: # only load the function once and put it in the resource cache function_file = f'{function_name}.py' - function_file = _overwrite_testing_function_file(function_file) + function_file = straxen.test_utils._overwrite_testing_function_file(function_file) function = get_resource(function_file, fmt='txt') # pylint: disable=exec-used exec(function) @@ -351,26 +288,8 @@ def pre_apply_function(data, run_id, target, function_name='pre_apply_function') def _overwrite_testing_function_file(function_file): - """For testing purposes allow this function file to be loaded from HOME/testing_folder""" - if not straxen._is_on_pytest(): - # If we are not on a pytest, never try using a local file. - return function_file - - home = os.environ.get('HOME') - if home is None: - # Impossible to load from non-existent folder - return function_file - - testing_file = os.path.join(home, function_file) - - if os.path.exists(testing_file): - # For testing purposes allow loading from 'home/testing_folder' - warn(f'Using local function: {function_file} from {testing_file}! ' - f'If you are not integrated testing on github you should ' - f'absolutely remove this file. (See #559)') - function_file = testing_file - - return function_file + warn('Use straxen.test_utils._overwrite_testing_function_file') + return straxen._overwrite_testing_function_file(function_file) @export @@ -613,6 +532,7 @@ def _swap_values_in_array(data_arr, buffer, items, replacements): break return buffer + ## # Old XENON1T Stuff ## diff --git a/straxen/contexts.py b/straxen/contexts.py index 2b0ef2e6c..82c53511c 100644 --- a/straxen/contexts.py +++ b/straxen/contexts.py @@ -2,10 +2,11 @@ import strax import straxen from copy import deepcopy +from .rucio import HAVE_ADMIX +import os common_opts = dict( register_all=[ - straxen.event_processing, straxen.double_scatter, ], # Register all peak/pulse processing by hand as 1T does not need to have @@ -17,11 +18,17 @@ straxen.MergedS2s, straxen.Peaks, straxen.PeakBasics, - straxen.PeakProximity], + straxen.PeakProximity, + straxen.Events, + straxen.EventBasics, + straxen.EventPositions, + straxen.CorrectedAreas, + straxen.EnergyEstimates, + ], check_available=('raw_records', 'peak_basics'), store_run_fields=( 'name', 'number', - 'start', 'end', 'livetime', 'mode')) + 'start', 'end', 'livetime', 'mode', 'source')) xnt_common_config = dict( n_tpc_pmts=straxen.n_tpc_pmts, @@ -68,11 +75,13 @@ straxen.PeakPositionsMLP, straxen.PeakPositionsGCN, straxen.PeakPositionsNT, + straxen.S2ReconPosDiff, straxen.PeakBasicsHighEnergy, straxen.PeaksHighEnergy, straxen.PeakletsHighEnergy, straxen.PeakletClassificationHighEnergy, straxen.MergedS2sHighEnergy, + straxen.PeakVetoTagging, straxen.EventInfo, ], 'register_all': common_opts['register_all'] + [straxen.veto_veto_regions, @@ -86,10 +95,11 @@ straxen.online_monitor, straxen.event_area_per_channel, straxen.event_patternfit, + straxen.event_processing, + straxen.event_shadow, ], 'use_per_run_defaults': False, }) - ## # XENONnT ## @@ -103,14 +113,14 @@ def xenonnt(cmt_version='global_ONLINE', **kwargs): def xenonnt_online(output_folder='./strax_data', - use_rucio=False, we_are_the_daq=False, + download_heavy=False, _minimum_run_number=7157, _maximum_run_number=None, _database_init=True, _forbid_creation_of=None, - _rucio_path='/dali/lgrandi/rucio/', _include_rucio_remote=False, + _rucio_path='/dali/lgrandi/rucio/', _raw_path='/dali/lgrandi/xenonnt/raw', _processed_path='/dali/lgrandi/xenonnt/processed', _add_online_monitor_frontend=False, @@ -121,8 +131,8 @@ def xenonnt_online(output_folder='./strax_data', :param output_folder: str, Path of the strax.DataDirectory where new data can be stored - :param use_rucio: bool, whether or not to use the rucio frontend :param we_are_the_daq: bool, if we have admin access to upload data + :param download_heavy: bool, whether or not to allow downloads of heavy data (raw_records*, less the aqmon) :param _minimum_run_number: int, lowest number to consider :param _maximum_run_number: Highest number to consider. When None (the default) consider all runs that are higher than the @@ -147,7 +157,7 @@ def xenonnt_online(output_folder='./strax_data', st = strax.Context( config=straxen.contexts.xnt_common_config, **context_options) - st.register([straxen.DAQReader, straxen.LEDCalibration]) + st.register([straxen.DAQReader, straxen.LEDCalibration, straxen.LEDAfterpulseProcessing]) st.storage = [ straxen.RunDB( @@ -169,21 +179,21 @@ def xenonnt_online(output_folder='./strax_data', readonly=True, )] if output_folder: - st.storage.append( - strax.DataDirectory(output_folder, - provide_run_metadata=True, - )) - + st.storage += [strax.DataDirectory(output_folder, + provide_run_metadata=True, + )] st.context_config['forbid_creation_of'] = straxen.daqreader.DAQReader.provides if _forbid_creation_of is not None: st.context_config['forbid_creation_of'] += strax.to_str_tuple(_forbid_creation_of) - # if we said so, add the rucio frontend to storage - if use_rucio: - st.storage.append(straxen.rucio.RucioFrontend( - include_remote=straxen.RUCIO_AVAILABLE, - staging_dir=output_folder, - )) + # Add the rucio frontend if we are able to + if HAVE_ADMIX: + rucio_frontend = straxen.rucio.RucioFrontend( + include_remote=_include_rucio_remote, + staging_dir=os.path.join(output_folder, 'rucio'), + download_heavy=download_heavy, + ) + st.storage += [rucio_frontend] # Only the online monitor backend for the DAQ if _database_init and (_add_online_monitor_frontend or we_are_the_daq): @@ -191,7 +201,10 @@ def xenonnt_online(output_folder='./strax_data', readonly=not we_are_the_daq, take_only=('veto_intervals', 'online_peak_monitor', - 'event_basics',))] + 'event_basics', + 'online_monitor_nv', + 'online_monitor_mv', + ))] # Remap the data if it is before channel swap (because of wrongly cabled # signal cable connectors) These are runs older than run 8797. Runs @@ -209,21 +222,32 @@ def xenonnt_online(output_folder='./strax_data', def xenonnt_led(**kwargs): st = xenonnt_online(**kwargs) - st.context_config['check_available'] = ('raw_records', 'led_calibration') + st.set_context_config( + {'check_available': ('raw_records', 'led_calibration'), + 'free_options': list(xnt_common_config.keys()) + }) # Return a new context with only raw_records and led_calibration registered st = st.new_context( replace=True, config=st.config, storage=st.storage, **st.context_config) - st.register([straxen.DAQReader, straxen.LEDCalibration]) + st.register([straxen.DAQReader, + straxen.LEDCalibration, + straxen.nVETORecorder, + straxen.nVETOPulseProcessing, + straxen.nVETOHitlets, + straxen.nVetoExtTimings, ]) + st.set_config({"coincidence_level_recorder_nv": 1}) return st + def xenonnt_simulation( output_folder='./strax_data', + wfsim_registry='RawRecordsFromFaxNT', cmt_run_id_sim=None, cmt_run_id_proc=None, - cmt_version='v3', + cmt_version='global_v5', fax_config='fax_config_nt_design.json', overwrite_from_fax_file_sim=False, overwrite_from_fax_file_proc=False, @@ -234,7 +258,7 @@ def xenonnt_simulation( drift_time_gate='electron_drift_time_gate', drift_velocity_liquid='electron_drift_velocity', electron_lifetime_liquid='elife_conf'), - **kwargs): + **kwargs): """ The most generic context that allows for setting full divergent settings for simulation purposes @@ -255,6 +279,7 @@ def xenonnt_simulation( 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 @@ -285,8 +310,15 @@ def xenonnt_simulation( check_raw_record_overlaps=True, **straxen.contexts.xnt_common_config,), **straxen.contexts.xnt_common_opts, **kwargs) - st.register(wfsim.RawRecordsFromFaxNT) - st.apply_cmt_version(f'global_{cmt_version}') + 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) @@ -331,7 +363,7 @@ def xenonnt_simulation( 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]) + 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 @@ -339,16 +371,16 @@ def xenonnt_simulation( 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(): + 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: + f'CMT to fax config!') + for fax_key, cmt_key in _config_overlap.items(): + if cmt_key == option: _name_index = 2 if 'cmt_run_id' in cmt_options[option] else 0 st.config['fax_config_override_from_cmt'][fax_key] = ( cmt_options[option][_name_index] + '_constant', cmt_option_overwrite_sim[option]) - del(_name_index) + del _name_index del(fax_key, cmt_key) # User customized for simulation for option in cmt_option_overwrite_proc: @@ -358,7 +390,7 @@ def xenonnt_simulation( _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) + del _name_index # Only for simulations st.set_config({"event_info_function": "disabled"}) @@ -504,7 +536,10 @@ def xenon1t_dali(output_folder='./strax_data', build_lowlevel=False, **kwargs): def xenon1t_led(**kwargs): st = xenon1t_dali(**kwargs) - st.context_config['check_available'] = ('raw_records', 'led_calibration') + st.set_context_config( + {'check_available': ('raw_records', 'led_calibration'), + 'free_options': list(x1t_context_config.keys()) + }) # Return a new context with only raw_records and led_calibration registered st = st.new_context( replace=True, @@ -525,4 +560,5 @@ def xenon1t_simulation(output_folder='./strax_data'): **x1t_common_config), **x1t_context_config) st.register(wfsim.RawRecordsFromFax1T) + st.deregister_plugins_with_missing_dependencies() return st diff --git a/straxen/corrections_services.py b/straxen/corrections_services.py index 3c844c4ad..fee4684b1 100644 --- a/straxen/corrections_services.py +++ b/straxen/corrections_services.py @@ -14,10 +14,10 @@ export, __all__ = strax.exporter() -corrections_w_file = ['mlp_model', 'gcn_model', 'cnn_model', - 's2_xy_map', 's1_xyz_map_mlp', 's1_xyz_map_cnn', - 's1_xyz_map_gcn', 'fdc_map_mlp', 'fdc_map_gcn', - 'fdc_map_cnn'] +corrections_w_file = ['mlp_model', 'cnn_model', 'gcn_model', + 's2_xy_map_mlp', 's2_xy_map_cnn', 's2_xy_map_gcn', 's2_xy_map', + 's1_xyz_map_mlp', 's1_xyz_map_cnn', 's1_xyz_map_gcn', + 'fdc_map_mlp', 'fdc_map_cnn', 'fdc_map_gcn'] single_value_corrections = ['elife_xenon1t', 'elife', 'baseline_samples_nv', 'electron_drift_velocity', 'electron_drift_time_gate'] @@ -27,12 +27,14 @@ # needed because we pass these names as strax options which then get paired with the default reconstruction algorithm # important for apply_cmt_version -posrec_corrections_basenames = ['s1_xyz_map', 'fdc_map'] +posrec_corrections_basenames = ['s1_xyz_map', 'fdc_map', 's2_xy_map'] class CMTVersionError(Exception): pass +class CMTnanValueError(Exception): + pass @export class CorrectionsManagementServices(): @@ -130,14 +132,22 @@ def _get_correction(self, run_id, correction, version): pmts = list(gains.keys()) for it_correction in pmts: # loop over all PMTs if correction in it_correction: - df = self.interface.read(it_correction) + df = self.interface.read_at(it_correction, when) + if df[version].isnull().values.any(): + raise CMTnanValueError(f"For {it_correction} there are NaN values, this means no correction available " + f"for {run_id} in version {version}, please check e-logbook for more info ") + if version in 'ONLINE': df = self.interface.interpolate(df, when, how='fill') else: df = self.interface.interpolate(df, when) values.append(df.loc[df.index == when, version].values[0]) else: - df = self.interface.read(correction) + df = self.interface.read_at(correction, when) + if df[version].isnull().values.any(): + raise CMTnanValueError(f"For {correction} there are NaN values, this means no correction available " + f"for {run_id} in version {version}, please check e-logbook for more info ") + if correction in corrections_w_file or correction in arrays_corrections or version in 'ONLINE': df = self.interface.interpolate(df, when, how='fill') else: @@ -189,12 +199,10 @@ def get_pmt_gains(self, run_id, model_type, version, to_pe = self._get_correction(run_id, target_detector, version) # be cautious with very early runs, check that not all are None - if np.isnan(to_pe).all(): + if np.isnan(to_pe).any(): raise ValueError( f"to_pe(PMT gains) values are NaN, no data available " - f"for {run_id} in the gain model with version " - f"{version}, please set constant values for " - f"{run_id}") + f"for {run_id} in the gain model with version") else: raise ValueError(f"{model_type} not implemented for to_pe values") diff --git a/straxen/get_corrections.py b/straxen/get_corrections.py index b12c5aa97..86cb6f6b5 100644 --- a/straxen/get_corrections.py +++ b/straxen/get_corrections.py @@ -131,6 +131,7 @@ def get_cmt_resource(run_id, conf, fmt=''): return straxen.get_resource(get_correction_from_cmt(run_id, conf), fmt=fmt) +@export def is_cmt_option(config): """ Check if the input configuration is cmt style. diff --git a/straxen/matplotlib_utils.py b/straxen/matplotlib_utils.py index 8e130f4e8..9e9dd8219 100644 --- a/straxen/matplotlib_utils.py +++ b/straxen/matplotlib_utils.py @@ -6,6 +6,7 @@ import strax import straxen + export, __all__ = strax.exporter() @@ -18,14 +19,13 @@ def plot_pmts( extend='neither', vmin=None, vmax=None, **kwargs): """Plot the PMT arrays side-by-side, coloring the PMTS with c. - :param c: Array of colors to use. Must have len() n_tpc_pmts :param label: Label for the color bar :param figsize: Figure size to use. :param extend: same as plt.colorbar(extend=...) :param vmin: Minimum of color scale :param vmax: maximum of color scale - + :param show_axis_labels: if True it will show x and y labels Other arguments are passed to plot_on_single_pmt_array. """ if vmin is None: @@ -36,15 +36,16 @@ def plot_pmts( # Single-valued array passed vmax += 1 if figsize is None: - figsize = (11, 4) if xenon1t else (13, 5.5) + figsize = (11.25, 4.25) if xenon1t else (13.25, 5.75) f, axes = plt.subplots(1, 2, figsize=figsize) + plot_result = None for array_i, array_name in enumerate(['top', 'bottom']): ax = axes[array_i] plt.sca(ax) plt.title(array_name.capitalize()) - plot_on_single_pmt_array( + plot_result = plot_on_single_pmt_array( c, xenon1t=xenon1t, array_name=array_name, @@ -52,13 +53,16 @@ def plot_pmts( vmin=vmin, vmax=vmax, **kwargs) + axes[0].set_xlabel('x [cm]') + axes[0].xaxis.set_label_coords(1.035, -0.075) + axes[0].set_ylabel('y [cm]') + axes[1].yaxis.tick_right() axes[1].yaxis.set_label_position('right') plt.tight_layout() plt.subplots_adjust(wspace=0) - plt.colorbar(ax=axes, extend=extend, label=label) - + plt.colorbar(mappable=plot_result, ax=axes, extend=extend, label=label) @export @@ -74,7 +78,6 @@ def plot_on_single_pmt_array( dead_pmts=None, dead_pmt_color='gray', **kwargs): """Plot one of the PMT arrays and color it by c. - :param c: Array of colors to use. Must be len() of the number of TPC PMTs :param label: Label for the color bar :param pmt_label_size: Fontsize for the PMT number labels. @@ -84,7 +87,6 @@ def plot_on_single_pmt_array( :param extend: same as plt.colorbar(extend=...) :param vmin: Minimum of color scale :param vmax: maximum of color scale - Other arguments are passed to plt.scatter. """ if vmin is None: @@ -129,7 +131,7 @@ def plot_on_single_pmt_array( ax.set_axis_off() if dead_pmts is not None: _dead_mask = [pi in dead_pmts for pi in pos['i']] - result = plt.scatter( + plt.scatter( pos[_dead_mask]['x'], pos[_dead_mask]['y'], c=dead_pmt_color, @@ -201,6 +203,7 @@ def draw_box(x, y, **kwargs): plt.gca().add_patch(matplotlib.patches.Rectangle( (x[0], y[0]), x[1] - x[0], y[1] - y[0], facecolor='none', **kwargs)) + @export def plot_single_pulse(records, run_id, pulse_i=''): """ diff --git a/straxen/mini_analysis.py b/straxen/mini_analysis.py index 21a5e4ddb..4ac9ee616 100644 --- a/straxen/mini_analysis.py +++ b/straxen/mini_analysis.py @@ -94,7 +94,7 @@ def wrapped_f(context: strax.Context, run_id: str, **kwargs): for dkind, dtypes in deps_by_kind.items(): if dkind in kwargs: # Already have data, just apply cuts - kwargs[dkind] = context.apply_selection( + kwargs[dkind] = strax.apply_selection( kwargs[dkind], selection_str=kwargs['selection_str'], time_range=kwargs['time_range'], diff --git a/straxen/misc.py b/straxen/misc.py index 6e41ef0be..a22c15dd3 100644 --- a/straxen/misc.py +++ b/straxen/misc.py @@ -7,9 +7,10 @@ import warnings import datetime import pytz -from os import environ as os_environ from importlib import import_module +from git import Repo, InvalidGitRepositoryError from configparser import NoSectionError +import typing as ty export, __all__ = strax.exporter() @@ -41,7 +42,9 @@ def do_round(x): @export -def print_versions(modules=('strax', 'straxen', 'cutax'), return_string=False): +def print_versions(modules=('strax', 'straxen', 'cutax'), + return_string=False, + include_git=True): """ Print versions of modules installed. @@ -50,6 +53,8 @@ def print_versions(modules=('strax', 'straxen', 'cutax'), return_string=False): 'cutax', 'pema')) :param return_string: optional. Instead of printing the message, return a string + :param include_git: Include the current branch and latest + commit hash :return: optional, the message that would have been printed """ message = (f'Working on {socket.getfqdn()} with the following ' @@ -59,29 +64,65 @@ def print_versions(modules=('strax', 'straxen', 'cutax'), return_string=False): for m in strax.to_str_tuple(modules): try: mod = import_module(m) - message += f'\n{m}' - if hasattr(mod, '__version__'): - message += f'\tv{mod.__version__}' - if hasattr(mod, '__path__'): - message += f'\t{mod.__path__[0]}' except (ModuleNotFoundError, ImportError): print(f'{m} is not installed') + continue + + message += f'\n{m}' + if hasattr(mod, '__version__'): + message += f'\tv{mod.__version__}' + if hasattr(mod, '__path__'): + module_path = mod.__path__[0] + message += f'\t{module_path}' + if include_git: + try: + repo = Repo(module_path, search_parent_directories=True) + except InvalidGitRepositoryError: + # not a git repo + pass + else: + try: + branch = repo.active_branch + except TypeError: + branch = 'unknown' + try: + commit_hash = repo.head.object.hexsha + except TypeError: + commit_hash = 'unknown' + message += f'\tgit branch:{branch} | {commit_hash[:7]}' if return_string: return message print(message) @export -def utilix_is_configured(header='RunDB', section='xent_database') -> bool: +def utilix_is_configured(header: str = 'RunDB', + section: str = 'xent_database', + warning_message: ty.Union[None, bool, str] = None, + ) -> bool: """ Check if we have the right connection to :return: bool, can we connect to the Mongo database? + + :param header: Which header to check in the utilix config file + :param section: Which entry in the header to check to exist + :param warning_message: If utilix is not configured, warn the user. + if None -> generic warning + if str -> use the string to warn + if False -> don't warn """ try: - return (hasattr(straxen.uconfig, 'get') and - straxen.uconfig.get(header, section) is not None) + is_configured = (hasattr(straxen.uconfig, 'get') and + straxen.uconfig.get(header, section) is not None) except NoSectionError: - return False + is_configured = False + + should_report = bool(warning_message) or warning_message is None + if not is_configured and should_report: + if warning_message is None: + warning_message = 'Utilix is not configured, cannot proceed' + warnings.warn(warning_message) + return is_configured @export @@ -203,6 +244,27 @@ def _convert_to_datetime(time_widget, time_zone): return time, time_ns +@strax.Context.add_method +def extract_latest_comment(self): + """ + Extract the latest comment in the runs-database. This just adds info to st.runs + + Example: + st.extract_latest_comment() + st.select_runs(available=('raw_records')) + """ + if self.runs is None or 'comments' not in self.runs.keys(): + self.scan_runs(store_fields=('comments',)) + latest_comments = _parse_to_last_comment(self.runs['comments']) + self.runs['comments'] = latest_comments + return self.runs + + +def _parse_to_last_comment(comments): + """Unpack to get the last comment (hence the -1) or give '' when there is none""" + return [(c[-1]['comment'] if hasattr(c, '__len__') else '') for c in comments] + + @export def convert_array_to_df(array: np.ndarray) -> pd.DataFrame: """ @@ -214,9 +276,3 @@ def convert_array_to_df(array: np.ndarray) -> pd.DataFrame: """ keys = [key for key in array.dtype.names if array[key].ndim == 1] return pd.DataFrame(array[keys]) - - -@export -def _is_on_pytest(): - """Check if we are on a pytest""" - return 'PYTEST_CURRENT_TEST' in os_environ diff --git a/straxen/mongo_storage.py b/straxen/mongo_storage.py index f0e3683f6..c63bdf58b 100644 --- a/straxen/mongo_storage.py +++ b/straxen/mongo_storage.py @@ -32,6 +32,7 @@ def __init__(self, file_database='files', config_identifier='config_name', collection=None, + _test_on_init=False, ): """ GridFsInterface @@ -46,7 +47,10 @@ def __init__(self, PyMongo DataName Collection to bypass normal initiation using utilix. Should be an object of the form: pymongo.MongoClient(..).DATABASE_NAME.COLLECTION_NAME + :param _test_on_init: Test if the collection is empty on init + (only deactivate if you are using a brand new database)! """ + if collection is None: if not readonly: # We want admin access to start writing data! @@ -79,7 +83,8 @@ def __init__(self, # Set collection and make sure it can at least do a 'find' operation self.collection = collection - self.test_find() + if _test_on_init: + self.test_find() # This is the identifier under which we store the files. self.config_identifier = config_identifier diff --git a/straxen/numbafied_scipy.py b/straxen/numbafied_scipy.py new file mode 100644 index 000000000..f81b5fd1e --- /dev/null +++ b/straxen/numbafied_scipy.py @@ -0,0 +1,36 @@ +import numba +import ctypes +import strax +from scipy.special.cython_special import loggamma, betainc, gammaln + +export, __all__ = strax.exporter() +# Numba versions of scipy functions: +# Based on http://numba.pydata.org/numba-doc/latest/extending/high-level.html but +# omitting the address step. +# Unfortunately this approach does not allow to cache the resulting numba functions. +# Any ideas in this regard would be welcome. +# All subsequent function will treat their inputs as 64bit floats: +double = ctypes.c_double + +functype = ctypes.CFUNCTYPE(double, double) +gammaln_float64 = functype(gammaln) +@export +@numba.njit +def numba_gammaln(x): + return gammaln_float64(x) + + +functype = ctypes.CFUNCTYPE(double, double, double, double) +betainc_float64 = functype(betainc) +@export +@numba.njit +def numba_betainc(x1, x2, x3): + return betainc_float64(x1, x2, x3) + + +functype = ctypes.CFUNCTYPE(double, double) +loggamma_float64 = functype(loggamma) +@export +@numba.njit +def numba_loggamma(x): + return loggamma_float64(x) diff --git a/straxen/plugins/__init__.py b/straxen/plugins/__init__.py index 934a4907f..29c0f51df 100644 --- a/straxen/plugins/__init__.py +++ b/straxen/plugins/__init__.py @@ -28,6 +28,9 @@ from . import event_processing from .event_processing import * +from . import afterpulse_processing +from .afterpulse_processing import * + from . import double_scatter from .double_scatter import * @@ -56,3 +59,6 @@ from . import online_monitor from .online_monitor import * + +from . import event_shadow +from .event_shadow import * diff --git a/straxen/plugins/afterpulse_processing.py b/straxen/plugins/afterpulse_processing.py new file mode 100644 index 000000000..ddf9ed16d --- /dev/null +++ b/straxen/plugins/afterpulse_processing.py @@ -0,0 +1,298 @@ +import numba +import numpy as np +import strax +import straxen +from straxen.get_corrections import is_cmt_option + +export, __all__ = strax.exporter() + + +@export +@strax.takes_config( + strax.Option('gain_model', + help='PMT gain model. Specify as (model_type, model_config)', + ), + strax.Option('n_tpc_pmts', + type=int, + help="Number of PMTs in TPC", + ), + strax.Option('LED_window_left', + default=50, + help='Left boundary of sample range for LED pulse integration', + ), + strax.Option('LED_window_right', + default=100, + help='Right boundary of sample range for LED pulse integration', + ), + strax.Option('baseline_samples', + default=40, + help='Number of samples to use at start of WF to determine the baseline', + ), + strax.Option('hit_min_amplitude', + track=True, + default=('hit_thresholds_tpc', 'ONLINE', True), + help='Minimum hit amplitude in ADC counts above baseline. ' + 'Specify as a tuple of length n_tpc_pmts, or a number,' + 'or a string like "pmt_commissioning_initial" which means calling' + 'hitfinder_thresholds.py' + 'or a tuple like (correction=str, version=str, nT=boolean),' + 'which means we are using cmt.', + ), + strax.Option('hit_min_height_over_noise', + default=4, + help='Minimum hit amplitude in numbers of baseline_rms above baseline.' + 'Actual threshold used is max(hit_min_amplitude, hit_min_' + 'height_over_noise * baseline_rms).', + ), + strax.Option('save_outside_hits', + default=(3, 20), + help='Save (left, right) samples besides hits; cut the rest', + ), +) +class LEDAfterpulseProcessing(strax.Plugin): + __version__ = '0.5.1' + depends_on = 'raw_records' + data_kind = 'afterpulses' + provides = 'afterpulses' + compressor = 'zstd' + parallel = 'process' + rechunk_on_save = True + + def infer_dtype(self): + dtype = dtype_afterpulses() + return dtype + + def setup(self): + self.to_pe = straxen.get_correction_from_cmt(self.run_id, self.config['gain_model']) + self.hit_left_extension, self.hit_right_extension = self.config['save_outside_hits'] + + # Check config of `hit_min_amplitude` and define hit thresholds + # if cmt config + if is_cmt_option(self.config['hit_min_amplitude']): + self.hit_thresholds = straxen.get_correction_from_cmt( + self.run_id, + self.config['hit_min_amplitude']) + # if hitfinder_thresholds config + elif isinstance(self.config['hit_min_amplitude'], str): + self.hit_thresholds = straxen.hit_min_amplitude( + self.config['hit_min_amplitude']) + else: # int or array + self.hit_thresholds = self.config['hit_min_amplitude'] + + def compute(self, raw_records): + + # Convert everything to the records data type -- adds extra fields. + records = strax.raw_to_records(raw_records) + del raw_records + + # calculate baseline and baseline rms + strax.baseline(records, + baseline_samples=self.config['baseline_samples'], + flip=True) + + # find all hits + hits = strax.find_hits(records, + min_amplitude=self.hit_thresholds, + min_height_over_noise=self.config['hit_min_height_over_noise'], + ) + + # sort hits by record_i and time, then find LED hit and afterpulse + # hits within the same record + hits_ap = find_ap(hits, + records, + LED_window_left=self.config['LED_window_left'], + LED_window_right=self.config['LED_window_right'], + hit_left_extension=self.hit_left_extension, + hit_right_extension=self.hit_right_extension, + ) + + hits_ap['area_pe'] = hits_ap['area'] * self.to_pe[hits_ap['channel']] + hits_ap['height_pe'] = hits_ap['height'] * self.to_pe[hits_ap['channel']] + + return hits_ap + + +@export +def find_ap(hits, records, LED_window_left, LED_window_right, hit_left_extension, + hit_right_extension): + buffer = np.zeros(len(hits), dtype=dtype_afterpulses()) + + if not len(hits): + return buffer + + # sort hits first by record_i, then by time + hits_sorted = np.sort(hits, order=('record_i', 'time')) + res = _find_ap(hits_sorted, records, LED_window_left, LED_window_right, + hit_left_extension, hit_right_extension, buffer=buffer) + return res + + +@numba.jit(nopython=True, nogil=True, cache=True) +def _find_ap(hits, records, LED_window_left, LED_window_right, hit_left_extension, + hit_right_extension, buffer=None): + # hits need to be sorted by record_i, then time! + offset = 0 + + is_LED = False + t_LED = None + + prev_record_i = hits[0]['record_i'] + record_data = records[prev_record_i]['data'] + record_len = records[prev_record_i]['length'] + baseline_fpart = records[prev_record_i]['baseline'] % 1 + + for h_i, h in enumerate(hits): + + if h['record_i'] > prev_record_i: + # start of a new record + is_LED = False + # only increment buffer if the old one is not empty! this happens + # when no (LED) hit is found in the previous record + if not buffer[offset]['time'] == 0: + offset += 1 + prev_record_i = h['record_i'] + record_data = records[prev_record_i]['data'] + baseline_fpart = records[prev_record_i]['baseline'] % 1 + + res = buffer[offset] + + if h['left'] < LED_window_left: + # if hit is before LED window: discard + continue + + if h['left'] < LED_window_right: + # hit is in LED window + if not is_LED: + # this is the first hit in the LED window + fill_hitpars(res, h, hit_left_extension, hit_right_extension, + record_data, record_len, baseline_fpart) + + # set the LED time in the current WF + t_LED = res['sample_10pc_area'] + is_LED = True + + continue + + # more hits in LED window: extend the first (merging all hits in the LED window) + fill_hitpars(res, h, hit_left_extension, hit_right_extension, + record_data, record_len, + baseline_fpart, extend=True) + + t_LED = res['sample_10pc_area'] + + continue + + # Here begins a new hit after the LED window + if (h['left'] >= LED_window_right) and not is_LED: + # no LED hit found: ignore and go to next hit (until new record begins) + continue + + # if a hit is completely inside the previous hit's right_extension, + # then skip it (because it's already included in the previous hit) + if h['right'] <= res['right_integration']: + continue + + # if a hit only partly overlaps with the previous hit's right_ + # extension, merge them (extend previous hit by this one) + if h['left'] <= res['right_integration']: + fill_hitpars(res, h, hit_left_extension, hit_right_extension, record_data, record_len, + baseline_fpart, extend=True) + + res['tdelay'] = res['sample_10pc_area'] - t_LED + + continue + + # an actual new hit increases the buffer index + offset += 1 + res = buffer[offset] + + fill_hitpars(res, h, hit_left_extension, hit_right_extension, record_data, record_len, + baseline_fpart) + + res['tdelay'] = res['sample_10pc_area'] - t_LED + + return buffer[:offset] + + +@export +@numba.jit(nopython=True, nogil=True, cache=True) +def get_sample_area_quantile(data, quantile, baseline_fpart): + """ + returns first sample index in hit where integrated area of hit is above total area + """ + + area = 0 + area_tot = data.sum() + len(data) * baseline_fpart + + for d_i, d in enumerate(data): + area += d + baseline_fpart + if area > (quantile * area_tot): + return d_i + if d_i == len(data) - 1: + # if no quantile was found, something went wrong + # (negative area due to wrong baseline, caused by real events that + # by coincidence fall in the first samples of the trigger window) + # print('no quantile found: set to 0') + return 0 + + # What happened here?! + return 0 + + +@numba.jit(nopython=True, nogil=True, cache=True) +def fill_hitpars(result, hit, hit_left_extension, hit_right_extension, record_data, record_len, + baseline_fpart, extend=False): + if not extend: # fill first time only + result['time'] = hit['time'] - hit_left_extension * hit['dt'] + result['dt'] = hit['dt'] + result['channel'] = hit['channel'] + result['left'] = hit['left'] + result['record_i'] = hit['record_i'] + result['threshold'] = hit['threshold'] + result['left_integration'] = hit['left'] - hit_left_extension + result['height'] = hit['height'] + + # fill always (if hits are merged, only these will be updated) + result['right'] = hit['right'] + result['right_integration'] = hit['right'] + hit_right_extension + if result['right_integration'] > record_len: + result['right_integration'] = record_len # cap right_integration at end of record + result['length'] = result['right_integration'] - result['left_integration'] + + hit_data = record_data[result['left_integration']:result['right_integration']] + result['area'] = hit_data.sum() + result['length'] * baseline_fpart + result['sample_10pc_area'] = result['left_integration'] + get_sample_area_quantile( + hit_data, 0.1, baseline_fpart) + result['sample_50pc_area'] = result['left_integration'] + get_sample_area_quantile( + hit_data, 0.5, baseline_fpart) + if len(hit_data): + result['max'] = result['left_integration'] + hit_data.argmax() + + if extend: # only when merging hits + result['height'] = max(result['height'], hit['height']) + + +@export +def dtype_afterpulses(): + # define new data type for afterpulse data + dtype_ap = [ + (('Channel/PMT number', 'channel'), '0) # If area equals or smaller than 0 - assume 0 + is_zero = areas <= 0 # If area equals or smaller than 0 - assume 0 res[is_zero] = 2.*mu[is_zero] # if zero channel has negative expectation, assume LLH to be 0 there # this happens in the normalization factor calculation when mu is received from area - neg_mu = mu<0.0 - res[is_zero*neg_mu] = 0.0 + neg_mu = mu < 0.0 + res[is_zero | neg_mu] = 0.0 return res + # continuous and discrete binomial test # https://github.com/poliastro/cephes/blob/master/src/bdtr.c -def bdtrc(k, n, p): - if (k < 0): - return (1.0) - if (k == n): - return (0.0) +@numba.vectorize([numba.float64(numba.float64, numba.float64, numba.float64)]) +def binom_pmf(k, n, p): + scale_log = numba_gammaln(n + 1) - numba_gammaln(n - k + 1) - numba_gammaln(k + 1) + ret_log = scale_log + k * np.log(p) + (n - k) * np.log(1 - p) + return np.exp(ret_log) + +@numba.njit +def binom_cdf(k, n, p): + if k < 0: + return np.nan + if k == n: + return 1.0 dn = n - k - if (k == 0): - if (p < .01): - dk = -np.expm1(dn * np.log1p(-p)) - else: - dk = 1.0 - np.exp(dn * np.log(1.0 - p)) + if k == 0: + dk = np.exp(dn * np.log(1.0 - p)) else: dk = k + 1 - dk = betainc(dk, dn, p) + dk = numba_betainc(dn, dk, 1.0 - p) return dk -def bdtr(k, n, p): - if (k < 0): - return np.nan - if (k == n): - return (1.0) +@numba.njit +def binom_sf(k, n, p): + if k < 0: + return 1.0 + if k == n: + return 0.0 dn = n - k - if (k == 0): - dk = np.exp(dn * np.log(1.0 - p)) + if k == 0: + if p < .01: + dk = -np.expm1(dn * np.log1p(-p)) + else: + dk = 1.0 - np.exp(dn * np.log(1.0 - p)) else: dk = k + 1 - dk = betainc(dn, dk, 1.0 - p) + dk = numba_betainc(dk, dn, p) return dk -# continuous binomial distribution -def binom_pmf(k, n, p): - scale_log = gammaln(n + 1) - gammaln(n - k + 1) - gammaln(k + 1) - ret_log = scale_log + k * np.log(p) + (n - k) * np.log(1 - p) - return np.exp(ret_log) - -def binom_cdf(k, n, p): - return bdtr(k, n, p) +@numba.njit +def _find_interval(n, p): + mu = n*p + sigma = np.sqrt(n*p*(1-p)) + if mu-2*sigma < 0: + s1_min_range = 0 + else: + s1_min_range = mu-2*sigma + if mu+2*sigma > n: + s1_max_range = n + else: + s1_max_range = mu+2*sigma + return s1_min_range, s1_max_range -def binom_sf(k, n, p): - return bdtrc(k, n, p) +@numba.njit +def _get_min_and_max(s1_min, s1_max, n, p, shift): + s1 = np.arange(s1_min, s1_max, 0.01) + ds1 = binom_pmf(s1, n, p) + s1argmax = np.argmax(ds1) + if np.argmax(ds1) - shift > 0: + minimum = s1[s1argmax - shift] + else: + minimum = s1[0] + maximum = s1[s1argmax + shift] + return minimum, maximum, s1[s1argmax] +@numba.njit def binom_test(k, n, p): """ The main purpose of this algorithm is to find the value j on the @@ -371,71 +398,82 @@ def binom_test(k, n, p): integrate the tails outward from k and j. In the case where either k or j are zero, only the non-zero tail is integrated. """ + # define the S1 interval for finding the maximum + _s1_min_range, _s1_max_range = _find_interval(n, p) + + # compute the binomial probability for each S1 and define the Binom range for later + minimum, maximum, s1max = _get_min_and_max(_s1_min_range, _s1_max_range, n, p, shift=2) + + # comments TODO + _d0 = binom_pmf(0, n, p) + _dmax = binom_pmf(s1max, n, p) + _dn = binom_pmf(n, n, p) d = binom_pmf(k, n, p) rerr = 1 + 1e-7 d = d * rerr + _d0 = _d0 * rerr + _dmax = _dmax * rerr + _dn = _dn * rerr # define number of interaction for finding the the value j # the exceptional case of n<=0, is avoid since n_iter is at least 2 if n > 0: - n_iter = int(max(np.round(np.log10(n)) + 1, 2)) + n_iter = int(np.round_(np.log10(n)) + 1) + n_iter = max(n_iter, 2) else: n_iter = 2 - - if k < n * p: - # if binom_pmf(n, n, p) > d, with d<<1e-3, means that we have - # to look for j value above n. It is likely that the binomial - # test for such j is extremely low such that we can stop - # the algorithm and return 0 - if binom_pmf(n, n, p) > d: - for n_ in np.arange(n, 2*n, 1): - if binom_pmf(n_, n, p) < d: - j_min, j_max = k, n_ - do_test = True - break - do_test = False + # comments TODO + if k < minimum: + if (_d0 >= d) and (_d0 > _dmax): + n_iter, j_min, j_max = -1, 0, 0 + elif _dn > d: + n_iter, j_min, j_max = -2, 0, 0 else: - j_min, j_max = k, n - do_test = True + j_min, j_max = s1max, n def _check_(d, y0, y1): return (d>y1) and (d<=y0) - else: - if binom_pmf(0, n, p) > d: - n_iter, j_min, j_max = 0, 0, 0 + # comments TODO + elif k>maximum: + if _d0 >= d: + n_iter, j_min, j_max = -1, 0, 0 else: - j_min, j_max = 0, k - do_test = True + j_min, j_max = 0, s1max def _check_(d, y0, y1): return (d>=y0) and (d=y0) and (d 0) + for s_i in [1, 2]: + for peak_type in ['', 'alt_']: + if event_i[f'{peak_type}s{s_i}_index'] == -1: + continue + + index = event_i[f'{peak_type}s{s_i}_index'] + result_i[f'{peak_type}s{s_i}_veto_tag'] = tags_i[index]['veto_tag'] + result_i[f'{peak_type}s{s_i}_dt_veto'] = tags_i[index]['time_to_closest_veto'] + + +@export @strax.takes_config( strax.Option( 's1_xyz_correction_map', @@ -544,20 +603,20 @@ def compute(self, events): (first_sr1_run, pax_file('XENON1T_s1_xyz_lce_true_kr83m_SR1_pax-680_fdc-3d_v0.json'))]), # noqa strax.Option( 's2_xy_correction_map', - help="S2 (x, y) correction map. Correct S2 position dependence " + help="S2 (x, y) correction map. Correct S2 position dependence, including S2 top, bottom, and total." "manly due to bending of anode/gate-grid, PMT quantum efficiency " "and extraction field distribution, as well as other geometric factors.", default_by_run=[ (0, pax_file('XENON1T_s2_xy_ly_SR0_24Feb2017.json')), (170118_1327, pax_file('XENON1T_s2_xy_ly_SR1_v2.2.json'))]), - strax.Option( + strax.Option( 'elife_conf', default=("elife", "ONLINE", True), help='Electron lifetime ' 'Specify as (model_type->str, model_config->str, is_nT->bool) ' 'where model_type can be "elife" or "elife_constant" ' 'and model_config can be a version.' - ), + ), *DEFAULT_POSREC_ALGO_OPTION ) class CorrectedAreas(strax.Plugin): @@ -571,66 +630,100 @@ class CorrectedAreas(strax.Plugin): Note: Please be aware that for both, the main and alternative S1, the area is corrected according to the xy-position of the main S2. + + There are now 3 components of cS2s: cs2_top, cS2_bottom and cs2. + cs2_top and cs2_bottom are corrected by the corresponding maps, + and cs2 is the sum of the two. """ - __version__ = '0.1.1' + __version__ = '0.1.3' depends_on = ['event_basics', 'event_positions'] - dtype = [('cs1', np.float32, 'Corrected S1 area [PE]'), - ('cs2', np.float32, 'Corrected S2 area [PE]'), - ('alt_cs1', np.float32, 'Corrected area of the alternate S1 [PE]'), - ('alt_cs2', np.float32, 'Corrected area of the alternate S2 [PE]') - ] + strax.time_fields + + def infer_dtype(self): + dtype = [] + dtype += strax.time_fields + + for peak_type, peak_name in zip(['', 'alt_'], ['main', 'alternate']): + dtype += [(f'{peak_type}cs1', np.float32, f'Corrected area of {peak_name} S1 [PE]'), + (f'{peak_type}cs2_wo_elifecorr', np.float32, + f'Corrected area of {peak_name} S2 before elife correction (s2 xy correction only) [PE]'), + (f'{peak_type}cs2_area_fraction_top', np.float32, + f'Fraction of area seen by the top PMT array for corrected {peak_name} S2'), + (f'{peak_type}cs2_bottom', np.float32, + f'Corrected area of {peak_name} S2 in the bottom PMT array [PE]'), + (f'{peak_type}cs2', np.float32, f'Corrected area of {peak_name} S2 [PE]'),] + + return dtype def setup(self): self.elife = get_correction_from_cmt(self.run_id, self.config['elife_conf']) - if isinstance(self.config['s1_xyz_correction_map'], str): - self.config['s1_xyz_correction_map'] = [self.config['s1_xyz_correction_map']] - if isinstance(self.config['s2_xy_correction_map'], str): - self.config['s2_xy_correction_map'] = [self.config['s2_xy_correction_map']] - self.s1_map = InterpolatingMap( get_cmt_resource(self.run_id, tuple(['suffix', self.config['default_reconstruction_algorithm'], - *self.config['s1_xyz_correction_map']]), + *strax.to_str_tuple(self.config['s1_xyz_correction_map'])]), fmt='text')) self.s2_map = InterpolatingMap( get_cmt_resource(self.run_id, - tuple([*self.config['s2_xy_correction_map']]), - fmt='text')) + tuple(['suffix', + self.config['default_reconstruction_algorithm'], + *strax.to_str_tuple(self.config['s2_xy_correction_map'])]), + fmt='json')) def compute(self, events): + result = dict( + time=events['time'], + endtime=strax.endtime(events) + ) + # S1 corrections depend on the actual corrected event position. # We use this also for the alternate S1; for e.g. Kr this is # fine as the S1 correction varies slowly. event_positions = np.vstack([events['x'], events['y'], events['z']]).T + for peak_type in ["", "alt_"]: + result[f"{peak_type}cs1"] = events[f'{peak_type}s1_area'] / self.s1_map(event_positions) + + # s2 corrections + # S2 top and bottom are corrected separately, and cS2 total is the sum of the two + # figure out the map name + if len(self.s2_map.map_names) > 1: + s2_top_map_name = "map_top" + s2_bottom_map_name = "map_bottom" + else: + s2_top_map_name = "map" + s2_bottom_map_name = "map" + + for peak_type in ["", "alt_"]: + # S2(x,y) corrections use the observed S2 positions + s2_positions = np.vstack([events[f'{peak_type}s2_x'], events[f'{peak_type}s2_y']]).T + + # corrected s2 with s2 xy map only, i.e. no elife correction + # this is for s2-only events which don't have drift time info + cs2_top_wo_elifecorr = (events[f'{peak_type}s2_area'] * events[f'{peak_type}s2_area_fraction_top'] / + self.s2_map(s2_positions, map_name=s2_top_map_name)) + cs2_bottom_wo_elifecorr = (events[f'{peak_type}s2_area'] * + (1 - events[f'{peak_type}s2_area_fraction_top']) / + self.s2_map(s2_positions, map_name=s2_bottom_map_name)) + result[f"{peak_type}cs2_wo_elifecorr"] = (cs2_top_wo_elifecorr + cs2_bottom_wo_elifecorr) + + # cs2aft doesn't need elife correction as it cancels + result[f"{peak_type}cs2_area_fraction_top"] = cs2_top_wo_elifecorr / result[f"{peak_type}cs2_wo_elifecorr"] + + # For electron lifetime corrections to the S2s, + # use drift time computed using the main S1. + el_string = peak_type + "s2_interaction_" if peak_type == "alt_" else peak_type + elife_correction = np.exp(events[f'{el_string}drift_time'] / self.elife) + + cs2_top = cs2_top_wo_elifecorr * elife_correction + result[f"{peak_type}cs2_bottom"] = cs2_bottom_wo_elifecorr * elife_correction + result[f"{peak_type}cs2"] = cs2_top + result[f"{peak_type}cs2_bottom"] - # For electron lifetime corrections to the S2s, - # use lifetimes computed using the main S1. - lifetime_corr = np.exp(events['drift_time'] / self.elife) - alt_lifetime_corr = ( - np.exp((events['alt_s2_interaction_drift_time']) - / self.elife)) - - # S2(x,y) corrections use the observed S2 positions - s2_positions = np.vstack([events['s2_x'], events['s2_y']]).T - alt_s2_positions = np.vstack([events['alt_s2_x'], events['alt_s2_y']]).T - - return dict( - time=events['time'], - endtime=strax.endtime(events), - - cs1=events['s1_area'] / self.s1_map(event_positions), - alt_cs1=events['alt_s1_area'] / self.s1_map(event_positions), - - cs2=(events['s2_area'] * lifetime_corr - / self.s2_map(s2_positions)), - alt_cs2=(events['alt_s2_area'] * alt_lifetime_corr - / self.s2_map(alt_s2_positions))) + return result +@export @strax.takes_config( strax.Option( 'g1', diff --git a/straxen/plugins/event_shadow.py b/straxen/plugins/event_shadow.py new file mode 100644 index 000000000..1e1e6b051 --- /dev/null +++ b/straxen/plugins/event_shadow.py @@ -0,0 +1,70 @@ +import numpy as np +import strax +import numba +export, __all__ = strax.exporter() + +@export +@strax.takes_config( + strax.Option('pre_s2_area_threshold', default=1000, + help='Only take S2s large than this into account when calculating EventShadow.'), + strax.Option('time_window_backward', default=int(3e9), + help='Search for S2s causing shadow in this time window')) +class EventShadow(strax.Plugin): + """ + This plugin can find and calculate the previous S2 shadow at event level, + with time window backward and previous S2 area as options. + It also gives the area and position infomation of these previous S2s. + """ + __version__ = '0.0.6' + depends_on = ('event_basics','peak_basics','peak_positions') + provides = 'event_shadow' + save_when = strax.SaveWhen.EXPLICIT + + + def infer_dtype(self): + dtype = [ + ('pre_s2_area', np.float32,'previous s2 area [PE]'), + ('shadow_dt', np.int64,'time diffrence to the previous s2 [ns]'), + ('shadow', np.float32,'previous s2 shadow [PE/ns]'), + ('pre_s2_x', np.float32,'x of previous s2 peak causing shadow [cm]'), + ('pre_s2_y', np.float32,'y of previous s2 peak causing shadow [cm]'), + ('shadow_distance', np.float32,'distance to the previous s2 peak causing the max shadow [cm]') + ] + dtype += strax.time_fields + return dtype + + def compute(self, events, peaks): + roi_dt = np.dtype([(('back in time', 'time'), int), + (('till it begin','endtime'), int)]) + roi = np.zeros(len(events), dtype=roi_dt) + n_seconds = self.config['time_window_backward'] + roi['time'] = events['time'] - n_seconds + roi['endtime'] = events['time'] + mask_s2 = peaks['type'] == 2 + mask_s2 &= peaks['area'] > self.config['pre_s2_area_threshold'] + split_peaks = strax.split_touching_windows(peaks[mask_s2], roi) + res = np.zeros(len(events), self.dtype) + compute_shadow(events, split_peaks, res) + res['shadow_distance'] = ((res['pre_s2_x'] - events['s2_x'])**2+(res['pre_s2_y'] - events['s2_y'])**2)**0.5 + res['time'] = events['time'] + res['endtime'] = strax.endtime(events) + + return res + +def compute_shadow(events, split_peaks, res): + if len(res): + return _compute_shadow(events, split_peaks, res) + + +@numba.njit(cache=True) +def _compute_shadow(events, split_peaks, res): + for event_i, event_a in enumerate(events): + new_shadow = 0 + for peak_i, peak_a in enumerate(split_peaks[event_i]): + new_shadow = peak_a['area']/(event_a['s2_center_time']-peak_a['center_time']) + if new_shadow > res['shadow'][event_i]: + res['pre_s2_area'][event_i] = peak_a['area'] + res['shadow_dt'][event_i] = event_a['s2_center_time']-peak_a['center_time'] + res['pre_s2_x'][event_i] = peak_a['x'] + res['pre_s2_y'][event_i] = peak_a['y'] + res['shadow'][event_i] = new_shadow diff --git a/straxen/plugins/led_calibration.py b/straxen/plugins/led_calibration.py index c4d96dca9..62f0f4bbd 100644 --- a/straxen/plugins/led_calibration.py +++ b/straxen/plugins/led_calibration.py @@ -3,6 +3,7 @@ if you want to complain please contact: chiara@physik.uzh.ch, gvolta@physik.uzh.ch, kazama@isee.nagoya-u.ac.jp ''' import datetime +from immutabledict import immutabledict import strax import numba import numpy as np @@ -141,3 +142,87 @@ def get_area(records, led_window): Area['area'] = Area['area']/float(len(end_pos)) return Area + + +@export +@strax.takes_config( + strax.Option('channel_map', track=False, type=immutabledict, + help="immutabledict mapping subdetector to (min, max) " + "channel number."), +) +class nVetoExtTimings(strax.Plugin): + """ + Plugin which computes the time difference `delta_time` from pulse timing + of `hitlets_nv` to start time of `raw_records` which belong the `hitlets_nv`. + They are used as the external trigger timings. + """ + depends_on = ('raw_records_nv', 'hitlets_nv') + provides = 'ext_timings_nv' + data_kind = 'hitlets_nv' + + compressor = 'zstd' + __version__ = '0.0.1' + + def infer_dtype(self): + dtype = [] + dtype += strax.time_dt_fields + dtype += [(('Delta time from trigger timing [ns]', 'delta_time'), np.int16), + (('Index to which pulse (not record) the hitlet belongs to.', 'pulse_i'), + np.int32),] + return dtype + + def setup(self): + self.nv_pmt_start = self.config['channel_map']['nveto'][0] + self.nv_pmt_stop = self.config['channel_map']['nveto'][1] + 1 + + def compute(self, hitlets_nv, raw_records_nv): + + rr_nv = raw_records_nv[raw_records_nv['record_i'] == 0] + pulses = np.zeros(len(rr_nv), dtype=self.pulse_dtype()) + pulses['time'] = rr_nv['time'] + pulses['endtime'] = rr_nv['time'] + rr_nv['pulse_length'] * rr_nv['dt'] + pulses['channel'] = rr_nv['channel'] + + ext_timings_nv = np.zeros_like(hitlets_nv, dtype=self.dtype) + ext_timings_nv['time'] = hitlets_nv['time'] + ext_timings_nv['length'] = hitlets_nv['length'] + ext_timings_nv['dt'] = hitlets_nv['dt'] + self.calc_delta_time(ext_timings_nv, pulses, hitlets_nv, + self.nv_pmt_start, self.nv_pmt_stop) + + return ext_timings_nv + + @staticmethod + def pulse_dtype(): + pulse_dtype = [] + pulse_dtype += strax.time_fields + pulse_dtype += [(('PMT channel', 'channel'), np.int16)] + return pulse_dtype + + @staticmethod + @numba.jit + def calc_delta_time(ext_timings_nv_delta_time, pulses, hitlets_nv, nv_pmt_start, nv_pmt_stop): + """ + numpy access with fancy index returns copy, not view + This for-loop is required to substitute in one by one + """ + hitlet_index = np.arange(len(hitlets_nv)) + pulse_index = np.arange(len(pulses)) + for ch in range(nv_pmt_start, nv_pmt_stop): + mask_hitlets_in_channel = hitlets_nv['channel'] == ch + hitlet_in_channel_index = hitlet_index[mask_hitlets_in_channel] + + mask_pulse_in_channel = pulses['channel'] == ch + pulse_in_channel_index = pulse_index[mask_pulse_in_channel] + + hitlets_in_channel = hitlets_nv[hitlet_in_channel_index] + pulses_in_channel = pulses[pulse_in_channel_index] + hit_in_pulse_index = strax.fully_contained_in(hitlets_in_channel, pulses_in_channel) + for h_i, p_i in zip(hitlet_in_channel_index, hit_in_pulse_index): + if p_i == -1: + continue + res = ext_timings_nv_delta_time[h_i] + + res['delta_time'] = hitlets_nv[h_i]['time'] + hitlets_nv[h_i]['time_amplitude'] \ + - pulses_in_channel[p_i]['time'] + res['pulse_i'] = pulse_in_channel_index[p_i] diff --git a/straxen/plugins/online_monitor.py b/straxen/plugins/online_monitor.py index dd1c640cc..236154d09 100644 --- a/straxen/plugins/online_monitor.py +++ b/straxen/plugins/online_monitor.py @@ -1,5 +1,6 @@ import strax import numpy as np +from immutabledict import immutabledict export, __all__ = strax.exporter() @@ -175,3 +176,127 @@ def area_width_hist(self, data): range=self.config['area_vs_width_bounds'], bins=self.config['area_vs_width_nbins']) return hist.T + + +@export +@strax.takes_config( + strax.Option( + 'channel_map', + track=False, + type=immutabledict, + help='immutabledict mapping subdetector to (min, max) ' + 'channel number.'), + strax.Option( + 'events_area_bounds', + type=tuple, default=(-0.5, 130.5), + help='Boundaries area histogram of events_nv_area_per_chunk [PE]'), + strax.Option( + 'events_area_nbins', + type=int, default=131, + help='Number of bins of histogram of events_nv_area_per_chunk, ' + 'defined value 1 PE/bin') +) +class OnlineMonitorNV(strax.Plugin): + """ + Plugin to write data of nVeto detector to the online-monitor. + Data that is written by this plugin should be small (~MB/chunk) + to not overload the runs-database. + + This plugin takes 'hitlets_nv' and 'events_nv'. Although they are + not strictly related, they are aggregated into a single data_type + in order to minimize the number of documents in the online monitor. + + Produces 'online_monitor_nv' with info on the hitlets_nv and events_nv + """ + depends_on = ('hitlets_nv', 'events_nv') + provides = 'online_monitor_nv' + data_kind = 'online_monitor_nv' + rechunk_on_save = False + + # Needed in case we make again an muVETO child. + ends_with = '_nv' + + __version__ = '0.0.4' + + def infer_dtype(self): + self.channel_range = self.config['channel_map']['nveto'] + self.n_channel = (self.channel_range[1] - self.channel_range[0]) + 1 + return veto_monitor_dtype(self.ends_with, self.n_channel, self.config['events_area_nbins']) + + def compute(self, hitlets_nv, events_nv, start, end): + # General setup + res = np.zeros(1, dtype=self.dtype) + res['time'] = start + res['endtime'] = end + + # Count number of hitlets_nv per PMT + hitlets_channel_count, _ = np.histogram(hitlets_nv['channel'], + bins=self.n_channel, + range=[self.channel_range[0], + self.channel_range[1] + 1]) + res[f'hitlets{self.ends_with}_per_channel'] = hitlets_channel_count + + # Count number of events_nv with coincidence cut + res[f'events{self.ends_with}_per_chunk'] = len(events_nv) + sel = events_nv['n_contributing_pmt'] >= 4 + res[f'events{self.ends_with}_4coinc_per_chunk'] = np.sum(sel) + sel = events_nv['n_contributing_pmt'] >= 5 + res[f'events{self.ends_with}_5coinc_per_chunk'] = np.sum(sel) + sel = events_nv['n_contributing_pmt'] >= 8 + res[f'events{self.ends_with}_8coinc_per_chunk'] = np.sum(sel) + sel = events_nv['n_contributing_pmt'] >= 10 + res[f'events{self.ends_with}_10coinc_per_chunk'] = np.sum(sel) + + # Get histogram of events_nv_area per chunk + events_area, bins_ = np.histogram(events_nv['area'], + bins=self.config['events_area_nbins'], + range=self.config['events_area_bounds']) + res[f'events{self.ends_with}_area_per_chunk'] = events_area + return res + + +def veto_monitor_dtype(veto_name: str = '_nv', + n_pmts: int = 120, + n_bins: int = 131) -> list: + dtype = [] + dtype += strax.time_fields # because mutable + dtype += [((f'hitlets{veto_name} per channel', f'hitlets{veto_name}_per_channel'), (np.int64, n_pmts)), + ((f'events{veto_name}_area per chunk', f'events{veto_name}_area_per_chunk'), np.int64, n_bins), + ((f'events{veto_name} per chunk', f'events{veto_name}_per_chunk'), np.int64), + ((f'events{veto_name} 4-coincidence per chunk', f'events{veto_name}_4coinc_per_chunk'), np.int64), + ((f'events{veto_name} 5-coincidence per chunk', f'events{veto_name}_5coinc_per_chunk'), np.int64), + ((f'events{veto_name} 8-coincidence per chunk', f'events{veto_name}_8coinc_per_chunk'), np.int64), + ((f'events{veto_name} 10-coincidence per chunk', f'events{veto_name}_10coinc_per_chunk'), np.int64) + ] + return dtype + + +@export +@strax.takes_config( + strax.Option( + 'adc_to_pe_mv', + type=int, default=170.0, + help='conversion factor from ADC to PE for muon Veto') +) +class OnlineMonitorMV(OnlineMonitorNV): + __doc__ = OnlineMonitorNV.__doc__.replace('_nv', '_mv').replace('nVeto', 'muVeto') + depends_on = ('hitlets_mv', 'events_mv') + provides = 'online_monitor_mv' + data_kind = 'online_monitor_mv' + rechunk_on_save = False + + # Needed in case we make again an muVETO child. + ends_with = '_mv' + child_plugin = True + + __version__ = '0.0.1' + + def infer_dtype(self): + self.channel_range = self.config['channel_map']['mv'] + self.n_channel = (self.channel_range[1] - self.channel_range[0]) + 1 + return veto_monitor_dtype(self.ends_with, self.n_channel, self.config['events_area_nbins']) + + def compute(self, hitlets_mv, events_mv, start, end): + events_mv = np.copy(events_mv) + events_mv['area'] *= 1./self.config['adc_to_pe_mv'] + return super().compute(hitlets_mv, events_mv, start, end) diff --git a/straxen/plugins/peak_processing.py b/straxen/plugins/peak_processing.py index 2707b2fc8..045da06d2 100644 --- a/straxen/plugins/peak_processing.py +++ b/straxen/plugins/peak_processing.py @@ -3,12 +3,13 @@ import tempfile import numpy as np import numba +from enum import IntEnum import strax import straxen from straxen.common import pax_file, get_resource, first_sr1_run export, __all__ = strax.exporter() -from .pulse_processing import HE_PREAMBLE +from .pulse_processing import HE_PREAMBLE @export @@ -27,7 +28,7 @@ class PeakBasics(strax.Plugin): arrays. NB: This plugin can therefore be loaded as a pandas DataFrame. """ - __version__ = "0.0.9" + __version__ = "0.1.0" parallel = True depends_on = ('peaks',) provides = 'peak_basics' @@ -46,6 +47,8 @@ class PeakBasics(strax.Plugin): 'max_pmt'), np.int16), (('Area of signal in the largest-contributing PMT (PE)', 'max_pmt_area'), np.float32), + (('Total number of saturated channels', + 'n_saturated_channels'), np.int16), (('Width (in ns) of the central 50% area of the peak', 'range_50p_area'), np.float32), (('Width (in ns) of the central 90% area of the peak', @@ -61,6 +64,8 @@ class PeakBasics(strax.Plugin): 'rise_time'), np.float32), (('Hits within tight range of mean', 'tight_coincidence'), np.int16), + (('PMT channel within tight range of mean', + 'tight_coincidence_channel'), np.int16), (('Classification of the peak(let)', 'type'), np.int8) ] @@ -77,6 +82,8 @@ def compute(self, peaks): r['max_pmt'] = np.argmax(p['area_per_channel'], axis=1) r['max_pmt_area'] = np.max(p['area_per_channel'], axis=1) r['tight_coincidence'] = p['tight_coincidence'] + r['tight_coincidence_channel'] = p['tight_coincidence_channel'] + r['n_saturated_channels'] = p['n_saturated_channels'] n_top = self.config['n_top_pmts'] area_top = p['area_per_channel'][:, :n_top].sum(axis=1) @@ -324,3 +331,154 @@ def find_n_competing(peaks, windows, fraction): n_tot[i] = n_left[i] + np.sum(areas[i + 1:right_i] > threshold) return n_left, n_tot + + +@export +class VetoPeakTags(IntEnum): + """Identifies by which detector peak was tagged. + """ + # Peaks are not inside any veto interval + NO_VETO = 0 + # Peaks are inside a veto interval issued by: + NEUTRON_VETO = 1 + MUON_VETO = 2 + BOTH = 3 + + +@export +class PeakVetoTagging(strax.Plugin): + """ + Plugin which tags S1 peaks according to muon and neutron-vetos. + Tagging S2s is does not make sense as they occur with a delay. + However, we compute for both S1/S2 the time delay to the closest veto + region. + + * untagged: 0 + * neutron-veto: 1 + * muon-veto: 2 + * both vetos: 3 + """ + __version__ = '0.0.1' + depends_on = ('peak_basics', 'veto_regions_nv', 'veto_regions_mv') + provides = ('peak_veto_tags') + save_when = strax.SaveWhen.TARGET + + dtype = strax.time_fields + [ + ('veto_tag', np.int8, + 'Veto tag for S1 peaks. unatagged: 0, nveto: 1, mveto: 2, both: 3'), + ('time_to_closest_veto', np.int64, 'Time to closest veto interval boundary in ns (can be ' + 'negative if closest boundary comes before peak.). ') + ] + + def get_time_difference(self, peaks, veto_regions_nv, veto_regions_mv): + """ + Computes time differences to closest nv/mv veto signal. + + It might be that neutron-veto and muon-veto signals overlap + Hence we compute first the individual time differences to the + corresponding vetos and keep afterwards the smallest ones. + """ + dt_nv = get_time_to_closest_veto(peaks, veto_regions_nv) + dt_mv = get_time_to_closest_veto(peaks, veto_regions_mv) + + dts = np.transpose([dt_nv, dt_mv]) + ind_axis1 = np.argmin(np.abs(dts), axis=1) + return self._get_smallest_value(dts, ind_axis1) + + @staticmethod + @numba.njit(cache=True, nogil=True) + def _get_smallest_value(time_differences, index): + res = np.zeros(len(time_differences), np.int64) + for res_ind, (ind, dt) in enumerate(zip(index, time_differences)): + res[res_ind] = dt[ind] + return res + + def compute(self, peaks, veto_regions_nv, veto_regions_mv): + touching_mv = strax.touching_windows(peaks, veto_regions_mv) + touching_nv = strax.touching_windows(peaks, veto_regions_nv) + + tags = np.zeros(len(peaks)) + tags = tag_peaks(tags, touching_nv, straxen.VetoPeakTags.NEUTRON_VETO) + tags = tag_peaks(tags, touching_mv, straxen.VetoPeakTags.MUON_VETO) + + dt = self.get_time_difference(peaks, veto_regions_nv, veto_regions_mv) + return {'time': peaks['time'], + 'endtime': strax.endtime(peaks), + 'veto_tag': tags, + 'time_to_closest_veto': dt, + } + + +@numba.njit(cache=True, nogil=True) +def tag_peaks(tags, touching_windows, tag_number): + """Tags every peak which are within the corresponding touching window + with the defined tag number. + + :param tags: numpy.array in which the tags should be stored. Should + be of length peaks. + :param touching_windows: Start/End index of tags to be set to tag + value. + :param tag_number: integer representing the tag. + :return: Updated tags. + """ + pre_tags = np.zeros(len(tags), dtype=np.int8) + for start, end in touching_windows: + pre_tags[start:end] = tag_number + tags += pre_tags + return tags + + +def get_time_to_closest_veto(peaks, veto_intervals): + """Computes time difference between peak and closest veto interval. + + The time difference is always computed from peaks-time field to + the time or endtime of the veto_interval depending on which distance + is smaller. + """ + vetos = np.zeros(len(veto_intervals)+2, strax.time_fields) + vetos[1:-1]['time'] = veto_intervals['time'] + vetos[1:-1]['endtime'] = strax.endtime(veto_intervals) + vetos[-1]['time'] = straxen.INFINITY_64BIT_SIGNED + vetos[-1]['endtime'] = straxen.INFINITY_64BIT_SIGNED + vetos[0]['time'] = -straxen.INFINITY_64BIT_SIGNED + vetos[0]['endtime'] = -straxen.INFINITY_64BIT_SIGNED + return _get_time_to_closest_veto(peaks, vetos) + + +@numba.njit(cache=True, nogil=True) +def _get_time_to_closest_veto(peaks, vetos): + res = np.zeros(len(peaks), dtype=np.int64) + veto_index = 0 + for ind, p in enumerate(peaks): + for veto_index in range(veto_index, len(vetos)): + if veto_index+1 == len(vetos): + # If we reach here all future peaks are closest to last veto: + res[ind] = np.abs(vetos[-1]['time'] - p['time']) + break + + # Current interval can be before or after current peak, hence + # we have to check which distance is smaller. + dt_current_veto = min(np.abs(vetos[veto_index]['time'] - p['time']), + np.abs(vetos[veto_index]['endtime'] - p['time']) + ) + # Next interval is always further in the future as the + # current one, hence we only have to compute distance with "time". + dt_next_veto = np.abs(vetos[veto_index+1]['time'] - p['time']) + + # Next veto is closer so we have to repeat in case next + 1 + # is even closer. + if dt_current_veto >= dt_next_veto: + veto_index += 1 + continue + + # Now compute time difference for real: + dt_time = vetos[veto_index]['time'] - p['time'] + dt_endtime = vetos[veto_index]['endtime'] - p['time'] + + if np.abs(dt_time) < np.abs(dt_endtime): + res[ind] = dt_time + else: + res[ind] = dt_endtime + break + + return res diff --git a/straxen/plugins/peaklet_processing.py b/straxen/plugins/peaklet_processing.py index d942e5811..661060f03 100644 --- a/straxen/plugins/peaklet_processing.py +++ b/straxen/plugins/peaklet_processing.py @@ -1,6 +1,7 @@ import numba import numpy as np import strax +from immutabledict import immutabledict from strax.processing.general import _touching_windows import straxen from .pulse_processing import HITFINDER_OPTIONS, HITFINDER_OPTIONS_he, HE_PREAMBLE @@ -63,6 +64,9 @@ "to correct saturated samples"), strax.Option('peaklet_max_duration', default=int(10e6), help="Maximum duration [ns] of a peaklet"), + strax.Option('channel_map', track=False, type=immutabledict, + help="immutabledict mapping subdetector to (min, max) " + "channel number."), *HITFINDER_OPTIONS, ) class Peaklets(strax.Plugin): @@ -92,7 +96,7 @@ class Peaklets(strax.Plugin): parallel = 'process' compressor = 'zstd' - __version__ = '0.3.10' + __version__ = '0.5.0' def infer_dtype(self): return dict(peaklets=strax.peak_dtype( @@ -121,6 +125,8 @@ def setup(self): self.config['hit_min_amplitude']) else: # int or array self.hit_thresholds = self.config['hit_min_amplitude'] + + self.channel_range = self.config['channel_map']['tpc'] def compute(self, records, start, end): r = records @@ -153,7 +159,8 @@ def compute(self, records, start, end): # Get hits outside peaklets, and store them separately. # fully_contained is OK provided gap_threshold > extension, # which is asserted inside strax.find_peaks. - lone_hits = hits[strax.fully_contained_in(hits, peaklets) == -1] + is_lone_hit = strax.fully_contained_in(hits, peaklets) == -1 + lone_hits = hits[is_lone_hit] strax.integrate_lone_hits( lone_hits, records, peaklets, save_outside_hits=(self.config['peak_left_extension'], @@ -161,14 +168,40 @@ def compute(self, records, start, end): n_channels=len(self.to_pe)) # Compute basic peak properties -- needed before natural breaks - strax.sum_waveform(peaklets, r, self.to_pe) + hits = hits[~is_lone_hit] + # Define regions outside of peaks such that _find_hit_integration_bounds + # is not extended beyond a peak. + outside_peaks = self.create_outside_peaks_region(peaklets, start, end) + strax.find_hit_integration_bounds( + hits, outside_peaks, records, + save_outside_hits=(self.config['peak_left_extension'], + self.config['peak_right_extension']), + n_channels=len(self.to_pe), + allow_bounds_beyond_records=True, + ) + + # Transform hits to hitlets for naming conventions. A hit refers + # to the central part above threshold a hitlet to the entire signal + # including the left and right extension. + # (We are not going to use the actual hitlet data_type here.) + hitlets = hits + del hits + + hitlet_time_shift = (hitlets['left'] - hitlets['left_integration']) * hitlets['dt'] + hitlets['time'] = hitlets['time'] - hitlet_time_shift + hitlets['length'] = (hitlets['right_integration'] - hitlets['left_integration']) + hitlets = strax.sort_by_time(hitlets) + rlinks = strax.record_links(records) + + strax.sum_waveform(peaklets, hitlets, r, rlinks, self.to_pe) + strax.compute_widths(peaklets) # Split peaks using low-split natural breaks; # see https://github.com/XENONnT/straxen/pull/45 # and https://github.com/AxFoundation/strax/pull/225 peaklets = strax.split_peaks( - peaklets, r, self.to_pe, + peaklets, hitlets, r, rlinks, self.to_pe, algorithm='natural_breaks', threshold=self.natural_breaks_threshold, split_low=True, @@ -187,7 +220,7 @@ def compute(self, records, start, end): if self.config['saturation_correction_on']: peak_list = peak_saturation_correction( - r, peaklets, self.to_pe, + r, rlinks, peaklets, hitlets, self.to_pe, reference_length=self.config['saturation_reference_length'], min_reference_length=self.config['saturation_min_reference_length']) @@ -200,27 +233,33 @@ def compute(self, records, start, end): # (b) increase strax memory usage / max_messages, # possibly due to its currently primitive scheduling. hit_max_times = np.sort( - hits['time'] - + hits['dt'] * hit_max_sample(records, hits)) + hitlets['time'] + + hitlets['dt'] * hit_max_sample(records, hitlets) + + hitlet_time_shift # add time shift again to get correct maximum + ) peaklet_max_times = ( peaklets['time'] + np.argmax(peaklets['data'], axis=1) * peaklets['dt']) - peaklets['tight_coincidence'] = get_tight_coin( + tight_coincidence, tight_coincidence_channel = get_tight_coin( hit_max_times, + hitlets['channel'], peaklet_max_times, self.config['tight_coincidence_window_left'], - self.config['tight_coincidence_window_right']) + self.config['tight_coincidence_window_right'], + self.channel_range) + + peaklets['tight_coincidence'] = tight_coincidence + peaklets['tight_coincidence_channel'] = tight_coincidence_channel if self.config['diagnose_sorting'] and len(r): assert np.diff(r['time']).min(initial=1) >= 0, "Records not sorted" - assert np.diff(hits['time']).min(initial=1) >= 0, "Hits not sorted" + assert np.diff(hitlets['time']).min(initial=1) >= 0, "Hits/Hitlets not sorted" assert np.all(peaklets['time'][1:] >= strax.endtime(peaklets)[:-1]), "Peaks not disjoint" # Update nhits of peaklets: - counts = strax.touching_windows(hits, peaklets) + counts = strax.touching_windows(hitlets, peaklets) counts = np.diff(counts, axis=1).flatten() - counts += 1 peaklets['n_hits'] = counts return dict(peaklets=peaklets, @@ -253,16 +292,43 @@ def clip_peaklet_times(peaklets, start, end): if strax.endtime(p) > end: p['length'] = (end - p['time']) // p['dt'] + @staticmethod + def create_outside_peaks_region(peaklets, start, end): + """ + Creates time intervals which are outside peaks. + + :param peaklets: Peaklets for which intervals should be computed. + :param start: Chunk start + :param end: Chunk end + :return: array of strax.time_fields dtype. + """ + if not len(peaklets): + return np.zeros(0, dtype=strax.time_fields) + + outside_peaks = np.zeros(len(peaklets) + 1, + dtype=strax.time_fields) + + outside_peaks[0]['time'] = start + outside_peaks[0]['endtime'] = peaklets[0]['time'] + outside_peaks[1:-1]['time'] = strax.endtime(peaklets[:-1]) + outside_peaks[1:-1]['endtime'] = peaklets['time'][1:] + outside_peaks[-1]['time'] = strax.endtime(peaklets[-1]) + outside_peaks[-1]['endtime'] = end + return outside_peaks + @numba.jit(nopython=True, nogil=True, cache=True) -def peak_saturation_correction(records, peaks, to_pe, +def peak_saturation_correction(records, rlinks, peaks, hitlets, to_pe, reference_length=100, min_reference_length=20, use_classification=False, ): """Correct the area and per pmt area of peaks from saturation :param records: Records + :param rlinks: strax.record_links of corresponding records. :param peaks: Peaklets / Peaks + :param hitlets: Hitlets found in records to build peaks. + (Hitlets are hits including the left/right extension) :param to_pe: adc to PE conversion (length should equal number of PMTs) :param reference_length: Maximum number of reference sample used to correct saturated samples @@ -334,7 +400,7 @@ def peak_saturation_correction(records, peaks, to_pe, peaks[peak_i]['length'] = p['length'] * p['dt'] / dt peaks[peak_i]['dt'] = dt - strax.sum_waveform(peaks, records, to_pe, peak_list) + strax.sum_waveform(peaks, hitlets, records, rlinks, to_pe, peak_list) return peak_list @@ -467,6 +533,8 @@ def setup(self): self.config['hit_min_amplitude_he']) else: # int or array self.hit_thresholds = self.config['hit_min_amplitude_he'] + + self.channel_range = self.config['channel_map']['he'] def compute(self, records_he, start, end): result = super().compute(records_he, start, end) @@ -547,16 +615,27 @@ def compute(self, peaklets_he): "where the gap size of the first point is the maximum gap to allow merging" "and the area of the last point is the maximum area to allow merging. " "The format is ((log10(area), max_gap), (..., ...), (..., ...))" - )) + ), + strax.Option('gain_model', + help='PMT gain model. Specify as ' + '(str(model_config), str(version), nT-->boolean'), + strax.Option('merge_without_s1', default=True, + help="If true, S1s will be igored during the merging. " + "It's now possible for a S1 to be inside a S2 post merging"), +) class MergedS2s(strax.OverlapWindowPlugin): """ Merge together peaklets if peak finding favours that they would form a single peak instead. """ - depends_on = ('peaklets', 'peaklet_classification') + depends_on = ('peaklets', 'peaklet_classification', 'lone_hits') data_kind = 'merged_s2s' provides = 'merged_s2s' - __version__ = '0.2.0' + __version__ = '0.4.0' + + def setup(self): + self.to_pe = straxen.get_correction_from_cmt(self.run_id, + self.config['gain_model']) def infer_dtype(self): return strax.unpack_dtype(self.deps['peaklets'].dtype_for('peaklets')) @@ -565,7 +644,10 @@ def get_window_size(self): return 5 * (int(self.config['s2_merge_gap_thresholds'][0][1]) + self.config['s2_merge_max_duration']) - def compute(self, peaklets): + def compute(self, peaklets, lone_hits): + if self.config['merge_without_s1']: + peaklets = peaklets[peaklets['type'] != 1] + if len(peaklets) <= 1: return np.zeros(0, dtype=peaklets.dtype) @@ -594,6 +676,16 @@ def compute(self, peaklets): max_buffer=int(self.config['s2_merge_max_duration'] // peaklets['dt'].min())) merged_s2s['type'] = 2 + + # Updated time and length of lone_hits and sort again: + lh = np.copy(lone_hits) + del lone_hits + lh_time_shift = (lh['left'] - lh['left_integration']) *lh['dt'] + lh['time'] = lh['time'] - lh_time_shift + lh['length'] = (lh['right_integration'] - lh['left_integration']) + lh = strax.sort_by_time(lh) + strax.add_lone_hits(merged_s2s, lh, self.to_pe) + strax.compute_widths(merged_s2s) return merged_s2s @@ -702,13 +794,20 @@ def infer_dtype(self): return strax.unpack_dtype(self.deps['peaklets_he'].dtype_for('peaklets_he')) def compute(self, peaklets_he): - return super().compute(peaklets_he) + # There are not any lone hits for the high energy channel, + # so create a dummy for the compute method. + lone_hits = np.zeros(0, dtype=strax.hit_dtype) + return super().compute(peaklets_he, lone_hits) @export @strax.takes_config( strax.Option('diagnose_sorting', track=False, default=False, - help="Enable runtime checks for sorting and disjointness")) + help="Enable runtime checks for sorting and disjointness"), + strax.Option('merge_without_s1', default=True, + help="If true, S1s will be igored during the merging. " + "It's now possible for a S1 to be inside a S2 post merging"), +) class Peaks(strax.Plugin): """ Merge peaklets and merged S2s such that we obtain our peaks @@ -719,9 +818,9 @@ class Peaks(strax.Plugin): data_kind = 'peaks' provides = 'peaks' parallel = True - save_when = strax.SaveWhen.NEVER + save_when = strax.SaveWhen.EXPLICIT - __version__ = '0.1.1' + __version__ = '0.1.2' def infer_dtype(self): return self.deps['peaklets'].dtype_for('peaklets') @@ -729,13 +828,24 @@ def infer_dtype(self): def compute(self, peaklets, merged_s2s): # Remove fake merged S2s from dirty hack, see above merged_s2s = merged_s2s[merged_s2s['type'] != FAKE_MERGED_S2_TYPE] - - peaks = strax.replace_merged(peaklets, merged_s2s) + + if self.config['merge_without_s1']: + is_s1 = peaklets['type'] == 1 + peaks = strax.replace_merged(peaklets[~is_s1], merged_s2s) + peaks = strax.sort_by_time(np.concatenate([peaklets[is_s1], + peaks])) + else: + peaks = strax.replace_merged(peaklets, merged_s2s) if self.config['diagnose_sorting']: assert np.all(np.diff(peaks['time']) >= 0), "Peaks not sorted" - assert np.all(peaks['time'][1:] - >= strax.endtime(peaks)[:-1]), "Peaks not disjoint" + if self.config['merge_without_s1']: + to_check = peaks['type'] != 1 + else: + to_check = peaks['type'] != FAKE_MERGED_S2_TYPE + + assert np.all(peaks['time'][to_check][1:] + >= strax.endtime(peaks)[to_check][:-1]), "Peaks not disjoint" return peaks @@ -756,33 +866,55 @@ def compute(self, peaklets_he, merged_s2s_he): @numba.jit(nopython=True, nogil=True, cache=True) -def get_tight_coin(hit_max_times, peak_max_times, left, right): - """Calculates the tight coincidence +def get_tight_coin(hit_max_times, hit_channel, peak_max_times, left, right, + channels=(0, 493)): + """Calculates the tight coincidence based on hits and PMT channels. Defined by number of hits within a specified time range of the the peak's maximum amplitude. Imitates tight_coincidence variable in pax: github.com/XENON1T/pax/blob/master/pax/plugins/peak_processing/BasicProperties.py + + :param hit_max_times: Time of the hit amplitude in ns. + :param hit_channel: PMT channels of the hits + :param peak_max_times: Time of the peaks maximum in ns. + :param left: Left boundary in which we search for the tight + coincidence in ns. + :param right: Right boundary in which we search for the tight + coincidence in ns. + :param channel_range: (min/max) channel for the corresponding detector. + + :returns: n_coin_hit, n_coin_channel of length peaks containing the + tight coincidence. """ left_hit_i = 0 - n_coin = np.zeros(len(peak_max_times), dtype=np.int16) + n_coin_hit = np.zeros(len(peak_max_times), dtype=np.int16) + n_coin_channel = np.zeros(len(peak_max_times), dtype=np.int16) + start_ch, end_ch = channels + channels_seen = np.zeros(end_ch-start_ch+1, dtype=np.bool_) # loop over peaks for p_i, p_t in enumerate(peak_max_times): - + channels_seen[:] = 0 # loop over hits starting from the last one we left at for left_hit_i in range(left_hit_i, len(hit_max_times)): # if the hit is in the window, its a tight coin d = hit_max_times[left_hit_i] - p_t if (-left <= d) & (d <= right): - n_coin[p_i] += 1 + n_coin_hit[p_i] += 1 + channels_seen[hit_channel[left_hit_i]-start_ch] = 1 # stop the loop when we know we're outside the range if d > right: + n_coin_channel[p_i] = np.sum(channels_seen) break + + # Add channel information in case there are no hits beyond + # the last peak: + n_coin_channel[p_i] = np.sum(channels_seen) - return n_coin + return n_coin_hit, n_coin_channel @numba.njit(cache=True, nogil=True) diff --git a/straxen/plugins/position_reconstruction.py b/straxen/plugins/position_reconstruction.py index 19cb51de3..975ab5083 100644 --- a/straxen/plugins/position_reconstruction.py +++ b/straxen/plugins/position_reconstruction.py @@ -14,6 +14,7 @@ default="mlp", )]) + @export @strax.takes_config( strax.Option('min_reconstruction_area', @@ -22,7 +23,6 @@ strax.Option('n_top_pmts', default=straxen.n_top_pmts, help="Number of top PMTs") ) - class PeakPositionsBaseNT(strax.Plugin): """ Base class for reconstructions. @@ -120,6 +120,7 @@ def _get_model_file_name(self): model_file = straxen.get_correction_from_cmt(self.run_id, model_from_config) return model_file + @export @strax.takes_config( strax.Option('mlp_model', @@ -203,3 +204,97 @@ def compute(self, peaks): for xy in ('x', 'y'): result[xy] = peaks[f'{xy}_{algorithm}'] return result + + +@export +@strax.takes_config( + strax.Option('recon_alg_included', help = 'The list of all reconstruction algorithm considered.', + default = ('_mlp', '_gcn', '_cnn'), + ) +) +class S2ReconPosDiff(strax.Plugin): + ''' + Plugin that provides position reconstruction difference for S2s in events, see note: + https://xe1t-wiki.lngs.infn.it/doku.php?id=xenon:shengchao:sr0:reconstruction_quality + ''' + + __version__ = '0.0.3' + parallel = True + depends_on = 'event_basics' + provides = 's2_recon_pos_diff' + save_when = strax.SaveWhen.EXPLICIT + + def infer_dtype(self): + dtype = [ + ('s2_recon_avg_x', np.float32, + 'Mean value of x for main S2'), + ('alt_s2_recon_avg_x', np.float32, + 'Mean value of x for alternatice S2'), + ('s2_recon_avg_y', np.float32, + 'Mean value of y for main S2'), + ('alt_s2_recon_avg_y', np.float32, + 'Mean value of y for alternatice S2'), + ('s2_recon_pos_diff', np.float32, + 'Reconstructed position difference for main S2'), + ('alt_s2_recon_pos_diff', np.float32, + 'Reconstructed position difference for alternative S2'), + ] + dtype += strax.time_fields + return dtype + + def compute(self, events): + + result = np.zeros(len(events), dtype = self.dtype) + result['time'] = events['time'] + result['endtime'] = strax.endtime(events) + # Computing position difference + self.compute_pos_diff(events, result) + return result + + def cal_avg_and_std(self, values, axis = 1): + average = np.mean(values, axis = axis) + std = np.std(values, axis = axis) + return average, std + + def eval_recon(self, data, name_x_list, name_y_list): + """ + This function reads the name list based on s2/alt_s2 and all recon algorithm registered + Each row consists the reconstructed x/y and their average and standard deviation is calculated + """ + x_avg, x_std = self.cal_avg_and_std(np.array(data[name_x_list].tolist())) #lazy fix to delete field name in array, otherwise np.mean will complain + y_avg, y_std = self.cal_avg_and_std(np.array(data[name_y_list].tolist())) + r_std = np.sqrt(x_std**2 + y_std**2) + res = x_avg, y_avg, r_std + return res + + def compute_pos_diff(self, events, result): + + alg_list = self.config['recon_alg_included'] + for peak_type in ['s2', 'alt_s2']: + # Selecting S2s for pos diff + # - must exist (index != -1) + # - must have positive AFT + # - must contain all alg info + cur_s2_bool = (events[peak_type + '_index'] !=- 1) + cur_s2_bool &= (events[peak_type + '_area_fraction_top'] > 0) + for name in self.config['recon_alg_included']: + cur_s2_bool &= ~np.isnan(events[peak_type+'_x'+name]) + cur_s2_bool &= ~np.isnan(events[peak_type+'_y'+name]) + + # default value is nan, it will be ovewrite if the event satisfy the requirments + result[peak_type + '_recon_pos_diff'][:] = np.nan + result[peak_type + '_recon_avg_x'][:] = np.nan + result[peak_type + '_recon_avg_y'][:] = np.nan + + if np.any(cur_s2_bool): + name_x_list = [] + name_y_list = [] + for alg in alg_list: + name_x_list.append(peak_type + '_x' + alg) + name_y_list.append(peak_type + '_y' + alg) + + # Calculating average x,y, and position difference + x_avg, y_avg, r_std = self.eval_recon(events[cur_s2_bool], name_x_list, name_y_list) + result[peak_type + '_recon_pos_diff'][cur_s2_bool] = r_std + result[peak_type + '_recon_avg_x'][cur_s2_bool] = x_avg + result[peak_type + '_recon_avg_y'][cur_s2_bool] = y_avg diff --git a/straxen/plugins/veto_events.py b/straxen/plugins/veto_events.py index e3a2ecbed..6783d2581 100644 --- a/straxen/plugins/veto_events.py +++ b/straxen/plugins/veto_events.py @@ -10,6 +10,7 @@ export, __all__ = strax.exporter() + @strax.takes_config( strax.Option('event_left_extension_nv', default=0, help="Extends events this many ns to the left"), @@ -128,6 +129,7 @@ def compute_nveto_event_properties(events: np.ndarray, ch = hit['channel'] - start_channel e['area_per_channel'][ch] += hit['area'] + @export def find_veto_events(hitlets: np.ndarray, coincidence_level: int, @@ -223,7 +225,8 @@ def _make_event(hitlets: np.ndarray, @strax.takes_config( strax.Option('position_max_time_nv', default=20, - help="Time [ns] within an evnet use to compute the azimuthal angle of the event."), + help="Time [ns] within an event use to compute the azimuthal angle of the " + "event."), strax.Option('nveto_pmt_position_map', help="nVeto PMT position mapfile", default='nveto_pmt_position.csv'), @@ -243,7 +246,7 @@ class nVETOEventPositions(strax.Plugin): # Needed in case we make again an muVETO child. ends_with = '_nv' - __version__ = '0.0.1' + __version__ = '0.1.0' def infer_dtype(self): return veto_event_positions_dtype() @@ -274,8 +277,8 @@ def compute(self, events_nv, hitlets_nv): event_angles['n_prompt_hitlets'] = n_prompt # Compute azimuthal angle and xyz positions: - angle = compute_average_angle(hits_in_events, - self.pmt_properties) + angle = get_average_angle(hits_in_events, + self.pmt_properties) event_angles['angle'] = angle compute_positions(event_angles, events_nv, hits_in_events, self.pmt_properties) strax.copy_to_buffer(events_nv, event_angles, f'_copy_events{self.ends_with}') @@ -303,9 +306,22 @@ def veto_event_positions_dtype() -> list: @numba.njit(cache=True, nogil=True) def compute_positions(event_angles: np.ndarray, events: np.ndarray, - contained_hitlets: np.ndarray, + contained_hitlets: numba.typed.typedlist.List, pmt_pos: np.ndarray, start_channel: int = 2000): + """ + Function which computes some artificial event position for a given + neutron/muon-veto event. The position is computed based on a simple + area weighted mean. Please note that the event position can be + reconstructed in unphysical regions like being within the TPC. + + :param event_angles: Result array of the veto_event_position dtype. + The result is updated inplace. + :param events: Events for which the position should be computed. + :param contained_hitlets: Hitlets contained in each event. + :param pmt_pos: Position of the veto PMTs + :param start_channel: Starting channel of the detector. + """ for e_angles, e, hitlets in zip(event_angles, events, contained_hitlets): if e['area']: ch = hitlets['channel'] - start_channel @@ -313,20 +329,26 @@ def compute_positions(event_angles: np.ndarray, pos_y = pmt_pos['y'][ch] pos_z = pmt_pos['z'][ch] - e_angles['pos_x'] = np.sum(pos_x * hitlets['area']) / e['area'] - e_angles['pos_y'] = np.sum(pos_y * hitlets['area']) / e['area'] - e_angles['pos_z'] = np.sum(pos_z * hitlets['area']) / e['area'] - if len(hitlets) and np.sum(hitlets['area']) > 0: - w = hitlets['area'] / e['area'] # normalized weights - e_angles['pos_x_spread'] = np.sqrt(np.sum(w * np.power(pos_x - e_angles['pos_x'], 2)) / np.sum(w)) - e_angles['pos_y_spread'] = np.sqrt(np.sum(w * np.power(pos_y - e_angles['pos_y'], 2)) / np.sum(w)) - e_angles['pos_z_spread'] = np.sqrt(np.sum(w * np.power(pos_z - e_angles['pos_z'], 2)) / np.sum(w)) + e_angles['pos_x'] = np.sum(pos_x * hitlets['area'])/e['area'] + e_angles['pos_y'] = np.sum(pos_y * hitlets['area'])/e['area'] + e_angles['pos_z'] = np.sum(pos_z * hitlets['area'])/e['area'] + w = hitlets['area'] / e['area'] # normalized weights + if len(hitlets) and np.sum(w) > 0: + e_angles['pos_x_spread'] = np.sqrt( + np.sum(w * (pos_x - e_angles['pos_x'])**2)/np.sum(w) + ) + e_angles['pos_y_spread'] = np.sqrt( + np.sum(w * (pos_y - e_angles['pos_y'])**2)/np.sum(w) + ) + e_angles['pos_z_spread'] = np.sqrt( + np.sum(w * (pos_z - e_angles['pos_z'])**2)/np.sum(w) + ) @numba.njit(cache=True, nogil=True) -def compute_average_angle(hitlets_in_event: np.ndarray, - pmt_properties: np.ndarray, - start_channel: int = 2000,) -> np.ndarray: +def get_average_angle(hitlets_in_event: numba.typed.typedlist.List, + pmt_properties: np.ndarray, + start_channel: int = 2000, ) -> np.ndarray: """ Computes azimuthal angle as an area weighted mean over all hitlets. diff --git a/straxen/plugins/veto_hitlets.py b/straxen/plugins/veto_hitlets.py index d83dc1ff8..8e09afa51 100644 --- a/straxen/plugins/veto_hitlets.py +++ b/straxen/plugins/veto_hitlets.py @@ -71,7 +71,7 @@ class nVETOHitlets(strax.Plugin): Note: Hitlets are getting chopped if extended in not recorded regions. """ - __version__ = '0.1.0' + __version__ = '0.1.1' parallel = 'process' rechunk_on_save = True @@ -126,7 +126,9 @@ def compute(self, records_nv, start, end): min_hitlet_sample=600) temp_hitlets = strax.split_peaks(temp_hitlets, + None, # Only needed for peak splitting records_nv, + None, # Only needed for peak splitting self.to_pe, data_type='hitlets', algorithm='local_minimum', diff --git a/straxen/rucio.py b/straxen/rucio.py index 9b41496c2..f376cea7f 100644 --- a/straxen/rucio.py +++ b/straxen/rucio.py @@ -1,54 +1,43 @@ -import socket -import re -from tqdm import tqdm -import json -from bson import json_util -import os import glob import hashlib -import time -from utilix import xent_collection +import json +import os +import re +import socket +from warnings import warn +import numpy as np import strax -import warnings +from bson import json_util +from utilix import xent_collection try: - from rucio.client.client import Client - from rucio.client.rseclient import RSEClient - from rucio.client.replicaclient import ReplicaClient - from rucio.client.downloadclient import DownloadClient - from rucio.client.didclient import DIDClient + import admix from rucio.common.exception import DataIdentifierNotFound + HAVE_ADMIX = True +except ImportError: + HAVE_ADMIX = False - rucio_client = Client() - rse_client = RSEClient() - replica_client = ReplicaClient() - did_client = DIDClient() - download_client = DownloadClient() - RUCIO_AVAILABLE = True -except (ModuleNotFoundError, RuntimeError): - warnings.warn("No installation of rucio-clients found. Can't use rucio remote backend.") - RUCIO_AVAILABLE = False export, __all__ = strax.exporter() -__all__ += ['RUCIO_AVAILABLE'] class TooMuchDataError(Exception): pass -class DownloadError(Exception): - pass - - @export class RucioFrontend(strax.StorageFrontend): """ Uses the rucio client for the data find. """ - local_rses = {'UC_DALI_USERDISK': r'.rcc.'} + local_rses = {'UC_DALI_USERDISK': r'.rcc.', + 'SDSC_USERDISK': r'.sdsc.' + } local_did_cache = None - local_rucio_path = None + path = None + local_prefixes = {'UC_DALI_USERDISK': '/dali/lgrandi/rucio/', + 'SDSC_USERDISK': '/expanse/lustre/projects/chi135/shockley/rucio', + } def __init__(self, include_remote=False, @@ -90,14 +79,19 @@ def __init__(self, # rucio backend to read from that path rucio_prefix = self.get_rse_prefix(local_rse) self.backends.append(RucioLocalBackend(rucio_prefix)) - self.local_rucio_path = rucio_prefix + self.path = rucio_prefix if include_remote: - self.backends.append(RucioRemoteBackend(staging_dir, download_heavy=download_heavy)) + if not HAVE_ADMIX: + self.log.warning("You passed use_remote=True to rucio fronted, " + "but you don't have access to admix/rucio! Using local backed only.") + else: + self.backends.append(RucioRemoteBackend(staging_dir, + download_heavy=download_heavy)) def __repr__(self): # List the relevant attributes - attributes = ('include_remote', 'exclude', 'take_only', 'readonly') + attributes = ('include_remote', 'readonly', 'path', 'exclude', 'take_only') representation = f'{self.__class__.__module__}.{self.__class__.__name__}' for attr in attributes: if hasattr(self, attr) and getattr(self, attr): @@ -105,17 +99,9 @@ def __repr__(self): return representation def find_several(self, keys, **kwargs): - if not len(keys): - return [] - - ret = [] - for key in keys: - did = key_to_rucio_did(key) - if self.did_is_local(did): - ret.append(('RucioLocalBackend', did)) - else: - ret.append(False) - return ret + # for performance, dont do find_several with this plugin + # we basically do the same query we would do in the RunDB plugin + return np.zeros_like(keys, dtype=bool).tolist() def _find(self, key: strax.DataKey, write, allow_incomplete, fuzzy_for, fuzzy_for_options): did = key_to_rucio_did(key) @@ -126,12 +112,10 @@ def _find(self, key: strax.DataKey, write, allow_incomplete, fuzzy_for, fuzzy_fo if self.did_is_local(did): return "RucioLocalBackend", did elif self.include_remote: - # only do this part if we include the remote backend try: - # check if the DID exists - scope, name = did.split(':') - did_client.get_did(scope, name) - return "RucioRemoteBackend", did + rules = admix.rucio.list_rules(did, state="OK") + if len(rules): + return "RucioRemoteBackend", did except DataIdentifierNotFound: pass @@ -145,15 +129,12 @@ def _find(self, key: strax.DataKey, write, allow_incomplete, fuzzy_for, fuzzy_fo raise strax.DataNotAvailable def get_rse_prefix(self, rse): - try: - rse_info = rse_client.get_rse(rse) - prefix = rse_info['protocols'][0]['prefix'] - except NameError as e: - if self.local_rse == 'UC_DALI_USERDISK': - # If we cannot load rucio, let's try the default - prefix = '/dali/lgrandi/rucio/' - else: - raise e + if HAVE_ADMIX: + prefix = admix.rucio.get_rse_prefix(rse) + elif self.local_rse in self.local_prefixes: + prefix = self.local_prefixes[self.local_rse] + else: + raise ValueError(f'We are not on dali nor expanse and thus cannot load rucio') return prefix def did_is_local(self, did): @@ -166,7 +147,7 @@ def did_is_local(self, did): """ try: md = self._get_backend("RucioLocalBackend").get_metadata(did) - except (strax.DataNotAvailable, strax.DataCorrupted): + except (strax.DataNotAvailable, strax.DataCorrupted, KeyError): return False return self._all_chunk_stored(md, did) @@ -180,7 +161,7 @@ def _all_chunk_stored(self, md: dict, did: str) -> bool: for chunk in md.get('chunks', []): if chunk.get('filename'): _did = f"{scope}:{chunk['filename']}" - ch_path = rucio_path(self.local_rucio_path, _did) + ch_path = rucio_path(self.path, _did) if not os.path.exists(ch_path): return False return True @@ -249,10 +230,13 @@ class RucioRemoteBackend(strax.FileSytemBackend): # datatypes we don't want to download since they're too heavy heavy_types = ['raw_records', 'raw_records_nv', 'raw_records_he'] + # for caching RSE locations + dset_cache = {} + def __init__(self, staging_dir, download_heavy=False, **kwargs): """ :param staging_dir: Path (a string) where to save data. Must be a writable location. - :param *args: Passed to strax.FileSystemBackend + :param download_heavy: Whether or not to allow downloads of the heaviest data (raw_records*, less aqmon and MV) :param **kwargs: Passed to strax.FileSystemBackend """ @@ -266,30 +250,22 @@ def __init__(self, staging_dir, download_heavy=False, **kwargs): except OSError: raise PermissionError(f"You told the rucio backend to download data to {staging_dir}, " f"but that path is not writable by your user") - super().__init__(**kwargs) self.staging_dir = staging_dir self.download_heavy = download_heavy - def get_metadata(self, dset_did, rse='UC_OSG_USERDISK', **kwargs): - base_dir = os.path.join(self.staging_dir, did_to_dirname(dset_did)) - - # define where the metadata will go (or where it already might be) - number, dtype, hsh = parse_did(dset_did) - metadata_file = f"{dtype}-{hsh}-metadata.json" - metadata_path = os.path.join(base_dir, metadata_file) - - # download if it doesn't exist - if not os.path.exists(metadata_path): - metadata_did = f'{dset_did}-metadata.json' - did_dict = dict(did=metadata_did, - base_dir=base_dir, - no_subdir=True, - rse=rse - ) - print(f"Downloading {metadata_did}") - self._download([did_dict]) - + def _get_metadata(self, dset_did, **kwargs): + if dset_did in self.dset_cache: + rse = self.dset_cache[dset_did] + else: + rses = admix.rucio.get_rses(dset_did) + rse = admix.downloader.determine_rse(rses) + self.dset_cache[dset_did] = rse + + metadata_did = f'{dset_did}-metadata.json' + downloaded = admix.download(metadata_did, rse=rse, location=self.staging_dir) + assert len(downloaded) == 1, f"{metadata_did} should be a single file. We found {len(downloaded)}." + metadata_path = downloaded[0] # check again if not os.path.exists(metadata_path): raise FileNotFoundError(f"No metadata found at {metadata_path}") @@ -297,10 +273,10 @@ def get_metadata(self, dset_did, rse='UC_OSG_USERDISK', **kwargs): with open(metadata_path, mode='r') as f: return json.loads(f.read()) - def _read_chunk(self, dset_did, chunk_info, dtype, compressor, rse="UC_OSG_USERDISK"): + def _read_chunk(self, dset_did, chunk_info, dtype, compressor): base_dir = os.path.join(self.staging_dir, did_to_dirname(dset_did)) chunk_file = chunk_info['filename'] - chunk_path = os.path.join(base_dir, chunk_file) + chunk_path = os.path.abspath(os.path.join(base_dir, chunk_file)) if not os.path.exists(chunk_path): number, datatype, hsh = parse_did(dset_did) if datatype in self.heavy_types and not self.download_heavy: @@ -309,16 +285,20 @@ def _read_chunk(self, dset_did, chunk_info, dtype, compressor, rse="UC_OSG_USERD "doing, pass download_heavy=True to the Rucio " "frontend. If not, check your context and/or ask " "someone if this raw data is needed locally.") - raise DownloadError(error_msg) + warn(error_msg) + raise strax.DataNotAvailable scope, name = dset_did.split(':') chunk_did = f"{scope}:{chunk_file}" - print(f"Downloading {chunk_did}") - did_dict = dict(did=chunk_did, - base_dir=base_dir, - no_subdir=True, - rse=rse, - ) - self._download([did_dict]) + if dset_did in self.dset_cache: + rse = self.dset_cache[dset_did] + else: + rses = admix.rucio.get_rses(dset_did) + rse = admix.downloader.determine_rse(rses) + self.dset_cache[dset_did] = rse + + downloaded = admix.download(chunk_did, rse=rse, location=self.staging_dir) + assert len(downloaded) == 1, f"{chunk_did} should be a single file. We found {len(downloaded)}." + assert chunk_path == downloaded[0] # check again if not os.path.exists(chunk_path): @@ -326,44 +306,25 @@ def _read_chunk(self, dset_did, chunk_info, dtype, compressor, rse="UC_OSG_USERD return strax.load_file(chunk_path, dtype=dtype, compressor=compressor) - def _saver(self, dirname, metadata): + def _saver(self, dirname, metadata, **kwargs): raise NotImplementedError("Cannot save directly into rucio (yet), upload with admix instead") - def _download(self, did_dict_list): - # need to pass a list of dicts - # let's try 3 times - success = False - _try = 1 - while _try <= 3 and not success: - if _try > 1: - for did_dict in did_dict_list: - did_dict['rse'] = None - try: - download_client.download_dids(did_dict_list) - success = True - except KeyboardInterrupt: - raise - except Exception: - sleep = 3**_try - print(f"Download try #{_try} failed. Sleeping for {sleep} seconds and trying again...") - time.sleep(sleep) - _try += 1 - if not success: - raise DownloadError(f"Error downloading from rucio.") - class RucioSaver(strax.Saver): """ TODO Saves data to rucio if you are the production user """ - def __init__(self): + def __init__(self, *args, **kwargs): raise NotImplementedError def rucio_path(root_dir, did): - """Convert target to path according to rucio convention. - See the __hash method here: https://github.com/rucio/rucio/blob/1.20.15/lib/rucio/rse/protocols/protocol.py""" + """ + Convert target to path according to rucio convention. + See the __hash method here: + https://github.com/rucio/rucio/blob/1.20.15/lib/rucio/rse/protocols/protocol.py + """ scope, filename = did.split(':') # disable bandit rucio_md5 = hashlib.md5(did.encode('utf-8')).hexdigest() # nosec @@ -394,17 +355,8 @@ def key_to_rucio_did(key: strax.DataKey) -> str: return f'xnt_{key.run_id}:{key.data_type}-{key.lineage_hash}' -def key_to_rucio_meta(key: strax.DataKey) -> str: - return f'{str(key.data_type)}-{key.lineage_hash}-metadata.json' - - def read_md(path: str) -> json: with open(path, mode='r') as f: md = json.loads(f.read(), object_hook=json_util.object_hook) return md - - -def list_datasets(scope): - datasets = [d for d in rucio_client.list_dids(scope, {'type': 'dataset'}, type='dataset')] - return datasets diff --git a/straxen/rundb.py b/straxen/rundb.py index 315f3d7c2..dff3e25ae 100644 --- a/straxen/rundb.py +++ b/straxen/rundb.py @@ -5,7 +5,7 @@ from tqdm import tqdm from copy import deepcopy import strax -from .rucio import key_to_rucio_did +from .rucio import key_to_rucio_did, RucioLocalBackend import warnings try: @@ -122,8 +122,7 @@ def __init__(self, self.available_query.append({'host': host_alias}) if self.rucio_path is not None: - # TODO replace with rucio backend in the rucio module - self.backends.append(strax.rucio(self.rucio_path)) + self.backends.append(RucioLocalBackend(self.rucio_path)) # When querying for rucio, add that it should be dali-userdisk self.available_query.append({'host': 'rucio-catalogue', 'location': 'UC_DALI_USERDISK', @@ -183,9 +182,8 @@ def _find(self, key: strax.DataKey, datum = doc['data'][0] error_message = f'Expected {rucio_key} got data on {datum["location"]}' assert datum.get('did', '') == rucio_key, error_message - backend_name, backend_key = ( - datum['protocol'], - f'{key.run_id}-{key.data_type}-{key.lineage_hash}') + backend_name = 'RucioLocalBackend' + backend_key = key_to_rucio_did(key) return backend_name, backend_key dq = self._data_query(key) @@ -268,20 +266,6 @@ def find_several(self, keys: typing.List[strax.DataKey], **kwargs): return [results_dict.get(k.run_id, False) for k in keys] - def _list_available(self, key: strax.DataKey, - allow_incomplete, fuzzy_for, fuzzy_for_options): - if fuzzy_for or fuzzy_for_options or allow_incomplete: - # The RunDB frontend can do neither fuzzy nor incomplete - warnings.warn('RunDB cannot do fuzzy or incomplete') - - q = self._data_query(key) - q.update(self.number_query()) - - cursor = self.collection.find( - q, - projection=[self.runid_field]) - return [x[self.runid_field] for x in cursor] - def _scan_runs(self, store_fields): query = self.number_query() projection = strax.to_str_tuple(list(store_fields)) @@ -292,8 +276,9 @@ def _scan_runs(self, store_fields): cursor = self.collection.find( filter=query, projection=projection) - for doc in tqdm(cursor, desc='Fetching run info from MongoDB', - total=cursor.count()): + for doc in strax.utils.tqdm( + cursor, desc='Fetching run info from MongoDB', + total=cursor.count()): del doc['_id'] if self.reader_ini_name_is_mode: doc['mode'] = \ @@ -302,7 +287,7 @@ def _scan_runs(self, store_fields): def run_metadata(self, run_id, projection=None): if run_id.startswith('_'): - # Superruns are currently not supprorted.. + # Superruns are currently not supported.. raise strax.DataNotAvailable if self.runid_field == 'name': diff --git a/straxen/scada.py b/straxen/scada.py index 4cd947b0a..dc12916ed 100644 --- a/straxen/scada.py +++ b/straxen/scada.py @@ -321,31 +321,25 @@ def _query_single_parameter(self, offset = 1 else: offset = 0 - ntries = 0 - max_tries = 40000 # This corresponds to ~23 years - while ntries < max_tries: - temp_df = self._query(query, - self.SCData_URL, - start=(start // 10**9) + offset, - end=(end // 10**9), - query_type_lab=query_type_lab, - seconds_interval=every_nth_value, - raise_error_message=False # No valid value in query range... - ) # +1 since it is end before exclusive - if temp_df.empty: - # In case WebInterface does not return any data, e.g. if query range too small - break - times = (temp_df['timestampseconds'].values * 10**9).astype(' 0 assert 'cs1' in df.columns assert df['cs1'].sum() > 0 assert not np.all(np.isnan(df['x'].values)) + def test_event_info_double(self): + df = self.st.get_df(self.run_id, 'event_info_double') + assert 'cs2_a' in df.columns + assert df['cs2_a'].sum() > 0 + assert len(df) > 0 + def test_get_livetime_sec(self): st = self.st - events = st.get_array(self.run_id, 'peaks') + events = st.get_array(self.run_id, 'events') straxen.get_livetime_sec(st, test_run_id_1T, things=events) def test_mini_analysis(self): @@ -60,3 +64,34 @@ def count_rr(raw_records): n = self.st.count_rr(self.run_id) assert n > 100 + + @staticmethod + def _extract_latest_comment(context, + test_for_target='raw_records', + **context_kwargs, + ): + if context == 'xenonnt_online' and not straxen.utilix_is_configured(): + return + st = getattr(straxen.contexts, context)(**context_kwargs) + assert hasattr(st, 'extract_latest_comment'), "extract_latest_comment not added to context?" + st.extract_latest_comment() + assert st.runs is not None, "No registry build?" + assert 'comments' in st.runs.keys() + st.select_runs(available=test_for_target) + if context == 'demo': + assert len(st.runs) + assert f'{test_for_target}_available' in st.runs.keys() + + def test_extract_latest_comment_nt(self, **opt): + """Run the test for nt (but only 2000 runs""" + self._extract_latest_comment(context='xenonnt_online', + _minimum_run_number=10_000, + _maximum_run_number=12_000, + **opt) + + def test_extract_latest_comment_demo(self): + self._extract_latest_comment(context='demo') + + def test_extract_latest_comment_lone_hits(self): + """Run the test for some target that is not in the default availability check""" + self.test_extract_latest_comment_nt(test_for_target='lone_hits') diff --git a/tests/test_cmt.py b/tests/test_cmt.py index ad8805871..e64696da9 100644 --- a/tests/test_cmt.py +++ b/tests/test_cmt.py @@ -6,65 +6,51 @@ import numpy as np from warnings import warn from .test_basics import test_run_id_1T -from .test_plugins import test_run_id_nT +from straxen.test_utils import nt_test_run_id as test_run_id_nT from straxen.common import aux_repo +import unittest + +@unittest.skipIf(not straxen.utilix_is_configured(), "No db access, cannot test!") def test_connect_to_db(): """ Test connection to db """ - if not straxen.utilix_is_configured(): - warn('Cannot do test becaus ' - 'no have access to the database.') - return - - username=None - password=None - mongo_url=None - is_nt=True - mongo_kwargs = {'url': mongo_url, - 'user': username, - 'password': password, - 'database': 'corrections'} - corrections_collection = utilix.rundb.xent_collection(**mongo_kwargs) + corrections_collection = utilix.rundb.xent_collection(database='corrections') client = corrections_collection.database.client cmt = strax.CorrectionsInterface(client, database_name='corrections') df = cmt.read('global_xenonnt') mes = 'Return empty dataframe when reading DB. Please check' assert not df.empty, mes + +@unittest.skipIf(not straxen.utilix_is_configured(), "No db access, cannot test!") def test_1T_elife(): """ Test elife from CMT DB against historical data(aux file) """ - if not straxen.utilix_is_configured(): - warn('Cannot do test becaus ' - 'no have access to the database.') - return - elife_conf = ('elife_xenon1t', 'ONLINE', False) elife_cmt = straxen.get_correction_from_cmt(test_run_id_1T, elife_conf) - elife_file = elife_conf=aux_repo + '3548132b55f81a43654dba5141366041e1daaf01/strax_files/elife.npy' + elife_file = aux_repo + '3548132b55f81a43654dba5141366041e1daaf01/strax_files/elife.npy' x = straxen.get_resource(elife_file, fmt='npy') run_index = np.where(x['run_id'] == int(test_run_id_1T))[0] elife = x[run_index[0]]['e_life'] mes = 'Elife values do not match. Please check' assert elife_cmt == elife, mes + +@unittest.skipIf(not straxen.utilix_is_configured(), "No db access, cannot test!") def test_cmt_conf_option(option='mlp_model', version='ONLINE', is_nT=True): """ Test CMT conf options If wrong conf is passed it would raise an error accordingly """ - if not straxen.utilix_is_configured(): - warn('Cannot do test becaus ' - 'no have access to the database.') - return - conf = option, version, is_nT correction = straxen.get_correction_from_cmt(test_run_id_nT, conf) assert isinstance(correction, (float, int, str, np.ndarray)) + +@unittest.skipIf(not straxen.utilix_is_configured(), "No db access, cannot test!") def test_mc_wrapper_elife(run_id='009000', cmt_id='016000', mc_id='mc_0', @@ -78,8 +64,6 @@ def test_mc_wrapper_elife(run_id='009000', and the test does not work). :return: None """ - if not straxen.utilix_is_configured(): - return assert np.abs(int(run_id) - int(cmt_id)) > 500, 'runs must be far apart' # First for the run-id let's get the value @@ -106,6 +90,7 @@ def test_mc_wrapper_elife(run_id='009000', assert elife == mc_elife_same +@unittest.skipIf(not straxen.utilix_is_configured(), "No db access, cannot test!") def test_mc_wrapper_gains(run_id='009000', cmt_id='016000', mc_id='mc_0', @@ -123,9 +108,8 @@ def test_mc_wrapper_gains(run_id='009000', the testing time due to faster CMT queries is reduced). :return: None """ - if not straxen.utilix_is_configured() or not execute: + if not execute: return - assert np.abs(int(run_id) - int(cmt_id)) > 500, 'runs must be far apart' # First for the run-id let's get the value @@ -156,16 +140,4 @@ def test_is_cmt_option(): The example dummy_option works at least before Jun 13 2021 """ dummy_option = ('hit_thresholds_tpc', 'ONLINE', True) - assert is_cmt_option(dummy_option), 'Structure of CMT options changed!' - - -def is_cmt_option(config): - """ - Check if the input configuration is cmt style. - """ - is_cmt = (isinstance(config, tuple) - and len(config) == 3 - and isinstance(config[0], str) - and isinstance(config[1], (str, int, float)) - and isinstance(config[2], bool)) - return is_cmt + assert straxen.is_cmt_option(dummy_option), 'Structure of CMT options changed!' diff --git a/tests/test_contexts.py b/tests/test_contexts.py index 049db13ed..089c5551a 100644 --- a/tests/test_contexts.py +++ b/tests/test_contexts.py @@ -6,7 +6,7 @@ import straxen import tempfile import os - +import unittest ## # XENONnT @@ -14,35 +14,27 @@ def test_xenonnt_online(): - st = xenonnt_online(_database_init=False, use_rucio=False) + st = xenonnt_online(_database_init=False) st.search_field('time') def test_xennonnt(): if straxen.utilix_is_configured(): - st = xenonnt(_database_init=False, use_rucio=False) - st.search_field('time') - - -def test_xennonnt_latest(cmt_version='latest'): - if straxen.utilix_is_configured(): - st = xenonnt(cmt_version, _database_init=False, use_rucio=False) + st = xenonnt(_database_init=False) st.search_field('time') def test_xenonnt_led(): - st = xenonnt_led(_database_init=False, use_rucio=False) + st = xenonnt_led(_database_init=False) st.search_field('time') +@unittest.skipIf(not straxen.utilix_is_configured(), "No db access, cannot test!") def test_nt_is_nt_online(): - if not straxen.utilix_is_configured(): - # Cannot contact CMT without the database - return # Test that nT and nT online are the same - st_online = xenonnt_online(_database_init=False, use_rucio=False) + st_online = xenonnt_online(_database_init=False) - st = xenonnt(_database_init=False, use_rucio=False) + st = xenonnt(_database_init=False) for plugin in st._plugin_class_registry.keys(): print(f'Checking {plugin}') nt_key = st.key_for('0', plugin) @@ -86,3 +78,9 @@ def test_fake_daq(): def test_xenon1t_led(): st = xenon1t_led() st.search_field('time') + + +@unittest.skipIf('ALLOW_WFSIM_TEST' not in os.environ, "if you want test wfsim context do `export 'ALLOW_WFSIM_TEST'=1`") +def test_sim_context(): + st = straxen.contexts.xenonnt_simulation(cmt_run_id_sim='008000', cmt_version='global_ONLINE') + st.search_field('time') diff --git a/tests/test_count_pulses.py b/tests/test_count_pulses.py index 8bc60bc68..f1b3eeaab 100644 --- a/tests/test_count_pulses.py +++ b/tests/test_count_pulses.py @@ -43,7 +43,7 @@ def _check_pulse_count(records): # Not sure how to check lone pulses other than duplicating logic # already in count_pulses, so just do a basic check: assert count['lone_pulse_area'][ch] <= count['pulse_area'][ch] - + # Check baseline values: # nan does not exist for integer: mean = straxen.NO_PULSE_COUNTS if np.isnan(np.mean(rc0['baseline'])) else int(np.mean(rc0['baseline'])) diff --git a/tests/test_database_frontends.py b/tests/test_database_frontends.py new file mode 100644 index 000000000..5fa5c0809 --- /dev/null +++ b/tests/test_database_frontends.py @@ -0,0 +1,149 @@ +import unittest +import strax +from strax.testutils import Records, Peaks +import straxen +import os +import shutil +import tempfile +import pymongo +import datetime + + +def mongo_uri_not_set(): + return 'TEST_MONGO_URI' not in os.environ + + +class TestRunDBFrontend(unittest.TestCase): + """ + Test the saving behavior of the context with the straxen.RunDB + + Requires write access to some pymongo server, the URI of witch is to be set + as an environment variable under: + + TEST_MONGO_URI + + At the moment this is just an empty database but you can also use some free + ATLAS mongo server. + """ + _run_test = True + + def setUp(self): + # Just to make sure we are running some mongo server, see test-class docstring + if 'TEST_MONGO_URI' not in os.environ: + return + self.test_run_ids = ['0', '1'] + self.all_targets = ('peaks', 'records') + + uri = os.environ.get('TEST_MONGO_URI') + db_name = 'test_rundb' + self.collection_name = 'test_rundb_coll' + client = pymongo.MongoClient(uri) + self.database = client[db_name] + collection = self.database[self.collection_name] + self.path = os.path.join(tempfile.gettempdir(), 'strax_data') + assert self.collection_name not in self.database.list_collection_names() + + if not straxen.utilix_is_configured(): + # Bit of an ugly hack but there is no way to get around this + # function even though we don't need it + straxen.rundb.utilix.rundb.xent_collection = lambda *args, **kwargs: collection + + self.rundb_sf = straxen.RunDB(readonly=False, + runid_field='number', + new_data_path=self.path, + minimum_run_number=-1, + ) + self.rundb_sf.client = client + self.rundb_sf.collection = collection + + self.st = strax.Context(register=[Records, Peaks], + storage=[self.rundb_sf], + use_per_run_defaults=False, + config=dict(bonus_area=0), + ) + for run_id in self.test_run_ids: + collection.insert_one(_rundoc_format(run_id)) + assert not self.is_all_targets_stored + + @unittest.skipIf(mongo_uri_not_set(), "No access to test database") + def tearDown(self): + self.database[self.collection_name].drop() + if os.path.exists(self.path): + print(f'rm {self.path}') + shutil.rmtree(self.path) + + @property + def is_all_targets_stored(self) -> bool: + """This should always be False as one of the targets (records) is not stored in mongo""" + return all([all( + [self.st.is_stored(r, t) for t in self.all_targets]) + for r in self.test_run_ids]) + + @unittest.skipIf(mongo_uri_not_set(), "No access to test database") + def test_finding_runs(self): + rdb = self.rundb_sf + col = self.database[self.collection_name] + assert col.find_one() is not None + query = rdb.number_query() + assert col.find_one(query) is not None + runs = self.st.select_runs() + assert len(runs) == len(self.test_run_ids) + + @unittest.skipIf(mongo_uri_not_set(), "No access to test database") + def test_write_and_load(self): + assert not self.is_all_targets_stored + + # Make ALL the data + # NB: the context writes to ALL the storage frontends that are susceptible + for t in self.all_targets: + self.st.make(self.test_run_ids, t) + + for r in self.test_run_ids: + print(self.st.available_for_run(r)) + assert self.is_all_targets_stored + + # Double check that we can load data from mongo even if we cannot make it + self.st.context_config['forbid_creation_of'] = self.all_targets + peaks = self.st.get_array(self.test_run_ids, self.all_targets[-1]) + assert len(peaks) + runs = self.st.select_runs(available=self.all_targets) + assert len(runs) == len(self.test_run_ids) + + # Insert a new run number and check that it's not marked as available + self.database[self.collection_name].insert_one(_rundoc_format(3)) + self.st.runs = None # Reset + all_runs = self.st.select_runs() + available_runs = self.st.select_runs(available=self.all_targets) + assert len(available_runs) == len(self.test_run_ids) + assert len(all_runs) == len(self.test_run_ids) + 1 + + @unittest.skipIf(mongo_uri_not_set(), "No access to test database") + def test_lineage_changes(self): + st = strax.Context(register=[Records, Peaks], + storage=[self.rundb_sf], + use_per_run_defaults=True, + ) + lineages = [st.key_for(r, 'peaks').lineage_hash for r in self.test_run_ids] + assert len(set(lineages)) > 1 + with self.assertRaises(ValueError): + # Lineage changing per run is not allowed! + st.select_runs(available='peaks') + + +def _rundoc_format(run_id): + start = datetime.datetime.fromtimestamp(0) + datetime.timedelta(days=int(run_id)) + end = start + datetime.timedelta(days=1) + doc = { + 'comments': [{'comment': 'some testdoc', + 'date': start, + 'user': 'master user'}], + 'data': [], + 'detectors': ['tpc'], + + 'mode': 'test', + 'number': int(run_id), + 'source': 'none', + 'start': start, + 'end': end, + 'user': 'master user'} + return doc diff --git a/tests/test_holoviews_utils.py b/tests/test_holoviews_utils.py index 7cf7aae23..187a24eb2 100644 --- a/tests/test_holoviews_utils.py +++ b/tests/test_holoviews_utils.py @@ -1,22 +1,14 @@ import strax import straxen from straxen.holoviews_utils import nVETOEventDisplay - import holoviews as hv import panel as pn import numpy as np - from tempfile import TemporaryDirectory import os -dummy_map = np.zeros(120, dtype=[('x', np.int32), - ('y', np.int32), - ('z', np.int32), - ('channel', np.int32),]) -dummy_map['x'] = np.arange(0, 120) -dummy_map['y'] = np.arange(0, 120) -dummy_map['z'] = np.arange(0, 120) -dummy_map['channel'] = np.arange(2000, 2120, 1, dtype=np.int32) +_dummy_map = straxen.test_utils._nveto_pmt_dummy_df.to_records() + def test_hitlets_to_hv_points(): hit = np.zeros(1, dtype=strax.hit_dtype) @@ -26,10 +18,11 @@ def test_hitlets_to_hv_points(): hit['channel'] = 2000 hit['area'] = 1 - nvd = nVETOEventDisplay(pmt_map=dummy_map) + nvd = nVETOEventDisplay(pmt_map=_dummy_map) points = nvd.hitlets_to_hv_points(hit, t_ref=0) - m = [hit[key] == points.data[key] for key in hit.dtype.names if key in points.data.columns.values] + m = [hit[key] == points.data[key] for key in hit.dtype.names if + key in points.data.columns.values] assert np.all(m), 'Data has not been converted corretly into hv.Points.' @@ -41,7 +34,7 @@ def test_hitlet_matrix(): hit['channel'] = 2000 hit['area'] = 1 - nvd = nVETOEventDisplay(pmt_map=dummy_map) + nvd = nVETOEventDisplay(pmt_map=_dummy_map) hit_m = nvd.plot_hitlet_matrix(hitlets=hit) with TemporaryDirectory() as d: @@ -54,7 +47,7 @@ def test_plot_nveto_pattern(): hit['channel'] = 2000 hit['area'] = 1 - nvd = nVETOEventDisplay(pmt_map=dummy_map) + nvd = nVETOEventDisplay(pmt_map=_dummy_map) pmt_plot = nvd.plot_nveto(hitlets=hit) with TemporaryDirectory() as d: # Have to store plot to make sure it is rendered @@ -76,7 +69,7 @@ def test_nveto_event_display(): event['endtime'] = hit['time'] + 40 event['area'] = hit['area'] - nvd = nVETOEventDisplay(event, hit, pmt_map=dummy_map, run_id='014986') + nvd = nVETOEventDisplay(event, hit, pmt_map=_dummy_map, run_id='014986') dispaly = nvd.plot_event_display() with TemporaryDirectory() as d: @@ -89,7 +82,7 @@ def test_array_to_df_and_make_sliders(): + straxen.veto_events.veto_event_positions_dtype()[2:]) evt = np.zeros(1, dtype) - nvd = nVETOEventDisplay(pmt_map=dummy_map) + nvd = nVETOEventDisplay(pmt_map=_dummy_map) df = straxen.convert_array_to_df(evt) nvd._make_sliders_and_tables(df) diff --git a/tests/test_led_calibration.py b/tests/test_led_calibration.py new file mode 100644 index 000000000..2bfe57099 --- /dev/null +++ b/tests/test_led_calibration.py @@ -0,0 +1,32 @@ +import strax +from strax.testutils import several_fake_records +import straxen +import numpy as np +from hypothesis import given, settings + + +@given(several_fake_records) +@settings(deadline=None) +def test_ext_timings_nv(records): + """ + Little test for nVetoExtTimings. + """ + if not straxen.utilix_is_configured(): + return + # several fake records do not have any pulse length + # and channel start at 0, convert to nv: + records['pulse_length'] = records['length'] + records['channel'] += 2000 + + st = straxen.contexts.xenonnt_led() + plugin = st.get_single_plugin('1', 'ext_timings_nv') + hits = strax.find_hits(records, min_amplitude=1) + + hitlets = np.zeros(len(hits), strax.hitlet_dtype()) + strax.copy_to_buffer(hits, hitlets, '_refresh_hit_to_hitlets') + result = plugin.compute(hitlets, records) + assert len(result) == len(hits) + assert np.all(result['time'] == hits['time']) + assert np.all(result['pulse_i'] == hits['record_i']) + true_dt = hits['time'] - records[hits['record_i']]['time'] + assert np.all(result['delta_time'] == true_dt) diff --git a/tests/test_mini_analyses.py b/tests/test_mini_analyses.py new file mode 100644 index 000000000..898ba00a6 --- /dev/null +++ b/tests/test_mini_analyses.py @@ -0,0 +1,227 @@ +import straxen +import pandas +import os +import numpy as np +import strax +from matplotlib.pyplot import clf as plt_clf +from straxen.test_utils import nt_test_context, nt_test_run_id +import unittest + + +def test_pmt_pos_1t(): + """ + Test if we can get the 1T PMT positions + """ + pandas.DataFrame(straxen.pmt_positions(True)) + + +def test_pmt_pos_nt(): + """ + Test if we can get the nT PMT positions + """ + pandas.DataFrame(straxen.pmt_positions(False)) + + +class TestMiniAnalyses(unittest.TestCase): + """ + + """ + # They were added on 25/10/2021 and may be outdated by now + _expected_test_results = { + 'peak_basics': 40, + 'n_s1': 19, + 'run_live_time': 4.7516763, + 'event_basics': 20, + } + + @classmethod + def setUpClass(cls) -> None: + cls.st = nt_test_context() + # For al the WF plotting, we might need records, let's make those + cls.st.make(nt_test_run_id, 'records') + cls.first_peak = cls.st.get_array(nt_test_run_id, 'peak_basics')[0] + cls.first_event = cls.st.get_array(nt_test_run_id, 'event_basics')[0] + + def tearDown(self): + plt_clf() + + def test_target_peaks(self, target='peak_basics', tol=2): + assert target in self._expected_test_results, f'No expectation for {target}?!' + data = self.st.get_array(nt_test_run_id, target) + message = (f'Got more/less data for {target}. If you changed something ' + f'on {target}, please update the numbers in ' + f'tests/test_mini_analyses.TestMiniAnalyses._expected_test_results') + if not straxen.utilix_is_configured(): + # If we do things with dummy maps, things might be slightly different + tol += 10 + assert np.abs(len(data) - self._expected_test_results[target]) < tol, message + + def test_target_events(self): + self.test_target_peaks(target='event_basics') + + def test_plot_waveform(self, deep=False): + self.st.plot_waveform(nt_test_run_id, time_within=self.first_peak, deep=deep) + + def test_plot_waveform_deep(self): + self.test_plot_waveform(deep=True) + + def test_plot_hit_pattern(self): + self.st.plot_hit_pattern(nt_test_run_id, time_within=self.first_peak, xenon1t=False) + + def test_plot_records_matrix(self): + self.st_attr_for_one_peak('plot_records_matrix') + + def test_raw_records_matrix(self): + self.st_attr_for_one_peak('raw_records_matrix') + + def test_event_display_simple(self): + plot_all_positions = straxen.utilix_is_configured() + self.st.event_display_simple(nt_test_run_id, + time_within=self.first_event, + xenon1t=False, + plot_all_positions=plot_all_positions, + ) + + def test_single_event_plot(self): + plot_all_positions = straxen.utilix_is_configured() + straxen.analyses.event_display.plot_single_event( + self.st, + nt_test_run_id, + events=self.st.get_array(nt_test_run_id, 'events'), + event_number=self.first_event['event_number'], + xenon1t=False, + plot_all_positions=plot_all_positions, + ) + + def test_event_display_interactive(self): + self.st.event_display_interactive(nt_test_run_id, + time_within=self.first_event, + xenon1t=False, + ) + + def test_plot_peaks_aft_histogram(self): + self.st.plot_peaks_aft_histogram(nt_test_run_id) + + def test_event_scatter(self): + self.st.event_scatter(nt_test_run_id) + + def test_energy_spectrum(self): + self.st.plot_energy_spectrum(nt_test_run_id) + + def test_peak_classification(self): + self.st.plot_peak_classification(nt_test_run_id) + + def st_attr_for_one_peak(self, function_name): + f = getattr(self.st, function_name) + f(nt_test_run_id, time_within=self.first_peak) + + def test_waveform_display(self): + self.st_attr_for_one_peak('waveform_display') + + def test_hvdisp_plot_pmt_pattern(self): + self.st_attr_for_one_peak('hvdisp_plot_pmt_pattern') + + def test_hvdisp_plot_peak_waveforms(self): + self.st_attr_for_one_peak('hvdisp_plot_peak_waveforms') + + def test_plot_pulses_tpc(self): + self.st.plot_pulses_tpc(nt_test_run_id, max_plots=2, plot_hits=True, + ignore_time_warning=True) + + def test_plot_pulses_mv(self): + self.st.plot_pulses_mv(nt_test_run_id, max_plots=2, plot_hits=True, + ignore_time_warning=True) + + def test_plot_pulses_nv(self): + self.st.plot_pulses_nv(nt_test_run_id, max_plots=2, plot_hits=True, + ignore_time_warning=True) + + @unittest.skipIf(not straxen.utilix_is_configured(), "No db access, cannot test!") + def test_event_display(self): + self.st.event_display(nt_test_run_id, time_within=self.first_event) + + def test_calc_livetime(self): + try: + live_time = straxen.get_livetime_sec(self.st, nt_test_run_id) + except strax.RunMetadataNotAvailable: + things = self.st.get_array(nt_test_run_id, 'peaks') + live_time = straxen.get_livetime_sec(self.st, nt_test_run_id, things=things) + assertion_statement = "Live-time calculation is wrong" + assert live_time == self._expected_test_results['run_live_time'], assertion_statement + + def test_df_wiki(self): + df = self.st.get_df(nt_test_run_id, 'peak_basics') + straxen.dataframe_to_wiki(df) + + def test_interactive_display(self): + fig = self.st.event_display_interactive(nt_test_run_id, + time_within=self.first_event, + xenon1t=False, + plot_record_matrix=True, + ) + save_as = 'test_display.html' + fig.save(save_as) + assert os.path.exists(save_as) + os.remove(save_as) + assert not os.path.exists(save_as), f'Should have removed {save_as}' + + def test_selector(self): + from straxen.analyses.bokeh_waveform_plot import DataSelectionHist + p = self.st.get_array(nt_test_run_id, 'peak_basics') + ds = DataSelectionHist('ds') + fig = ds.histogram2d(p, + p['area'], + p['area'], + bins=50, + hist_range=((0, 200), (0, 2000)), + log_color_scale=True, + clim=(10, None), + undeflow_color='white') + + import bokeh.plotting as bklt + save_as = 'test_data_selector.html' + bklt.save(fig, save_as) + assert os.path.exists(save_as) + os.remove(save_as) + assert not os.path.exists(save_as), f'Should have removed {save_as}' + + @unittest.skipIf(not straxen.utilix_is_configured(), "No db access, cannot test!") + def test_nt_daq_plot(self): + self.st.daq_plot(nt_test_run_id, + time_within=self.first_peak, + vmin=0.1, + vmax=1, + ) + + @unittest.skipIf(not straxen.utilix_is_configured(), "No db access, cannot test!") + def test_nt_daq_plot_grouped(self): + self.st.plot_records_matrix(nt_test_run_id, + time_within=self.first_peak, + vmin=0.1, + vmax=1, + group_by='ADC ID', + ) + + def test_records_matrix_downsample(self): + self.st.records_matrix(nt_test_run_id, + time_within=self.first_event, + max_samples=20 + ) + + @unittest.skipIf(not straxen.utilix_is_configured(), "No db access, cannot test!") + def test_load_corrected_positions(self): + self.st.load_corrected_positions(nt_test_run_id, time_within=self.first_peak) + + @unittest.skipIf(not straxen.utilix_is_configured(), "No db access, cannot test!") + def test_nv_event_display(self): + self.st.make(nt_test_run_id, 'events_nv') + self.st.make(nt_test_run_id, 'event_positions_nv') + with self.assertRaises(ValueError): + self.st.plot_nveto_event_display(nt_test_run_id, time_within=self.first_peak) + + +def test_plots(): + """Make some plots""" + c = np.ones(straxen.n_tpc_pmts) + straxen.plot_pmts(c) + straxen.plot_pmts(c, log_scale=True) diff --git a/tests/test_misc.py b/tests/test_misc.py index c217f96cf..08123ce21 100644 --- a/tests/test_misc.py +++ b/tests/test_misc.py @@ -1,4 +1,6 @@ -from straxen.misc import TimeWidgets +from straxen.misc import TimeWidgets, print_versions +import straxen +import unittest def test_widgets(): @@ -14,10 +16,12 @@ def test_widgets(): start_utc, end_utc = tw.get_start_end() h_in_ns_unix = 60*60*10**9 - unix_conversion_worked = start_utc - start == h_in_ns_unix or start_utc - start == 2 * h_in_ns_unix - assert unix_conversion_worked - unix_conversion_worked = start_utc - end == h_in_ns_unix or start_utc - end == 2 * h_in_ns_unix - assert unix_conversion_worked + assert (start_utc - start == h_in_ns_unix + or start_utc - start == 2 * h_in_ns_unix + or start_utc - start == 0 * h_in_ns_unix) + assert (start_utc - end == h_in_ns_unix + or start_utc - end == 2 * h_in_ns_unix + or start_utc - end == 0 * h_in_ns_unix) def test_change_in_fields(): @@ -41,3 +45,24 @@ def test_change_in_fields(): start00, _ = tw.get_start_end() assert start20 - start00 == minutes, 'Time field did not update its value!' + + +def test_print_versions(modules=('numpy', 'straxen', 'non_existing_module')): + for return_string in [True, False]: + for include_git in [True, False]: + res = print_versions(modules, + return_string=return_string, + include_git=include_git) + if return_string: + assert res is not None + + +class HitAmplitude(unittest.TestCase): + def test_non_existing(self): + with self.assertRaises(ValueError): + straxen.hit_min_amplitude('non existing key') + + @staticmethod + def test_get_hit_amplitude(): + straxen.hit_min_amplitude('pmt_commissioning_initial') + straxen.hit_min_amplitude('pmt_commissioning_initial_he') diff --git a/tests/test_mongo_downloader.py b/tests/test_mongo_downloader.py new file mode 100644 index 000000000..b3f465e40 --- /dev/null +++ b/tests/test_mongo_downloader.py @@ -0,0 +1,74 @@ +import unittest +import straxen +import os +import pymongo + + +def mongo_uri_not_set(): + return 'TEST_MONGO_URI' not in os.environ + + +class TestMongoDownloader(unittest.TestCase): + """ + Test the saving behavior of the context with the mogno downloader + + Requires write access to some pymongo server, the URI of witch is to be set + as an environment variable under: + + TEST_MONGO_URI + + At the moment this is just an empty database but you can also use some free + ATLAS mongo server. + """ + _run_test = True + + def setUp(self): + # Just to make sure we are running some mongo server, see test-class docstring + if 'TEST_MONGO_URI' not in os.environ: + self._run_test = False + return + uri = os.environ.get('TEST_MONGO_URI') + db_name = 'test_rundb' + collection_name = 'fs.files' + client = pymongo.MongoClient(uri) + database = client[db_name] + collection = database[collection_name] + self.downloader = straxen.MongoDownloader(collection=collection, + readonly=True, + file_database=None, + _test_on_init=False, + ) + self.uploader = straxen.MongoUploader(collection=collection, + readonly=False, + file_database=None, + _test_on_init=False, + ) + self.collection = collection + + @unittest.skipIf(mongo_uri_not_set(), "No access to test database") + def tearDown(self): + self.collection.drop() + + @unittest.skipIf(mongo_uri_not_set(), "No access to test database") + def test_upload(self): + with self.assertRaises(ConnectionError): + # Should be empty! + self.downloader.test_find() + file_name = 'test.txt' + file_content = 'This is a test' + with open(file_name, 'w') as f: + f.write(file_content) + assert os.path.exists(file_name) + self.uploader.upload_from_dict({file_name: os.path.abspath(file_name)}) + assert self.uploader.md5_stored(file_name) + assert self.downloader.config_exists(file_name) + download_path = self.downloader.download_single(file_name) + assert os.path.exists(download_path) + read_file = straxen.get_resource(download_path) + assert file_content == read_file + os.remove(file_name) + assert not os.path.exists(file_name) + self.downloader.test_find() + + with self.assertRaises(NotImplementedError): + self.downloader.download_all() diff --git a/tests/test_mongo_interactions.py b/tests/test_mongo_interactions.py index e547dc558..5427d53ee 100644 --- a/tests/test_mongo_interactions.py +++ b/tests/test_mongo_interactions.py @@ -4,13 +4,12 @@ work e.g. on travis jobs and therefore the tests failing locally will not show up in Pull Requests. """ - import straxen import os -from warnings import warn -from .test_plugins import test_run_id_nT +import unittest +@unittest.skipIf(not straxen.utilix_is_configured(), "No db access, cannot test!") def test_select_runs(check_n_runs=2): """ Test (if we have a connection) if we can perform strax.select_runs @@ -18,11 +17,6 @@ def test_select_runs(check_n_runs=2): :param check_n_runs: int, the number of runs we want to check """ - - if not straxen.utilix_is_configured(): - warn('Makes no sense to test the select runs because we do not ' - 'have access to the database.') - return assert check_n_runs >= 1 st = straxen.contexts.xenonnt_online(use_rucio=False) run_col = st.storage[0].collection @@ -38,12 +32,9 @@ def test_select_runs(check_n_runs=2): st.select_runs() +@unittest.skipIf(not straxen.utilix_is_configured(), "Cannot download because utilix is not configured") def test_downloader(): """Test if we can download a small file from the downloader""" - if not straxen.utilix_is_configured(): - warn('Cannot download because utilix is not configured') - return - downloader = straxen.MongoDownloader() path = downloader.download_single('to_pe_nt.npy') assert os.path.exists(path) @@ -61,6 +52,7 @@ def _patch_om_init(take_only): return straxen.OnlineMonitor(uri=uri, take_only=take_only) +@unittest.skipIf(not straxen.utilix_is_configured(), "No db access, cannot test!") def test_online_monitor(target='online_peak_monitor', max_tries=3): """ See if we can get some data from the online monitor before max_tries @@ -69,10 +61,8 @@ def test_online_monitor(target='online_peak_monitor', max_tries=3): :param max_tries: number of queries max allowed to get a non-failing run """ - if not straxen.utilix_is_configured(): - warn('Cannot test online monitor because utilix is not configured') - return st = straxen.contexts.xenonnt_online(use_rucio=False) + straxen.get_mongo_uri() om = _patch_om_init(target) st.storage = [om] max_run = None diff --git a/tests/test_nveto_events.py b/tests/test_nveto_events.py index 2ab59f306..d5c9b4491 100644 --- a/tests/test_nveto_events.py +++ b/tests/test_nveto_events.py @@ -170,7 +170,7 @@ def _test_ambiguity(hitlets_ids_in_event): time_range=(0, 15), channel_range=(2000, 2010), length_range=(1, 4), ), - hnp.arrays(np.float32, elements=hst.floats(0, 10, width=32), shape=7), + hnp.arrays(np.float32, elements=hst.floats(-1, 10, width=32), shape=7), ) def test_nveto_event_plugin(hitlets, area): hitlets['area'] = area @@ -229,9 +229,9 @@ def test_nveto_event_plugin(hitlets, area): npmt_pos, start_channel=2000) - angle = straxen.plugins.veto_events.compute_average_angle(split_hitlets, - npmt_pos, - start_channel=2000) + angle = straxen.plugins.veto_events.get_average_angle(split_hitlets, + npmt_pos, + start_channel=2000) # Compute truth angles: truth_angle = np.angle(events_angle['pos_x']+events_angle['pos_y']*1j) # Replace not defined angles, into zeros to match np.angles return diff --git a/tests/test_nveto_recorder.py b/tests/test_nveto_recorder.py index 162aeedc4..4c4d86d44 100644 --- a/tests/test_nveto_recorder.py +++ b/tests/test_nveto_recorder.py @@ -6,7 +6,6 @@ class TestMergeIntervals(unittest.TestCase): - def setUp(self): self.intervals = np.zeros(4, dtype=strax.time_fields) self.intervals['time'] = [2, 3, 7, 20] @@ -34,7 +33,6 @@ def test_merge_overlapping_intervals(self): class TestCoincidence(unittest.TestCase): - def setUp(self): self.intervals = np.zeros(8, dtype=strax.time_fields) self.intervals['time'] = [3, 6, 9, 12, 15, 18, 21, 38] @@ -134,3 +132,12 @@ def _test_coincidence(self, resolving_time, coincidence, pre_trigger, endtime_is_correct = np.all(coincidence['endtime'] == endtime_truth) print(coincidence['endtime'], endtime_truth) assert endtime_is_correct, 'Coincidence does not have the correct endtime' + + +def test_nv_for_dummy_rr(): + """Basic test to run the nv rr for dummy raw-records""" + st = straxen.test_utils.nt_test_context(deregister=()) + st.context_config['forbid_creation_of'] = tuple() + st.register(straxen.test_utils.DummyRawRecords) + st.make(straxen.test_utils.nt_test_run_id, 'hitlets_nv') + st.make(straxen.test_utils.nt_test_run_id, 'events_tagged') diff --git a/tests/test_peaklet_processing.py b/tests/test_peaklet_processing.py new file mode 100644 index 000000000..f79ff895f --- /dev/null +++ b/tests/test_peaklet_processing.py @@ -0,0 +1,83 @@ +import numpy as np +from hypothesis import given, settings +import hypothesis.strategies as strat + +import strax +from strax.testutils import fake_hits +import straxen +from straxen.plugins.peaklet_processing import get_tight_coin + + +@settings(deadline=None) +@given(strat.lists(strat.integers(min_value=0, max_value=10), + min_size=8, max_size=8, unique=True), + ) +def test_create_outside_peaks_region(time): + time = np.sort(time) + time_intervals = np.zeros(len(time)//2, strax.time_dt_fields) + time_intervals['time'] = time[::2] + time_intervals['length'] = time[1::2] - time[::2] + time_intervals['dt'] = 1 + + st = straxen.contexts.demo() + p = st.get_single_plugin('0', 'peaklets') + outside = p.create_outside_peaks_region(time_intervals, 0, np.max(time)) + + touching = strax.touching_windows(outside, time_intervals, window=0) + + for tw in touching: + print(tw) + assert np.diff(tw) == 0, 'Intervals overlap although they should not!' + + +def test_n_hits(): + if not straxen.utilix_is_configured(): + return + records = np.zeros(2, dtype=strax.record_dtype()) + records['length'] = 5 + records['pulse_length'] = 5 + records['dt'] = 1 + records['channel'] = [0, 1] + records['data'][0, :5] = [0, 1, 1, 0, 1] + records['data'][1, :5] = [0, 1, 0, 0, 0] + + st = straxen.contexts.xenonnt_online() + st.set_config({'hit_min_amplitude': 1}) + p = st.get_single_plugin('0', 'peaklets') + res = p.compute(records, 0, 999) + peaklets = res['peaklets'] + assert peaklets['n_hits'] == 3, f"Peaklet has the wrong number of hits!" + + +@given(fake_hits, + strat.lists(elements=strat.integers(0, 9), min_size=20)) +@settings(deadline=None) +def test_tight_coincidence(hits, channel): + hits['area'] = 1 + hits['channel'] = channel[:len(hits)] # In case there are less channel then hits (unlikely) + gap_threshold = 10 + peaks = strax.find_peaks(hits, + adc_to_pe=np.ones(10), + right_extension=0, left_extension=0, + gap_threshold=gap_threshold, + min_channels=1, + min_area=0) + + peaks_max_time = peaks['time'] + peaks['length']//2 + hits_max_time = hits['time'] + hits['length']//2 + + left = 5 + right = 5 + tight_coin, tight_coin_channel = get_tight_coin(hits_max_time, + hits['channel'], + peaks_max_time, + left, + right, + ) + for ind, p_max_t in enumerate(peaks_max_time): + m_hits_in_peak = (hits_max_time >= (p_max_t - left)) + m_hits_in_peak &= (hits_max_time <= (p_max_t + right)) + n_hits = np.sum(m_hits_in_peak) + n_channel = len(np.unique(hits[m_hits_in_peak]['channel'])) + assert n_hits == tight_coin[ind], 'Wrong number of tight hits' + assert n_channel == tight_coin_channel[ind], f'Wrong number of tight channel got {tight_coin_channel[ind]}, but expectd {n_channel}' # noqa diff --git a/tests/test_peaks.py b/tests/test_peaks.py index 9e25873d3..c08bfc5ae 100644 --- a/tests/test_peaks.py +++ b/tests/test_peaks.py @@ -10,13 +10,13 @@ TEST_DATA_LENGTH = 3 R_TOL_DEFAULT = 1e-5 + def _not_close_to_0_or_1(x, rtol=R_TOL_DEFAULT): return not (np.isclose(x, 1, rtol=rtol) or np.isclose(x, 0, rtol=rtol)) class TestComputePeakBasics(unittest.TestCase): """Tests for peak basics plugin""" - def setUp(self, context=straxen.contexts.demo): self.st = context() self.n_top = self.st.config.get('n_top_pmts', 2) @@ -72,3 +72,90 @@ def get_test_peaks(n_top, length=2, sum_wf_samples=10): test_data['dt'] = 1 test_data['length'] = length return test_data + + +def create_unique_intervals(size, time_range=(0, 40), allow_zero_length=True): + """ + Hypothesis stragtegy which creates unqiue time intervals. + + :param size: Number of intervals desired. Can be less if non-unique + intervals are found. + :param time_range: Time range in which intervals should be. + :param allow_zero_length: If true allow zero length intervals. + """ + strat = strategies.lists(elements=strategies.integers(*time_range), + min_size=size * 2, + max_size=size * 2 + ).map(lambda x: _convert_to_interval(x, allow_zero_length)) + return strat + + +def _convert_to_interval(time_stamps, allow_zero_length): + time_stamps = np.sort(time_stamps) + intervals = np.zeros(len(time_stamps) // 2, strax.time_dt_fields) + intervals['dt'] = 1 + intervals['time'] = time_stamps[::2] + intervals['length'] = time_stamps[1::2] - time_stamps[::2] + + if not allow_zero_length: + intervals = intervals[intervals['length'] > 0] + return np.unique(intervals) + + +@settings(deadline=None) +@given(create_unique_intervals(10, time_range=(0, 30), allow_zero_length=False), + create_unique_intervals(5, time_range=(5, 25), allow_zero_length=False) + ) +def test_tag_peaks(peaks, veto_intervals): + peaks_in_vetos = strax.touching_windows(peaks, veto_intervals) + + tags = np.zeros(len(peaks)) + straxen.plugins.peak_processing.tag_peaks(tags, peaks_in_vetos, 1) + + # Make an additional dummy array to test if function worked: + dtype = [] + dtype += strax.time_dt_fields + dtype += [(('peak tag', 'tag'), np.int8)] + tagged_peaks = np.zeros(len(peaks), dtype) + tagged_peaks['time'] = peaks['time'] + tagged_peaks['length'] = peaks['length'] + tagged_peaks['dt'] = 1 + tagged_peaks['tag'] = tags + + split_tagged_peaks = strax.split_touching_windows(tagged_peaks, veto_intervals) + + for split_peaks in split_tagged_peaks: + if not len(split_peaks): + continue + assert np.all(split_peaks['tag'] == 1), f'Not all peaks were tagged properly {split_peaks}' + + +@settings(deadline=None) +@given(create_unique_intervals(10, time_range=(50, 80), allow_zero_length=False), + create_unique_intervals(5, time_range=(55, 75), allow_zero_length=False) + ) +def test_get_time_to_clostest_veto(peaks, veto_intervals): + time_differences = straxen.plugins.peak_processing.get_time_to_closest_veto(peaks, + veto_intervals + ) + + if not len(peaks): + assert not len(time_differences), 'Input is empty but output is not?' + + for ind, p in enumerate(peaks): + if len(veto_intervals): + dt = np.concatenate([veto_intervals['time'] - p['time'], + strax.endtime(veto_intervals) - p['time']]) + # If the distance to the curren/next veto interval is identical + # the function favors positive values. Hence sort and reverse + # order for the test. + dt = np.sort(dt)[::-1] + ind_argmin = np.argmin(np.abs(dt)) + dt = dt[ind_argmin] + else: + # If there are not any veto intervalls function will compute clostest + # difference to +/- 64 bit infinity. + dt = np.abs(straxen.INFINITY_64BIT_SIGNED - p['time']) + + mes = f'Wrong time difference to closest event. expected: "{dt}" got: "{time_differences[ind]}"' + assert dt == time_differences[ind], mes diff --git a/tests/test_plugins.py b/tests/test_plugins.py index 0284a4bc9..c71f065fd 100644 --- a/tests/test_plugins.py +++ b/tests/test_plugins.py @@ -1,135 +1,14 @@ import tempfile import strax -import numpy as np -from immutabledict import immutabledict import straxen -from straxen.common import pax_file, aux_repo -## -# Tools -## -# Let's make a dummy map for NVeto -nveto_pmt_dummy_df = {'channel': list(range(2000, 2120)), - 'x': list(range(120)), - 'y': list(range(120)), - 'z': list(range(120))} - -# Some configs are better obtained from the strax_auxiliary_files repo. -# Let's use small files, we don't want to spend a lot of time downloading -# some file. -testing_config_nT = dict( - nn_architecture= - aux_repo + 'f0df03e1f45b5bdd9be364c5caefdaf3c74e044e/fax_files/mlp_model.json', - nn_weights= - aux_repo + 'f0df03e1f45b5bdd9be364c5caefdaf3c74e044e/fax_files/mlp_model.h5', - gain_model=("to_pe_placeholder", True), - s2_xy_correction_map=pax_file('XENON1T_s2_xy_ly_SR0_24Feb2017.json'), - elife_conf=('elife_constant', 1e6), - baseline_samples_nv=10, - fdc_map=pax_file('XENON1T_FDC_SR0_data_driven_3d_correction_tf_nn_v0.json.gz'), - gain_model_nv=("adc_nv", True), - gain_model_mv=("adc_mv", True), - nveto_pmt_position_map=nveto_pmt_dummy_df, - s1_xyz_correction_map=pax_file('XENON1T_s1_xyz_lce_true_kr83m_SR0_pax-680_fdc-3d_v0.json'), - electron_drift_velocity=("electron_drift_velocity_constant", 1e-4), - s1_aft_map=aux_repo + 'ffdadba3439ae7922b19f5dd6479348b253c09b0/strax_files/s1_aft_UNITY_xyz_XENONnT.json', - s2_optical_map=aux_repo + '8a6f0c1a4da4f50546918cd15604f505d971a724/strax_files/s2_map_UNITY_xy_XENONnT.json', - s1_optical_map=aux_repo + '8a6f0c1a4da4f50546918cd15604f505d971a724/strax_files/s1_lce_UNITY_xyz_XENONnT.json', - electron_drift_time_gate=("electron_drift_time_gate_constant", 2700), - hit_min_amplitude='pmt_commissioning_initial', - hit_min_amplitude_nv=20, - hit_min_amplitude_mv=80, - hit_min_amplitude_he='pmt_commissioning_initial_he' -) - -testing_config_1T = dict( - hev_gain_model=('1T_to_pe_placeholder', False), - gain_model=('1T_to_pe_placeholder', False), - elife_conf=('elife_constant', 1e6), - electron_drift_velocity=("electron_drift_velocity_constant", 1e-4), - electron_drift_time_gate=("electron_drift_time_gate_constant", 1700), -) - -test_run_id_nT = '008900' -test_run_id_1T = '180423_1021' - - -@strax.takes_config( - strax.Option('secret_time_offset', default=0, track=False), - strax.Option('recs_per_chunk', default=10, track=False), - strax.Option('n_chunks', default=2, track=False, - help='Number of chunks for the dummy raw records we are writing here'), - strax.Option('channel_map', track=False, type=immutabledict, - help="frozendict mapping subdetector to (min, max) " - "channel number.") -) -class DummyRawRecords(strax.Plugin): - """ - Provide dummy raw records for the mayor raw_record types - """ - provides = ('raw_records', - 'raw_records_he', - 'raw_records_nv', - 'raw_records_aqmon', - 'raw_records_aux_mv', - 'raw_records_mv' - ) - parallel = 'process' - depends_on = tuple() - data_kind = immutabledict(zip(provides, provides)) - rechunk_on_save = False - dtype = {p: strax.raw_record_dtype() for p in provides} - - def setup(self): - self.channel_map_keys = {'he': 'he', - 'nv': 'nveto', - 'aqmon': 'aqmon', - 'aux_mv': 'aux_mv', - 's_mv': 'mv', - } # s_mv otherwise same as aux in endswith - - def source_finished(self): - return True - - def is_ready(self, chunk_i): - return chunk_i < self.config['n_chunks'] - - def compute(self, chunk_i): - t0 = chunk_i + self.config['secret_time_offset'] - if chunk_i < self.config['n_chunks'] - 1: - # One filled chunk - r = np.zeros(self.config['recs_per_chunk'], self.dtype['raw_records']) - r['time'] = t0 - r['length'] = r['dt'] = 1 - r['channel'] = np.arange(len(r)) - else: - # One empty chunk - r = np.zeros(0, self.dtype['raw_records']) - - res = {} - for p in self.provides: - rr = np.copy(r) - # Add detector specific channel offset: - for key, channel_key in self.channel_map_keys.items(): - if channel_key not in self.config['channel_map']: - # Channel map for 1T is different. - continue - if p.endswith(key): - s, e = self.config['channel_map'][channel_key] - rr['channel'] += s - res[p] = self.chunk(start=t0, end=t0 + 1, data=rr, data_type=p) - return res - - -# Don't concern ourselves with rr_aqmon et cetera -forbidden_plugins = tuple([p for p in - straxen.daqreader.DAQReader.provides - if p not in DummyRawRecords.provides]) +from straxen.test_utils import nt_test_run_id, DummyRawRecords, testing_config_1T, test_run_id_1T def _run_plugins(st, make_all=False, - run_id=test_run_id_nT, - **proces_kwargs): + run_id=nt_test_run_id, + from_scratch=False, + **process_kwargs): """ Try all plugins (except the DAQReader) for a given context (st) to see if we can really push some (empty) data from it and don't have any nasty @@ -137,16 +16,23 @@ def _run_plugins(st, """ with tempfile.TemporaryDirectory() as temp_dir: - st.storage = [strax.DataDirectory(temp_dir)] - - # As we use a temporary directory we should have a clean start - assert not st.is_stored(run_id, 'raw_records'), 'have RR???' + if from_scratch: + st.storage = [strax.DataDirectory(temp_dir)] + # As we use a temporary directory we should have a clean start + assert not st.is_stored(run_id, 'raw_records'), 'have RR???' + + # Don't concern ourselves with rr_aqmon et cetera + _forbidden_plugins = tuple([p for p in + straxen.daqreader.DAQReader.provides + if p not in + st._plugin_class_registry['raw_records'].provides]) + st.set_context_config({'forbid_creation_of': _forbidden_plugins}) # Create event info target = 'event_info' st.make(run_id=run_id, targets=target, - **proces_kwargs) + **process_kwargs) # The stuff should be there assert st.is_stored(run_id, target), f'Could not make {target}' @@ -155,11 +41,13 @@ def _run_plugins(st, return end_targets = set(st._get_end_targets(st._plugin_class_registry)) - for p in end_targets-set(forbidden_plugins): + for p in end_targets - set(_forbidden_plugins): + if 'raw' in p: + continue st.make(run_id, p) # Now make sure we can get some data for all plugins all_datatypes = set(st._plugin_class_registry.keys()) - for p in all_datatypes-set(forbidden_plugins): + for p in all_datatypes - set(_forbidden_plugins): should_be_stored = (st._plugin_class_registry[p].save_when == strax.SaveWhen.ALWAYS) if should_be_stored: @@ -168,23 +56,12 @@ def _run_plugins(st, print("Wonderful all plugins work (= at least they don't fail), bye bye") -def _update_context(st, max_workers, fallback_gains=None, nt=True): - # Change config to allow for testing both multiprocessing and lazy mode - st.set_context_config({'forbid_creation_of': forbidden_plugins}) +def _update_context(st, max_workers, nt=True): # Ignore strax-internal warnings st.set_context_config({'free_options': tuple(st.config.keys())}) - st.register(DummyRawRecords) - if nt and not straxen.utilix_is_configured(): - st.set_config(testing_config_nT) - del st._plugin_class_registry['peak_positions_mlp'] - del st._plugin_class_registry['peak_positions_cnn'] - del st._plugin_class_registry['peak_positions_gcn'] - st.register(straxen.PeakPositions1T) - print(f"Using {st._plugin_class_registry['peak_positions']} for posrec tests") - st.set_config({'gain_model': fallback_gains}) - - elif not nt: - if straxen.utilix_is_configured(): + if not nt: + st.register(DummyRawRecords) + if straxen.utilix_is_configured(warning_message=False): # Set some placeholder gain as this takes too long for 1T to load from CMT st.set_config({k: v for k, v in testing_config_1T.items() if k in ('hev_gain_model', 'gain_model')}) @@ -241,11 +118,6 @@ def _test_child_options(st, run_id): f'"{option_name}"!') -## -# Tests -## - - def test_1T(ncores=1): if ncores == 1: print('-- 1T lazy mode --') @@ -259,7 +131,7 @@ def test_1T(ncores=1): _plugin_class.save_when = strax.SaveWhen.ALWAYS # Run the test - _run_plugins(st, make_all=True, max_wokers=ncores, run_id=test_run_id_1T) + _run_plugins(st, make_all=True, max_workers=ncores, run_id=test_run_id_1T, from_scratch=True) # Test issue #233 st.search_field('cs1') @@ -274,25 +146,21 @@ def test_1T(ncores=1): def test_nT(ncores=1): if ncores == 1: print('-- nT lazy mode --') - st = straxen.contexts.xenonnt_online(_database_init=straxen.utilix_is_configured(), - use_rucio=False) - offline_gain_model = ("to_pe_placeholder", True) - _update_context(st, ncores, fallback_gains=offline_gain_model, nt=True) + init_database = straxen.utilix_is_configured(warning_message=False) + st = straxen.test_utils.nt_test_context( + _database_init=init_database, + use_rucio=False, + ) + _update_context(st, ncores, nt=True) # Lets take an abandoned run where we actually have gains for in the CMT - _run_plugins(st, make_all=True, max_wokers=ncores, run_id=test_run_id_nT) + _run_plugins(st, make_all=True, max_workers=ncores, run_id=nt_test_run_id) # Test issue #233 st.search_field('cs1') # Test of child plugins: - _test_child_options(st, test_run_id_nT) + _test_child_options(st, nt_test_run_id) print(st.context_config) def test_nT_mutlticore(): print('nT multicore') test_nT(2) - -# Disable the test below as it saves some time in travis and gives limited new -# information as most development is on nT-plugins. -# def test_1T_mutlticore(): -# print('1T multicore') -# test_1T(2) diff --git a/tests/test_rucio.py b/tests/test_rucio.py new file mode 100644 index 000000000..7caf02833 --- /dev/null +++ b/tests/test_rucio.py @@ -0,0 +1,78 @@ +import straxen +import unittest +import strax +import socket + + +class TestBasics(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + """ + For testing purposes, slightly alter the RucioFrontend such that + we can run tests outside of dali too + """ + if not straxen.utilix_is_configured(): + return + if 'rcc' not in socket.getfqdn(): + # If we are not on RCC, for testing, add some dummy site + straxen.RucioFrontend.local_rses = {'UC_DALI_USERDISK': r'.rcc.', + 'test_rucio': f'{socket.getfqdn()}'} + straxen.RucioFrontend.get_rse_prefix = lambda *x: 'test_rucio' + + # Some non-existing keys that we will try finding in the test cases. + cls.test_keys = [ + strax.DataKey(run_id=run_id, + data_type='dtype', + lineage={'dtype': ['Plugin', '0.0.0.', {}], } + ) + for run_id in ('-1', '-2') + ] + + @unittest.skipIf(not straxen.utilix_is_configured(), "No db access, cannot test!") + def test_load_context_defaults(self): + """Don't fail immediately if we start a context due to Rucio""" + st = straxen.contexts.xenonnt_online(_minimum_run_number=10_000, + _maximum_run_number=10_010, + ) + st.select_runs() + + @unittest.skipIf(not straxen.utilix_is_configured(), "No db access, cannot test!") + def test_find_local(self): + """Make sure that we don't find the non existing data""" + rucio = straxen.RucioFrontend(include_remote=False,) + self.assertRaises(strax.DataNotAvailable, + rucio.find, + self.test_keys[0] + ) + + @unittest.skipIf(not straxen.utilix_is_configured(), "No db access, cannot test!") + def test_find_several_local(self): + """Let's try finding some keys (won't be available)""" + rucio = straxen.RucioFrontend(include_remote=False,) + print(rucio) + found = rucio.find_several(self.test_keys) + # We shouldn't find any of these + assert found == [False for _ in self.test_keys] + + @unittest.skipIf(not straxen.utilix_is_configured(), "No db access, cannot test!") + def test_find_several_remote(self): + """ + Let's try running a find_several with the include remote. + This should fail but when no rucio is installed or else it + shouldn't find any data. + """ + try: + rucio = straxen.RucioFrontend(include_remote=True,) + except ImportError: + pass + else: + found = rucio.find_several(self.test_keys) + # We shouldn't find any of these + assert found == [False for _ in self.test_keys] + + @unittest.skipIf(not straxen.utilix_is_configured(), "No db access, cannot test!") + def test_find_local(self): + """Make sure that we don't find the non existing data""" + run_db = straxen.RunDB(rucio_path='./rucio_test') + with self.assertRaises(strax.DataNotAvailable): + run_db.find(self.test_keys[0]) diff --git a/tests/test_scada.py b/tests/test_scada.py index a23e97f20..0d2c25216 100644 --- a/tests/test_scada.py +++ b/tests/test_scada.py @@ -1,18 +1,14 @@ import numpy as np import straxen -import warnings - +import unittest +@unittest.skipIf(not straxen.utilix_is_configured('scada','scdata_url',), + "Cannot test scada since we have no access to xenon secrets.')") def test_query_sc_values(): """ Unity test for the SCADAInterface. Query a fixed range and check if return is correct. """ - - if not straxen.utilix_is_configured('scada', 'scdata_url'): - warnings.warn('Cannot test scada since we have no access to xenon secrets.') - return - print('Testing SCADAInterface') sc = straxen.SCADAInterface(use_progress_bar=False) @@ -22,14 +18,14 @@ def test_query_sc_values(): start = 1609682275000000000 # Add micro-second to check if query does not fail if inquery precsion > SC precision start += 10**6 - end = start + 5 * 10**9 + end = start + 5*10**9 parameters = {'SomeParameter': 'XE1T.CTPC.Board06.Chan011.VMon'} df = sc.get_scada_values(parameters, start=start, end=end, every_nth_value=1, - query_type_lab=False,) + query_type_lab=False, ) assert df['SomeParameter'][0] // 1 == 1253, 'First values returned is not corrrect.' assert np.all(np.isnan(df['SomeParameter'][1:])), 'Subsequent values are not correct.' @@ -53,7 +49,7 @@ def test_query_sc_values(): end=end, fill_gaps='forwardfill', every_nth_value=1, - query_type_lab=False,) + query_type_lab=False, ) df = sc.get_scada_values(parameters, start=start, @@ -83,3 +79,9 @@ def test_query_sc_values(): query_type_lab=True,) assert np.all(df['SomeParameter'] // 1 == -96), 'Not all values are correct for query type lab.' + + +def test_average_scada(): + t = np.arange(0, 100, 10) + t_t, t_a = straxen.scada._average_scada(t / 1e9, t, 1) + assert len(t_a) == len(t), 'Scada deleted some of my 10 datapoints!' diff --git a/tests/test_several.py b/tests/test_several.py deleted file mode 100644 index 0f5bcac6f..000000000 --- a/tests/test_several.py +++ /dev/null @@ -1,280 +0,0 @@ -"""Test several functions distibuted over common.py, misc.py, scada.py""" -import straxen -import pandas -import os -import tempfile -from .test_basics import test_run_id_1T -import numpy as np -import strax -from matplotlib.pyplot import clf as plt_clf - - -def test_pmt_pos_1t(): - """ - Test if we can get the 1T PMT positions - """ - pandas.DataFrame(straxen.pmt_positions(True)) - - -def test_pmt_pos_nt(): - """ - Test if we can get the nT PMT positions - """ - pandas.DataFrame(straxen.pmt_positions(False)) - - -def test_secret(): - """ - Check something in the sectets. This should not work because we - don't have any. - """ - try: - straxen.get_secret('somethingnonexistent') - except ValueError: - # Good we got some message we cannot load something that does - # not exist, - pass - - -# If one of the test below fail, perhaps these values need to be updated. -# They were added on 27/11/2020 and may be outdated by now -EXPECTED_OUTCOMES_TEST_SEVERAL = { - 'n_peaks': 128, - 'n_s1': 6, - 'run_live_time': 0.17933107, - 'n_events': 2, -} - - -def test_several(): - """ - Test several other functions in straxen. Is kind of messy but saves - time as we won't load data many times - :return: - """ - with tempfile.TemporaryDirectory() as temp_dir: - try: - print("Temporary directory is ", temp_dir) - os.chdir(temp_dir) - - print("Downloading test data (if needed)") - st = straxen.contexts.demo() - st.make(test_run_id_1T, 'records') - # Ignore strax-internal warnings - st.set_context_config({'free_options': tuple(st.config.keys())}) - st.make(test_run_id_1T, 'records') - - print("Get peaks") - p = st.get_array(test_run_id_1T, 'peaks') - - # Do checks on there number of peaks - assertion_statement = ("Got /more peaks than expected, perhaps " - "the test is outdated or clustering has " - "really changed") - assert np.abs(len(p) - - EXPECTED_OUTCOMES_TEST_SEVERAL['n_peaks']) < 5, assertion_statement - - events = st.get_array(test_run_id_1T, 'event_info') - print('plot wf') - peak_i = 0 - st.plot_waveform(test_run_id_1T, time_range=(p[peak_i]['time'], strax.endtime(p[peak_i]))) - plt_clf() - - print('plot hit pattern') - peak_i = 1 - st.plot_hit_pattern(test_run_id_1T, time_range=(p[peak_i]['time'], strax.endtime(p[peak_i])), xenon1t=True) - plt_clf() - - print('plot (raw)records matrix') - peak_i = 2 - assert st.is_stored(test_run_id_1T, 'records'), "no records" - assert st.is_stored(test_run_id_1T, 'raw_records'), "no raw records" - st.plot_records_matrix(test_run_id_1T, time_range=(p[peak_i]['time'], - strax.endtime(p[peak_i]))) - - st.raw_records_matrix(test_run_id_1T, time_range=(p[peak_i]['time'], - strax.endtime(p[peak_i]))) - st.plot_waveform(test_run_id_1T, - time_range=(p[peak_i]['time'], - strax.endtime(p[peak_i])), - deep=True) - plt_clf() - - straxen.analyses.event_display.plot_single_event(st, - test_run_id_1T, - events, - xenon1t=True, - event_number=0, - records_matrix='raw') - st.event_display_simple(test_run_id_1T, - time_range=(events[0]['time'], - events[0]['endtime']), - xenon1t=True) - plt_clf() - - st.event_display_interactive(test_run_id_1T, time_range=(events[0]['time'], - events[0]['endtime']), - xenon1t=True) - plt_clf() - - print('plot aft') - st.plot_peaks_aft_histogram(test_run_id_1T) - plt_clf() - - print('plot event scatter') - st.event_scatter(test_run_id_1T) - plt_clf() - - print('plot event scatter') - st.plot_energy_spectrum(test_run_id_1T) - plt_clf() - - print('plot peak clsassification') - st.plot_peak_classification(test_run_id_1T) - plt_clf() - - print("plot holoviews") - peak_i = 3 - st.waveform_display(test_run_id_1T, - time_range=(p[peak_i]['time'], - strax.endtime(p[peak_i]))) - st.hvdisp_plot_pmt_pattern(test_run_id_1T, - time_range=(p[peak_i]['time'], - strax.endtime(p[peak_i]))) - st.hvdisp_plot_peak_waveforms(test_run_id_1T, - time_range=(p[peak_i]['time'], - strax.endtime(p[peak_i]))) - - - print('Plot single pulse:') - st.plot_pulses_tpc(test_run_id_1T, max_plots=2, plot_hits=True, ignore_time_warning=True) - - print("Check live-time") - live_time = straxen.get_livetime_sec(st, test_run_id_1T, things=p) - assertion_statement = "Live-time calculation is wrong" - assert live_time == EXPECTED_OUTCOMES_TEST_SEVERAL['run_live_time'], assertion_statement - - print('Check the peak_basics') - df = st.get_df(test_run_id_1T, 'peak_basics') - assertion_statement = ("Got less/more S1s than expected, perhaps " - "the test is outdated or classification " - "has really changed.") - assert np.abs(np.sum(df['type'] == 1) - - EXPECTED_OUTCOMES_TEST_SEVERAL['n_s1']) < 2, assertion_statement - df = df[:10] - - print("Check that we can write nice wiki dfs") - straxen.dataframe_to_wiki(df) - - print("Abuse the peaks to show that _average_scada works") - p = p[:10] - p_t, p_a = straxen.scada._average_scada( - p['time']/1e9, - p['time'], - 1) - assert len(p_a) == len(p), 'Scada deleted some of my 10 peaks!' - - print('Check the number of events') - events = st.get_array(test_run_id_1T, 'event_info_double') - assertion_statement = ("Got less/ore events than expected, " - "perhaps the test is outdated or something " - "changed in the processing.") - assert len(events) == EXPECTED_OUTCOMES_TEST_SEVERAL['n_events'], assertion_statement - - print("Plot bokkeh:") - fig = st.event_display_interactive(test_run_id_1T, - time_range=(events[0]['time'], - events[0]['endtime']), - xenon1t=True, - plot_record_matrix=True, - ) - fig.save('test_display.html') - - # Test data selector: - from straxen.analyses.bokeh_waveform_plot import DataSelectionHist - ds = DataSelectionHist('ds') - fig = ds.histogram2d(p, - p['area'], - p['area'], - bins=50, - hist_range=((0, 200), (0, 2000)), - log_color_scale=True, - clim=(10, None), - undeflow_color='white') - - import bokeh.plotting as bklt - bklt.save(fig, 'test_data_selector.html') - - # On windows, you cannot delete the current process' - # working directory, so we have to chdir out first. - finally: - os.chdir('..') - - -def test_plots(): - """Make some plots""" - c = np.ones(straxen.n_tpc_pmts) - straxen.plot_pmts(c) - straxen.plot_pmts(c, log_scale=True) - - -def test_print_version(): - straxen.print_versions(['strax', 'something_that_does_not_exist']) - - -def test_nt_minianalyses(): - """Number of tests to be run on nT like configs""" - if not straxen.utilix_is_configured(): - return - with tempfile.TemporaryDirectory() as temp_dir: - try: - print("Temporary directory is ", temp_dir) - os.chdir(temp_dir) - from .test_plugins import DummyRawRecords, testing_config_nT, test_run_id_nT - st = straxen.contexts.xenonnt_online(use_rucio=False) - rundb = st.storage[0] - rundb.readonly = True - st.storage = [rundb, strax.DataDirectory(temp_dir)] - - # We want to test the FDC map that only works with CMT - test_conf = testing_config_nT.copy() - del test_conf['fdc_map'] - - st.set_config(test_conf) - st.set_context_config(dict(forbid_creation_of=())) - st.register(DummyRawRecords) - - rr = st.get_array(test_run_id_nT, 'raw_records') - st.make(test_run_id_nT, 'records') - st.make(test_run_id_nT, 'peak_basics') - - st.daq_plot(test_run_id_nT, - time_range=(rr['time'][0], strax.endtime(rr)[-1]), - vmin=0.1, - vmax=1, - ) - - st.plot_records_matrix(test_run_id_nT, - time_range=(rr['time'][0], - strax.endtime(rr)[-1]), - vmin=0.1, - vmax=1, - group_by='ADC ID', - ) - plt_clf() - - st.make(test_run_id_nT, 'event_info') - st.load_corrected_positions(test_run_id_nT, - time_range=(rr['time'][0], - strax.endtime(rr)[-1]), - - ) - # This would be nice to add but with empty events it does not work - # st.event_display(test_run_id_nT, - # time_range=(rr['time'][0], - # strax.endtime(rr)[-1]), - # ) - # On windows, you cannot delete the current process'git p - # working directory, so we have to chdir out first. - finally: - os.chdir('..') diff --git a/tests/test_veto_hitlets.py b/tests/test_veto_hitlets.py index 8ce0c6c25..ec5ba9b03 100644 --- a/tests/test_veto_hitlets.py +++ b/tests/test_veto_hitlets.py @@ -6,7 +6,6 @@ class TestRemoveSwtichedOffChannels(unittest.TestCase): - def setUp(self): self.channel_range = (10, 19) self.to_pe = np.zeros(20) diff --git a/tests/test_veto_veto_regions.py b/tests/test_veto_veto_regions.py index 51f30743c..ca055f9c3 100644 --- a/tests/test_veto_veto_regions.py +++ b/tests/test_veto_veto_regions.py @@ -32,7 +32,7 @@ def test_concatenate_overlapping_intervals(self): right_extension=right_extension) assert len(vetos) == 2, 'Got the wrong number of veto intervals!' - time_is_correct = vetos[0]['time'] == self.events['time'][0]-left_extension + time_is_correct = vetos[0]['time'] == self.events['time'][0] - left_extension assert time_is_correct, 'First veto event has the wrong time!' time_is_correct = vetos[0]['endtime'] == self.events['endtime'][2] + right_extension assert time_is_correct, 'First veto event has the wrong endtime!' @@ -61,12 +61,12 @@ def _test_threshold_type(events, field, threshold_type, threshold): left_extension=0, right_extension=0) print(events[field], thresholds, vetos) - assert len(vetos) == 0, f'Vetos for {threshold_type} threshold should be empty since it is below threshold!' + assert len(vetos) == 0, f'Vetos for {threshold_type} threshold should be empty since it is below threshold!' # noqa events[field] = threshold vetos = straxen.plugins.veto_veto_regions.create_veto_intervals(events, **thresholds, left_extension=0, right_extension=0) - assert len(vetos) == 1, f'{threshold_type} threshold did not work, have a wrong number of vetos!' + assert len(vetos) == 1, f'{threshold_type} threshold did not work, have a wrong number of vetos!' # noqa events[field] = 1