ddale · Dec 16, 2015 · Mar 23, 2014 · Mar 24, 2014 · Mar 24, 2014 · Mar 24, 2014
diff --git a/quantities/quantity.py b/quantities/quantity.py
@@ -428,6 +428,15 @@ def sum(self, axis=None, dtype=None, out=None):
             copy=False
         )
 
+    @with_doc(np.nansum)
+    def nansum(self, axis=None, dtype=None, out=None):
+        import numpy as np
+        return Quantity(
+            np.nansum(self.magnitude, axis, dtype, out),
+            self.dimensionality,
+            copy=False
+        )
+
     @with_doc(np.ndarray.fill)
     def fill(self, value):
         self.magnitude.fill(value)
@@ -477,6 +486,14 @@ def max(self, axis=None, out=None):
             copy=False
         )
 
+    @with_doc(np.nanmax)
+    def nanmax(self, axis=None, out=None):
+        return Quantity(
+            np.nanmax(self.magnitude),
+            self.dimensionality,
+            copy=False
+        )
+
     @with_doc(np.ndarray.min)
     def min(self, axis=None, out=None):
         return Quantity(
@@ -485,10 +502,30 @@ def min(self, axis=None, out=None):
             copy=False
         )
 
+    @with_doc(np.nanmin)
+    def nanmin(self, axis=None, out=None):
+        return Quantity(
+            np.nanmin(self.magnitude),
+            self.dimensionality,
+            copy=False
+        )
+
     @with_doc(np.ndarray.argmin)
     def argmin(self,axis=None, out=None):
         return self.magnitude.argmin()
 
+    @with_doc(np.ndarray.argmax)
+    def argmax(self,axis=None, out=None):
+        return self.magnitude.argmax()
+
+    @with_doc(np.nanargmin)
+    def nanargmin(self,axis=None, out=None):
+        return np.nanargmin(self.magnitude)
+
+    @with_doc(np.nanargmax)
+    def nanargmax(self,axis=None, out=None):
+        return np.nanargmax(self.magnitude)
+
     @with_doc(np.ndarray.ptp)
     def ptp(self, axis=None, out=None):
         return Quantity(
@@ -549,6 +586,14 @@ def mean(self, axis=None, dtype=None, out=None):
             self.dimensionality,
             copy=False)
 
+    @with_doc(np.nanmean)
+    def nanmean(self, axis=None, dtype=None, out=None):
+        import numpy as np
+        return Quantity(
+            np.nanmean(self.magnitude, axis, dtype, out),
+            self.dimensionality,
+            copy=False)
+
     @with_doc(np.ndarray.var)
     def var(self, axis=None, dtype=None, out=None, ddof=0):
         return Quantity(
@@ -565,6 +610,14 @@ def std(self, axis=None, dtype=None, out=None, ddof=0):
             copy=False
         )
 
+    @with_doc(np.nanstd)
+    def nanstd(self, axis=None, dtype=None, out=None, ddof=0):
+        return Quantity(
+            np.nanstd(self.magnitude, axis, dtype, out, ddof),
+            self._dimensionality,
+            copy=False
+        )
+
     @with_doc(np.ndarray.prod)
     def prod(self, axis=None, dtype=None, out=None):
         if axis == None:

diff --git a/quantities/tests/test_methods.py b/quantities/tests/test_methods.py
@@ -2,7 +2,7 @@
 
 from .. import units as pq
 from .common import TestCase
-
+import numpy as np
 
 class TestQuantityMethods(TestCase):
 
@@ -17,6 +17,12 @@ def test_sum(self):
         self.assertQuantityEqual(self.q.sum(0), [4, 6]*pq.m)
         self.assertQuantityEqual(self.q.sum(1), [3, 7]*pq.m)
 
+    def test_nansum(self):
+        import numpy as np
+        qnan = [[1,2], [3,4], [np.nan,np.nan]] * pq.m
+        self.assertQuantityEqual(qnan.nansum(), 10*pq.m )
+        self.assertQuantityEqual(qnan.nansum(0), [4,6]*pq.m )
+
     def test_fill(self):
         self.q.fill(6 * pq.ft)
         self.assertQuantityEqual(self.q, [[6, 6], [6, 6]] * pq.ft)
@@ -103,15 +109,31 @@ def test_nonzero(self):
     def test_max(self):
         self.assertQuantityEqual(self.q.max(), 4*pq.m)
 
