<commit_before># coding=utf-8
# Filename: k40.py
# pylint: disable=locally-disabled
"""
A collection of k40 related functions and modules.

"""
from __future__ import division, absolute_import, print_function

import os
from itertools import combinations
from collections import defaultdict
from functools import partial
from datetime import datetime
import io

from scipy import optimize
import numpy as np
import h5py
import pickle
import numba as nb

import km3pipe as kp
from km3pipe.io.daq import TMCHData

__author__ = "<NAME>"
__email__ = "<EMAIL>"
__status__ = "Development"

log = kp.logger.logging.getLogger(__name__)  # pylint: disable=C0103
# log.setLevel(logging.DEBUG)

TIMESLICE_LENGTH = 0.1  # [s]


class K40BackgroundSubtractor(kp.Module):
    """Subtracts random coincidence background from K40 data

    Required Services
    -----------------
    'MedianPMTRates()': dict (key=dom_id, value=list of pmt rates)

    Output Keys
    -----------
    'K40Counts': dict, Corrected K40 counts

    """
    def configure(self):
        self.combs = list(combinations(range(31), 2))

    def process(self, blob):
        print('Subtracting random background calculated from single rates')
        counts = self.services['TwofoldCounts']
        dom_ids = list(counts.keys())
        mean_rates = self.services['GetMedianPMTRates']()
        corrected_counts = {}
        livetime = self.services['GetLivetime']()
        for dom_id in dom_ids:
            try:
                pmt_rates = mean_rates[dom_id]
            except KeyError:
                log.warning("Skipping BG correction for DOM {}."
                            .format(dom_id))
                corrected_counts[dom_id] = counts[dom_id]
                continue

            k40_rates = counts[dom_id] / livetime
            bg_rates = []
            for c in self.combs:
                bg_rates.append(pmt_rates[c[0]]*pmt_rates[c[1]]*1e-9)
            corrected_counts[dom_id] = (k40_rates.T -
                                        np.array(bg_rates)).T * livetime
        blob["CorrectedTwofoldCounts"] = corrected_counts
        pickle.dump({'data': corrected_counts,
                     'livetime': livetime},
                    open("k40_counts_bg_sub.p", "wb"))
        pickle.dump(mean_rates, open('mean_rates.p', 'wb'))
        return blob


class IntraDOMCalibrator(kp.Module):
    """Intra DOM calibrator which performs the calibration from K40Counts.

    Parameters
    ----------
    det_id: int
      Detector ID [default: 14]
    ctmin: float
      Minimum cos(angle)
    mode: str ('offline' | 'online')
      Calibration mode [default: 'online']

    Input Keys
    ----------
    'TwofoldCounts': dict (key=dom_id,
                           value=matrix of k40 counts 465x(dt*2+1))
    'CorrectedTwofoldCounts': dict (key=dom_id,
                                    value=matrix of k40 counts 465x(dt*2+1))

    Output Keys
    -----------
    'IntraDOMCalibration': dict (key=dom_id, value=calibration)

    """
    def configure(self):
        det_id = self.get("det_id") or 14
        self.detector = kp.hardware.Detector(det_id=det_id)
        self.ctmin = self.require("ctmin")
        self.mode = self.get("mode", default="online")
        self.calib_filename = self.get("calib_filename", default="k40_cal.npy")

    def process(self, blob):
        if self.mode != 'online':
            return blob

        if 'CorrectedTwofoldCounts' in blob:
            log.info("Using corrected twofold counts")
            fit_background = False
            twofold_counts = blob['CorrectedTwofoldCounts']
        else:
            log.info("No corrected twofold counts found, fitting background.")
            twofold_counts = self.services['TwofoldCounts']
            fit_background = True

        blob['IntraDOMCalibration'] = self.calibrate(twofold_counts,
                                                     fit_background)
        return blob

    def calibrate(self, twofold_counts, fit_background=False):
        print("Starting calibration:")
        calibration = {}

        for dom_id, data in twofold_counts.items():
            print(" calibrating DOM '{0}'".format(dom_id))
            try:
                calib = calibrate_dom(dom_id, data,
                                      self.detector,
                                      livetime=self.services['GetLivetime'](),
                                      fit_background=fit_background,
                                      ad_fit_shape='exp', ctmin=self.ctmin)
            except RuntimeError:
                log.error(" skipping DOM '{0}'.".format(dom_id))
            else:
                calibration[dom_id] = calib

        return calibration

    def finish(self):
        if self.mode == 'offline':
            print("Starting offline calibration")
            twofold_counts = self.services['TwofoldCounts']
            calibration = self.calibrate(twofold_counts, fit_background=True)
            print("Dumping calibration to '{}'.".format(self.calib_filename))
            np.save(self.calib_filename, calibration)

<commit_msg>Uses new summaryslice background subtractor for offline calibration<commit_after># coding=utf-8
# Filename: k40.py
# pylint: disable=locally-disabled
"""
A collection of k40 related functions and modules.

"""
from __future__ import division, absolute_import, print_function

import os
from itertools import combinations
from collections import defaultdict
from functools import partial
from datetime import datetime
import io

from scipy import optimize
import numpy as np
import h5py
import pickle
import numba as nb

import km3pipe as kp
from km3pipe.io.daq import TMCHData

__author__ = "<NAME>"
__email__ = "<EMAIL>"
__status__ = "Development"

