Skip to content
This repository was archived by the owner on Jan 7, 2023. It is now read-only.

Commit 377db79

Browse files
ndawendawe
committed
Merge pull request #206 from ndawe/master
array2tree: support strings and fixed-size subarrays
2 parents b4f9a19 + 19e488c commit 377db79

File tree

11 files changed

+13926
-11756
lines changed

11 files changed

+13926
-11756
lines changed

‎root_numpy/_tree.py‎

Lines changed: 55 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -335,9 +335,8 @@ def tree2rec(tree,
335335
def array2tree(arr, name='tree', tree=None):
336336
"""Convert a numpy structured array into a ROOT TTree.
337337
338-
.. warning::
339-
This function is experimental. Please report problems.
340-
Not all data types are supported (``np.object`` and ``np.float16``).
338+
Fields of basic types, strings, and fixed-size subarrays of basic types are
339+
supported. ``np.object`` and ``np.float16`` are currently not supported.
341340
342341
Parameters
343342
----------
@@ -359,6 +358,26 @@ def array2tree(arr, name='tree', tree=None):
359358
--------
360359
array2root
361360
361+
Examples
362+
--------
363+
364+
>>> from root_numpy import array2tree
365+
>>> import numpy as np
366+
>>>
367+
>>> a = np.array([(1, 2.5, 3.4),
368+
... (4, 5, 6.8)],
369+
... dtype=[('a', np.int32),
370+
... ('b', np.float32),
371+
... ('c', np.float64)])
372+
>>> tree = array2tree(a)
373+
>>> tree.Scan()
374+
************************************************
375+
* Row * a * b * c *
376+
************************************************
377+
* 0 * 1 * 2.5 * 3.4 *
378+
* 1 * 4 * 5 * 6.8 *
379+
************************************************
380+
362381
"""
363382
import ROOT
364383
if tree is not None:
@@ -374,9 +393,8 @@ def array2tree(arr, name='tree', tree=None):
374393
def array2root(arr, filename, treename='tree', mode='update'):
375394
"""Convert a numpy array into a ROOT TTree and save it in a ROOT TFile.
376395
377-
.. warning::
378-
This function is experimental. Please report problems.
379-
Not all data types are supported (``np.object`` and ``np.float16``).
396+
Fields of basic types, strings, and fixed-size subarrays of basic types are
397+
supported. ``np.object`` and ``np.float16`` are currently not supported.
380398
381399
Parameters
382400
----------
@@ -396,5 +414,36 @@ def array2root(arr, filename, treename='tree', mode='update'):
396414
--------
397415
array2tree
398416
417+
Examples
418+
--------
419+
420+
>>> from root_numpy import array2root, root2array
421+
>>> import numpy as np
422+
>>>
423+
>>> a = np.array([(1, 2.5, 3.4),
424+
... (4, 5, 6.8)],
425+
... dtype=[('a', np.int32),
426+
... ('b', np.float32),
427+
... ('c', np.float64)])
428+
>>> array2root(a, 'test.root', mode='recreate')
429+
>>> root2array('test.root')
430+
array([(1, 2.5, 3.4), (4, 5.0, 6.8)],
431+
dtype=[('a', '<i4'), ('b', '<f4'), ('c', '<f8')])
432+
>>>
433+
>>> a = np.array(['', 'a', 'ab', 'abc', 'xyz', ''],
434+
... dtype=[('string', 'S3')])
435+
>>> array2root(a, 'test.root', mode='recreate')
436+
>>> root2array('test.root')
437+
array([('',), ('a',), ('ab',), ('abc',), ('xyz',), ('',)],
438+
dtype=[('string', 'S3')])
439+
>>>
440+
>>> a = np.array([([1, 2, 3],),
441+
... ([4, 5, 6],)],
442+
... dtype=[('array', np.int32, (3,))])
443+
>>> array2root(a, 'test.root', mode='recreate')
444+
>>> root2array('test.root')
445+
array([([1, 2, 3],), ([4, 5, 6],)],
446+
dtype=[('array', '<i4', (3,))])
447+
399448
"""
400449
_librootnumpy.array2root(arr, filename, treename, mode)

‎root_numpy/info.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,5 +6,5 @@
66
|_| \___/ \___/ \__|___|_| |_|\__,_|_| |_| |_| .__/ \__, | {0}
77
|_____| |_| |___/
88
"""
9-
__version__ = '4.2.1.dev0'
9+
__version__ = '4.3.0.dev0'
1010
__doc__ = __doc__.format(__version__) # pylint:disable=redefined-builtin

‎root_numpy/src/Column.h‎

Lines changed: 102 additions & 101 deletions
Original file line numberDiff line numberDiff line change
@@ -9,124 +9,125 @@
99
class Column
1010
{
1111
public:
12-
virtual ~Column() {}
13-
virtual int GetLen() = 0;
14-
virtual int GetCountLen() = 0;
15-
virtual int GetSize() = 0;
16-
virtual void* GetValuePointer() = 0;
17-
virtual const char* GetTypeName() = 0;
18-
19-
// Column name
20-
std::string name;
21-
// Name of the ROOT type
22-
std::string type;
12+
13+
virtual ~Column() {}
14+
virtual int GetLen() = 0;
15+
virtual int GetCountLen() = 0;
16+
virtual int GetSize() = 0;
17+
virtual void* GetValuePointer() = 0;
18+
virtual const char* GetTypeName() = 0;
19+
20+
// Column name
21+
std::string name;
22+
// Name of the ROOT type
23+
std::string type;
2324
};
2425

2526

2627
class FormulaColumn: public Column
2728
{
2829
public:
2930

30-
FormulaColumn(std::string _name, TTreeFormula* _formula)
31-
{
32-
name = _name;
33-
formula = _formula;
34-
type = "Double_t";
35-
value = new double[1];
36-
}
37-
38-
~FormulaColumn()
39-
{
40-
delete[] value;
41-
}
42-
43-
int GetLen()
44-
{
45-
return 1;
46-
}
47-
48-
int GetCountLen()
49-
{
50-
return 1;
51-
}
52-
53-
int GetSize()
54-
{
55-
return sizeof(double) * GetLen();
56-
}
57-
58-
void* GetValuePointer()
59-
{
60-
// required, as in TTreePlayer
61-
formula->GetNdata();
62-
value[0] = formula->EvalInstance(0);
63-
return value;
64-
}
65-
66-
const char* GetTypeName()
67-
{
68-
return "double";
69-
}
70-
71-
TTreeFormula* formula;
72-
double* value;
31+
FormulaColumn(std::string _name, TTreeFormula* _formula)
32+
{
33+
name = _name;
34+
formula = _formula;
35+
type = "Double_t";
36+
value = new double[1];
37+
}
38+
39+
~FormulaColumn()
40+
{
41+
delete[] value;
42+
}
43+
44+
int GetLen()
45+
{
46+
return 1;
47+
}
48+
49+
int GetCountLen()
50+
{
51+
return 1;
52+
}
53+
54+
int GetSize()
55+
{
56+
return sizeof(double) * GetLen();
57+
}
58+
59+
void* GetValuePointer()
60+
{
61+
// required, as in TTreePlayer
62+
formula->GetNdata();
63+
value[0] = formula->EvalInstance(0);
64+
return value;
65+
}
66+
67+
const char* GetTypeName()
68+
{
69+
return "double";
70+
}
71+
72+
TTreeFormula* formula;
73+
double* value;
7374
};
7475

7576

7677
class BranchColumn: public Column
7778
{
7879
public:
7980

80-
BranchColumn(std::string& _name, TLeaf* _leaf)
81-
{
82-
name = _name;
83-
leaf = _leaf;
84-
type = leaf->GetTypeName();
85-
}
86-
87-
void SetLeaf(TLeaf* newleaf, bool check=false)
81+
BranchColumn(std::string& _name, TLeaf* _leaf)
82+
{
83+
name = _name;
84+
leaf = _leaf;
85+
type = leaf->GetTypeName();
86+
}
87+
88+
void SetLeaf(TLeaf* newleaf, bool check=false)
89+
{
90+
leaf = newleaf;
91+
if (check)
8892
{
89-
leaf = newleaf;
90-
if (check)
91-
{
92-
assert(leaf->GetTypeName() == type);
93-
// TODO: compare shape
94-
}
93+
assert(leaf->GetTypeName() == type);
94+
// TODO: compare shape
9595
}
96-
97-
int GetLen()
96+
}
97+
98+
int GetLen()
99+
{
100+
// get len of this block (in unit of element)
101+
return leaf->GetLen();
102+
}
103+
104+
int GetCountLen()
105+
{
106+
// get count leaf value
107+
TLeaf* count_leaf = leaf->GetLeafCount();
108+
if (count_leaf != NULL)
98109
{
99-
// get len of this block (in unit of element)
100-
return leaf->GetLen();
110+
return int(count_leaf->GetValue());
101111
}
102-
103-
int GetCountLen()
104-
{
105-
// get count leaf value
106-
TLeaf* count_leaf = leaf->GetLeafCount();
107-
if (count_leaf != NULL)
108-
{
109-
return int(count_leaf->GetValue());
110-
}
111-
return 1;
112-
}
113-
114-
int GetSize()
115-
{
116-
// get size of this block in bytes
117-
return leaf->GetLenType() * leaf->GetLen();
118-
}
119-
120-
void* GetValuePointer()
121-
{
122-
return leaf->GetValuePointer();
123-
}
124-
125-
const char* GetTypeName()
126-
{
127-
return leaf->GetTypeName();
128-
}
129-
130-
TLeaf* leaf;
112+
return 1;
113+
}
114+
115+
int GetSize()
116+
{
117+
// get size of this block in bytes
118+
return leaf->GetLenType() * leaf->GetLen();
119+
}
120+
121+
void* GetValuePointer()
122+
{
123+
return leaf->GetValuePointer();
124+
}
125+
126+
const char* GetTypeName()
127+
{
128+
return leaf->GetTypeName();
129+
}
130+
131+
TLeaf* leaf;
131132
};
132133
#endif

‎root_numpy/src/ROOT.pxi‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,8 @@ cdef extern from "TLeaf.h":
5252
TLeaf* GetLeafCount()
5353
TLeaf* GetLeafCounter(int&)
5454
TBranch* GetBranch()
55+
int GetLen()
56+
int GetLenStatic()
5557

5658
cdef extern from "TTree.h":
5759
cdef cppclass TTree:

0 commit comments

Comments
 (0)