previous commit (min,max,nan..) works w/ 2D arrays

mankoff · mankoff · commit 6b4d1692530a · 2014-03-24T14:34:33.000-04:00
diff --git a/quantities/tests/test_uncertainty.py b/quantities/tests/test_uncertainty.py
@@ -90,29 +90,30 @@ def test_uncertainty_nansum(self):
         self.assertQuantityEqual(np.sum(uq), np.nansum(uq_nan))
 
     def test_uncertainty_minmax_nan_arg(self):
-        q = [1,2] * pq.m
-        # quantity
+        q = [[1, 2], [3, 4]] * pq.m        # quantity
         self.assertQuantityEqual(q.min(), 1*pq.m) # min
-        self.assertQuantityEqual(q.max(), 2*pq.m) # max
+        self.assertQuantityEqual(q.max(), 4*pq.m) # max
         self.assertQuantityEqual(q.argmin(), 0) # argmin
-        self.assertQuantityEqual(q.argmax(), 1) # argmax
+        self.assertQuantityEqual(q.argmax(), 3) # argmax
         # uncertain quantity
-        uq = UncertainQuantity([1,2], pq.m, [1,1])
+        uq = UncertainQuantity([[1,2],[3,4]], pq.m, [[1,1],[1,1]])
         self.assertQuantityEqual(uq.min(), 1*pq.m) # min
-        self.assertQuantityEqual(uq.max(), 2*pq.m) # max
+        self.assertQuantityEqual(uq.max(), 4*pq.m) # max
         self.assertQuantityEqual(uq.argmin(), 0) # argmin
-        self.assertQuantityEqual(uq.argmax(), 1) # argmax
+        self.assertQuantityEqual(uq.argmax(), 3) # argmax
         # now repeat the above with NaNs
-        nanq = [1,2, np.nan] * pq.m
-        nanuq = UncertainQuantity([1,2, np.nan], pq.m, [1,1,np.nan])
+        nanq = [[1, 2], [3, 4], [np.nan,np.nan]] * pq.m        # quantity
+        nanuq = UncertainQuantity([[1,2],[3,4],[np.nan,np.nan]],
+                               pq.m,
+                               [[1,1],[1,1],[np.nan,np.nan]])
         self.assertQuantityEqual(nanq.nanmin(), 1*pq.m) # min
-        self.assertQuantityEqual(nanq.nanmax(), 2*pq.m) # max
+        self.assertQuantityEqual(nanq.nanmax(), 4*pq.m) # max
         self.assertQuantityEqual(nanq.nanargmin(), 0) # argmin
-        self.assertQuantityEqual(nanq.nanargmax(), 1) # argmax
+        self.assertQuantityEqual(nanq.nanargmax(), 3) # argmax
         self.assertQuantityEqual(nanuq.nanmin(), 1*pq.m) # min
-        self.assertQuantityEqual(nanuq.nanmax(), 2*pq.m) # max
+        self.assertQuantityEqual(nanuq.nanmax(), 4*pq.m) # max
         self.assertQuantityEqual(nanuq.nanargmin(), 0) # argmin
-        self.assertQuantityEqual(nanuq.nanargmax(), 1) # argmax
+        self.assertQuantityEqual(nanuq.nanargmax(), 3) # argmax
         
             
     
diff --git a/quantities/uncertainquantity.py b/quantities/uncertainquantity.py
@@ -243,22 +243,22 @@ def sqrt(self, out=None):
 
     @with_doc(np.ndarray.max)
     def max(self, axis=None, out=None):
-        idx = np.argmax(self.magnitude)
+        idx = np.unravel_index(np.argmax(self.magnitude), self.shape)
         return self[idx]
 
     @with_doc(np.nanmax)
     def nanmax(self, axis=None, out=None):
-        idx = np.nanargmax(self.magnitude)
+        idx = np.unravel_index(np.nanargmax(self.magnitude), self.shape)
         return self[idx]
 
     @with_doc(np.ndarray.min)
     def min(self, axis=None, out=None):
-        idx = np.argmin(self.magnitude)
+        idx = np.unravel_index(np.argmin(self.magnitude), self.shape)
         return self[idx]
 
     @with_doc(np.nanmin)
     def nanmin(self, axis=None, out=None):
-        idx = np.nanargmin(self.magnitude)
+        idx = np.unravel_index(np.nanargmin(self.magnitude), self.shape)
         return self[idx]
 
     @with_doc(np.ndarray.argmin)
