Skip to content

MAINT: cleanup np.average #7382

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Mar 7, 2016
Merged

MAINT: cleanup np.average #7382

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 13 additions & 8 deletions numpy/lib/function_base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -898,15 +898,19 @@ def average(a, axis=None, weights=None, returned=False):
TypeError: Axis must be specified when shapes of a and weights differ.

"""
if not isinstance(a, np.matrix):
a = np.asarray(a)
a = np.asanyarray(a)

if weights is None:
avg = a.mean(axis)
scl = avg.dtype.type(a.size/avg.size)
else:
a = a + 0.0
wgt = np.asarray(weights)
wgt = np.asanyarray(weights)

if issubclass(a.dtype.type, (np.integer, np.bool_)):
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should this be an or or and with the same for wgt? Hmmm, actually, if, which one, hmmm :), I guess or?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the type of wgt doesn't matter. The current code does as follows for the four combinations to think about (float means f4 or f8):

a     wgt        result_dtype
---   ---        ---
int   int        f8
int   float      f8
float int        f8
float float2     biggest(float, float2)

I suppose in the 2nd line we could cast to whatever wgt's floating type was (f4 or f8) but I'm not sure whether that's helpful or unnecessarily complicated.

(Edit: Actually the 2nd line will give back f8 no matter what we do, since int+f4 coerces to f8 anyway)

result_dtype = np.result_type(a.dtype, wgt.dtype, 'f8')
else:
result_dtype = np.result_type(a.dtype, wgt.dtype)

# Sanity checks
if a.shape != wgt.shape:
if axis is None:
Expand All @@ -921,17 +925,18 @@ def average(a, axis=None, weights=None, returned=False):
"Length of weights not compatible with specified axis.")

# setup wgt to broadcast along axis
wgt = np.array(wgt, copy=0, ndmin=a.ndim).swapaxes(-1, axis)
wgt = np.broadcast_to(wgt, (a.ndim-1)*(1,) + wgt.shape)
wgt = wgt.swapaxes(-1, axis)

scl = wgt.sum(axis=axis, dtype=np.result_type(a.dtype, wgt.dtype))
scl = wgt.sum(axis=axis, dtype=result_dtype)
if (scl == 0.0).any():
raise ZeroDivisionError(
"Weights sum to zero, can't be normalized")

avg = np.multiply(a, wgt).sum(axis)/scl
avg = np.multiply(a, wgt, dtype=result_dtype).sum(axis)/scl

if returned:
scl = np.multiply(avg, 0) + scl
scl = np.broadcast_to(scl, avg.shape)
return avg, scl
else:
return avg
Expand Down
23 changes: 23 additions & 0 deletions numpy/lib/tests/test_function_base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -167,6 +167,29 @@ def test_returned(self):
avg, scl = average(y, weights=w2, axis=1, returned=True)
assert_array_equal(scl, np.array([1., 6.]))

def test_subclasses(self):
class subclass(np.ndarray):
pass
a = np.array([[1,2],[3,4]]).view(subclass)
w = np.array([[1,2],[3,4]]).view(subclass)

assert_equal(type(np.average(a, weights=w)), subclass)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should that not be better as assert_(type(np.average(a, weights=w)) is subclass)?

EDIT: Although that only works for new style classes. Hmm, assert_equal does seem to work properly there.
EDIT: But not for old style classes either.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's not clear to me what assert_equal does for types, so maybe is ok, if not obvious.


# also test matrices
a = np.matrix([[1,2],[3,4]])
w = np.matrix([[1,2],[3,4]])

r = np.average(a, axis=0, weights=w)
assert_equal(type(r), np.matrix)
assert_equal(r, [[2.5, 10.0/3]])

def test_upcasting(self):
types = [('i4', 'i4', 'f8'), ('i4', 'f4', 'f8'), ('f4', 'i4', 'f8'),
('f4', 'f4', 'f4'), ('f4', 'f8', 'f8')]
for at, wt, rt in types:
a = np.array([[1,2],[3,4]], dtype=at)
w = np.array([[1,2],[3,4]], dtype=wt)
assert_equal(np.average(a, weights=w).dtype, np.dtype(rt))

class TestSelect(TestCase):
choices = [np.array([1, 2, 3]),
Expand Down