MetadataUtils.py

Go to the documentation of this file.

#!/bin/env python
from __future__ import print_function
from builtins import str
from builtins import range
from past.builtins import basestring
from builtins import object
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", "NOVA1800a00000"]
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 list(in_md.items()):
        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 = int(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(b"{"):])

        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(b"{"):])

        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(b"\x1b[?1034h", b""))

        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(object):

    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 list(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  list(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