Improve code quality of example scripts (#1828)

Jürgen Knödlseder
1 parent 256eb549
Showing 5 changed files with 128 additions and 128 deletions
cscripts/ioutils.py
cscripts/obsutils.py
examples/generate_prod3_irfs.py
examples/make_pointings.py
examples/make_ts_distributions.py
@@ -18,6 +18,7 @@
 #
 # ==========================================================================
 import csv
+import sys
 import gammalib
 import ctools
@@ -62,3 +63,59 @@ def write_csv_row(outfile, row, colnames, colvalues):
     # Return
     return
+
+
+# ============================== #
+# Get options from argument list #
+# ============================== #
+def get_arg_options(options, usage):
+    """
+    Get options from argument list
+
+    Parameters
+    ----------
+    options : list of dict
+        List of possible options and default values
+    usage : str
+        Usage string to be shown in case of a problem
+
+    Returns
+    -------
+    options : list of dict
+        Update list of options
+    """
+    # If there are no command line arguments then show usage string and
+    # exit
+    if len(sys.argv) < 1:
+        print(usage)
+        sys.exit()
+
+    # First possible option is element 1
+    i = 1
+
+    # Loop over all arguments
+    while i < len(sys.argv):
+
+        # Search for options
+        for option in options:
+            if sys.argv[i] == option['option']:
+
+                # If an option has been found then check if there is a
+                # following parameter and extract that parameter as a
+                # string
+                if len(sys.argv) > i+1:
+                    i += 1
+                    try:
+                        option['value'] = str(sys.argv[i])
+                    except:
+                        print(usage)
+                        sys.exit()
+                else:
+                    print(usage)
+                    sys.exit()
+
+        # Next item
+        i += 1
+
+    # Return options
+    return options
@@ -31,9 +31,9 @@ def sim(obs, log=False, debug=False, chatter=2, edisp=False, seed=0,
     Simulate events for all observations in the container
     Simulate events for all observations using ctobssim. If the number of
-    energy bins is positive, the events are then binned in a counts cube
-    using ctbin. If multiple observations are simulated, the counts cube is
-    a stacked cube and the corresponding response cubes are computed using
+    energy bins is positive, the events are binned into a counts cube using
+    ctbin. If multiple observations are simulated, the counts cube is a
+    stacked cube and the corresponding response cubes are computed using
     ctexpcube, ctpsfcube, ctbkgcube and optionally ctedispcube. The response
     cubes are attached to the first observation in the container, which
     normally is the observation with the counts cube.
@@ -203,8 +203,8 @@ def set_obs(pntdir, tstart=0.0, duration=1800.0, deadc=0.95, \
     Set a single CTA observation
     The function sets a single CTA observation containing an empty CTA
-    event list. By looping over this function you can add CTA observations
-    to the observation container.
+    event list. By looping over this function CTA observations can be
+    added to the observation container.
     Parameters
     ----------
@@ -235,19 +235,19 @@ def set_obs(pntdir, tstart=0.0, duration=1800.0, deadc=0.95, \
         CTA observation
     """
     # Allocate CTA observation
-    obs_cta = gammalib.GCTAObservation()
+    obs = gammalib.GCTAObservation()
-    # Set calibration database
+    # Set CTA calibration database
     db = gammalib.GCaldb()
     if (gammalib.dir_exists(caldb)):
         db.rootdir(caldb)
     else:
         db.open('cta', caldb)
-    # Set pointing direction
+    # Set pointing direction for CTA observation
     pnt = gammalib.GCTAPointing()
     pnt.dir(pntdir)
-    obs_cta.pointing(pnt)
+    obs.pointing(pnt)
     # Set ROI
     roi     = gammalib.GCTARoi()
@@ -266,22 +266,26 @@ def set_obs(pntdir, tstart=0.0, duration=1800.0, deadc=0.95, \
     # Allocate event list
     events = gammalib.GCTAEventList()
+
+    # Set ROI, GTI and energy boundaries for event list
     events.roi(roi)
     events.gti(gti)
     events.ebounds(ebounds)
-    obs_cta.events(events)
-    # Set instrument response
-    obs_cta.response(irf, db)
+    # Set the event list as the events for CTA observation
+    obs.events(events)
+
+    # Set instrument response for CTA observation
+    obs.response(irf, db)
-    # Set ontime, livetime, and deadtime correction factor
-    obs_cta.ontime(duration)
-    obs_cta.livetime(duration*deadc)
-    obs_cta.deadc(deadc)
-    obs_cta.id(obsid)
+    # Set ontime, livetime, and deadtime correction factor for CTA observation
+    obs.ontime(duration)
+    obs.livetime(duration*deadc)
+    obs.deadc(deadc)
+    obs.id(obsid)
     # Return CTA observation
-    return obs_cta
+    return obs
 # ============================ #
 #! /usr/bin/env python
 # ==========================================================================
-# Generates Prod3 IRFs from ROOT performance files.
+# Generates Prod3 IRFs from ROOT performance files
 #
 # Copyright (C) 2016 Juergen Knoedlseder
 #
@@ -30,7 +30,7 @@ def set_prod3_desy(datadir):
     Set Prod3 DESY database
     This function sets the calibration database. By adding dictionaries to
-    the db list a number of calibrations can be defined for which IRFs
+    the "db" list a number of calibrations can be defined for which IRFs
     should be generated. All calibrations will be added to the calibration
     database.
@@ -44,22 +44,10 @@ def set_prod3_desy(datadir):
     db : list of dict
         Database
     """
-    # Set database attributes
-    path       = datadir
-    file       = 'irf.root'
-    psftype    = 'Gauss'
-    oversample = 3
-    #norm1d     = True   # Will normalise the on-axis 2D IRFs on the 1D IRFs
-    norm1d     = False
-    bgdinfill  = False
-    bgdethres  = 10.0
-
     # Set database content
-    db = [{'inst': 'prod3', 'id': 'South_50h', 'psftype': psftype,
-           'oversample': oversample, 'norm1d': norm1d,
-           'bgdinfill': bgdinfill, 'bgdethres': bgdethres,
-           'path': path, 'file': file}
-         ]
+    db = [{'inst': 'prod3', 'id': 'South_50h', 'analysis': 'DESY',
+           'oversample': 3, 'norm1d': False, 'bgdinfill': False,
+           'bgdethres': 10.0, 'path': datadir, 'file': 'irf.root'}]
     # Return database
     return db
@@ -82,36 +70,6 @@ if __name__ == &#39;__main__&#39;:
     # Loop over entries
     for entry in entries:
-        # Set optional attributes
-        if 'rebin' in entry:
-            rebin = entry['rebin']
-        else:
-            rebin = False
-        if 'eascale' in entry:
-            eascale = entry['eascale']
-        else:
-            eascale = 1.0
-        if 'bgdscale' in entry:
-            bgdscale = entry['bgdscale']
-        else:
-            bgdscale = 1.0
-        if 'bgdethres' in entry:
-            bgdethres = entry['bgdethres']
-        else:
-            bgdethres = 1000.0
-        if 'bgdinfill' in entry:
-            bgdinfill = entry['bgdinfill']
-        else:
-            bgdinfill = False
-        if 'oversample' in entry:
-            oversample = entry['oversample']
-        else:
-            oversample = 1
-        if 'norm1d' in entry:
-            norm1d = entry['norm1d']
-        else:
-            norm1d = False
-
         # Set filename
         filename = entry['path']+'/'+entry['file']
@@ -122,15 +80,15 @@ if __name__ == &#39;__main__&#39;:
         caldb['inst']          = entry['inst']
         caldb['id']            = entry['id']
         caldb['version']       = entry['inst']
-        caldb['analysis']      = 'DESY'
-        caldb['psftype']       = entry['psftype']
-        caldb['norm1d']        = norm1d
-        caldb['rebin']         = rebin
-        caldb['eascale']       = eascale
-        caldb['bgdscale']      = bgdscale
-        caldb['bgdoversample'] = oversample
-        caldb['bgdethres']     = bgdethres
-        caldb['bgdinfill']     = bgdinfill
+        caldb['analysis']      = entry.setdefault('analysis', 'DESY')
+        caldb['psftype']       = entry.setdefault('psftype', 'Gauss')
+        caldb['norm1d']        = entry.setdefault('norm1d', False)
+        caldb['rebin']         = entry.setdefault('rebin', False)
+        caldb['eascale']       = entry.setdefault('eascale', 1.0)
+        caldb['bgdscale']      = entry.setdefault('bgdscale', 1.0)
+        caldb['bgdoversample'] = entry.setdefault('oversample', 1)
+        caldb['bgdethres']     = entry.setdefault('bgdethres', 1000.0)
+        caldb['bgdinfill']     = entry.setdefault('bgdinfill', False)
         caldb['logfile']       = 'generate_prod3_irfs.log'
         # Add CALDB entry
@@ -56,7 +56,7 @@ def get_positions(xmin, xmax, ymin, ymax, step):
     # Initialise positions
     positions = []
-    # Determine ystep and number of rows
+    # Determine ystep and number of pointing rows
     ystep   = step * math.sqrt(0.75)
     ny      = int((ymax-ymin)/ystep+1.5)
     rescale = (ymax-ymin)/(ystep*(ny-1))
@@ -67,13 +67,16 @@ def get_positions(xmin, xmax, ymin, ymax, step):
     y = ymin
     for row in range(ny):
-        # Determine xstep and number of pointings in row
-        xstep = step / math.cos(gammalib.deg2rad * y)
-        nx    = int((xmax-xmin)/xstep+1.5)
+        # Determine xstep and number of pointings in row. The xstep increases
+        # with the cosine of the latitude so that the distance of the step
+        # is invariant of latitude.
+        xstep   = step / math.cos(gammalib.deg2rad * y)
+        nx      = int((xmax-xmin)/xstep+1.5)
         rescale = (xmax-xmin)/(xstep*(nx-1))
         xstep  *= rescale
-        # Set x offset
+        # Set x offset. For every second row the x position is displace by
+        # half a step size.
         if row % 2 == 0:
             x = xmin
         else:
@@ -151,7 +154,9 @@ def set_patch(lmin=-30.0, lmax=30.0, bmin=-1.5, bmax=+1.5,
         # Optionally add-in site dependent IRF
         if site != None:
-            # Handle automatic site switching
+            # If automatic site switching is requested then set the North
+            # IRF for declinations greater or equal than "autodec" and the
+            # South IRF for declination smaller than "autodec".
             if site == 'Automatic':
                 pos = gammalib.GSkyDir()
                 pos.lb_deg(pnt['x'], pnt['y'])
@@ -160,10 +165,13 @@ def set_patch(lmin=-30.0, lmax=30.0, bmin=-1.5, bmax=+1.5,
                     obs_site = 'North'
                 else:
                     obs_site = 'South'
+
+            # ... otherwise use the specified site, North or South
             else:
                 obs_site = site
-            # Set site IRF
+            # Set site dependent Prod2 50 hour IRF and add up exposure for
+            # a given site
             caldb = 'prod2'
             if obs_site == 'North':
                 if lst:
@@ -87,53 +87,26 @@ def create_ts(datadir, loge, emin, emax, ntrials=100, duration=180000.0,
 # ======================== #
 if __name__ == '__main__':
-    # Get input arguments
-    usage = 'make_ts_distributions.py [-n ntrials] [-e enumbins] [-m max_threads]'
-    if len(sys.argv) < 1:
-        print(usage)
-        sys.exit()
-
-    # Set default parameters
-    ntrials     = 100
-    enumbins    = 0
-    duration    = 180000.0
-    max_threads = 1
-    datadir     = 'data'
-
-    # Parameter dictionnary
-    pars = [{'option': '-n', 'value': ntrials},
-            {'option': '-e', 'value': enumbins},
-            {'option': '-d', 'value': duration},
-            {'option': '-m', 'value': max_threads},
-            {'option': '-p', 'value': datadir}]
-
-    # Gather parameters from command line
-    i = 1
-    while i < len(sys.argv):
-
-        # Search for options
-        for par in pars:
-            if sys.argv[i] == par['option']:
-                if len(sys.argv) > i+1:
-                    i += 1
-                    try:
-                        par['value'] = int(sys.argv[i])
-                    except:
-                        print(usage)
-                        sys.exit()
-                else:
-                    print(usage)
-                    sys.exit()
-
-        # Next item
-        i += 1
-
-    # Recover parameters
-    ntrials     = pars[0]['value']
-    enumbins    = pars[1]['value']
-    duration    = pars[2]['value']
-    max_threads = pars[3]['value']
-    datadir     = pars[4]['value']
+    # Set usage string
+    usage = 'make_ts_distributions.py [-n ntrials] [-e enumbins]'\
+            ' [-d duration] [-m max_threads] [-p datadir]'
+
+    # Set default options
+    options = [{'option': '-n', 'value': '100'},
+               {'option': '-e', 'value': '0'},
+               {'option': '-d', 'value': '180000.0'},
+               {'option': '-m', 'value': '1'},
+               {'option': '-p', 'value': 'data'}]
+
+    # Get options from command line arguments
+    options = cscripts.ioutils.get_arg_options(options, usage)
+
+    # Extract script parameters from options
+    ntrials     = int(options[0]['value'])
+    enumbins    = int(options[1]['value'])
+    duration    = float(options[2]['value'])
+    max_threads = int(options[3]['value'])
+    datadir     = options[4]['value']
     # Loop over energy bands. The energy bands are those that are also
     # used for sensitivity computation.