MetadataUtils.py

Go to the documentation of this file.

#!/bin/env python
import os, sys, stat, pwd
import datetime
from time import sleep
import re
import subprocess
import json
from samweb_client.utility import fileEnstoreChecksum

detectors = dict()
detectors["fd"]   = "fardet"
detectors["nd"]   = "neardet"
detectors["tb"]   = "testbeam"
detectors["ndos"] = "ndos"

##allow this to potentially expand beyond just generators we use now
neutrinoGenerators=["genie", "gibuu"]
neutrinoGenieTunes=["default", "G1802a0211a", "G1810b0211a", "G1810i00000", "G1801b0211a", "G1802a00000", "G1801a0211a", "G1810a0211a", "G0000b00000", "G1802b0211a", "G1801a00000", "G1801b00000", "G1810a00000", "G1810j00000", "G1802b00000", "G1810b00000","N1810j0211a"]
cosmicGenerators=["cry"]

run_pattern = re.compile(r"^.*?_r([0-9]+)_s([0-9]+).+")
sim_pattern = re.compile(r"^.*_(20\d\d\d\d\d\d_\d\d\d\d\d\d)")
stream_pattern = re.compile(r"^.*?_r[0-9]+_s[0-9]+_(.*?)[\.,_].*")

parent_pattern = re.compile(r'(,\n.*)?(?P<parent_str>"parents":.*)', re.MULTILINE)

def isNone(value):
    if value ==   None : return True
    if value == "none" : return True
    return False

def appendField(base,extension):

    if   isNone(     base) and isNone(extension): return "none"
    elif isNone(     base): return extension
    elif isNone(extension): return base
    elif base == extension: return base

    return base +"_" + extension


def addMetadataToFCL(fclFile,parName,parValue):
    fclFile.write("\nphysics.analyzers.metadata.params.")
    fclFile.write(parName)
    fclFile.write(": ")
    fclFile.write(parValue)
    fclFile.write(" \n")

def unCamelCase(s):
    """ Convert CamelCase to camel_case """
    # \B matches an empty string which is NOT at the beginning of a word
    # so requiring this means no _ will be inserted at the start
    return re.sub( r'\B([A-Z])', r'_\1', s).lower()

def cleanup_metadata(in_md):
    """ Numerous keys are stored in the files in ways that we need to 'translate'
        so that they can be given to SAM. """

    md = {}
    for k, v in in_md.iteritems():
        k = str(k)
        if len(k) == 0:
            continue
        if isinstance(v, basestring) and len(v) == 0:
            v = "none"

        # thanks, sam_metadata_dumper
        if '.' not in k:
            k = unCamelCase(k)
        else:
            k = k.lower()
            v = str(v)

        # fix for poorly formatted metadata from FileReducer pre-r24852
        if k.lower() == "parents" and '","file_name":' in v:
            v = v.replace('","file_name":', '"},{"file_name":')

        # note that for correct JSON serialization,
        # any lists or maps need to be processed into Python lists/maps.
        # if that fails, just fall back to the string we had.
        if isinstance(v, basestring) and any(c in v for c in "[]{}"):
            try:
                v = eval(v)
            except:
                pass

        # specific bugfix for run list duplicates
        # (fixed upstream in r24859)
        if k == "runs":
            v = list(set([tuple(run) for run in v]))

        # sam_metadata_dumper spits out [ run, subrun, event ].
        # but the field itself should just be the run number...
        if k in ('first_event', 'last_event'):
            try:
                v = int(v)
            except TypeError:
                if len(v) == 3:
                    v = v[2]
                else:
                    continue

        # this is historical, but I can't guarantee it's not needed any more,
        # even though every file I can see now has this as an int...
        if k == 'simulated.cycle':
            v = eval(v)

        # more type coersion
        if k in ('start_date', 'end_date'):
            v = long(v)

        # ummm... yeah.
        if k == 'process_name' and md.get('application',[]).get('name') is None:
            k = 'application_name'
        elif k == 'stream_name':
            k = 'data_stream'

        if k.startswith('application_'):
            if 'application' not in md:
                md['application'] = {}
            md['application'][k[12:]] = v
        elif k in ('file_format_era', 'file_format_version'):
            pass
        else:
            md[k] = v

    return md

