Merge pull request #144 from ndawe/master

[MRG] New general evaluate() function

docs/reference/index.rst

@@ -26,6 +26,7 @@ what they are and what they do.
    fill_hist
    fill_graph
    random_sample
+   evaluate
    list_trees
    list_branches
    list_structures

root_numpy/__init__.py

@@ -8,6 +8,7 @@
 from _sample import random_sample
 from _array import array
 from _matrix import matrix
+from _evaluate import evaluate
 from _warnings import RootNumpyWarning, RootNumpyUnconvertibleWarning
 from _utils import (
     stretch, blockwise_inner_join,
@@ -26,6 +27,7 @@
     'random_sample',
     'array',
     'matrix',
+    'evaluate',
     'list_trees',
     'list_branches',
     'list_structures',

root_numpy/_evaluate.py

@@ -0,0 +1,99 @@
+import numpy as np
+import _librootnumpy
+
+
+__all__ = [
+    'evaluate',
+]
+
+
+def evaluate(root_object, array):
+    """
+    Evaluate a ROOT histogram, function, graph, or spline at all values
+    in a NumPy array and return the resulting array.
+
+    Parameters
+    ----------
+    root_object : TH[1|2|3], TF[1|2|3], TGraph, TSpline
+        A ROOT histogram, function, graph, or spline
+    array : ndarray
+        An array containing the values to evaluate the ROOT object on.
+        The shape must match the dimensionality of the ROOT object.
+
+    Returns
+    -------
+    y : array
+        An array containing the values of the ROOT object evaluated at each
+        value in the input array.
+
+    Raises
+    ------
+    TypeError
+        If the ROOT object is not a histogram, function, graph, or spline
+    ValueError
+        If the shape of the array is not compatible with the dimensionality
+        of the ROOT object being evaluated.
+
+    Examples
+    --------
+    >>> from root_numpy import evaluate
+    >>> from ROOT import TF1, TF2
+    >>> func = TF1("f1", "x*x")
+    >>> evaluate(func, [1, 2, 3, 4])
+    array([  1.,   4.,   9.,  16.])
+    >>> func = TF2("f2", "x*y")
+    >>> evaluate(func, [[1, 1], [1, 2], [3, 1]])
+    array([ 1.,  2.,  3.])
+    """
+    import ROOT
+    array = np.asarray(array, dtype=np.double)
+    if isinstance(root_object, ROOT.TH1):
+        if isinstance(root_object, ROOT.TH3):
+            if array.ndim != 2:
+                raise ValueError("array must be 2-dimensional")
+            if array.shape[1] != 3:
+                raise ValueError(
+                    "length of the second dimension must equal "
+                    "the dimension of the histogram")
+            return _librootnumpy.evaluate_h3(ROOT.AsCObject(root_object), array)
+        elif isinstance(root_object, ROOT.TH2):
+            if array.ndim != 2:
+                raise ValueError("array must be 2-dimensional")
+            if array.shape[1] != 2:
+                raise ValueError(
+                    "length of the second dimension must equal "
+                    "the dimension of the histogram")
+            return _librootnumpy.evaluate_h2(ROOT.AsCObject(root_object), array)
+        if array.ndim != 1:
+            raise ValueError("array must be 1-dimensional")
+        return _librootnumpy.evaluate_h1(ROOT.AsCObject(root_object), array)
+    elif isinstance(root_object, ROOT.TF1):
+        if isinstance(root_object, ROOT.TF3):
+            if array.ndim != 2:
+                raise ValueError("array must be 2-dimensional")
+            if array.shape[1] != 3:
+                raise ValueError(
+                    "length of the second dimension must equal "
+                    "the dimension of the function")
+            return _librootnumpy.evaluate_f3(ROOT.AsCObject(root_object), array)
+        elif isinstance(root_object, ROOT.TF2):
+            if array.ndim != 2:
+                raise ValueError("array must be 2-dimensional")
+            if array.shape[1] != 2:
+                raise ValueError(
+                    "length of the second dimension must equal "
+                    "the dimension of the function")
+            return _librootnumpy.evaluate_f2(ROOT.AsCObject(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, ROOT.TGraph):
+        if array.ndim != 1:
+            raise ValueError("array must be 1-dimensional")
+        return _librootnumpy.evaluate_graph(ROOT.AsCObject(root_object), array)
+    elif isinstance(root_object, ROOT.TSpline):
+        if array.ndim != 1:
+            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_numpy/_hist.py

@@ -68,6 +68,5 @@ def fill_hist(hist, array, weights=None, return_indices=False):
             raise ValueError("array must be 1-dimensional")
         return _librootnumpy.fill_h1(
             ROOT.AsCObject(hist), array, weights, return_indices)
-    else:
-        raise TypeError(
-            "hist must be an instance of ROOT.TH1, ROOT.TH2, or ROOT.TH3")
+    raise TypeError(
+        "hist must be an instance of ROOT.TH1, ROOT.TH2, or ROOT.TH3")

root_numpy/info.py

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

root_numpy/src/ROOT.pxi

@@ -85,43 +85,57 @@ cdef extern from "TClassEdit.h" namespace "TClassEdit":
 cdef extern from "TF1.h":
     cdef cppclass TF1:
         double GetRandom()
+        double Eval(double x)
 
 cdef extern from "TF2.h":
     cdef cppclass TF2:
         double GetRandom2(double& x, double& y)
+        double Eval(double x, double y)
 
 cdef extern from "TF3.h":
     cdef cppclass TF3:
         double GetRandom3(double& x, double& y, double& z)
+        double Eval(double x, double y, double z)
 
 cdef extern from "TH1.h":
     cdef cppclass TH1:
         double GetRandom()
         int Fill(double x)
         int Fill(double x, double w)
+        int FindBin(double x)
+        double GetBinContent(int bin)
 
 cdef extern from "TH2.h":
     cdef cppclass TH2:
         double GetRandom2(double& x, double& y)
         int Fill(double x, double y)
         int Fill(double x, double y, double w)
+        int FindBin(double x, double y)
+        double GetBinContent(int bin)
 
 cdef extern from "TH3.h":
     cdef cppclass TH3:
         double GetRandom3(double& x, double& y, double& z)
         int Fill(double x, double y, double z)
         int Fill(double x, double y, double z, double w)
+        int FindBin(double x, double y, double z)
+        double GetBinContent(int bin)
 
 cdef extern from "TGraph.h":
     cdef cppclass TGraph:
         void Set(int n)
         void SetPoint(int i, double x, double y)
+        double Eval(double x)
 
 cdef extern from "TGraph2D.h":
     cdef cppclass TGraph2D:
         void Set(int n)
         void SetPoint(int i, double x, double y, double z)
 
+cdef extern from "TSpline.h":
+    cdef cppclass TSpline:
+        double Eval(double x)
+
 cdef extern from "TArrayD.h":
     cdef cppclass TArrayD:
         int GetSize()