Merge pull request #171 from CoffeaTeam/docs

More docstrings

coffea/lumi_tools/__init__.py

@@ -1,6 +1,8 @@
 """Tools to parse CMS luminosity non-event data
 
-Details...
+These tools are currently tailored to the CMS experiment
+data formats, however they could be generalized and/or compartmentalized
+into a standalone package.
 """
 from .lumi_tools import (
     LumiData,

coffea/lumi_tools/lumi_tools.py

@@ -7,12 +7,18 @@
 
 
 class LumiData(object):
-    """
-    Class to hold and parse per-lumiSection integrated lumi values
-    as returned by brilcalc, e.g. with a command such as:
-    $ brilcalc lumi -c /cvmfs/cms.cern.ch/SITECONF/local/JobConfig/site-local-config.xml \
-                    -b "STABLE BEAMS" --normtag=/cvmfs/cms-bril.cern.ch/cms-lumi-pog/Normtags/normtag_PHYSICS.json \
-                    -u /pb --byls --output-style csv -i Cert_294927-306462_13TeV_PromptReco_Collisions17_JSON.txt > lumi2017.csv
+    r"""Holds per-lumiSection integrated lumi values
+
+    Parameters
+    ----------
+        lumi_csv : str
+            The path the the luminosity csv output file
+
+    The values are extracted from the csv output as returned by brilcalc, e.g. with a command such as::
+
+        brilcalc lumi -c /cvmfs/cms.cern.ch/SITECONF/local/JobConfig/site-local-config.xml \
+                 -b "STABLE BEAMS" --normtag=/cvmfs/cms-bril.cern.ch/cms-lumi-pog/Normtags/normtag_PHYSICS.json \
+                 -u /pb --byls --output-style csv -i Cert_294927-306462_13TeV_PromptReco_Collisions17_JSON.txt > lumi2017.csv
     """
     def __init__(self, lumi_csv):
         self._lumidata = np.loadtxt(lumi_csv, delimiter=',', usecols=(0, 1, 6, 7), converters={
@@ -26,32 +32,41 @@ def __init__(self, lumi_csv):
         self.build_lumi_table()
 
     def build_lumi_table(self):
+        """Build index for numba-compiled functions
+
+        This needs to be executed upon unpickling, it should be part of
+        a custom deserialize function.
+        """
         runs = self._lumidata[:, 0].astype('u4')
         lumis = self._lumidata[:, 1].astype('u4')
-        LumiData.build_lumi_table_kernel(runs, lumis, self._lumidata, self.index)
+        LumiData._build_lumi_table_kernel(runs, lumis, self._lumidata, self.index)
 
     @staticmethod
     @numba.njit(parallel=False, fastmath=False)
-    def build_lumi_table_kernel(runs, lumis, lumidata, index):
+    def _build_lumi_table_kernel(runs, lumis, lumidata, index):
         for i in range(len(runs)):
             run = runs[i]
             lumi = lumis[i]
             index[(run, lumi)] = float(lumidata[i, 2])
 
     def get_lumi(self, runlumis):
-        """
-        Return integrated lumi
-        runlumis: 2d numpy array of [[run,lumi], [run,lumi], ...] or LumiList object
+        """Calculate integrated lumi
+
+        Parameters
+        ----------
+            runlumis : numpy.ndarray or LumiList
+                A 2d numpy array of ``[[run,lumi], [run,lumi], ...]`` or `LumiList` object
+                of the lumiSections to integrate over.
         """
         if isinstance(runlumis, LumiList):
             runlumis = runlumis.array
         tot_lumi = np.zeros((1, ), dtype=np.float64)
-        LumiData.get_lumi_kernel(runlumis[:, 0], runlumis[:, 1], self.index, tot_lumi)
+        LumiData._get_lumi_kernel(runlumis[:, 0], runlumis[:, 1], self.index, tot_lumi)
         return tot_lumi[0]
 
     @staticmethod
     @numba.njit(parallel=False, fastmath=False)
-    def get_lumi_kernel(runs, lumis, index, tot_lumi):
+    def _get_lumi_kernel(runs, lumis, index, tot_lumi):
         ks_done = set()
         for iev in range(len(runs)):
             run = np.uint32(runs[iev])
@@ -63,10 +78,14 @@ def get_lumi_kernel(runs, lumis, index, tot_lumi):
 
 
 class LumiMask(object):
-    """
-    Class that parses a 'golden json' into an efficient valid lumiSection lookup table
-    Instantiate with the json file, and call with an array of runs and lumiSections, to
-    return a boolean array, where valid lumiSections are marked True
+    """Holds a luminosity mask index, and provides vectorized lookup
+
+    Parameters
+    ----------
+        jsonfile : str
+            Path the the 'golden json' file or other valid lumiSection database in json format.
+
+    This class parses a CMS lumi json into an efficient valid lumiSection lookup table
     """
     def __init__(self, jsonfile):
         with open(jsonfile) as fin:
@@ -83,18 +102,33 @@ def __init__(self, jsonfile):
             self._masks[np.uint32(run)] = mask
 
     def __call__(self, runs, lumis):
+        """Check if run and lumi are valid
+
+        Parameters
+        ----------
+            runs : numpy.ndarray
+                Vectorized list of run numbers
+            lumis : numpy.ndarray
+                Vectorized list of lumiSection numbers
+
+        Returns
+        -------
+            mask_out : numpy.ndarray
+                An array of dtype `bool` where valid (run, lumi) tuples
+                will have their corresponding entry set ``True``.
+        """
         mask_out = np.zeros(dtype='bool', shape=runs.shape)
-        LumiMask.apply_run_lumi_mask(self._masks, runs, lumis, mask_out)
+        LumiMask._apply_run_lumi_mask(self._masks, runs, lumis, mask_out)
         return mask_out
 
     @staticmethod
-    def apply_run_lumi_mask(masks, runs, lumis, mask_out):
-        LumiMask.apply_run_lumi_mask_kernel(masks, runs, lumis, mask_out)
+    def _apply_run_lumi_mask(masks, runs, lumis, mask_out):
+        LumiMask._apply_run_lumi_mask_kernel(masks, runs, lumis, mask_out)
 
     # This could be run in parallel, but windows does not support it
     @staticmethod
     @numba.njit(parallel=False, fastmath=True)
-    def apply_run_lumi_mask_kernel(masks, runs, lumis, mask_out):
+    def _apply_run_lumi_mask_kernel(masks, runs, lumis, mask_out):
         for iev in numba.prange(len(runs)):
             run = np.uint32(runs[iev])
             lumi = np.uint32(lumis[iev])
@@ -107,9 +141,16 @@ def apply_run_lumi_mask_kernel(masks, runs, lumis, mask_out):
 
 
 class LumiList(object):
-    """
-    Mergeable (using +=) list of unique (run,lumiSection) values
-    The member array can be passed to LumiData.get_lumi()
+    """Mergeable list of unique (run, lumiSection) values
+
+    This list can be merged with another via ``+=``.
+
+    Parameters
+    ----------
+        runs : numpy.ndarray
+            Vectorized list of run numbers
+        lumis : numpy.ndarray
+            Vectorized list of lumiSection values
     """
     def __init__(self, runs=None, lumis=None):
         self.array = np.zeros(shape=(0, 2))
@@ -125,4 +166,5 @@ def __iadd__(self, other):
         return self
 
     def clear(self):
+        """Clear current lumi list"""
         self.array = np.zeros(shape=(0, 2))

coffea/processor/__init__.py

@@ -1,6 +1,6 @@
 """A framework for analysis scale-out
 
-Details...
+
 """
 from .processor import ProcessorABC
 from .dataframe import (

coffea/util.py

@@ -1,6 +1,5 @@
 """Utility functions
 
-Details...
 """
 try:
     import awkward.numba as awkward
@@ -19,18 +18,17 @@
 
 
 def load(filename):
-    '''
-    Load a coffea file from disk
+    '''Load a coffea file from disk
     '''
     with lz4.frame.open(filename) as fin:
         output = cloudpickle.load(fin)
     return output
 
 
 def save(output, filename):
-    '''
-    Save a coffea object or collection thereof to disk
-    Suggested suffix: .coffea
+    '''Save a coffea object or collection thereof to disk
+
+    This function can accept any picklable object.  Suggested suffix: ``.coffea``
     '''
     with lz4.frame.open(filename, 'wb') as fout:
         cloudpickle.dump(output, fout)

tests/test_hist_plot.py

@@ -68,7 +68,7 @@ def test_plot1d():
     lepton_kinematics = fill_lepton_kinematics()
 
     # looking at lepton pt for all eta
-    lepton_pt = lepton_kinematics.project("eta", overflow='under')
+    lepton_pt = lepton_kinematics.integrate("eta", overflow='under')
 
     fig, ax, primitives = hist.plot1d(lepton_pt, overlay="flavor", stack=True,
                                       fill_opts={'alpha': .5, 'edgecolor': (0,0,0,0.3)})
@@ -93,7 +93,7 @@ def test_plot2d():
     lepton_kinematics = fill_lepton_kinematics()
 
     # looking at lepton pt for all eta
-    muon_kinematics = lepton_kinematics.project("flavor","muon")
+    muon_kinematics = lepton_kinematics.integrate("flavor","muon")
 
     fig, ax, primitives = hist.plot2d(muon_kinematics, "eta")
 

tests/test_lumi_tools.py

@@ -24,13 +24,13 @@ def test_lumidata():
     )
     pyruns = lumidata._lumidata[:, 0].astype('u4')
     pylumis = lumidata._lumidata[:, 1].astype('u4')
-    LumiData.build_lumi_table_kernel.py_func(pyruns, pylumis, lumidata._lumidata, py_index)
+    LumiData._build_lumi_table_kernel.py_func(pyruns, pylumis, lumidata._lumidata, py_index)
 
     assert(len(py_index) == len(lumidata.index))
 
     # test get_lumi_kernel
     py_tot_lumi = np.zeros((1, ), dtype=np.float64)
-    LumiData.get_lumi_kernel.py_func(runslumis[:, 0], runslumis[:, 1], py_index, py_tot_lumi)
+    LumiData._get_lumi_kernel.py_func(runslumis[:, 0], runslumis[:, 1], py_index, py_tot_lumi)
 
     assert(abs(py_tot_lumi[0] - l) < 1e-4)
 
@@ -45,7 +45,7 @@ def test_lumimask():
 
     # test underlying py_func
     py_mask = np.zeros(dtype='bool', shape=runs.shape)
-    LumiMask.apply_run_lumi_mask_kernel.py_func(lumimask._masks,
+    LumiMask._apply_run_lumi_mask_kernel.py_func(lumimask._masks,
                                                 runs, lumis,
                                                 py_mask)