def createMetadata(inFile):
    filesize = os.path.getsize(inFile)
    filename = os.path.basename(inFile)
    if filename.endswith("caf.root"):
        md = {'file_name': filename, 'group': 'nova', 'file_size': filesize, 'file_format': 'unknown'}
        try:
            # data = subprocess.check_output(["extractCAFMetadata", str(inFile)])
            data = subprocess.check_output(["extractCAFMetadata", str(inFile)])
        except subprocess.CalledProcessError:
            print "extractCAFMetadata failed!"
            return None

        dumperDict = json.loads(data[data.find("{"):])

        md.update(cleanup_metadata(dumperDict))
        md['crc'] = fileEnstoreChecksum(inFile)

        return md
    elif filename.endswith(".h5"):
        md = {'file_name': filename, 'group': 'nova', 'file_size': filesize, 'file_format': 'unknown'}
        try:
            data = subprocess.check_output(["extractHDF5Metadata", str(inFile)])
        except subprocess.CalledProcessError:
            print "extractHDF5Metadata failed!"
            return None

        dumperDict = json.loads(data[data.find("{"):])

        md.update(cleanup_metadata(dumperDict))
        md['crc'] = fileEnstoreChecksum(inFile)

        return md

    else:
        md = {'file_name': filename, 'group': 'nova', 'file_size': filesize, 'file_format': 'artroot'}
        try:
            data = subprocess.check_output(["sam_metadata_dumper", str(inFile)])
        except subprocess.CalledProcessError:
            print "sam_metadata_dumper failed!"
            return None

        # hilariously, sam_metadata_dumper prints out the 'parents' metadata twice:
        # once, with format like '[ "<f1>", "<f2>", ... ]' (incorrect?)
        # and again, with format like '[ {"filename": "<f1>"}, {"filename": "<f2>"}, ...]' (correct?)
        # eliminate the first one.
        matched = False
        for matches in parent_pattern.findall(data):
            match = matches[-1]
            parent_dict = eval("{%s}" % match)
            if isinstance(parent_dict["parents"][0], basestring):
                matched = True
                break
        if matched:
            data = data.replace("".join(matches), "")

        dumperDict = json.loads(data.replace("\x1b[?1034h", ""))

        for dumperKey in dumperDict:
            if str(inFile) in str(dumperKey):
                tmpMD = dumperDict[dumperKey]
                md.update(cleanup_metadata(tmpMD))
                break
        md['crc'] = fileEnstoreChecksum(inFile)

        return md

