Skip to content

Commit e86419a

Browse files
ddaleddale
committed
Merge pull request python-quantities#79 from mankoff/master
Calculate mean w/ uncertainties
2 parents fe3effb + 6b4d169 commit e86419a

4 files changed

Lines changed: 215 additions & 2 deletions

File tree

‎quantities/quantity.py‎

Lines changed: 53 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)
@@ -477,6 +486,14 @@ def max(self, axis=None, out=None):
477486
copy=False
478487
)
479488

489+
@with_doc(np.nanmax)
490+
def nanmax(self, axis=None, out=None):
491+
return Quantity(
492+
np.nanmax(self.magnitude),
493+
self.dimensionality,
494+
copy=False
495+
)
496+
480497
@with_doc(np.ndarray.min)
481498
def min(self, axis=None, out=None):
482499
return Quantity(
@@ -485,10 +502,30 @@ def min(self, axis=None, out=None):
485502
copy=False
486503
)
487504

505+
@with_doc(np.nanmin)
506+
def nanmin(self, axis=None, out=None):
507+
return Quantity(
508+
np.nanmin(self.magnitude),
509+
self.dimensionality,
510+
copy=False
511+
)
512+
488513
@with_doc(np.ndarray.argmin)
489514
def argmin(self,axis=None, out=None):
490515
return self.magnitude.argmin()
491516

517+
@with_doc(np.ndarray.argmax)
518+
def argmax(self,axis=None, out=None):
519+
return self.magnitude.argmax()
520+
521+
@with_doc(np.nanargmin)
522+
def nanargmin(self,axis=None, out=None):
523+
return np.nanargmin(self.magnitude)
524+
525+
@with_doc(np.nanargmax)
526+
def nanargmax(self,axis=None, out=None):
527+
return np.nanargmax(self.magnitude)
528+
492529
@with_doc(np.ndarray.ptp)
493530
def ptp(self, axis=None, out=None):
494531
return Quantity(
@@ -549,6 +586,14 @@ def mean(self, axis=None, dtype=None, out=None):
549586
self.dimensionality,
550587
copy=False)
551588

589+
@with_doc(np.nanmean)
590+
def nanmean(self, axis=None, dtype=None, out=None):
591+
import numpy as np
592+
return Quantity(
593+
np.nanmean(self.magnitude, axis, dtype, out),
594+
self.dimensionality,
595+
copy=False)
596+
552597
@with_doc(np.ndarray.var)
553598
def var(self, axis=None, dtype=None, out=None, ddof=0):
554599
return Quantity(
@@ -565,6 +610,14 @@ def std(self, axis=None, dtype=None, out=None, ddof=0):
565610
copy=False
566611
)
567612

613+
@with_doc(np.nanstd)
614+
def nanstd(self, axis=None, dtype=None, out=None, ddof=0):
615+
return Quantity(
616+
np.nanstd(self.magnitude, axis, dtype, out, ddof),
617+
self._dimensionality,
618+
copy=False
619+
)
620+
568621
@with_doc(np.ndarray.prod)
569622
def prod(self, axis=None, dtype=None, out=None):
570623
if axis == None:

‎quantities/tests/test_methods.py‎

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
from .. import units as pq
44
from .common import TestCase
5-
5+
import numpy as np
66

77
class TestQuantityMethods(TestCase):
88

@@ -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)
@@ -103,15 +109,31 @@ def test_nonzero(self):
103109
def test_max(self):
104110
self.assertQuantityEqual(self.q.max(), 4*pq.m)
105111

112+
def test_nanmax(self):
113+
q = np.append(self.q, np.nan) * self.q.units
114+
self.assertQuantityEqual(q.nanmax(), 4*pq.m)
115+
106116
def test_argmax(self):
107117
self.assertEqual(self.q.argmax(), 3)
108118

119+
def test_nanargmax(self):
120+
q = np.append(self.q, np.nan) * self.q.units
121+
self.assertEqual(self.q.nanargmax(), 3)
122+
109123
def test_min(self):
110124
self.assertEqual(self.q.min(), 1 * pq.m)
111125

126+
def test_nanmin(self):
127+
q = np.append(self.q, np.nan) * self.q.units
128+
self.assertQuantityEqual(q.nanmin(), 1*pq.m)
129+
112130
def test_argmin(self):
113131
self.assertEqual(self.q.argmin(), 0)
114132

