Skip to content

Commit 7bd93c7

Browse files
committed
nansum for quantities and uncertain quantities
1 parent de880b7 commit 7bd93c7

4 files changed

Lines changed: 32 additions & 3 deletions

File tree

‎quantities/quantity.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -428,6 +428,15 @@ def sum(self, axis=None, dtype=None, out=None):
428428
copy=False
429429
)
430430

431+
@with_doc(np.nansum)
432+
def nansum(self, axis=None, dtype=None, out=None):
433+
import numpy as np
434+
return Quantity(
435+
np.nansum(self.magnitude, axis, dtype, out),
436+
self.dimensionality,
437+
copy=False
438+
)
439+
431440
@with_doc(np.ndarray.fill)
432441
def fill(self, value):
433442
self.magnitude.fill(value)

‎quantities/tests/test_methods.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,12 @@ def test_sum(self):
1717
self.assertQuantityEqual(self.q.sum(0), [4, 6]*pq.m)
1818
self.assertQuantityEqual(self.q.sum(1), [3, 7]*pq.m)
1919

20+
def test_nansum(self):
21+
import numpy as np
22+
qnan = [[1,2], [3,4], [np.nan,np.nan]] * pq.m
23+
self.assertQuantityEqual(qnan.nansum(), 10*pq.m )
24+
self.assertQuantityEqual(qnan.nansum(0), [4,6]*pq.m )
25+
2026
def test_fill(self):
2127
self.q.fill(6 * pq.ft)
2228
self.assertQuantityEqual(self.q, [[6, 6], [6, 6]] * pq.ft)

‎quantities/tests/test_uncertainty.py‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
from .. import units as pq
44
from ..uncertainquantity import UncertainQuantity
55
from .common import TestCase
6+
import numpy as np
67

78

89
class TestUncertainty(TestCase):
@@ -62,20 +63,17 @@ def test_uncertainquantity_divide(self):
6263
self.assertQuantityEqual((1/a).uncertainty, [0.1, 0.05]/pq.m)
6364

6465
def test_uncertaintity_mean(self):
65-
import numpy as np
6666
a = UncertainQuantity([1,2], 'm', [.1,.2])
6767
mean0 = np.sum(a)/np.size(a) # calculated traditionally
6868
mean1 = a.mean() # calculated using this code
6969
self.assertQuantityEqual(mean0, mean1)
7070

7171
def test_uncertaintity_nanmean(self):
72-
import numpy as np
7372
a = UncertainQuantity([1,2], 'm', [.1,.2])
7473
b = UncertainQuantity([1,2,np.nan], 'm', [.1,.2,np.nan])
7574
self.assertQuantityEqual(a.mean(),b.nanmean())
7675

7776
def test_uncertainquantity_subtract(self):
78-
import numpy as np
7977
a = UncertainQuantity(1, 'm', 1)
8078
b = a.copy() # different object
8179
self.assertQuantityEqual(a-a, UncertainQuantity(0, 'm', 0))
@@ -84,3 +82,9 @@ def test_uncertainquantity_subtract(self):
8482
def test_uncertainty_sqrt(self):
8583
a = UncertainQuantity([1,2], 'm', [.1,.2])
8684
self.assertQuantityEqual(a**0.5, a.sqrt())
85+
86+
def test_uncertainty_nansum(self):
87+
uq = UncertainQuantity([1,2], 'm', [1,1])
88+
uq_nan = UncertainQuantity([1,2,np.nan], 'm', [1,1,np.nan])
89+
self.assertQuantityEqual(np.sum(uq), np.nansum(uq))
90+
self.assertQuantityEqual(np.sum(uq), np.nansum(uq_nan))

‎quantities/uncertainquantity.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -210,6 +210,16 @@ def sum(self, axis=None, dtype=None, out=None):
210210
copy=False
211211
)
212212

213+
@with_doc(np.nansum)
214+
def nansum(self, axis=None, dtype=None, out=None):
215+
import numpy as np
216+
return UncertainQuantity(
217+
np.nansum(self.magnitude, axis, dtype, out),
218+
self.dimensionality,
219+
(np.nansum(self.uncertainty.magnitude**2, axis))**0.5,
220+
copy=False
221+
)
222+
213223
@with_doc(np.ndarray.mean)
214224
def mean(self, axis=None, dtype=None, out=None):
215225
return UncertainQuantity(

0 commit comments

Comments
 (0)