MCenterSpills.py

Go to the documentation of this file.

#!/usr/bin/env python

#..................................................
# INCLUDES
#..................................................
import sys, os
import urllib2, csv
import getopt
from datetime import date, datetime, timedelta
import psycopg2
import psycopg2.extras

# Change to output directory if you want to specify where the root file goes
BASE_DIR = os.getcwd()
filename = 'MCenterSpills.root'

DETECTOR = 'testbeam'
DB_CURSOR     = None
DB_QUERY_TIME = timedelta(0)

#..................................................
def usage():
    """ Print out usage.
    """
    print "Usage:"
    print "python MCenterSpills.py -s YYYY-mm-dd -e YYYY-mm-dd -n #"
    print ""
    print "Where:"
    print "-s : start date"
    print "-e : end date"
    print "-n : number of days"
    print ""
    print "You can specify any 2 of the above 3 options. The third will be calculated for you."
    print "There is no default, you must supply options"
    print ""
    print "Creates root file with date, first and last run, number of recorded spills, and recorded POT [counts] for each day."

#..................................................
def options():
    """ Parse the command line options
    Returns:
    t1 (datetime) : First day to report
    t2 (datetime) : Last day to report
    n  (int)      : Number of days to report
    """
    date_start = None
    date_end   = None
    ndays      = None
    
    global DETECTOR

    try:
        (opts, args) = getopt.getopt(sys.argv[1:],\
                                         "hs:e:n:",\
                                         ["help",\
                                          "start=",\
                                          "end=",\
                                          "ndays="])
    except getopt.GetoptError as err:
        print str(err)
        sys.exit(2)
    for opt, arg in opts:
        if opt in ("-h", "--help"):
            print "help opt"
            usage()
            exit(0)
        elif opt in ("-s", "--start"):
            date_start = arg
        elif opt in ("-e", "--end"):
            date_end = arg
        elif opt in ("-n", "--ndays"):
            ndays = int(arg)
        else:
            assert False, "unknown option"


    if (date_start!=None):
        t0 = datetime.strptime(date_start,'%Y-%m-%d')
    if (date_end!=None):
        t1 = datetime.strptime(date_end,'%Y-%m-%d')

    # missing start date - use end date and ndays to construct
    if ((date_start==None)&(date_end!=None)&(ndays!=None)):
        t0 = t1
        t0 -= timedelta(days=(ndays-1))

    # missing end date - use start date and ndays to construct
    if ((date_start!=None)&(date_end==None)&(ndays!=None)):
        t1 = t0
        t1 += timedelta(days=ndays)

    # missing ndays - use start and end dates to construct
    if ((date_start!=None)&(date_end!=None)&(ndays==None)):
        n = (t1-t0)+timedelta(days=1)
        ndays = int(n.days)

    if ((date_start!=None)&(date_end!=None)&(ndays!=None)):
        #
        # Have everything we need, but are they consistent?
        #
        d = (t1-t0)+timedelta(days=1)
        n = int(d.days)
        if (ndays != n):
            print 'Inconsistent time range'
            exit(2)

    return (t0,t1,ndays)

#..................................................

def runs_by_day(day0):
    """ Get a complete list of run numbers for a specified day

    Any run that starts or stops on the specified date is contaned in
    the list

    Args:
    day0: A datetime object for the requested day

    Returns:
    An list of SQL results for run,tstart,tstop ordered by start time
    """
    day1 = day0+timedelta(days=1)
    return runs_for_period(day0, day1)

#..................................................

def runs_for_period(day0, day1):
    """ Get all runs that start or stop within a particular date range

    Args:
    day0 : A datetime object for first day
    day1 : A datetime object for the last day
    """
    sql0 = "select run,tstart,tstop from "+DETECTOR+".runs";
    cut1 = " where partition=1";
    cut2 = " tstop is not null"
    cut3 = "tstart >= '{}'".format(day0.strftime("%Y-%m-%dT%H:%M:%S"))
    cut4 = "tstart < '{}'".format(day1.strftime("%Y-%m-%dT%H:%M:%S"))
    
    cut5 = "tstop >= '{}'".format(day0.strftime("%Y-%m-%dT%H:%M:%S"))
    cut6 = "tstop < '{}'".format(day1.strftime("%Y-%m-%dT%H:%M:%S"))

    sql1 = "order by tstart asc"
    
    SQL = \
        sql0 + \
        cut1 + " and " + \
        cut2 + " and " + \
        "((" + cut3+" and "+cut4+") or ("+cut5+" and "+cut6+")) " + \
        sql1

    
    DB_CURSOR.execute(SQL);
    runs = DB_CURSOR.fetchall();

    return runs