133+
def test_nanargmax(self):
134+
q = np.append(self.q, np.nan) * self.q.units
135+
self.assertEqual(self.q.nanargmin(), 0)
136+
115137
def test_ptp(self):
116138
self.assertQuantityEqual(self.q.ptp(), 3 * pq.m)
117139

@@ -146,12 +168,23 @@ def test_cumsum(self):
146168
def test_mean(self):
147169
self.assertQuantityEqual(self.q.mean(), 2.5 * pq.m)
148170

171+
def test_nanmean(self):
172+
import numpy as np
173+
q = [[1,2], [3,4], [np.nan,np.nan]] * pq.m
174+
self.assertQuantityEqual(q.nanmean(), self.q.mean())
175+
149176
def test_var(self):
150177
self.assertQuantityEqual(self.q.var(), 1.25*pq.m**2)
151178

152179
def test_std(self):
153180
self.assertQuantityEqual(self.q.std(), 1.11803*pq.m, delta=1e-5)
154181

182+
def test_nanstd(self):
183+
import numpy as np
184+
q0 = [[1,2], [3,4]] * pq.m
185+
q1 = [[1,2], [3,4], [np.nan,np.nan]] * pq.m
186+
self.assertQuantityEqual(q0.std(), q1.nanstd())
187+
155188
def test_prod(self):
156189
self.assertQuantityEqual(self.q.prod(), 24 * pq.m**4)
157190

‎quantities/tests/test_uncertainty.py‎

Lines changed: 57 additions & 0 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):
@@ -60,3 +61,59 @@ def test_uncertainquantity_divide(self):
6061
self.assertQuantityEqual((a/2).uncertainty, [0.05, 0.1 ]*pq.m)
6162
self.assertQuantityEqual(1/a, [1., 0.5]/pq.m)
6263
self.assertQuantityEqual((1/a).uncertainty, [0.1, 0.05]/pq.m)
64+
65+
def test_uncertaintity_mean(self):
66+
a = UncertainQuantity([1,2], 'm', [.1,.2])
67+
mean0 = np.sum(a)/np.size(a) # calculated traditionally
68+
mean1 = a.mean() # calculated using this code
69+
self.assertQuantityEqual(mean0, mean1)
70+
71+
def test_uncertaintity_nanmean(self):
72+
a = UncertainQuantity([1,2], 'm', [.1,.2])
73+
b = UncertainQuantity([1,2,np.nan], 'm', [.1,.2,np.nan])
74+
self.assertQuantityEqual(a.mean(),b.nanmean())
75+
76+
def test_uncertainquantity_subtract(self):
77+
a = UncertainQuantity(1, 'm', 1)
78+
b = a.copy() # different object
79+
self.assertQuantityEqual(a-a, UncertainQuantity(0, 'm', 0))
80+
self.assertQuantityEqual(a-b, UncertainQuantity(0, 'm', np.sqrt(2)))
81+
82+
def test_uncertainty_sqrt(self):
83+
a = UncertainQuantity([1,2], 'm', [.1,.2])
84+
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))
91+
92+
def test_uncertainty_minmax_nan_arg(self):
93+
q = [[1, 2], [3, 4]] * pq.m # quantity
94+
self.assertQuantityEqual(q.min(), 1*pq.m) # min
95+
self.assertQuantityEqual(q.max(), 4*pq.m) # max
96+
self.assertQuantityEqual(q.argmin(), 0) # argmin
97+
self.assertQuantityEqual(q.argmax(), 3) # argmax
98+
# uncertain quantity
99+
uq = UncertainQuantity([[1,2],[3,4]], pq.m, [[1,1],[1,1]])
100+
self.assertQuantityEqual(uq.min(), 1*pq.m) # min
101+
self.assertQuantityEqual(uq.max(), 4*pq.m) # max
102+
self.assertQuantityEqual(uq.argmin(), 0) # argmin
103+
self.assertQuantityEqual(uq.argmax(), 3) # argmax
104+
# now repeat the above with NaNs
105+
nanq = [[1, 2], [3, 4], [np.nan,np.nan]] * pq.m # quantity
106+
nanuq = UncertainQuantity([[1,2],[3,4],[np.nan,np.nan]],
107+
pq.m,
108+
[[1,1],[1,1],[np.nan,np.nan]])
109+
self.assertQuantityEqual(nanq.nanmin(), 1*pq.m) # min
110+
self.assertQuantityEqual(nanq.nanmax(), 4*pq.m) # max
111+
self.assertQuantityEqual(nanq.nanargmin(), 0) # argmin
112+
self.assertQuantityEqual(nanq.nanargmax(), 3) # argmax
113+
self.assertQuantityEqual(nanuq.nanmin(), 1*pq.m) # min
114+
self.assertQuantityEqual(nanuq.nanmax(), 4*pq.m) # max
115+
self.assertQuantityEqual(nanuq.nanargmin(), 0) # argmin
116+
self.assertQuantityEqual(nanuq.nanargmax(), 3) # argmax
117+
118+
119+

