FilterAlgs.py

Go to the documentation of this file.

#!/usr/bin/env python
#
# Data Filtering algorithms
# 
# Useful algorithms for filtering data written to output file
#

__all__ = [ "FilterAlg", "CoincidenceFilterAlg" ]

# Load GaudiPython
from DybPython.DybPythonAlg import DybPythonAlg
from GaudiPython import SUCCESS, FAILURE
from GaudiPython import gbl
import GaudiKernel.SystemOfUnits as units

# Shortcuts to DayaBay classes
TimeStamp = gbl.TimeStamp
Detector = gbl.DayaBay.Detector
ReadoutHeader = gbl.DayaBay.ReadoutHeader
RecHeader = gbl.DayaBay.RecHeader
RegistrationSequence = gbl.DayaBay.RegistrationSequence
Rndm = gbl.Rndm 

# Basic Filtering algorithm
class FilterAlg(DybPythonAlg):
    "Filter the reconstructed data from a general data file"
    def __init__(self,name):
        DybPythonAlg.__init__(self,name)
        # Properties
        self.RegSeqLocation = RegistrationSequence.defaultLocation()
        self.ClearStore = True
        self.StorePath = ""
        self.Prescale = None
        self.DetectorNames = []
        # Internal variables
        self.detectors = []
        self.uniform_random = Rndm.Numbers()
        return

    def initialize(self):
        status = DybPythonAlg.initialize(self)
        self.info("initializing")
        for detName in self.DetectorNames:
            self.detectors.append(Detector(detName))
        # Initialize random number service if applying prescale 
        if self.Prescale != None:
            self.randSvc = self.svc('IRndmGenSvc','RndmGenSvc')
            if self.randSvc == None:
                self.error("Failed to initialize random number service.")
                return FAILURE
            status= self.uniform_random.initialize(self.randSvc, Rndm.Flat(0,1))
            if not status.isSuccess():
                self.error("Failed to initialize uniform random numbers.")
                return status
        return status

    def execute(self):
        self.info("executing")
        evt = self.evtSvc()
        # Get Registration Sequence
        regSequence = evt[self.RegSeqLocation]
        if regSequence == None:
            self.error("Failed to retrieve Registration Sequence from"
                       + self.regseqLocation)
            return FAILURE
        # Get List of TES objects registered this execution cycle
        registrations = regSequence.registrations()
        for registration in registrations:
            hdrObj = registration.object()
            if hdrObj == None:
                self.error("Failed to retrieve object at "+registration.path())
                return FAILURE
            if self.ClearStore:
                # Start by turning off storage of all items
                regSequence.registration(hdrObj).setStore(False)
            if len(self.detectors)>0:
                # Filter by detector if a detector list is defined
                currentDetector = Detector(hdrObj.context().GetSite(),
                                           hdrObj.context().GetDetId())
                if currentDetector not in self.detectors:
                    continue
            if registration.path() == self.StorePath:
                # Store if it is at the selected TES path, check prescale
                if self.Prescale == None:
                    regSequence.registration(hdrObj).setStore(True)
                    self.debug("Storing "+registration.path()+" in output file")
                elif self.Prescale >= self.uniform_random():
                    regSequence.registration(hdrObj).setStore(True)
                    self.debug("Storing "+registration.path()+" in output file")
        return SUCCESS
        
    def finalize(self):
        self.info("finalizing")
        status = DybPythonAlg.finalize(self)
        return status


