Merge pull request #265 from ndawe/master

[MRG] fix handling of fixed and variable length array expressions
scikit-hep · Aug 17, 2016 · 640beb7 · 640beb7
2 parents 8022e9e + 0bc0b62
commit 640beb7
Show file tree

Hide file tree

Showing 10 changed files with 4,586 additions and 3,864 deletions.
diff --git a/docs/index.rst b/docs/index.rst
@@ -1,9 +1,3 @@
-.. raw:: html
-
-   <a href="https://github.com/rootpy/root_numpy">
-      <img style="position: absolute; top: 0; right: 0; border: 0;"
-      src="_static/fork_me.png">
-   </a>
 
 ==========
 root_numpy

diff --git a/docs/start.rst b/docs/start.rst
@@ -3,19 +3,19 @@
 Getting Started
 ===============
 
-Try root_numpy on `CERN's LXPLUS <http://information-technology.web.cern.ch/services/lxplus-service>`_
-======================================================================================================
+If you have access to CERN's CVMFS then you can activate an environment with
+compatible builds of python, ROOT, numpy and root_numpy with the following::
 
-First set up an environment with consistent GCC, ROOT and Python builds::
-
-   export LCGENV_PATH=/afs/cern.ch/sw/lcg/releases
-   /cvmfs/sft.cern.ch/lcg/releases/lcgenv/latest/lcgenv -p LCG_84 x86_64-slc6-gcc49-opt ROOT > lcgenv.sh
+   export LCGENV_PATH=/cvmfs/sft.cern.ch/lcg/releases
+   /cvmfs/sft.cern.ch/lcg/releases/lcgenv/latest/lcgenv -p LCG_85swan2 --ignore Grid x86_64-slc6-gcc49-opt root_numpy > lcgenv.sh
    echo 'export PATH=$HOME/.local/bin:$PATH' >> lcgenv.sh
    source lcgenv.sh
 
 In new terminal sessions, only the last line above will be required.
 
-Install pip and virtualenv::
+If you want to instead use your own installation of root_numpy along with any
+other packages you need, then continue with setting up a virtualenv. First
+install pip and virtualenv::
 
    curl -O https://bootstrap.pypa.io/get-pip.py
    python get-pip.py --user

diff --git a/root_numpy/_tree.py b/root_numpy/_tree.py
@@ -224,7 +224,7 @@ def root2rec(filenames,
              include_weight=False,
              weight_name='weight',
              cache_size=-1,
-             warn_missing_tree=False):
+             warn_missing_tree=False):  # pragma: no cover
     """View the result of :func:`root2array` as a record array.
 
     .. warning:: ``root2rec`` is deprecated and will be removed in
@@ -385,7 +385,7 @@ def tree2rec(tree,
              step=None,
              include_weight=False,
              weight_name='weight',
-             cache_size=-1):
+             cache_size=-1):  # pragma: no cover
     """View the result of :func:`tree2array` as a record array.
 
     .. warning:: ``tree2rec`` is deprecated and will be removed in

diff --git a/root_numpy/src/Column.h b/root_numpy/src/Column.h
@@ -10,6 +10,9 @@ class Column
 {
     public:
 
+    Column(std::string _name, std::string _type):
+        name(_name),
+        type(_type){}
     virtual ~Column() {}
     virtual int GetLen() = 0;
     virtual int GetCountLen() = 0;
@@ -24,26 +27,24 @@ class Column
 };
 
 
