8000 BUG: Fixed masked array behavior for scalar inputs to np.ma.atleast_*d · numpy/numpy@54b68dd · GitHub
[go: up one dir, main page]
Skip to content

Commit 54b68dd

Browse files
madphysicistmadphysicist
committed
BUG: Fixed masked array behavior for scalar inputs to np.ma.atleast_*d
1 parent d92c75b commit 54b68dd

File tree

< 8000 span class="prc-TooltipV2-Tooltip-cYMVY" data-direction="se" aria-hidden="true" id=":R4atdab:">Collapse file tree

4 files changed

+133
-52
lines changed

4 files changed

+133
-52
lines changed

‎numpy/core/shape_base.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -146,7 +146,7 @@ def atleast_3d(*arys):
146146
>>> x = np.arange(12.0).reshape(4,3)
147147
>>> np.atleast_3d(x).shape
148148
(4, 3, 1)
149-
>>> np.atleast_3d(x).base is x
149+
>>> np.atleast_3d(x).base is x.base # x is a reshape, so not base itself
150150
True
151151
152152
>>> for arr in np.atleast_3d([1, 2], [[1, 2]], [[[1, 2]]]):

‎numpy/lib/info.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -67,9 +67,9 @@
6767
------------------
6868
================ ===================
6969
squeeze Return a with length-one dimensions removed.
70-
atleast_1d Force arrays to be > 1D
71-
atleast_2d Force arrays to be > 2D
72-
atleast_3d Force arrays to be > 3D
70+
atleast_1d Force arrays to be >= 1D
71+
atleast_2d Force arrays to be >= 2D
72+
atleast_3d Force arrays to be >= 3D
7373
vstack Stack arrays vertically (row on row)
7474
hstack Stack arrays horizontally (column on column)
7575
column_stack Stack 1D arrays as columns into 2D array

‎numpy/ma/extras.py

Lines changed: 99 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -45,9 +45,8 @@ def issequence(seq):
4545
Is seq a sequence (ndarray, list or tuple)?
4646
4747
"""
48-
if isinstance(seq, (ndarray, tuple, list)):
49-
return True
50-
return False
48+
return isinstance(seq, (ndarray, tuple, list))
49+
5150

5251
def count_masked(arr, axis=None):
5352
"""
@@ -103,6 +102,7 @@ def count_masked(arr, axis=None):
103102
m = getmaskarray(arr)
104103
return m.sum(axis)
105104

105+
106106
def masked_all(shape, dtype=float):
107107
"""
108108
Empty masked array with all elements masked.
@@ -154,6 +154,7 @@ def masked_all(shape, dtype=float):
154154
mask=np.ones(shape, make_mask_descr(dtype)))
155155
return a
156156

157+
157158
def masked_all_like(arr):
158159
"""
159160
Empty masked array with the properties of an existing array.
@@ -221,6 +222,9 @@ class _fromnxfunction:
221222
as the wrapped NumPy function. The docstring of `newfunc` is adapted from
222223
the wrapped function as well, see `getdoc`.
223224
225+
This class should not be used directly. Instead, one of its extensions that
226+
provides support for a specific type of input should be used.
227+
224228
Parameters
225229
----------
226230
funcname : str
@@ -261,48 +265,100 @@ def getdoc(self):
261265
return
262266

