ecl_packer_unpacker.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
 
 
 
# Inspired by cdc/tests/cdc_packer_unpacker.py
 
# This test checks if ECLDigits that are packed and then unpacked are identical to the originals.
# ECLDigits are obtained in 2 ways:
# 1) From simulated muons
# 2) "Custom" ECLDigits that have different combinations of parameters spanning the allowed range.
 
import basf2 as b2
from ROOT import Belle2
from unittest import TestCase
from ROOT import gRandom
import simulation
import itertools
 
logLevel = b2.LogLevel.INFO  # LogLevel.DEBUG #
b2.set_random_seed(42)
 
eclDigitsDatastoreName = 'ECLDigits'
unpackerOutputDatastoreName = 'someECLUnpackerDatastoreName'
 
 
class addECLDigitsModule(b2.Module):
    """
    Adds ECLDigits with very large/small time/amplitude values
    """
 
    def __init__(self):
        """constructor"""
        super().__init__()
        amps = [0, 1, 100000, 262015]
        times = [-2048, -100, 0, 100, 2047]
        qualitys = [0, 1, 2, 3]
        chis = [0, 1, 254, 511]
        paramNames = ["amp", "time", "quality", "chi"]
        
        self.digitParamsdigitParams = [dict(zip(paramNames, params)) for params in itertools.product(amps, times, qualitys, chis)]
 
    def event(self):
        """
        event function
        """
        eclDigitsFromSimulation = Belle2.PyStoreArray(eclDigitsDatastoreName)
 
        # Check for used cellIds
        usedCellIds = list(map(lambda eclDigit: eclDigit.getCellId(), eclDigitsFromSimulation))
 
        # Create new ECLDigits and add them to datastore
        for digitParam in self.digitParamsdigitParams:
            # Skip combination of quality != 2 and chi != 0. Electronics can't output this
            if (digitParam['quality'] != 2) and (digitParam['chi'] != 0):
                continue
 
            # Skip combination of quality == 2 and amp != 0. Electronics can't output this
            if (digitParam['quality'] == 2) and (digitParam['time'] != 0):
                continue
 
            # Choose cellId that's not already used
            cellId = int(gRandom.Uniform(1, 8736))
            while cellId in usedCellIds:
                cellId = int(gRandom.Uniform(1, 8736))
 
            usedCellIds.append(cellId)
            # Create new ECLDigit
            eclDigit = Belle2.ECLDigit()
 
            # Fill ECLDigit
            eclDigit.setCellId(cellId)
            eclDigit.setAmp(digitParam['amp'])
            eclDigit.setTimeFit(digitParam['time'])
            eclDigit.setQuality(digitParam['quality'])
            eclDigit.setChi(digitParam['chi'])
 
            # Add ECLDigit to datastore
            newDigit = eclDigitsFromSimulation.appendNew()
            newDigit.__assign__(eclDigit)
 
 
class ECLPackerUnpackerTestModule(b2.Module):
    """
    module which checks if two collections of ECLDigits are equal
    """
 
    def sortECLDigits(self, unsortedPyStoreArray):
        """ use a some information to sort the ECLDigits list
            Returns a python-list containing the ECLDigits
        """
        # first convert to a python-list to be able to sort
        py_list = [x for x in unsortedPyStoreArray]
 
        # sort via a hierachy of sort keys
        return sorted(
            py_list,
            key=lambda x: (
                x.getCellId()))
 
    def event(self):
        """
        event function
        """
 
        # load the ECLDigits that were packed and unpacked
        eclDigitsPackedUnpacked_unsorted = Belle2.PyStoreArray(unpackerOutputDatastoreName)
        # ECLDigits direct from simulation
        eclDigitsFromSimulation_unsorted = Belle2.PyStoreArray(eclDigitsDatastoreName)
 
        if not len(eclDigitsFromSimulation_unsorted) == len(eclDigitsPackedUnpacked_unsorted):
            b2.B2FATAL("Different number of ECLDigits from simulation and after packing+unpacking")
 
        eclDigitsPackedUnpacked = self.sortECLDigitssortECLDigits(eclDigitsPackedUnpacked_unsorted)
        eclDigitsFromSimulation = self.sortECLDigitssortECLDigits(eclDigitsFromSimulation_unsorted)
 
        tc = TestCase('__init__')
 
        # check all quantities between the direct and the packed/unpacked ECLDigits
        for idx in range(len(eclDigitsPackedUnpacked)):
            digit = eclDigitsFromSimulation[idx]
            digitPackedUnpacked = eclDigitsPackedUnpacked[idx]
 
            b2.B2DEBUG(5, 'MC digit: cellid = ' +
                       str(digit.getCellId()) +
                       ', amp = ' +
                       str(digit.getAmp()) +
                       ', time = ' +
                       str(digit.getTimeFit()) +
                       ', quality = ' +
                       str(digit.getQuality()) +
                       ', chi = ' +
                       str(digit.getChi()) +
                       '\nUnpackedDigit: cellid = ' +
                       str(digitPackedUnpacked.getCellId()) +
                       ', amp = ' +
                       str(digitPackedUnpacked.getAmp()) +
                       ', time = ' +
                       str(digitPackedUnpacked.getTimeFit()) +
                       ', quality = ' +
                       str(digitPackedUnpacked.getQuality()) +
                       ', chi = ' +
                       str(digitPackedUnpacked.getChi()))
 
            tc.assertEqual(digit.getCellId(), digitPackedUnpacked.getCellId())
            tc.assertEqual(digit.getAmp(), digitPackedUnpacked.getAmp())
            tc.assertEqual(digit.getTimeFit(), digitPackedUnpacked.getTimeFit())
            tc.assertEqual(digit.getQuality(), digitPackedUnpacked.getQuality())
            tc.assertEqual(digit.getChi(), digitPackedUnpacked.getChi())
 
 
main = b2.create_path()
# Create Event information
eventinfosetter = b2.register_module('EventInfoSetter')
eventinfosetter.param({'evtNumList': [10]})
main.add_module(eventinfosetter)
 
# to run the framework the used modules need to be registered
particlegun = b2.register_module('ParticleGun')
particlegun.param('pdgCodes', [13, -13])
particlegun.param('nTracks', 10)
main.add_module(particlegun)
 
# add simulation for ECL only
simulation.add_simulation(main, components=['ECL'])
b2.set_module_parameters(main, type="Geometry", useDB=False, components=["ECL"])
 
# Add ECLDigits with wide range of amp, time, quality, chi2
main.add_module(addECLDigitsModule())
 
# add the packer which packs the ECLDigits from the simulation
ecl_packer = b2.register_module('ECLPacker')
main.add_module(ecl_packer)
 
# add the unpacker which unpacks the RawECL into ECLDigits and store them in a seperate datstore container
ecl_unpacker = b2.register_module('ECLUnpacker')
# set the unpacker output datastore name, so that we don't overwrite the "normal" ECLDigits
ecl_unpacker.param('ECLDigitsName', unpackerOutputDatastoreName)
main.add_module(ecl_unpacker)
 
# Run test module to check if the original ECLDigits and the ones that were packed and unpacked are identical
eclPackUnpackerChecker = ECLPackerUnpackerTestModule()
main.add_module(eclPackUnpackerChecker, logLevel=logLevel)
 
 
# Process events
b2.process(main)