-class MultiFormulaColumn: public Column
+template <typename T>
+class FormulaArrayColumn: public Column
 {
     public:
 
-    MultiFormulaColumn(std::string _name, TTreeFormula* _formula)
-    {
-        name = _name;
-        formula = _formula;
-        type = "Double_t";
-        value = NULL;
-    }
+    FormulaArrayColumn(std::string _name, std::string _type, TTreeFormula* _formula):
+        Column(_name, _type),
+        formula(_formula),
+        value(NULL){}
 
-    ~MultiFormulaColumn()
+    ~FormulaArrayColumn()
     {
         delete[] value;
     }
 
     const char* GetTypeName()
     {
-        return "double";
+        return type.c_str();
     }
 
     int GetLen()
@@ -58,33 +59,70 @@ class MultiFormulaColumn: public Column
 
     int GetSize()
     {
-        return sizeof(double) * GetLen();
+        return sizeof(T) * GetLen();
     }
 
     void* GetValuePointer()
     {
         delete[] value;
-        value = new double[formula->GetNdata()];
+        value = new T[formula->GetNdata()];
         for (int i = 0; i < formula->GetNdata(); ++i)
         {
-            value[i] = formula->EvalInstance(i);
+            value[i] = (T) formula->EvalInstance(i);
         }
         return value;
     }
 
     TTreeFormula* formula;
-    double* value;
+    T* value;
+};
+
+
+template <typename T>
+class FormulaFixedArrayColumn: public FormulaArrayColumn<T>
+{
+    public:
+
+    FormulaFixedArrayColumn(std::string _name, std::string _type, TTreeFormula* _formula):
+        FormulaArrayColumn<T>(_name, _type, _formula)
+    {
+        length = _formula->GetNdata();
+        this->value = new T[length];
+    }
+
+    int GetLen()
+    {
+        return length;
+    }
+
+    int GetCountLen()
+    {
+        return length;
+    }
+
+    void* GetValuePointer()
+    {
+        // Call to GetNdata() again required to update leaves
+        for (int i = 0; i < this->formula->GetNdata(); ++i)
+        {
+            this->value[i] = (T) this->formula->EvalInstance(i);
+        }
+        return this->value;
+    }
+
+    int length;
 };
 
 
-class FormulaColumn: public MultiFormulaColumn
+template <typename T>
+class FormulaColumn: public FormulaArrayColumn<T>
 {
     public:
 
-    FormulaColumn(std::string _name, TTreeFormula* _formula):
-        MultiFormulaColumn(_name, _formula)
+    FormulaColumn(std::string _name, std::string _type, TTreeFormula* _formula):
+        FormulaArrayColumn<T>(_name, _type, _formula)
     {
-        value = new double[1];
+        this->value = new T[1];
     }
 
     int GetLen()
@@ -99,9 +137,9 @@ class FormulaColumn: public MultiFormulaColumn
 
     void* GetValuePointer()
     {
-        formula->GetNdata(); // required, as in TTreePlayer
-        value[0] = formula->EvalInstance(0);
-        return value;
+        this->formula->GetNdata(); // required, as in TTreePlayer
+        this->value[0] = (T) this->formula->EvalInstance(0);
+        return this->value;
     }
 };
 
@@ -110,12 +148,9 @@ class BranchColumn: public Column
 {
     public:
 
-    BranchColumn(std::string& _name, TLeaf* _leaf)
-    {
-        name = _name;
-        leaf = _leaf;
-        type = leaf->GetTypeName();
-    }
+    BranchColumn(std::string _name, TLeaf* _leaf):
+        Column(_name, _leaf->GetTypeName()),
+        leaf(_leaf){}
 
     void SetLeaf(TLeaf* newleaf, bool check=false)
     {

diff --git a/root_numpy/src/ROOT.pxi b/root_numpy/src/ROOT.pxi
@@ -132,6 +132,7 @@ cdef extern from "TTreeFormula.h":
         int GetNdata()
         int GetMultiplicity()
         double EvalInstance(int)
+        bool IsInteger(bool)
 
 cdef extern from "TClassEdit.h" namespace "TClassEdit":
     string ResolveTypedef(const_char*, bool)

diff --git a/root_numpy/src/TreeChain.h b/root_numpy/src/TreeChain.h
@@ -302,7 +302,12 @@ class TreeChain
             ncodes = (*fit)->GetNcodes();
             for (n = 0; n < ncodes; ++n)
             {
-                branch = (*fit)->GetLeaf(n)->GetBranch();
+                leaf = (*fit)->GetLeaf(n);
+                if (leaf == NULL)
+                {
+                    continue;
+                }
+                branch = leaf->GetBranch();
                 // Branch may be a TObject split across multiple
                 // subbranches. These must be activated recursively.
                 activate_branch_recursive(branch);