263267
def __call__(self, *args, **params):
268+
pass
269+
270+
271+
class _fromnxfunction_single(_fromnxfunction):
272+
"""
273+
A version of `_fromnxfunction` that is called with a single array
274+
argument followed by auxiliary args that are passed verbatim for
275+
both the data and mask calls.
276+
"""
277+
def __call__(self, x, *args, **params):
264278
func = getattr(np, self.__name__)
265-
if len(args) == 1:
266-
x = args[0]
267-
if isinstance(x, ndarray):
268-
_d = func(x.__array__(), **params)
269-
_m = func(getmaskarray(x), **params)
270-
return masked_array(_d, mask=_m)
271-
elif isinstance(x, tuple) or isinstance(x, list):
272-
_d = func(tuple([np.asarray(a) for a in x]), **params)
273-
_m = func(tuple([getmaskarray(a) for a in x]), **params)
274-
return masked_array(_d, mask=_m)
275-
else:
276-
_d = func(np.asarray(x), **params)
277-
_m = func(getmaskarray(x), **params)
278-
return masked_array(_d, mask=_m)
279+
if isinstance(x, ndarray):
280+
_d = func(x.__array__(), *args, **params)
281+
_m = func(getmaskarray(x), *args, **params)
282+
return masked_array(_d, mask=_m)
279283
else:
280-
arrays = []
281-
args = list(args)
282-
while len(args) > 0 and issequence(args[0]):
283-
arrays.append(args.pop(0))
284-
res = []
285-
for x in arrays:
286-
_d = func(np.asarray(x), *args, **params)
287-
_m = func(getmaskarray(x), *args, **params)
288-
res.append(masked_array(_d, mask=_m))
289-
return res
290-
291-
atleast_1d = _fromnxfunction('atleast_1d')
292-
atleast_2d = _fromnxfunction('atleast_2d')
293-
atleast_3d = _fromnxfunction('atleast_3d')
294-
#atleast_1d = np.atleast_1d
295-
#atleast_2d = np.atleast_2d
296-
#atleast_3d = np.atleast_3d
297-
298-
vstack = row_stack = _fromnxfunction('vstack')
299-
hstack = _fromnxfunction('hstack')
300-
column_stack = _fromnxfunction('column_stack')
301-
dstack = _fromnxfunction('dstack')
302-
303-
hsplit = _fromnxfunction('hsplit')
304-
305-
diagflat = _fromnxfunction('diagflat')
284+
_d = func(np.asarray(x), *args, **params)
285+
_m = func(getmaskarray(x), *args, **params)
286+
return masked_array(_d, mask=_m)
287+
288+
289+
class _fromnxfunction_seq(_fromnxfunction):
290+
"""
291+
A version of `_fromnxfunction` that is called with a single sequence
292+
of arrays followed by auxiliary args that are passed verbatim for
293+
both the data and mask calls.
294+
"""
295+
def __call__(self, x, *args, **params):
296+
func = getattr(np, self.__name__)
297+
_d = func(tuple([np.asarray(a) for a in x]), *args, **params)
298+
_m = func(tuple([getmaskarray(a) for a in x]), *args, **params)
299+
return masked_array(_d, mask=_m)
300+
301+
302+
class _fromnxfunction_args(_fromnxfunction):
303+
"""
304+
A version of `_fromnxfunction` that is called with multiple array
305+
arguments. The first non-array-like input marks the beginning of the
306+
arguments that are passed verbatim for both the data and mask calls.
307+
Array arguments are processed independently and the results are
308+
returned in a list. If only one array is found, the return value is
309+
just the processed array instead of a list.
310+
"""
311+
def __call__(self, *args, **params):
312+
func = getattr(np, self.__name__)
313+
arrays = []
314+
args = list(args)
315+
while len(args) > 0 and issequence(args[0]):
316+
arrays.append(args.pop(0))
317+
res = []
318+
for x in arrays:
319+
_d = func(np.asarray(x), *args, **params)
320+
_m = func(getmaskarray(x), *args, **params)
321+
res.append(masked_array(_d, mask=_m))
322+
if len(arrays) == 1:
323+
return res[0]
324+
return res
325+
326+
327+
class _fromnxfunction_allargs(_fromnxfunction):
328+
"""
329+
A version of `_fromnxfunction` that is called with multiple array
330+
arguments. Similar to `_fromnxfunction_args` except that all args
331+
are converted to arrays even if they are not so already. This makes
332+
it possible to process scalars as 1-D arrays. Only keyword arguments
333+
are passed through verbatim for the data and mask calls. Arrays
334+
arguments are processed independently and the results are returned
335+
in a list. If only one arg is present, the return value is just the
336+
processed array instead of a list.
337+
"""
338+
def __call__(self, *args, **params):
339+
func = getattr(np, self.__name__)
340+
res = []
341+
for x in args:
342+
_d = func(np.asarray(x), **params)
343+
_m = func(getmaskarray(x), **params)
344+
res.append(masked_array(_d, mask=_m))
345+
if len(args) == 1:
346+
return res[0]
347+
return res
348+
349+
350+
atleast_1d = _fromnxfunction_allargs('atleast_1d')
351+
atleast_2d = _fromnxfunction_allargs('atleast_2d')
352+
atleast_3d = _fromnxfunction_allargs('atleast_3d')
353+
354+
vstack = row_stack = _fromnxfunction_seq('vstack')
355+
hstack = _fromnxfunction_seq('hstack')
356+
column_stack = _fromnxfunction_seq('column_stack')
357+
dstack = _fromnxfunction_seq('dstack')
358+
359+
hsplit = _fromnxfunction_single('hsplit')
360+
361+
diagflat = _fromnxfunction_single('diagflat')
306362

