Merge pull request #147 from ndawe/master

[MRG] evaluate: handle TFormula and string input

root_numpy/_evaluate.py

@@ -1,3 +1,4 @@
+import uuid
 import numpy as np
 import _librootnumpy
 
@@ -12,8 +13,9 @@ def evaluate(root_object, array):
 
     Parameters
     ----------
-    root_object : TH[1|2|3], TF[1|2|3], TGraph, TSpline
-        A ROOT histogram, function, graph, or spline
+    root_object : TH[1|2|3], TF[1|2|3], TFormula, TGraph, TSpline, or string
+        A ROOT histogram, function, formula, graph, spline, or string.
+        If a string is specified, a TFormula is created.
     array : ndarray
         An array containing the values to evaluate the ROOT object on.
         The shape must match the dimensionality of the ROOT object.
@@ -31,6 +33,8 @@ def evaluate(root_object, array):
     ValueError
         If the shape of the array is not compatible with the dimensionality
         of the ROOT object being evaluated.
+        If the string expression does not compile to a valid TFormula
+        expression.
 
     Examples
     --------
@@ -42,6 +46,8 @@ def evaluate(root_object, array):
     >>> func = TF2("f2", "x*y")
     >>> evaluate(func, [[1, 1], [1, 2], [3, 1]])
     array([ 1.,  2.,  3.])
+    >>> evaluate("x*y", [[1, 1], [1, 2], [3, 1]])
+    array([ 1.,  2.,  3.])
 
     """
     import ROOT
@@ -86,6 +92,33 @@ def evaluate(root_object, array):
         if array.ndim != 1:
             raise ValueError("array must be 1-dimensional")
         return _librootnumpy.evaluate_f1(ROOT.AsCObject(root_object), array)
+    elif isinstance(root_object, (basestring, ROOT.TFormula)):
+        if isinstance(root_object, basestring):
+            # attempt to create a formula
+            root_object = ROOT.TFormula(uuid.uuid4().hex, root_object)
+        ndim = root_object.GetNdim()
+        if ndim == 0:
+            raise ValueError("invalid formula expression")
+        if ndim == 1:
+            if array.ndim != 1:
+                raise ValueError("array must be 1-dimensional")
+            return _librootnumpy.evaluate_f1(
+                ROOT.AsCObject(root_object), array)
+        if array.ndim != 2:
+            raise ValueError("array must be 2-dimensional")
+        if array.shape[1] != ndim:
+            raise ValueError(
+                "length of the second dimension must equal "
+                "the dimension of the function")
+        if ndim == 2:
+            return _librootnumpy.evaluate_f2(
+                ROOT.AsCObject(root_object), array)
+        elif ndim == 3:
+            return _librootnumpy.evaluate_f3(
+                ROOT.AsCObject(root_object), array)
+        # 4d
+        return _librootnumpy.evaluate_f4(
+            ROOT.AsCObject(root_object), array)
     elif isinstance(root_object, ROOT.TGraph):
         if array.ndim != 1:
             raise ValueError("array must be 1-dimensional")
@@ -95,4 +128,5 @@ def evaluate(root_object, array):
             raise ValueError("array must be 1-dimensional")
         return _librootnumpy.evaluate_spline(ROOT.AsCObject(root_object), array)
     raise TypeError(
-        "root_object is not a ROOT histogram, function, graph, or spline")
+        "root_object is not a ROOT histogram, function, formula, "
+        "graph, spline or string")

root_numpy/info.py

@@ -6,5 +6,5 @@
 |_|  \___/ \___/ \__|___|_| |_|\__,_|_| |_| |_| .__/ \__, |  {0}
                    |_____|                    |_|    |___/
 """
-__version__ = '3.2.1.dev'
+__version__ = '3.3.0.dev'
 __doc__ = __doc__.format(__version__)

root_numpy/src/ROOT.pxi

@@ -72,6 +72,13 @@ cdef extern from "TList.h":
         TObject* At(int idx)
         int GetEntries()
 
+cdef extern from "TFormula.h":
+    cdef cppclass TFormula:
+        double Eval(double x)
+        double Eval(double x, double y)
+        double Eval(double x, double y, double z)
+        double Eval(double x, double y, double z, double t)
+
 cdef extern from "TTreeFormula.h":
     cdef cppclass TTreeFormula:
         TTreeFormula(const_char*, const_char*, TTree*)