add histogram2d to cythonfunctions

tests included

postpic/cythonfunctions.pyx

@@ -71,3 +71,57 @@ def histogram(np.ndarray[np.double_t, ndim=1] data, range=None, int bins=20,
                     if (xr > 1.0):
                         ret[xr - 1] += (0.5 - x + xr) * weights[i]
     return ret, bin_edges
+
+
+def histogram2d(np.ndarray[np.double_t, ndim=1] datax, np.ndarray[np.double_t, ndim=1] datay,
+                np.ndarray[np.double_t, ndim=1] weights=None,
+                range=None, bins=(20, 20), int order=0):
+    '''
+    Mimics numpy.histogram2d.
+    Additional Arguments:
+        - order = 0:
+            sets the order of the particle shapes.
+            order = 0 returns a normal histogram.
+            order = 1 uses top hat particle shape.
+    '''
+    cdef int n = len(datax)
+    if n != len(datay):
+            raise ValueError('datax and datay must be of equal length')
+    cdef np.ndarray[np.double_t, ndim=2] ret = np.zeros(bins, dtype=np.double)
+    cdef double xmin, xmax, ymin, ymax
+    if range is None:
+        xmin = np.min(datax)
+        xmax = np.max(datax)
+        ymin = np.min(datay)
+        ymax = np.max(datay)
+    else:
+        xmin = range[0][0]
+        xmax = range[0][1]
+        ymin = range[1][0]
+        ymax = range[1][1]
+    cdef int xbins = bins[0]
+    cdef int ybins = bins[1]
+    xedges = np.linspace(xmin, xmax, xbins+1)
+    yedges = np.linspace(ymin, ymax, ybins+1)
+    cdef double xtmp = 1.0 / (xmax - xmin) * xbins
+    cdef double ytmp = 1.0 / (ymax - ymin) * ybins
+    cdef double x
+    cdef double y
+    if order == 0:
+        # normal Histogram
+        for i in xrange(n):
+            x = (datax[i] - xmin) * xtmp
+            y = (datay[i] - ymin) * ytmp
+            if x > 0.0 and y > 0.0 and x < xbins and y < ybins:
+                if weights is None:
+                    ret[<int>x, <int>y] += 1.0
+                else:
+                    ret[<int>x, <int>y] += weights[i]
+    elif order == 1:
+        pass
+
+    return ret, xedges, yedges
+
+
+
+

test/test_cythonfunctions.py

@@ -22,7 +22,7 @@
 import postpic.cythonfunctions as cf
 import numpy as np
 
-class TestCythonfunctions(unittest.TestCase):
+class TestHistogram(unittest.TestCase):
 
     def setUp(self):
         self.data = np.random.random(1e3)
@@ -61,5 +61,37 @@ def test_histogramo1w(self):
         self.assertAlmostEqual(diff, 0)
         self.assertListEqual(list(npe), list(cfe))
 
+
+class TestHistogram2d(unittest.TestCase):
+
+    def setUp(self):
+        self.datax = np.random.random(1e3)
+        self.datay = 2 * np.random.random(1e3)
+        self.weights = np.random.random(1e3)
+
+    def test_histogram2d(self):
+        # in this special case, cf.histogram2d and np.histogram2d must yield equal results
+        # check hist and edges
+        cfh, cfex, cfey = cf.histogram2d(self.datax, self.datay,
+                                         bins=(20,30), range=((0,1), (0,2)), order=0)
+        nph, npex, npey = np.histogram2d(self.datax, self.datay,
+                                          bins=(20, 30), range=((0,1), (0,2)))
+        self.assertAlmostEqual(0.0, np.sum(np.abs(nph - cfh)))
+        self.assertListEqual(list(npex), list(cfex))
+        self.assertListEqual(list(npey), list(cfey))
+
+    def test_histogram2dw(self):
+        # in this special case, cf.histogram2d and np.histogram2d must yield equal results
+        # check hist and edges
+        cfh, cfex, cfey = cf.histogram2d(self.datax, self.datay, weights=self.weights,
+                                         bins=(20, 30), range=((0,1), (0,2)), order=0)
+        nph, npex, npey = np.histogram2d(self.datax, self.datay, weights=self.weights,
+                                         bins=(20, 30), range=((0,1), (0,2)))
+        self.assertAlmostEqual(0.0, np.sum(np.abs(nph - cfh)))
+        self.assertListEqual(list(npex), list(cfex))
+        self.assertListEqual(list(npey), list(cfey))
+
+
+
 if __name__ == '__main__':
     unittest.main()