307363

308364
#####--------------------------------------------------------------------------

‎numpy/ma/tests/test_extras.py

Lines changed: 30 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1202,29 +1202,54 @@ def test_setdiff1d_char_array(self):
12021202

12031203
class TestShapeBase(TestCase):
12041204

1205-
def test_atleast2d(self):
1205+
def test_atleast_2d(self):
12061206
# Test atleast_2d
12071207
a = masked_array([0, 1, 2], mask=[0, 1, 0])
12081208
b = atleast_2d(a)
12091209
assert_equal(b.shape, (1, 3))
12101210
assert_equal(b.mask.shape, b.data.shape)
12111211
assert_equal(a.shape, (3,))
12121212
assert_equal(a.mask.shape, a.data.shape)
1213+
assert_equal(b.mask.shape, b.data.shape)
12131214

12141215
def test_shape_scalar(self):
12151216
# the atleast and diagflat function should work with scalars
12161217
# GitHub issue #3367
1218+
# Additionally, the atleast functions should accept multiple scalars
1219+
# correctly
12171220
b = atleast_1d(1.0)
1218-
assert_equal(b.shape, (1, ))
1219-
assert_equal(b.mask.shape, b.data.shape)
1221+
assert_equal(b.shape, (1,))
1222+
assert_equal(b.mask.shape, b.shape)
1223+
assert_equal(b.data.shape, b.shape)
1224+
1225+
b = atleast_1d(1.0, 2.0)
1226+
for a in b:
1227+
assert_equal(a.shape, (1,))
1228+
assert_equal(a.mask.shape, a.shape)
1229+
assert_equal(a.data.shape, a.shape)
12201230

12211231
b = atleast_2d(1.0)
12221232
assert_equal(b.shape, (1, 1))
1223-
assert_equal(b.mask.shape, b.data.shape)
1233+
assert_equal(b.mask.shape, b.shape)
1234+
assert_equal(b.data.shape, b.shape)
1235+
1236+
b = atleast_2d(1.0, 2.0)
1237+
for a in b:
1238+
assert_equal(a.shape, (1, 1))
1239+
assert_equal(a.mask.shape, a.shape)
1240+
assert_equal(a.data.shape, a.shape)
12241241

12251242
b = atleast_3d(1.0)
12261243
assert_equal(b.shape, (1, 1, 1))
1227-
assert_equal(b.mask.shape, b.data.shape)
1244+
assert_equal(b.mask.shape, b.shape)
1245+
assert_equal(b.data.shape, b.shape)
1246+
1247+
b = atleast_3d(1.0, 2.0)
1248+
for a in b:
1249+
assert_equal(a.shape, (1, 1, 1))
1250+
assert_equal(a.mask.shape, a.shape)
1251+
assert_equal(a.data.shape, a.shape)
1252+
12281253

12291254
b = diagflat(1.0)
12301255
assert_equal(b.shape, (1, 1))

0 commit comments

Comments
 (0)
0