+    def test_nanmax(self):
+        q = np.append(self.q, np.nan) * self.q.units
+        self.assertQuantityEqual(q.nanmax(), 4*pq.m)
+
     def test_argmax(self):
         self.assertEqual(self.q.argmax(), 3)
 
+    def test_nanargmax(self):
+        q = np.append(self.q, np.nan) * self.q.units
+        self.assertEqual(self.q.nanargmax(), 3)
+
     def test_min(self):
         self.assertEqual(self.q.min(), 1 * pq.m)
 
+    def test_nanmin(self):
+        q = np.append(self.q, np.nan) * self.q.units
+        self.assertQuantityEqual(q.nanmin(), 1*pq.m)
+
     def test_argmin(self):
         self.assertEqual(self.q.argmin(), 0)
 
+    def test_nanargmax(self):
+        q = np.append(self.q, np.nan) * self.q.units
+        self.assertEqual(self.q.nanargmin(), 0)
+
     def test_ptp(self):
         self.assertQuantityEqual(self.q.ptp(), 3 * pq.m)
 
@@ -146,12 +168,23 @@ def test_cumsum(self):
     def test_mean(self):
         self.assertQuantityEqual(self.q.mean(), 2.5 * pq.m)
 
+    def test_nanmean(self):
+        import numpy as np    
+        q = [[1,2], [3,4], [np.nan,np.nan]] * pq.m
+        self.assertQuantityEqual(q.nanmean(), self.q.mean())
+
     def test_var(self):
         self.assertQuantityEqual(self.q.var(), 1.25*pq.m**2)
 
     def test_std(self):
         self.assertQuantityEqual(self.q.std(), 1.11803*pq.m, delta=1e-5)
 
+    def test_nanstd(self):
+        import numpy as np    
+        q0 = [[1,2], [3,4]] * pq.m
+        q1 = [[1,2], [3,4], [np.nan,np.nan]] * pq.m
+        self.assertQuantityEqual(q0.std(), q1.nanstd())
+
     def test_prod(self):
         self.assertQuantityEqual(self.q.prod(), 24 * pq.m**4)
 

diff --git a/quantities/tests/test_uncertainty.py b/quantities/tests/test_uncertainty.py
@@ -3,6 +3,7 @@
 from .. import units as pq
 from ..uncertainquantity import UncertainQuantity
 from .common import TestCase
+import numpy as np
 
 
 class TestUncertainty(TestCase):
@@ -60,3 +61,59 @@ def test_uncertainquantity_divide(self):
         self.assertQuantityEqual((a/2).uncertainty, [0.05, 0.1 ]*pq.m)
         self.assertQuantityEqual(1/a, [1., 0.5]/pq.m)
         self.assertQuantityEqual((1/a).uncertainty, [0.1, 0.05]/pq.m)
+
+    def test_uncertaintity_mean(self):
+        a = UncertainQuantity([1,2], 'm', [.1,.2])
+        mean0 = np.sum(a)/np.size(a) # calculated traditionally
+        mean1 = a.mean()        # calculated using this code
+        self.assertQuantityEqual(mean0, mean1)
+
+    def test_uncertaintity_nanmean(self):
+        a = UncertainQuantity([1,2], 'm', [.1,.2])
+        b = UncertainQuantity([1,2,np.nan], 'm', [.1,.2,np.nan])
+        self.assertQuantityEqual(a.mean(),b.nanmean())
+
+    def test_uncertainquantity_subtract(self):
+        a = UncertainQuantity(1, 'm', 1)
+        b = a.copy() # different object
+        self.assertQuantityEqual(a-a, UncertainQuantity(0, 'm', 0))
+        self.assertQuantityEqual(a-b, UncertainQuantity(0, 'm', np.sqrt(2)))
+
+    def test_uncertainty_sqrt(self):
+        a = UncertainQuantity([1,2], 'm', [.1,.2])
+        self.assertQuantityEqual(a**0.5, a.sqrt())
+
+    def test_uncertainty_nansum(self):
+        uq = UncertainQuantity([1,2], 'm', [1,1])
+        uq_nan = UncertainQuantity([1,2,np.nan], 'm', [1,1,np.nan])
+        self.assertQuantityEqual(np.sum(uq), np.nansum(uq))
+        self.assertQuantityEqual(np.sum(uq), np.nansum(uq_nan))
+
+    def test_uncertainty_minmax_nan_arg(self):
+        q = [[1, 2], [3, 4]] * pq.m        # quantity
+        self.assertQuantityEqual(q.min(), 1*pq.m) # min
+        self.assertQuantityEqual(q.max(), 4*pq.m) # max
+        self.assertQuantityEqual(q.argmin(), 0) # argmin
+        self.assertQuantityEqual(q.argmax(), 3) # argmax
+        # uncertain quantity
+        uq = UncertainQuantity([[1,2],[3,4]], pq.m, [[1,1],[1,1]])
+        self.assertQuantityEqual(uq.min(), 1*pq.m) # min
+        self.assertQuantityEqual(uq.max(), 4*pq.m) # max
+        self.assertQuantityEqual(uq.argmin(), 0) # argmin
+        self.assertQuantityEqual(uq.argmax(), 3) # argmax
+        # now repeat the above with NaNs
+        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(), 4*pq.m) # max
+        self.assertQuantityEqual(nanq.nanargmin(), 0) # argmin
+        self.assertQuantityEqual(nanq.nanargmax(), 3) # argmax
+        self.assertQuantityEqual(nanuq.nanmin(), 1*pq.m) # min
+        self.assertQuantityEqual(nanuq.nanmax(), 4*pq.m) # max
+        self.assertQuantityEqual(nanuq.nanargmin(), 0) # argmin
+        self.assertQuantityEqual(nanuq.nanargmax(), 3) # argmax
+
+
+
diff --git a/quantities/uncertainquantity.py b/quantities/uncertainquantity.py
@@ -102,7 +102,10 @@ def __radd__(self, other):
     @scale_other_units
     def __sub__(self, other):
         res = super(UncertainQuantity, self).__sub__(other)