‎quantities/uncertainquantity.py‎

Lines changed: 71 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -102,7 +102,10 @@ def __radd__(self, other):
102102
@scale_other_units
103103
def __sub__(self, other):
104104
res = super(UncertainQuantity, self).__sub__(other)
105-
u = (self.uncertainty**2+other.uncertainty**2)**0.5
105+
if self is not other:
106+
u = (self.uncertainty**2+other.uncertainty**2)**0.5
107+
else:
108+
u = self.uncertainty*0
106109
return UncertainQuantity(res, uncertainty=u, copy=False)
107110

108111
@with_doc(Quantity.__rsub__, use_header=False)
@@ -207,6 +210,73 @@ def sum(self, axis=None, dtype=None, out=None):
207210
copy=False
208211
)
209212

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+
223+
@with_doc(np.ndarray.mean)
224+
def mean(self, axis=None, dtype=None, out=None):
225+
return UncertainQuantity(
226+
self.magnitude.mean(axis, dtype, out),
227+
self.dimensionality,
228+
((1.0/np.size(self,axis))**2 * np.sum(self.uncertainty.magnitude**2, axis))**0.5,
229+
copy=False)
230+
231+
@with_doc(np.nanmean)
232+
def nanmean(self, axis=None, dtype=None, out=None):
233+
size = np.sum(~np.isnan(self),axis)
234+
return UncertainQuantity(
235+
np.nanmean(self.magnitude, axis, dtype, out),
236+
self.dimensionality,
237+
((1.0/size)**2 * np.nansum(np.nan_to_num(self.uncertainty.magnitude)**2, axis))**0.5,
238+
copy=False)
239+
240+
@with_doc(np.sqrt)
241+
def sqrt(self, out=None):
242+
return self**0.5
243+
244+
@with_doc(np.ndarray.max)
245+
def max(self, axis=None, out=None):
246+
idx = np.unravel_index(np.argmax(self.magnitude), self.shape)
247+
return self[idx]
248+
249+
@with_doc(np.nanmax)
250+
def nanmax(self, axis=None, out=None):
251+
idx = np.unravel_index(np.nanargmax(self.magnitude), self.shape)
252+
return self[idx]
253+
254+
@with_doc(np.ndarray.min)
255+
def min(self, axis=None, out=None):
256+
idx = np.unravel_index(np.argmin(self.magnitude), self.shape)
257+
return self[idx]
258+
259+
@with_doc(np.nanmin)
260+
def nanmin(self, axis=None, out=None):
261+
idx = np.unravel_index(np.nanargmin(self.magnitude), self.shape)
262+
return self[idx]
263+
264+
@with_doc(np.ndarray.argmin)
265+
def argmin(self,axis=None, out=None):
266+
return self.magnitude.argmin()
267+
268+
@with_doc(np.ndarray.argmax)
269+
def argmax(self,axis=None, out=None):
270+
return self.magnitude.argmax()
271+
272+
@with_doc(np.nanargmin)
273+
def nanargmin(self,axis=None, out=None):
274+
return np.nanargmin(self.magnitude)
275+
276+
@with_doc(np.nanargmax)
277+
def nanargmax(self,axis=None, out=None):
278+
return np.nanargmax(self.magnitude)
279+
210280
def __getstate__(self):
211281
"""
212282
Return the internal state of the quantity, for pickling

0 commit comments

Comments
 (0)