class metaDataMgr:

    def __init__(self,inFileName, metadata, release, systematic, skim, cosmicsPolarity, subversion, specialName):

        #make all lower case keys
        self.metadata = {k.lower():v for k,v in metadata.items()}
        self.inFileName=inFileName
        self.release = release
        self.setRunInfo()
        self.setDatasetTag()
        if not self.isSam4Users():
            self.setDetector()
            self.setDataFlag()
            if self.dataFlag == "data":
                self.setDataInfo()
            else:
                self.setSimInfo()
            self.setSkim(skim)
            self.cosmicsPolarity = cosmicsPolarity
            self.setSpecialInfo(subversion,systematic, specialName)

    def getMetadataParameter(self,parName):
        retVal=None
        if parName in self.metadata:
            retVal = self.metadata[parName]
        return retVal

    def setDatasetTag(self):
        mytag=self.getMetadataParameter("dataset.tag")
        if mytag!= None:
            self.mytag = mytag.lower()
        else:
            self.mytag = "none"

    def isSam4Users(self):
        return not isNone(self.mytag)

    def  setDetector(self):
        detector=self.getMetadataParameter("online.detector")
        if detector != None:
            self.detector = detector.lower()
        else:
            detID = self.getMetadataParameter("nova.detectorid")
            if detID == None:
                print "Error!!  No detector name found!!"
                exit(1)
            detID=detID.lower()
            if detID in  detectors.keys():
                self.detector=detectors[detID]
            else:
                print "No detector name found for detectorid:", detId
                exit(1)


    def setDataFlag(self):

        self.fileType=self.getMetadataParameter("file_type")

        if self.fileType == None:
            print "No file_type set!"
            exit(1)

        elif self.fileType == "importedDetector":
            self.dataFlag = "data"

        elif self.fileType == "importedSimulated":
            self.dataFlag = "sim"

        elif self.fileType == "text":
            self.dataFlag = "text"

        else:
            print "Unrecognized file_type:", file_type
            exit(1)


    def setRunInfo(self):

        #if new format "runs" metadata, use that, otherwise, get from file name
        if "runs" in self.metadata and len(self.metadata["runs"][0]) == 3:
            self.runNum = int(self.metadata["runs"][0][0])
            self.subNum = int(self.metadata["runs"][0][1])
        else:
            if "dataset.tag" in self.metadata != "none":
                nameofThetag = self.getMetadataParameter("dataset.tag")

                TOTrunNum = subprocess.Popen( [ "samweb", "count-files", "dataset.tag " + str(nameofThetag) ], shell=False, stdout=subprocess.PIPE ).stdout.read()
                Tot = int(TOTrunNum)
                for i in range(0,Tot):
                    self.runNum = int(TOTrunNum)
                    self.subNum = int(i)

            else:
                inFileBase=os.path.basename(self.inFileName)
                run_match = run_pattern.match(inFileBase)
                if run_match != None:
                    self.runNum = int(run_match.groups()[0])
                    self.subNum = int(run_match.groups()[1])

                else:
                    print "No run number/subrun number found!"
                    exit(1)

    def setDataInfo(self):

        #get data_stream from metadata.  If it is "all", try extracting from the file name.
        tmpStream = 'all'
        if "data_stream" in self.metadata:
            tmpStream = self.metadata["data_stream"]
        elif "online.stream" in self.metadata:
            tmpStream = self.metadata["online.stream"]

        if self.getMetadataParameter("nova.hornpolarity"):
            self.horn=self.getMetadataParameter("nova.hornpolarity")
        else:
            self.horn=""

        if tmpStream == 'all':
            stream_match = stream_pattern.match(inFileBase)
            if stream_match != None:
                tmpStream = stream_match.groups()[0]
            if tmpStream[0] == 't':
                tmpStream = tmpStream[1:]

        try:
            self.stream = int(tmpStream)
            self.streamEntry = 't{0:02d}'.format(self.stream)
        except ValueError:
            self.stream = tmpStream
            self.streamEntry = self.stream
        except:
            print "No data stream found!"
            exit(1)

    def setSimInfo(self):

        self.generator = self.getMetadataParameter("simulated.generator")

        # Some early gibuu files accidentally have generator set to singlep
        # (fixed in r32723). If that's the case, fix it up here. In particular
        # because it impacts the correct setting of flavorset.
        if self.generator == 'singlep' and self.getMetadataParameter('simulated.gibuuused') == 'true':
            self.generator = 'gibuu'

        if self.getMetadataParameter("nova.hornpolarity"):
            self.horn=self.getMetadataParameter("nova.hornpolarity")
        else:
            self.horn=""

        if self.generator in neutrinoGenerators:

            self.flavorset=self.getMetadataParameter("simulated.genieflavorset")
            if self.flavorset=="swap": self.flavorset="fluxswap"

            self.genierw=self.getMetadataParameter("simulated.genierw")
            if self.genierw==None: self.genierw = "nogenierw"

            self.fluxVersion=self.getMetadataParameter("nova.flux_version")
            if self.fluxVersion == None:
                self.fluxVersion=os.environ["NOVA_FLUX_VERSION"].replace("nova_","")

            if self.generator == 'genie' and self.getMetadataParameter('simulated.genieused') == 'true':
                self.NPPFX=self.getMetadataParameter("simulated.NPPFX")
                self.genietune=self.getMetadataParameter("simulated.genietune")
                if self.genietune not in neutrinoGenieTunes:
                    # If no genietune metadata field is found, the default tune
                    # is assumed
                    if self.genietune==None:
                        self.genietune="default"
                    else:
                        print "ERROR!!  Invalid GENIE tune: ", self.genietune
                        exit(1)

        elif self.generator in cosmicGenerators:
            self.flavorset=self.getMetadataParameter("simulated.cryflavorset")

        self.nevt = self.getMetadataParameter("simulated.number_of_spills")
        self.cycle = self.getMetadataParameter("simulated.cycle")

        matches=sim_pattern.match(self.inFileName)
        matchGroups=matches.groups()
        if len(matchGroups)<1:
            print "ERROR!!  Couldn't find timestamp in" , self.inFileName
            exit(1)
        self.timestamp = matchGroups[0]

    def setSpecialInfo(self,subversion,systematic,specialName):

        self.systematic = "none"
        mdSystematic=self.getMetadataParameter("nova.systematic")
        if  (not isNone(mdSystematic)) and (not isNone(systematic)):
            print "ERROR:  you specified a systematic, but this is already a systematic file"
            exit(1)
        elif not isNone(mdSystematic):
            self.systematic = mdSystematic
        elif not isNone(systematic):
            self.systematic = systematic

        if subversion != None:
            self.subversion = subversion
        else:
            self.subversion = self.getMetadataParameter("nova.subversion")
            if self.subversion == None:
                self.subversion="1"

        mdSpecialName = self.getMetadataParameter("nova.special")
        self.special = appendField(mdSpecialName,specialName)

    def setSkim(self,skim):
        mdSkim = self.getMetadataParameter("nova.skim")
        self.skim = appendField(mdSkim,skim)

    def  getOutputFileName(self,tier):
        outFileName=""
        if self.isSam4Users():
            outFileName = os.path.basename(self.inFileName.strip(tier+'.root'))
            outFileName += "_r%.8d_s%.2d" % (self.runNum, self.subNum)
        else:
            outFileName  = self.detector

            if self.dataFlag == "data":
                # data
                # <det>_<run>_<subrun>_<horn>_<stream>_<tag>_<npass>_data_<syst>_<skim>_<spec>.<tier>.root

                print "SETTING FILE NAME DATA"
                outFileName += "_r%.8d_s%.2d" % (self.runNum, self.subNum)

                if not isNone(self.cosmicsPolarity):
                    outFileName += "_%s" % self.cosmicsPolarity

                if self.horn != "":
                    outFileName += "_%s" % self.horn


                outFileName += "_%s" % self.streamEntry
                outFileName += "_%s" % self.release
                outFileName += "_v%s" % self.subversion
                outFileName += "_data"

            elif self.dataFlag == "sim":
                # sim
                #<det>_<gen>_<tune>_<flavor>_<genierw>_<horn>_<flux>_<nevt>_<run>_<subrun>_<cycle>_<tag>_<npass>_<timestamp>_sim_<syst>_<skim>_<spec>.<tier>.root
                outFileName += "_%s" % self.generator
                if self.generator in neutrinoGenerators:
                    outFileName += "_%s" % self.genietune
                outFileName += "_%s" % self.flavorset
                if self.generator in neutrinoGenerators: #these parameters do not apply for cosmic generators
                    outFileName += "_%s" % self.genierw
                if not isNone(self.cosmicsPolarity):
                    outFileName += "_%s" % self.cosmicsPolarity
                if self.horn != "":
                    outFileName += "_%s" % self.horn
                if self.generator in neutrinoGenerators: #these parameters do not apply for cosmic generators
                    outFileName += "_%s_%s" % ( self.fluxVersion, self.nevt )

                if self.cycle == None: # cosmic MC files with cycle missing
                    outFileName += "_r%.8d_s%.2d" % (self.runNum, self.subNum)
                else:
                    outFileName += "_r%.8d_s%.2d_c%.3d" % (self.runNum, self.subNum, self.cycle)

                outFileName += "_%s" % self.release
                outFileName += "_v%s" % self.subversion
                outFileName += "_%s" % self.timestamp
                outFileName += "_sim"

            if not isNone(self.systematic):
                outFileName += "_%s" % self.systematic
            if not isNone(self.skim):
                outFileName += "_%s" % self.skim
            if not isNone(self.special):
                outFileName += "_%s" % self.special

        if tier == "artdaq" :
            if self.dataFlag == "sim":
                outFileName += ".daq.root"
            elif self.dataFlag == "data" :
                outFileName += ".%s.root" % tier
        elif tier == "h5caf":
            outFileName += ".h5caf"
        else :
            outFileName += ".%s.root" % tier

        return outFileName