#..................................................

def run_minmax(runs):
    """ Find the highest and lowest run number in a list

    Args:
    runs: A list of SQL results (run,tstart,tstop)

    Returns:
    (runlo,runhi), the lowest and highest run numbers in the list
    """
    runmin = runs[0]['run']
    runmax = runs[0]['run']
    for r in runs:
        if (r['run']<runmin):
            runmin = r['run']
        if (r['run']>runmax):
            runmax = r['run']
    return (runmin,runmax)

#..................................................

def mc7_spills_between(t1, t2):
    """
    Return spill information (unix time and POT) for a datetime range

    Args:
    t1 (datetime): Start time
    t2 (datetime): End time

    Returns:
    List of (time,pot) tuples
    """
    global DB_QUERY_TIME

    spills = [];    

    tf1 = in_unix(t1)
    tf2 = in_unix(t2)
    webaddress = 'http://ifb-data.fnal.gov:8099'\
        '/ifbeam/data/data?'\
        'v=F:MC7SC1&e=e,36&t0={}&t1={}&&f=csv'.\
        format(tf1,tf2)

    dd = None
    #
    # Make three tries and then bail...
    #
    t0 =  datetime.utcnow()
    try:
        dd = urllib2.urlopen(webaddress,None,600)
    except:
        try:
            dd = urllib2.urlopen(webaddress,None,600)
        except:
            try:
                dd = urllib2.urlopen(webaddress,None,600)
            except:
                return spills
    t1 =  datetime.utcnow()

    DB_QUERY_TIME += (t1-t0)
    
    if (dd==None):
        return spills
    
    csvdata = csv.reader(dd)

    for row in csvdata:
        if (row.count('F:MC7SC1')>0):
            if (row[4]!='null'):
                spills.append( (float(row[2])/1000., float(row[4])) )

    return spills

#..................................................

def compute_pot(day,runs,spills):
    """ Tally the protons on target delivered and recorded in a day

    Args:
    day                       : datetime of day of interest
    runs   (sql)              : SQL set of runs (run, tstart, tend)
    spills (list[(time,pot)]) : Complete list of spills for time period

    Returns:
    (pot_recorded,pot_delivered) (float,float)
    """
    daystart = day
    dayend   = day + timedelta(days=1)

    pot_delivered = pot_between(daystart,dayend,spills)

    pot_recorded = 0.0
    spills_recorded = 0
    for r in runs:
        t1  = r['tstart']
        t2  = r['tstop']
        #
        # Handle end points. Runs that carried into this day and runs
        # that extedned beyond the end of this day should be pinned to
        # the start and end of the day.
        #
        if (t1<daystart):
            t1 = daystart
        if (t2>dayend):
            t2 = dayend

        pot_recorded += pot_between(t1,t2,spills)
        spills_recorded += spills_between(t1,t2,spills)

    return (pot_recorded,pot_delivered,spills_recorded)

#..................................................

def pot_between(t1,t2,spills):
    """ Compute the POT recorded between two time stamps.

    Args:
    t1 (datetime)              : Start time
    t2 (datetime)              : End time
    spills ( list[(time,pot)]) : Complete list of spills for time period

    Returns:
    pot (float) : Scintillator counts
    """

    tf1 = in_unix(t1)
    tf2 = in_unix(t2)

    sum = 0.0
    for (t,q) in spills:
        if (t>=tf1)&(t<=tf2):
            sum += q

    return sum

#..................................................

def spills_between(t1,t2,spills):
    """ Compute the POT recorded between two time stamps.

    Args:
    t1 (datetime)              : Start time
    t2 (datetime)              : End time
    spills ( list[(time,pot)]) : Complete list of spills for time period

    Returns:
    pot (float) : scintillator counts
    """

    tf1 = in_unix(t1)
    tf2 = in_unix(t2)

    nspill = 0
    for (t,q) in spills:
        if (t>=tf1)&(t<=tf2):
            nspill += 1

    return nspill

