Add histogramnd and fix histogram2d

2 files changed, 50 insertions, 6 deletions

diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py
index b548ce386..a8ccf5eaa 100644
--- a/numpy/lib/tests/test_function_base.py
+++ b/numpy/lib/tests/test_function_base.py
@@ -353,6 +353,31 @@ class test_histogram(NumpyTestCase):
         (a,b)=histogram(linspace(0,10,100))
         assert(all(a==10))
 
+class test_histogramnd(NumpyTestCase):
+    def check_simple(self):
+        x = array([[-.5, .5, 1.5], [-.5, 1.5, 2.5], [-.5, 2.5, .5], \
+        [.5, .5, 1.5], [.5, 1.5, 2.5], [.5, 2.5, 2.5]])
+        H, edges = histogramnd(x, (2,3,3), range = [[-1,1], [0,3], [0,3]])
+        answer = asarray([[[0,1,0], [0,0,1], [1,0,0]], [[0,1,0], [0,0,1], [0,0,1]]])
+        assert(all(H == answer))
+        # Check normalization
+        ed = [[-2,0,2], [0,1,2,3], [0,1,2,3]]
+        H, edges = histogramnd(x, bins = ed, normed = True)
+        assert(all(H == answer/12.))
+        # Check that H has the correct shape.
+        H, edges = histogramnd(x, (2,3,4), range = [[-1,1], [0,3], [0,4]], normed=True)
+        answer = asarray([[[0,1,0,0], [0,0,1,0], [1,0,0,0]], [[0,1,0,0], [0,0,1,0], [0,0,1,0]]])
+        assert_array_almost_equal(H, answer/6., 4)
+        # Check that a sequence of arrays is accepted and H has the correct shape.
+        z = [squeeze(y) for y in split(x,3,axis=1)]
+        H, edges = histogramnd(z, bins=(4,3,2),range=[[-2,2], [0,3], [0,2]])
+        answer = asarray([[[0,0],[0,0],[0,0]], 
+                          [[0,1], [0,0], [1,0]], 
+                          [[0,1], [0,0],[0,0]], 
+                          [[0,0],[0,0],[0,0]]])
+        assert_array_equal(H, answer)
+                
+
 class test_unique(NumpyTestCase):
     def check_simple(self):
         x = array([4,3,2,1,1,2,3,4, 0])
diff --git a/numpy/lib/tests/test_twodim_base.py b/numpy/lib/tests/test_twodim_base.py
index cfe51bb9c..3c6edfd24 100644
--- a/numpy/lib/tests/test_twodim_base.py
+++ b/numpy/lib/tests/test_twodim_base.py
@@ -5,7 +5,8 @@
 from numpy.testing import *
 set_package_path()
 from numpy import arange, rot90, add, fliplr, flipud, zeros, ones, eye, \
-     array, diag
+     array, diag, histogram2d
+import numpy as np
 restore_path()
 
 ##################################################
@@ -133,13 +134,12 @@ class test_rot90(NumpyTestCase):
 
 class test_histogram2d(NumpyTestCase):
     def check_simple(self):
-        import numpy as np
         x = array([ 0.41702200,  0.72032449,  0.00011437481, 0.302332573,  0.146755891])
         y = array([ 0.09233859,  0.18626021,  0.34556073,  0.39676747,  0.53881673])
         xedges = np.linspace(0,1,10)
         yedges = np.linspace(0,1,10)
-        H = np.histogram2d(x, y, (xedges, yedges))[0]
-        answer = np.array([[0, 0, 0, 1, 0, 0, 0, 0, 0],
+        H = histogram2d(x, y, (xedges, yedges))[0]
+        answer = array([[0, 0, 0, 1, 0, 0, 0, 0, 0],
                            [0, 0, 0, 0, 0, 0, 1, 0, 0],
                            [0, 0, 0, 0, 0, 0, 0, 0, 0],
                            [1, 0, 1, 0, 0, 0, 0, 0, 0],
@@ -148,7 +148,26 @@ class test_histogram2d(NumpyTestCase):
                            [0, 0, 0, 0, 0, 0, 0, 0, 0],
                            [0, 0, 0, 0, 0, 0, 0, 0, 0],
                            [0, 0, 0, 0, 0, 0, 0, 0, 0]])
-        assert_equal(H, answer)
-
+        assert_array_equal(H.T, answer)
+    def check_asym(self):
+        x = array([1, 1, 2, 3, 4, 4, 4, 5])
+        y = array([1, 3, 2, 0, 1, 2, 3, 4])
+        H, xed, yed = histogram2d(x,y, (6, 5), range = [[0,6],[0,5]], normed=True)
+        answer = array([[0.,0,0,0,0],
+                        [0,1,0,1,0],
+                        [0,0,1,0,0],
+                        [1,0,0,0,0],
+                        [0,1,1,1,0],
+                        [0,0,0,0,1]])
+        assert_array_almost_equal(H, answer/8., 3)
+    def check_norm(self):
+        x = array([1,2,3,1,2,3,1,2,3])
+        y = array([1,1,1,2,2,2,3,3,3]) 
+        H, xed, yed = histogram2d(x,y,[[1,2,3,5], [1,2,3,5]], normed=True)
+        answer=array([[1,1,.5],
+                     [1,1,.5],
+                     [.5,.5,.25]])/9.
+        assert_array_almost_equal(H, answer, 3)
+        
 if __name__ == "__main__":
     NumpyTest().run()