-        u = (self.uncertainty**2+other.uncertainty**2)**0.5
+        if self is not other:
+            u = (self.uncertainty**2+other.uncertainty**2)**0.5
+        else:
+            u = self.uncertainty*0
         return UncertainQuantity(res, uncertainty=u, copy=False)
 
     @with_doc(Quantity.__rsub__, use_header=False)
@@ -207,6 +210,73 @@ def sum(self, axis=None, dtype=None, out=None):
             copy=False
         )
 
+    @with_doc(np.nansum)
+    def nansum(self, axis=None, dtype=None, out=None):
+        import numpy as np
+        return UncertainQuantity(
+            np.nansum(self.magnitude, axis, dtype, out),
+            self.dimensionality,
+            (np.nansum(self.uncertainty.magnitude**2, axis))**0.5,
+            copy=False
+        )
+
+    @with_doc(np.ndarray.mean)
+    def mean(self, axis=None, dtype=None, out=None):
+        return UncertainQuantity(
+            self.magnitude.mean(axis, dtype, out),
+            self.dimensionality,
+            ((1.0/np.size(self,axis))**2 * np.sum(self.uncertainty.magnitude**2, axis))**0.5,
+            copy=False)
+
+    @with_doc(np.nanmean)
+    def nanmean(self, axis=None, dtype=None, out=None):
+        size = np.sum(~np.isnan(self),axis)
+        return UncertainQuantity(
+            np.nanmean(self.magnitude, axis, dtype, out),
+            self.dimensionality,
+            ((1.0/size)**2 * np.nansum(np.nan_to_num(self.uncertainty.magnitude)**2, axis))**0.5,
+            copy=False)
+
+    @with_doc(np.sqrt)
+    def sqrt(self, out=None):
+        return self**0.5
+
+    @with_doc(np.ndarray.max)
+    def max(self, axis=None, out=None):
+        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.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.unravel_index(np.argmin(self.magnitude), self.shape)
+        return self[idx]
+
+    @with_doc(np.nanmin)
+    def nanmin(self, axis=None, out=None):
+        idx = np.unravel_index(np.nanargmin(self.magnitude), self.shape)
+        return self[idx]
+
+    @with_doc(np.ndarray.argmin)
+    def argmin(self,axis=None, out=None):
+        return self.magnitude.argmin()
+
+    @with_doc(np.ndarray.argmax)
+    def argmax(self,axis=None, out=None):
+        return self.magnitude.argmax()
+
+    @with_doc(np.nanargmin)
+    def nanargmin(self,axis=None, out=None):
+        return np.nanargmin(self.magnitude)
+
+    @with_doc(np.nanargmax)
+    def nanargmax(self,axis=None, out=None):
+        return np.nanargmax(self.magnitude)
+
     def __getstate__(self):
         """
         Return the internal state of the quantity, for pickling