Merge pull request #206 from ndawe/master

array2tree: support strings and fixed-size subarrays
scikit-hep · Jun 19, 2015 · 377db79 · 377db79
2 parents b4f9a19 + 19e488c
commit 377db79
Show file tree

Hide file tree

Showing 11 changed files with 13,926 additions and 11,756 deletions.
diff --git a/root_numpy/_tree.py b/root_numpy/_tree.py
@@ -335,9 +335,8 @@ def tree2rec(tree,
 def array2tree(arr, name='tree', tree=None):
     """Convert a numpy structured array into a ROOT TTree.
 
-    .. warning::
-       This function is experimental. Please report problems.
-       Not all data types are supported (``np.object`` and ``np.float16``).
+    Fields of basic types, strings, and fixed-size subarrays of basic types are
+    supported. ``np.object`` and ``np.float16`` are currently not supported.
 
     Parameters
     ----------
@@ -359,6 +358,26 @@ def array2tree(arr, name='tree', tree=None):
     --------
     array2root
 
+    Examples
+    --------
+
+    >>> from root_numpy import array2tree
+    >>> import numpy as np
+    >>>
+    >>> a = np.array([(1, 2.5, 3.4),
+    ...               (4, 5, 6.8)],
+    ...              dtype=[('a', np.int32),
+    ...                     ('b', np.float32),
+    ...                     ('c', np.float64)])
+    >>> tree = array2tree(a)
+    >>> tree.Scan()
+    ************************************************
+    *    Row   *         a *         b *         c *
+    ************************************************
+    *        0 *         1 *       2.5 *       3.4 *
+    *        1 *         4 *         5 *       6.8 *
+    ************************************************
+
     """
     import ROOT
     if tree is not None:
@@ -374,9 +393,8 @@ def array2tree(arr, name='tree', tree=None):
 def array2root(arr, filename, treename='tree', mode='update'):
     """Convert a numpy array into a ROOT TTree and save it in a ROOT TFile.
 
-    .. warning::
-       This function is experimental. Please report problems.
-       Not all data types are supported (``np.object`` and ``np.float16``).
+    Fields of basic types, strings, and fixed-size subarrays of basic types are
+    supported. ``np.object`` and ``np.float16`` are currently not supported.
 
     Parameters
     ----------
@@ -396,5 +414,36 @@ def array2root(arr, filename, treename='tree', mode='update'):
     --------
     array2tree
 
+    Examples
+    --------
+
+    >>> from root_numpy import array2root, root2array
+    >>> import numpy as np
+    >>>
+    >>> a = np.array([(1, 2.5, 3.4),
+    ...               (4, 5, 6.8)],
+    ...              dtype=[('a', np.int32),
+    ...                     ('b', np.float32),
+    ...                     ('c', np.float64)])
+    >>> array2root(a, 'test.root', mode='recreate')
+    >>> root2array('test.root')
+    array([(1, 2.5, 3.4), (4, 5.0, 6.8)],
+          dtype=[('a', '<i4'), ('b', '<f4'), ('c', '<f8')])
+    >>>
+    >>> a = np.array(['', 'a', 'ab', 'abc', 'xyz', ''],
+    ...              dtype=[('string', 'S3')])
+    >>> array2root(a, 'test.root', mode='recreate')
+    >>> root2array('test.root')
+    array([('',), ('a',), ('ab',), ('abc',), ('xyz',), ('',)],
+          dtype=[('string', 'S3')])
+    >>>
+    >>> a = np.array([([1, 2, 3],),
+    ...               ([4, 5, 6],)],
+    ...              dtype=[('array', np.int32, (3,))])
+    >>> array2root(a, 'test.root', mode='recreate')
+    >>> root2array('test.root')
+    array([([1, 2, 3],), ([4, 5, 6],)],
+          dtype=[('array', '<i4', (3,))])
+
     """
     _librootnumpy.array2root(arr, filename, treename, mode)
diff --git a/root_numpy/info.py b/root_numpy/info.py
@@ -6,5 +6,5 @@
 |_|  \___/ \___/ \__|___|_| |_|\__,_|_| |_| |_| .__/ \__, |  {0}
                    |_____|                    |_|    |___/
 """
-__version__ = '4.2.1.dev0'
+__version__ = '4.3.0.dev0'
 __doc__ = __doc__.format(__version__)  # pylint:disable=redefined-builtin
diff --git a/root_numpy/src/Column.h b/root_numpy/src/Column.h
@@ -9,124 +9,125 @@
 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;
-
-        // Column name
-        std::string name;
-        // Name of the ROOT type
-        std::string type;
+
+    virtual ~Column() {}
+    virtual int GetLen() = 0;
+    virtual int GetCountLen() = 0;
+    virtual int GetSize() = 0;
+    virtual void* GetValuePointer() = 0;
+    virtual const char* GetTypeName() = 0;
+
+    // Column name
+    std::string name;
+    // Name of the ROOT type
+    std::string type;
 };
 
 
 class FormulaColumn: public Column
 {
     public:
 
-        FormulaColumn(std::string _name, TTreeFormula* _formula)
-        {
-            name = _name;
-            formula = _formula;
-            type = "Double_t";
-            value = new double[1];
-        }
-
-        ~FormulaColumn()
-        {
-            delete[] value;
-        }
-
-        int GetLen()
-        {
-            return 1;
-        }
-
-        int GetCountLen()
-        {
-            return 1;
-        }
-
-        int GetSize()
-        {
-            return sizeof(double) * GetLen();
-        }
-
-        void* GetValuePointer()
-        {
-            // required, as in TTreePlayer
-            formula->GetNdata();
-            value[0] = formula->EvalInstance(0);
-            return value;
-        }
-
-        const char* GetTypeName()
-        {
-            return "double";
-        }
-
-        TTreeFormula* formula;
-        double* value;
+    FormulaColumn(std::string _name, TTreeFormula* _formula)
+    {
+        name = _name;
+        formula = _formula;
+        type = "Double_t";
+        value = new double[1];
+    }
+
+    ~FormulaColumn()
+    {
+        delete[] value;
+    }
+
+    int GetLen()
+    {
+        return 1;
+    }
+
+    int GetCountLen()
+    {
+        return 1;
+    }
+
+    int GetSize()
+    {
+        return sizeof(double) * GetLen();
+    }
+
+    void* GetValuePointer()
+    {
+        // required, as in TTreePlayer
+        formula->GetNdata();
+        value[0] = formula->EvalInstance(0);
+        return value;
+    }
+
+    const char* GetTypeName()
+    {
+        return "double";
+    }
+
+    TTreeFormula* formula;
+    double* value;
 };
 
 
 class BranchColumn: public Column
 {
     public:
 
-        BranchColumn(std::string& _name, TLeaf* _leaf)
-        {
-            name = _name;
-            leaf = _leaf;
-            type = leaf->GetTypeName();
-        }
-
-        void SetLeaf(TLeaf* newleaf, bool check=false)
+    BranchColumn(std::string& _name, TLeaf* _leaf)
+    {
+        name = _name;
+        leaf = _leaf;
+        type = leaf->GetTypeName();
+    }
+
+    void SetLeaf(TLeaf* newleaf, bool check=false)
+    {
+        leaf = newleaf;
+        if (check)
         {
-            leaf = newleaf;
-            if (check)
-            {
-                assert(leaf->GetTypeName() == type);
-                // TODO: compare shape
-            }
+            assert(leaf->GetTypeName() == type);
+            // TODO: compare shape
         }
-
-        int GetLen()
+    }
+
+    int GetLen()
+    {
+        // get len of this block (in unit of element)
+        return leaf->GetLen();
+    }
+
+    int GetCountLen()
+    {
+        // get count leaf value
+        TLeaf* count_leaf = leaf->GetLeafCount();
+        if (count_leaf != NULL)
         {
-            // get len of this block (in unit of element)
-            return leaf->GetLen();
+            return int(count_leaf->GetValue());
         }
-
-        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
-            return leaf->GetLenType() * leaf->GetLen();
-        }
-
-        void* GetValuePointer()
-        {
-            return leaf->GetValuePointer();
-        }
-
-        const char* GetTypeName()
-        {
-            return leaf->GetTypeName();
-        }
-
-        TLeaf* leaf;
+        return 1;
+    }
+
+    int GetSize()
+    {
+        // get size of this block in bytes
+        return leaf->GetLenType() * leaf->GetLen();
+    }
+
+    void* GetValuePointer()
+    {
+        return leaf->GetValuePointer();
+    }
+
+    const char* GetTypeName()
+    {
+        return leaf->GetTypeName();
+    }
+
+    TLeaf* leaf;
 };
 #endif
diff --git a/root_numpy/src/ROOT.pxi b/root_numpy/src/ROOT.pxi
@@ -52,6 +52,8 @@ cdef extern from "TLeaf.h":
         TLeaf* GetLeafCount()
         TLeaf* GetLeafCounter(int&)
         TBranch* GetBranch()
+        int GetLen()
+        int GetLenStatic()
 
 cdef extern from "TTree.h":
     cdef cppclass TTree: