diff options
author | Duke Vijitbenjaronk <warut.vijitbenjaronk@gmail.com> | 2017-03-07 11:11:53 -0600 |
---|---|---|
committer | Eric Wieser <wieser.eric@gmail.com> | 2017-03-07 20:49:33 +0000 |
commit | 1588ae39ffb51ea916f03510671aab711fdfb568 (patch) | |
tree | c08267636b61856aa8d2726a0ac0bb92a327c96d /numpy/lib | |
parent | df9e134d1e40de4d3865597a239c5a1d68624f5d (diff) | |
download | numpy-1588ae39ffb51ea916f03510671aab711fdfb568.tar.gz |
BUG: Fix np.average with object array weights
Fixes #8696
Diffstat (limited to 'numpy/lib')
-rw-r--r-- | numpy/lib/function_base.py | 2 | ||||
-rw-r--r-- | numpy/lib/tests/test_function_base.py | 7 |
2 files changed, 8 insertions, 1 deletions
diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py index c54512c21..0903790bd 100644 --- a/numpy/lib/function_base.py +++ b/numpy/lib/function_base.py @@ -1135,7 +1135,7 @@ def average(a, axis=None, weights=None, returned=False): wgt = wgt.swapaxes(-1, axis) scl = wgt.sum(axis=axis, dtype=result_dtype) - if (scl == 0.0).any(): + if np.any(scl == 0.0): raise ZeroDivisionError( "Weights sum to zero, can't be normalized") diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py index 4fb0dba51..188c1c2ea 100644 --- a/numpy/lib/tests/test_function_base.py +++ b/numpy/lib/tests/test_function_base.py @@ -3,6 +3,7 @@ from __future__ import division, absolute_import, print_function import operator import warnings import sys +import decimal import numpy as np from numpy.testing import ( @@ -341,6 +342,12 @@ class TestAverage(TestCase): w = np.array([[1,2],[3,4]], dtype=wt) assert_equal(np.average(a, weights=w).dtype, np.dtype(rt)) + def test_object_dtype(self): + a = np.array([decimal.Decimal(x) for x in range(10)]) + w = np.array([decimal.Decimal(1) for _ in range(10)]) + w /= w.sum() + assert_almost_equal(a.mean(0), average(a, weights=w)) + class TestSelect(TestCase): choices = [np.array([1, 2, 3]), np.array([4, 5, 6]), |