10000 MAINT: make np.ma.apply_along_axis consistent with np.apply_along_axis by eric-wieser · Pull Request #8511 · numpy/numpy · GitHub
[go: up one dir, main page]
Skip to content

MAINT: make np.ma.apply_along_axis consistent with np.apply_along_axis #8511

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

Closed
wants to merge 2 commits into from
Closed

MAINT: make np.ma.apply_along_axis consistent with np.apply_along_axis #8511

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
72f88cd
MAINT: Rewrite np.ma.apply_along_axis in terms of apply_along_axis
eric-wieser Jan 21, 2017
39d105e
fixup last commit
eric-wieser May 29, 2018
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
82 changes: 5 additions & 77 deletions numpy/ma/extras.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -379,83 +379,11 @@ def apply_along_axis(func1d, axis, arr, *args, **kwargs):
"""
(This docstring should be overwritten)
"""
arr = array(arr, copy=False, subok=True)
nd = arr.ndim
axis = normalize_axis_index(axis, nd)
ind = [0] * (nd - 1)
i = np.zeros(nd, 'O')
indlist = list(range(nd))
indlist.remove(axis)
i[axis] = slice(None, None)
outshape = np.asarray(arr.shape).take(indlist)
i.put(indlist, ind)
j = i.copy()
res = func1d(arr[tuple(i.tolist())], *args, **kwargs)
# if res is a number, then we have a smaller output array
asscalar = np.isscalar(res)
if not asscalar:
try:
len(res)
except TypeError:
asscalar = True
# Note: we shouldn't set the dtype of the output from the first result
# so we force the type to object, and build a list of dtypes. We'll
# just take the largest, to avoid some downcasting
dtypes = []
if asscalar:
dtypes.append(np.asarray(res).dtype)
outarr = zeros(outshape, object)
outarr[tuple(ind)] = res
Ntot = np.product(outshape)
k = 1
while k < Ntot:
# increment the index
ind[-1] += 1
n = -1
while (ind[n] >= outshape[n]) and (n > (1 - nd)):
ind[n - 1] += 1
ind[n] = 0
n -= 1
i.put(indlist, ind)
res = func1d(arr[tuple(i.tolist())], *args, **kwargs)
outarr[tuple(ind)] = res
dtypes.append(asarray(res).dtype)
k += 1
else:
res = array(res, copy=False, subok=True)
j = i.copy()
j[axis] = ([slice(None, None)] * res.ndim)
j.put(indlist, ind)
Ntot = np.product(outshape)
holdshape = outshape
outshape = list(arr.shape)
outshape[axis] = res.shape
dtypes.append(asarray(res).dtype)
outshape = flatten_inplace(outshape)
outarr = zeros(outshape, object)
outarr[tuple(flatten_inplace(j.tolist()))] = res
k = 1
while k < Ntot:
# increment the index
ind[-1] += 1
n = -1
while (ind[n] >= holdshape[n]) and (n > (1 - nd)):
ind[n - 1] += 1
ind[n] = 0
n -= 1
i.put(indlist, ind)
j.put(indlist, ind)
res = func1d(arr[tuple(i.tolist())], *args, **kwargs)
outarr[tuple(flatten_inplace(j.tolist()))] = res
dtypes.append(asarray(res).dtype)
k += 1
max_dtypes = np.dtype(np.asarray(dtypes).max())
if not hasattr(arr, '_mask'):
result = np.asarray(outarr, dtype=max_dtypes)
else:
result = asarray(outarr, dtype=max_dtypes)
result.fill_value = ma.default_fill_value(result)
return result
def wrapped_func(a, *args, **kwargs):
res = func1d(a, *args, **kwargs)
return np.ma.asanyarray(res)

return np.apply_along_axis(wrapped_func, axis, arr, *args, **kwargs)
apply_along_axis.__doc__ = np.apply_along_axis.__doc__


Expand Down
15 changes: 15 additions & 0 deletions numpy/ma/tests/test_extras.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -640,6 +640,21 @@ def myfunc(b, offset=0):
xa = apply_along_axis(myfunc, 2, a, offset=1)
assert_equal(xa, [[2, 5], [8, 11]])

def test_mask(self):
x = np.arange(20).reshape(4, 5)
m = np.ma.masked_where(x%2 == 0, x)

# note that this lambda sometime returns np.ma.masked
# works fine with normal apply_along_axis
row0 = np.ma.apply_along_axis(lambda x: x[0], 0, m)

# fails with normal apply_along_axis, which doesn't use a masked_array
row1 = np.ma.apply_along_axis(lambda x: x[1], 0, m)

# note the first of these is a float, because that's the type of np.ma.masked...
assert_equal(row0, np.ma.masked_array([-1, 1.0, -1, 3.0, -1], [1, 0, 1, 0, 1]))
assert_equal(row1, np.ma.masked_array([5, -1, 7, -1, 9], [0, 1, 0, 1, 0]))


class TestApplyOverAxes(TestCase):
# Tests apply_over_axes
Expand Down
0