#..................................................

def in_unix(input):
    """ Convert a datetime to UNIX time

    Args:
    input: A datetime object

    Returns:
    A floating points number of seconds since Jan. 1, 1970 (UNIX standard)
    
    """
    start = datetime(year=1970,month=1,day=1,hour=0,minute=0,second=0)
    diff = input - start
    return diff.total_seconds()

#..................................................

def main():

    (t0, t1, ndays) = options()

    # setup PyRoot and other root stuff
    # Yes, these make more sense at the top of the file,
    # but if you import ROOT before parsing the arguments,
    # it takes over sys.argv and tries printing its own help
    # function instead of the one in this file.
    from ROOT import TFile, TTree
    from ROOT import gROOT, AddressOf
    
    file = TFile(BASE_DIR+"/"+filename,'RECREATE')

    gROOT.ProcessLine("struct BeamStruct {Double_t Date; Int_t spills_rec; Double_t pot_rec; Int_t run_low; Int_t run_high;};")
    from ROOT import BeamStruct
    beamStruct = BeamStruct()
    tree = TTree('BeamMetrics','POTtree')
    tree.Branch('Date',AddressOf(beamStruct,'Date'),'Date/D')
    tree.Branch('spills_rec',AddressOf(beamStruct,'spills_rec'),'spills_rec/I')
    tree.Branch('pot_rec', AddressOf(beamStruct,'pot_rec'),'pot_rec/D')
    tree.Branch('run_low',AddressOf(beamStruct,'run_low'),'run_low/I')
    tree.Branch('run_high',AddressOf(beamStruct,'run_high'),'run_high/I')


    # DB connection info for run info
    DB_NAME = os.environ['NOVADBNAME']  # 'nova_prod'
    DB_HOST = os.environ['NOVADBHOST']  # 'ifdbrep.fnal.gov'
    DB_USER = os.environ['NOVADBUSER']  # 'nova_reader'
    DB_PASS = open(os.environ['NOVADBPWDFILE'], 'r').readlines()[0].strip()
    DB_PORT = os.environ['NOVANEARDAQDBPORT']  # '5434'

    try:
        conn = psycopg2.connect(\
            "dbname=%s host=%s user=%s password=%s port=%s" % \
                (DB_NAME, DB_HOST, DB_USER, DB_PASS, DB_PORT))
    except:
        print "I am unable to connect to the database"

    global DB_CURSOR
    DB_CURSOR = conn.cursor(cursor_factory=psycopg2.extras.DictCursor)


    date0     = date(    year=t0.year, month=t0.month, day=t0.day)
    datetime0 = datetime(year=t0.year, month=t0.month, day=t0.day)


    # Loop over days to get spill info
    for i in range(0,ndays):
        datei     = date0     + timedelta(days=i)
        datetimei = datetime0 + timedelta(days=i)
        datetimef = datetimei + timedelta(days=1)
        # get unix timestamp for day - add 1 second to ensure binning is correct
        unixi     = (datetimei - datetime(1970,1,1)).total_seconds()+1
        
        #
        # Get the complete list of runs for the day
        #
        runs = runs_by_day(datei)

        #
        # Get the complete list of spills for the day
        #
        spills = mc7_spills_between(datetimei, datetimef)

        r1 = 0
        r2 = 0
        if (runs):
            (r1,r2) = run_minmax(runs);
        #
        # Compute recorded and delivered POT
        #
        (pot_rec,pot_del,num_spill_rec) = compute_pot(datetimei, runs, spills)


        print datei,",",\
            unixi,",",\
            r1,",",\
            r2,",",\
            len(spills),",",\
            num_spill_rec,",",\
            pot_del,",",\
            pot_rec;

        beamStruct.Date = unixi
        beamStruct.spills_rec = num_spill_rec
        beamStruct.pot_rec = pot_rec
        beamStruct.run_low = r1
        beamStruct.run_high = r2
        tree.Fill()

    file.Write()
    file.Close()

#..................................................

if __name__ == "__main__":
    main()