# Time-correlation Filtering algorithm
class CoincidenceFilterAlg(DybPythonAlg):
    "Filter stored data based on time between readouts"
    def __init__(self,name):
        DybPythonAlg.__init__(self,name)
        # Properties
        self.ReadoutLocation = ReadoutHeader.defaultLocation()
        self.RegSeqLocation = RegistrationSequence.defaultLocation()
        self.StorePaths = [ ReadoutHeader.defaultLocation() ]
        self.CoincidenceWindow = 1.0 * units.millisecond
        self.DetectorNames = []
        self.AdOnly = False
        self.SameDetectorOnly = False
        self.SaveIntermediate = False
        self.SaveNeighbors = False
        self.NeighborWindow = 1000.0 * units.nanosecond
        # Internal variables
        self.detectors = []
        return

    def initialize(self):
        status = DybPythonAlg.initialize(self)
        self.info("initializing")
        for detName in self.DetectorNames:
            self.detectors.append(Detector(detName))
        return status

    def execute(self):
        self.info("executing")
        # Get the readout headers in the archive
        readoutArchive = self.getAES(self.ReadoutLocation)
        if readoutArchive == None: return FAILURE
        if len(readoutArchive)<1: return SUCCESS # No ReadoutHeaders in archive
        # Get current readout (first in archive list)
        currentReadout = readoutArchive[0].readout()
        if currentReadout == None: return SUCCESS # No readout this cycle
        currentDet = currentReadout.detector()
        if self.AdOnly and not currentDet.isAD(): # Select AD coincidence only
            return SUCCESS
        if len(self.detectors)>0: # Use list of selected detectors
            if currentDet not in self.detectors:
                return SUCCESS  # Detector not in selected list 
        # Get previous readouts from archive
        previousReadouts = []
        for readoutHdr in readoutArchive[1:]:
            readout = readoutHdr.readout()
            if readout != None: previousReadouts.append(readout)
        if len(previousReadouts)<1: return SUCCESS # No preceeding readouts
        # Look for coincidences with the current readout
        promptReadouts = self.findPromptReadouts(currentReadout,
                                                 previousReadouts)
        if len(promptReadouts)==0: return SUCCESS # No prompts
        # Make the list of readouts we want to store in the output file
        storageReadoutHeaders = self.makeStorageList(currentReadout,
                                                     promptReadouts,
                                                     previousReadouts)
        # Mark the readouts (and associated data) for storage
        status = self.markForStorage(storageReadoutHeaders)
        return status
        
    def finalize(self):
        self.info("finalizing")
        status = DybPythonAlg.finalize(self)
        return status

    def findPromptReadouts(self, currentReadout, previousReadouts):
        # Use the coincidence time to find the prompt readouts
        promptReadouts = []
        for previousReadout in previousReadouts:
            if ( (self.SameDetectorOnly or self.AdOnly) 
                and currentReadout.detector() != previousReadout.detector()):
                continue  # Limited to the same detector?            
            if len(self.detectors)>0: # Use list of selected detectors
                if currentReadout.detector() not in self.detectors: continue
            # Check time between readouts
            deltaTime = TimeStamp(currentReadout.triggerTime().GetSec()
                                  - previousReadout.triggerTime().GetSec(),
                                  currentReadout.triggerTime().GetNanoSec()
                                  - previousReadout.triggerTime().GetNanoSec())
            self.verbose("dT readout: "+str(deltaTime.GetSeconds()))
            if deltaTime.GetSeconds()*units.second <= self.CoincidenceWindow:
                # Found coincidence
                self.debug("Found coincidence dT [s]: "
                           +str(deltaTime.GetSeconds()) )
                promptReadouts.append(previousReadout)
        return promptReadouts

    def markForStorage(self, storageHeaders):
        # Mark the headers (and associated data) for storage
        # First, recursively add input headers to storage list
        for header in storageHeaders:
            self.addInputHeaders( header, storageHeaders )
        # Second, loop over RegistrationSequences, and find header Registrations
        registrations = []
        status = self.findRegistrations( storageHeaders, registrations )
        if not status.isSuccess(): return status
        for header, registration in zip( storageHeaders, registrations ):
            if registration == None:
                self.error("Failed to find registration for header "
                           +header.name())
                return FAILURE
            if registration.path() in self.StorePaths:
                self.debug("Storing header at "+registration.path())
                registration.setStore(True) # Mark header for storage
        return SUCCESS

    def findRegistrations(self, headers, registrations):
        # Find the ObjectReg registrations for each header in list
        # Initialize registration list with 'None'
        for header in headers:
            registrations.append(None)
        # Get Registration Archive
        regSeqArchive = self.getAES(self.RegSeqLocation)
        if regSeqArchive == None: return FAILURE
        for regSequence in regSeqArchive:
            for registration in regSequence.registrations():
                hdrObj = registration.object()
                if hdrObj == None:
                    self.error("Failed to retrieve object at "
                               +registration.path())
                    return FAILURE
                if hdrObj in headers:
                    # Found the registration for this header
                    hdrIndex = headers.index(hdrObj)
                    registrations[hdrIndex] = regSequence.registration(hdrObj)
        return SUCCESS
    
    def addInputHeaders(self, header, headerList):
        # Recurse up the header.inputHeaders and add to headerList
        for inputHeader in header.inputHeaders():
            if inputHeader not in headerList:
                headerList.append(inputHeader)
                self.addInputHeaders(inputHeader, headerList)
        return

    def makeStorageList(self, delayedReadout, promptReadouts, previousReadouts):
        # Generate the list of readouts for storage
        storageHeaders = [delayedReadout.header()]
        storageHeaders += [readout.header() for readout in promptReadouts]
        if not self.SaveIntermediate and not self.SaveNeighbors:
            return storageHeaders
        # Save all readouts between or near prompt/delayed pair
        delayedTime = delayedReadout.triggerTime()
        earliestTime = TimeStamp(promptReadouts[0].triggerTime())
        # First, find earliest prompt readout
        for promptReadout in promptReadouts:
            if promptReadout.triggerTime() < earliestTime:
                earliestTime = promptReadout.triggerTime()
        # Add time window to catch readouts just before prompt?
        if self.SaveNeighbors:
            earliestTime.Add(TimeStamp(0,
                                       int(-self.NeighborWindow
                                           /units.nanosecond)))
        # Catch intermediate and neigboring readouts
        for previousReadout in previousReadouts:
            prevTime = previousReadout.triggerTime() 
            deltaT_earliest = TimeStamp( prevTime.GetSec()
                                         - earliestTime.GetSec(),
                                         prevTime.GetNanoSec()
                                         - earliestTime.GetNanoSec())
            deltaT_delayed = TimeStamp( prevTime.GetSec()
                                        - delayedTime.GetSec(),
                                        prevTime.GetNanoSec()
                                        - delayedTime.GetNanoSec())
            if (deltaT_earliest.GetSeconds() >= 0
                and deltaT_delayed.GetSeconds() <= 0):
                if previousReadout.header() not in storageHeaders:
                    storageHeaders.append(previousReadout.header())
        return storageHeaders