rewrite tree2array internals for improved speed (some tests show x2 speedup)

and support for arrays of arbitrary dimensionality

Author: ndawe

README.rst

@@ -46,10 +46,10 @@ At the core of root_numpy are powerful and flexible functions for converting
 or `structured arrays <http://docs.scipy.org/doc/numpy/user/basics.rec.html>`_
 as well as converting NumPy arrays back into ROOT TTrees. root_numpy can
 convert branches of strings and basic types such as bool, int, float, double,
-etc. as well as variable and fixed-length 1D arrays and 1D or 2D vectors of
-basic types and strings. root_numpy can also create columns in the output array
-that are expressions involving the TTree branches (i.e. ``'vect.Pt() / 1000'``)
-similar to ``TTree::Draw()``.
+etc. as well as variable-length and fixed-length multidimensional arrays and 1D
+or 2D vectors of basic types and strings. root_numpy can also create columns in
+the output array that are expressions involving the TTree branches (i.e.
+``'vect.Pt() / 1000'``) similar to ``TTree::Draw()``.
 
 For example, get a NumPy structured or record array from a TTree
 (copy and paste the following examples into your Python prompt):

root_numpy/_tree.py

@@ -114,11 +114,10 @@ def root2array(filenames,
     treename : str, optional (default=None)
         Name of the tree to convert (optional if each file contains exactly one
         tree).
-    branches : list of str, optional (default=None)
-        List of branch names to include as columns of the array. If None or
-        empty then include all branches than can be converted in the first
-        tree. If branches contains duplicate branches, only the first one is
-        used.
+    branches : list of strings or single string, optional (default=None)
+        List of branch names to include as columns of the array or a single
+        branch name to convert into a one-dimensional array. If None then
+        include all branches that can be converted.
     selection : str, optional (default=None)
         Only include entries fulfilling this condition.
     start, stop, step: int, optional (default=None)
@@ -189,16 +188,27 @@ def root2array(filenames,
         elif not trees:
             raise IOError(
                 "no trees present in {0}".format(filenames[0]))
-        else:
-            treename = trees[0]
+        treename = trees[0]
+
+    if isinstance(branches, string_types):
+        # single branch selected
+        branches = [branches]
+        flatten = True
+    else:
+        flatten = False
 
-    return _librootnumpy.root2array_fromFname(
+    arr = _librootnumpy.root2array_fromFname(
         filenames, treename, branches,
         selection,
         start, stop, step,
         include_weight,
         weight_name)
 
+    if flatten:
+        # select single column
+        return arr[branches[0]]
+    return arr
+
 
 def root2rec(filenames,
              treename=None,
@@ -245,11 +255,10 @@ def tree2array(tree,
     ----------
     tree : ROOT TTree instance
         The ROOT TTree to convert into an array.
-    branches : list of str, optional (default=None)
-        List of branch names to include as columns of the array. If None or
-        empty then include all branches than can be converted in the first
-        tree. If branches contains duplicate branches, only the first one is
-        used.
+    branches : list of strings or single string, optional (default=None)
+        List of branch names to include as columns of the array or a single
+        branch name to convert into a one-dimensional array. If None then
+        include all branches that can be converted.
     selection : str, optional (default=None)
         Only include entries fulfilling this condition.
     start, stop, step: int, optional (default=None)
@@ -271,13 +280,24 @@ def tree2array(tree,
     import ROOT
     if not isinstance(tree, ROOT.TTree):
         raise TypeError("tree must be a ROOT.TTree")
-    # will need AsCapsule for Python 3
     cobj = ROOT.AsCObject(tree)
+
+    if isinstance(branches, string_types):
+        # single branch selected
+        branches = [branches]
+        flatten = True
+    else:
+        flatten = False
+
     arr = _librootnumpy.root2array_fromCObj(
         cobj, branches, selection,
         start, stop, step,
         include_weight,
         weight_name)
+
+    if flatten:
+        # select single column
+        return arr[branches[0]]
     return arr
 
 

root_numpy/src/Column.h

@@ -6,45 +6,32 @@
 #include <string>
 
 
-enum ColumnType
-{
-    SINGLE = 1,
-    FIXED = 2,
-    VARY = 3
-};
-
-
 class Column
 {
     public:
         virtual ~Column() {}
         virtual int GetLen() = 0;
+        virtual int GetCountLen() = 0;
         virtual int GetSize() = 0;
         virtual void* GetValuePointer() = 0;
         virtual const char* GetTypeName() = 0;
 
-        bool skipped;
-        // single fixed vary?
-        ColumnType coltype;
-        // column name
-        std::string colname;
-        // useful in case of fixed element
-        int countval;
-        // name of the roottype
-        std::string rttype;
+        // Column name
+        std::string name;
+        // Name of the ROOT type
+        std::string type;
 };
 
 
 class FormulaColumn: public Column
 {
     public:
 
-        FormulaColumn(std::string _colname, TTreeFormula* _formula)
+        FormulaColumn(std::string _name, TTreeFormula* _formula)
         {
-            colname = _colname;
-            coltype = SINGLE;
+            name = _name;
             formula = _formula;
-            rttype = "Double_t";
+            type = "Double_t";
             value = new double[1];
         }
 
@@ -58,6 +45,11 @@ class FormulaColumn: public Column
             return 1;
         }
 
+        int GetCountLen()
+        {
+            return 1;
+        }
+
         int GetSize()
         {
             return sizeof(double) * GetLen();
@@ -81,73 +73,24 @@ class FormulaColumn: public Column
 };
 
 
-// This describes the structure of the tree
-// Converter should take this and make the appropriate data structure
 class BranchColumn: public Column
 {
     public:
 
-        static int find_coltype(TLeaf* leaf,
-                                ColumnType& coltype,
-                                int& countval)
-        {
-            // Check whether it's array if so of which type
-            TLeaf* len_leaf = leaf->GetLeafCounter(countval);
-            if (countval == 1)
-            {
-                if (len_leaf == 0)
-                { // single element
-                    coltype = SINGLE;
-                }
-                else
-                { // variable length
-                    coltype = VARY;
-                }
-            }
-            else if (countval > 0)
-            {
-                // fixed multiple array
-                coltype = FIXED;
-            }
-            else
-            {
-                // negative
-                std::string msg("Unable to understand the structure of leaf ");
-                msg += leaf->GetName();
-                PyErr_SetString(PyExc_IOError, msg.c_str());
-                return 0;
-            }
-            return 1;
-        }
-
-        static BranchColumn* build(TLeaf* leaf, const std::string& colname)
+        BranchColumn(std::string& _name, TLeaf* _leaf)
         {
-            BranchColumn* ret = new BranchColumn();
-            ret->leaf = leaf;
-            ret->colname = colname;
-            ret->skipped = false;
-            if (!find_coltype(leaf, ret->coltype, ret->countval))
-            {
-                delete ret;
-                return NULL;
-            }
-            ret->rttype = leaf->GetTypeName();
-            return ret;
+            name = _name;
+            leaf = _leaf;
+            type = leaf->GetTypeName();
         }
 
         void SetLeaf(TLeaf* newleaf, bool check=false)
         {
             leaf = newleaf;
             if (check)
             {
-                assert(leaf->GetTypeName() == rttype);
-                int cv;
-                ColumnType ct;
-                if (find_coltype(leaf, ct, cv) == 0)
-                    abort();
-                if (ct != coltype)
-                    abort();
-                //if(ct==FIXED){assert(cv==countval);}
+                assert(leaf->GetTypeName() == type);
+                // TODO: compare shape
             }
         }
 
@@ -157,6 +100,17 @@ class BranchColumn: public Column
             return leaf->GetLen();
         }
 
+        int GetCountLen()
+        {
+            // get count leaf value
+            TLeaf* count_leaf = leaf->GetLeafCount();
+            if (count_leaf != NULL)
+            {
+                return int(count_leaf->GetValue());
+            }
+            return 1;
+        }
+
         int GetSize()
         {
             // get size of this block in bytes

root_numpy/src/ROOT.pxi

@@ -46,16 +46,19 @@ cdef extern from "TBranch.h":
 
 cdef extern from "TLeaf.h":
     cdef cppclass TLeaf:
+        const_char* GetName()
+        const_char* GetTitle()
         const_char* GetTypeName()
+        TLeaf* GetLeafCount()
         TLeaf* GetLeafCounter(int&)
-        const_char* GetName()
-
+        TBranch* GetBranch()
+        
 cdef extern from "TTree.h":
     cdef cppclass TTree:
         TTree()
         TTree(const_char*,  const_char*)
-        void GetEntry(int i)
-        int GetEntries()
+        int GetEntry(long_long entry)
+        long_long GetEntries()
         void SetBranchAddress(const_char* bname, void* addr)
         void SetBranchStatus(const_char* bname, bool status)
         void Print()
@@ -66,15 +69,16 @@ cdef extern from "TTree.h":
         int Fill()
         int Scan()
         void Delete(void*)
-        long SetEntries(long)
+        long_long SetEntries(long_long)
+        double GetWeight()
         int Write()
         int Write(const_char* name, int option)
 
 cdef extern from "TChain.h":
     cdef cppclass TChain(TTree):
         TChain()
         TChain(const_char*)
-        int Add(const_char*, long)
+        int Add(const_char*, long_long)
         void Print()
 
 cdef extern from "TFormula.h":