log = kp.logger.logging.getLogger(__name__)  # pylint: disable=C0103
# log.setLevel(logging.DEBUG)

TIMESLICE_LENGTH = 0.1  # [s]


class K40BackgroundSubtractor(kp.Module):
    """Subtracts random coincidence background from K40 data

    Required Services
    -----------------
    'MedianPMTRates()': dict (key=dom_id, value=list of pmt rates)

    Output Keys
    -----------
    'K40Counts': dict, Corrected K40 counts

    """
    def configure(self):
        self.combs = list(combinations(range(31), 2))
        self.mode = self.get("mode", default='online')
        self.expose(self.get_corrected_counts, 'GetCorrectedTwofoldCounts')
        self.corrected_counts = None

    def process(self, blob):
        if self.mode != 'online':
            return blob
        print('Subtracting random background calculated from single rates')
        corrected_counts = self.subtract_background()
        blob['CorrectedTwofoldCounts'] = corrected_counts

        return blob

    def get_corrected_counts(self):
        return self.corrected_counts

    def subtract_background(self):
        counts = self.services['TwofoldCounts']
        dom_ids = list(counts.keys())
        mean_rates = self.services['GetMedianPMTRates']()
        corrected_counts = {}
        livetime = self.services['GetLivetime']()
        for dom_id in dom_ids:
            try:
                pmt_rates = mean_rates[dom_id]
            except KeyError:
                log.warning("Skipping BG correction for DOM {}."
                            .format(dom_id))
                corrected_counts[dom_id] = counts[dom_id]
                continue

            k40_rates = counts[dom_id] / livetime
            bg_rates = []
            for c in self.combs:
                bg_rates.append(pmt_rates[c[0]]*pmt_rates[c[1]]*1e-9)
            corrected_counts[dom_id] = (k40_rates.T -
                                        np.array(bg_rates)).T * livetime
        return corrected_counts

    def finish(self):
        if self.mode == 'offline':
            print('Subtracting background calculated from summaryslices.')
            self.corrected_counts = self.subtract_background()

    def dump(self, mean_rates, corrected_counts, livetime):
        pickle.dump(mean_rates, open('mean_rates.p', 'wb'))
        pickle.dump({'data': corrected_counts,
                     'livetime': livetime},
                    open("k40_counts_bg_sub.p", "wb"))


class IntraDOMCalibrator(kp.Module):
    """Intra DOM calibrator which performs the calibration from K40Counts.

    Parameters
    ----------
    det_id: int
      Detector ID [default: 14]
    ctmin: float
      Minimum cos(angle)
    mode: str ('offline' | 'online')
      Calibration mode [default: 'online']

    Input Keys
    ----------
    'TwofoldCounts': dict (key=dom_id,
                           value=matrix of k40 counts 465x(dt*2+1))
    'CorrectedTwofoldCounts': dict (key=dom_id,
                                    value=matrix of k40 counts 465x(dt*2+1))

    Output Keys
    -----------
    'IntraDOMCalibration': dict (key=dom_id, value=calibration)

    """
    def configure(self):
        det_id = self.get("det_id") or 14
        self.detector = kp.hardware.Detector(det_id=det_id)
        self.ctmin = self.require("ctmin")
        self.mode = self.get("mode", default="online")
        self.calib_filename = self.get("calib_filename", default="k40_cal.npy")

    def process(self, blob):
        if self.mode != 'online':
            return blob

        if 'CorrectedTwofoldCounts' in blob:
            log.info("Using corrected twofold counts")
            fit_background = False
            twofold_counts = blob['CorrectedTwofoldCounts']
        else:
            log.info("No corrected twofold counts found, fitting background.")
            twofold_counts = self.services['TwofoldCounts']
            fit_background = True

        blob['IntraDOMCalibration'] = self.calibrate(twofold_counts,
                                                     fit_background)
        return blob

    def calibrate(self, twofold_counts, fit_background=False):
        print("Starting calibration:")
        calibration = {}

        for dom_id, data in twofold_counts.items():
            print(" calibrating DOM '{0}'".format(dom_id))
            try:
                calib = calibrate_dom(dom_id, data,
                                      self.detector,
                                      livetime=self.services['GetLivetime'](),
                                      fit_background=fit_background,
                                      ad_fit_shape='exp', ctmin=self.ctmin)
            except RuntimeError:
                log.error(" skipping DOM '{0}'.".format(dom_id))
            else:
                calibration[dom_id] = calib

        return calibration

    def finish(self):
        if self.mode == 'offline':
            print("Starting offline calibration")
            if 'GetCorrectedTwofoldCounts' in self.services:
                print("Using corrected twofold counts")
                twofold_counts = self.services['GetCorrectedTwofoldCounts']()
                fit_background = False
            else:
                print("Using uncorrected twofold counts")
                twofold_counts = self.services['TwofoldCounts']
                fit_background = True
            calibration = self.calibrate(twofold_counts,
                                         fit_background=fit_background)
            print("Dumping calibration to '{}'.".format(self.calib_filename))
            np.save(self.calib_filename, calibration)

{'fraction_non_alphanumeric': 0.10561027084066127, 'fraction_numerical': 0.01037636299683433, 'mean_word_length': 2.090489130434783, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 20, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '870025', 'n_tokens_mistral': 3410, 'n_tokens_neox': 3165, 'n_words': 731}