8000 ENH: Added atleast_nd · numpy/numpy@8720da5 · GitHub
[go: up one dir, main page]

Skip to content

Commit 8720da5

Browse files
committed
ENH: Added atleast_nd
Function only accepts one array and an indexed placement spec. Masked arrays fully supported. atleast_1d and atleast_2d re-implemented in terms of new func. Tests included.
1 parent 5f2a5ae commit 8720da5

File tree

10 files changed

+228
-58
lines changed

10 files changed

+228
-58
lines changed

doc/release/1.16.0-notes.rst

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,9 @@ Highlights
1919
New functions
2020
=============
2121

22+
* `atleast_nd` generalizes ``atleast_1d``, ``atleast_2d`` and ``atleast_3d`` to
23+
arbitrary numbers of dimensions.
24+
2225

2326
Deprecations
2427
============

doc/source/reference/routines.array-manipulation.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ Changing number of dimensions
4040
atleast_1d
4141
atleast_2d
4242
atleast_3d
43+
atleast_nd
4344
broadcast
4445
broadcast_to
4546
broadcast_arrays

doc/source/reference/routines.ma.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,7 @@ Changing the number of dimensions
121121
ma.atleast_1d
122122
ma.atleast_2d
123123
ma.atleast_3d
124+
ma.atleast_nd
124125
ma.expand_dims
125126
ma.squeeze
126127

doc/source/user/quickstart.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -914,6 +914,7 @@ Conversions
914914
`atleast_1d`,
915915
`atleast_2d`,
916916
`atleast_3d`,
917+
`atleast_nd`,
917918
`mat`
918919
Manipulations
919920
`array_split`,

numpy/core/shape_base.py

Lines changed: 130 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,15 @@
11
from __future__ import division, absolute_import, print_function
22

3-
__all__ = ['atleast_1d', 'atleast_2d', 'atleast_3d', 'block', 'hstack',
4-
'stack', 'vstack']
3+
__all__ = ['atleast_1d', 'atleast_2d', 'atleast_3d', 'atleast_nd', 'block',
4+
'hstack', 'stack', 'vstack']
55

66

77
from . import numeric as _nx
88
from .numeric import array, asanyarray, newaxis
99
from .multiarray import normalize_axis_index
1010
from ._internal import recursive
1111

12+
1213
def atleast_1d(*arys):
1314
"""
1415
Convert inputs to arrays with at least one dimension.
@@ -29,7 +30,7 @@ def atleast_1d(*arys):
2930
3031
F438 See Also
3132
--------
32-
atleast_2d, atleast_3d
33+
atleast_2d, atleast_3d, atleast_nd
3334
3435
Examples
3536
--------
@@ -48,18 +49,10 @@ def atleast_1d(*arys):
4849
[array([1]), array([3, 4])]
4950
5051
"""
51-
res = []
52-
for ary in arys:
53-
ary = asanyarray(ary)
54-
if ary.ndim == 0:
55-
result = ary.reshape(1)
56-
else:
57-
result = ary
58-
res.append(result)
59-
if len(res) == 1:
60-
return res[0]
61-
else:
62-
return res
52+
if len(arys) == 1:
53+
return atleast_nd(arys[0], 1)
54+
return [atleast_nd(x, 1) for x in arys]
55+
6356

6457
def atleast_2d(*arys):
6558
"""
@@ -68,9 +61,9 @@ def atleast_2d(*arys):
6861
Parameters
6962
----------
7063
arys1, arys2, ... : array_like
71-
One or more array-like sequences. Non-array inputs are converted
72-
to arrays. Arrays that already have two or more dimensions are
73-
preserved.
64+
One or more array-like sequences. Non-array inputs are
65+
converted to arrays. Arrays that already have two or more
66+
dimensions are preserved.
7467
7568
Returns
7669
-------
@@ -81,7 +74,7 @@ def atleast_2d(*arys):
8174
8275
See Also
8376
--------
84-
atleast_1d, atleast_3d
77+
atleast_1d, atleast_3d, atleast_nd
8578
8679
Examples
8780
--------
@@ -98,20 +91,10 @@ def atleast_2d(*arys):
9891
[array([[1]]), array([[1, 2]]), array([[1, 2]])]
9992
10093
"""
101-
res = []
102-
for ary in arys:
103-
ary = asanyarray(ary)
104-
if ary.ndim == 0:
105-
result = ary.reshape(1, 1)
106-
elif ary.ndim == 1:
107-
result = ary[newaxis,:]
108-
else:
109-
result = ary
110-
res.append(result)
111-
if len(res) == 1:
112-
return res[0]
113-
else:
114-
return res
94+
if len(arys) == 1:
95+
return atleast_nd(arys[0], 2)
96+
return [atleast_nd(x, 2) for x in arys]
97+
11598

11699
def atleast_3d(*arys):
117100
"""
@@ -120,22 +103,31 @@ def atleast_3d(*arys):
120103
Parameters
121104
----------
122105
arys1, arys2, ... : array_like
123-
One or more array-like sequences. Non-array inputs are converted to
124-
arrays. Arrays that already have three or more dimensions are
125-
preserved.
106+
One or more array-like sequences. Non-array inputs are
107+
converted to arrays. Arrays that already have three or more
108+
dimensions are preserved.
126109
127110
Returns
128111
-------
129112
res1, res2, ... : ndarray
130-
An array, or list of arrays, each with ``a.ndim >= 3``. Copies are
131-
avoided where possible, and views with three or more dimensions are
132-
returned. For example, a 1-D array of shape ``(N,)`` becomes a view
133-
of shape ``(1, N, 1)``, and a 2-D array of shape ``(M, N)`` becomes a
134-
view of shape ``(M, N, 1)``.
113+
An array, or list of arrays, each with ``a.ndim >= 3``. Copies
114+
are avoided where possible, and views with three or more
115+
dimensions are returned. For example, a 1-D array of shape
116+
``(N,)`` becomes a view of shape ``(1, N, 1)``, and a 2-D array
117+
of shape ``(M, N)`` becomes a view of shape ``(M, N, 1)``.
135118
136119
See Also
137120
--------
138-
atleast_1d, atleast_2d
121+
atleast_1d, atleast_2d, atleast_nd
122+
123+
Notes
124+
-----
125+
As mentioned in the `Returns` section, the results of this fuction
126+
are not compatible with any of the other `atleast*` functions.
127+
`atleast_2d` prepends the unit dimension to a 1D array while
128+
`atleast_3d` appends it to a 2D array. The 1D array case both
129+
appends and prepends a dimension, while `atleast_nd` can only add
130+
dimensions to one end at a time.
139131
140132
Examples
141133
--------
@@ -168,9 +160,9 @@ def atleast_3d(*arys):
168160
if ary.ndim == 0:
169161
result = ary.reshape(1, 1, 1)
170162
elif ary.ndim == 1:
171-
result = ary[newaxis,:, newaxis]
163+
result = ary[newaxis, :, newaxis]
172164
elif ary.ndim == 2:
173-
result = ary[:,:, newaxis]
165+
result = ary[:, :, newaxis]
174166
else:
175167
result = ary
176168
res.append(result)
@@ -180,6 +172,100 @@ def atleast_3d(*arys):
180172
return res
181173

182174

175+
def atleast_nd(ary, ndim, pos=0):
176+
"""
177+
View input as array with at least `ndim` dimensions.
178+
179+
New unit dimensions are inserted at the index given by `pos` if
180+
necessary.
181+
182+
Parameters
183+
----------
184+
ary : array_like
185+
The input array. Non-array inputs are converted to arrays.
186+
Arrays that already have `ndim` or more dimensions are
187+
preserved.
188+
ndim : scalar
189+
The minimum number of dimensions required.
190+
pos : int, optional
191+
The index to insert the new dimensions. May range from
192+
``-ary.ndim - 1`` to ``+ary.ndim`` (inclusive). Non-negative
193+
indices indicate locations before the corresponding axis:
194+
``pos=0`` means to insert at the very beginning. Negative
195+
indices indicate locations after the corresponding axis:
196+
``pos=-1`` means to insert at the very end. 0 and -1 are always
197+
guaranteed to work. Any other number will depend on the
198+
dimensions of the existing array. Default is 0.
199+
200+
Returns
201+
-------
202+
res : ndarray
203+
An array with ``res.ndim >= ndim``. A view is returned for array
204+
inputs. Dimensions are prepended if `pos` is 0, so for example,
205+
a 1-D array of shape ``(N,)`` with ``ndim=4`` becomes a view of
206+
shape ``(1, 1, 1, N)``. Dimensions are appended if `pos` is -1,
207+
so for example a 2-D array of shape ``(M, N)`` becomes a view of
208+
shape ``(M, N, 1, 1)`` when ``ndim=4``.
209+
210+
See Also
211+
--------
212+
atleast_1d, atleast_2d, atleast_3d
213+
214+
Notes
215+
-----
216+
This function does not follow the convention of the other atleast_*d
217+
functions in numpy in that it only accepts a single array argument.
218+
To process multiple arrays, use a comprehension or loop around the
219+
function call. See examples below.
220+
221+
Setting ``pos=0`` is equivalent to how the array would be
222+
interpreted by numpy's broadcasting rules. There is no need to call
223+
this function for simple broadcasting. This is also roughly
224+
(but not exactly) equivalent to
225+
``np.array(ary, copy=False, subok=True, ndmin=ndim)``.
226+
227+
It is easy to create functions for specific dimensions similar to
228+
the other atleast_*d functions using Python's `functools.partial`
229+
function. An example is shown below.
230+
231+
Examples
232+
--------
233+
>>> np.atleast_nd(3.0, 4)
234+
array([[[[ 3.]]]])
235+
236+
>>> x = np.arange(3.0)
237+
>>> np.atleast_nd(x, 2).shape
238+
(1, 3)
239+
240+
>>> x = np.arange(12.0).reshape(4, 3)
241+
>>> np.atleast_nd(x, 5).shape
242+
(1, 1, 1, 4, 3)
243+
>>> np.atleast_nd(x, 5).base is x.base
244+
True
245+
246+
>>> [np.atleast_nd(x) for x in ((1, 2), [[1, 2]], [[[1, 2]]])]:
247+
[array([[1, 2]]), array([[1, 2]]), array([[[1, 2]]])]
248+
249+
>>> np.atleast_nd((1, 2), 5, pos=0).shape
250+
(1, 1, 1, 1, 2)
251+
>>> np.atleast_nd((1, 2), 5, pos=-1).shape
252+
(2, 1, 1, 1, 1)
253+
254+
>>> from functools import partial
255+
>>> atleast_4d = partial(np.atleast_nd, ndim=4)
256+
>>> atleast_4d([1, 2, 3])
257+
[[[[1, 2, 3]]]]
258+
"""
259+
ary = array(ary, copy=False, subok=True)
260+
if ary.ndim:
261+
pos = normalize_axis_index(pos, ary.ndim + 1)
262+
extra = ndim - ary.ndim
263+
if extra > 0:
264+
ind = pos * (slice(None),) + extra * (None,) + (Ellipsis,)
265+
ary = ary[ind]
266+
return ary
267+
268+
183269
def vstack(tup):
184270
"""
185271
Stack arrays in sequence vertically (row wise).

numpy/core/tests/test_shape_base.py

Lines changed: 55 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,8 @@
33
import warnings
44
import numpy as np
55
from numpy.core import (
6-
array, arange, atleast_1d, atleast_2d, atleast_3d, block, vstack, hstack,
7-
newaxis, concatenate, stack
6+
array, arange, atleast_1d, atleast_2d, atleast_3d, atleast_nd, block,
7+
concatenate, hstack, newaxis, stack, vstack,
88
)
99
from numpy.testing import (
1010
assert_, assert_raises, assert_array_equal, assert_equal,
@@ -125,6 +125,59 @@ def test_3D_array(self):
125125
assert_array_equal(res, desired)
126126

127127

128+
class TestAtleastNd(object):
129+
def test_0D_arrays(self):
130+
arrays = [array(1), array(2)]
131+
dims = [3, 2]
132+
expected = [array([[[1]]]), array([[2]])]
133+
134+
for x, y, d in zip(arrays, expected, dims):
135+
assert_array_equal(atleast_nd(x, d), y)
136+
assert_array_equal(atleast_nd(x, d, -1), y)
137+
138+
def test_nD_arrays(self):
139+
a = array([1])
140+
b = array([4, 5, 6])
141+
c = array([[2, 3]])
142+
d = array([[[2], [3]], [[2], [3]]])
143+
e = ((((1, 2), (3, 4)), ((5, 6), (7, 8))))
144+
arrays = (a, b, c, d, e)
145+
expected_before = (array([[[1]]]),
146+
array([[[4, 5, 6]]]),
147+
array([[[2, 3]]]),
148+
d,
149+
array(e))
150+
expected_after = (array([[[1]]]),
151+
array([[[4]], [[5]], [[6]]]),
152+
array([[[2], [3]]]),
153+
d,
154+
array(e))
155+
156+
for x, y in zip(arrays, expected_before):
157+
assert_array_equal(atleast_nd(x, 3), y)
158+
for x, y in zip(arrays, expected_after):
159+
assert_array_equal(atleast_nd(x, 3, pos=-1), y)
160+
161+
def test_nocopy(self):
162+
a = arange(12.0).reshape(4, 3)
163+
res = atleast_nd(a, 5)
164+
desired_shape = (1, 1, 1, 4, 3)
165+
desired_base = a.base # a was reshaped
166+
assert_equal(res.shape, desired_shape)
167+
assert_(res.base is desired_base)
168+
169+
def test_passthough(self):
170+
a = array([1, 2, 3])
171+
assert_(atleast_nd(a, 0) is a)
172+
assert_(atleast_nd(a, 1) is a)
173+
174+
def test_other_pos(self):
175+
a = arange(12.0).reshape(4, 3)
176+
res = atleast_nd(a, 4, pos=1)
177+
assert_equal(res.shape, (4, 1, 1, 3))
178+
assert_raises(ValueError, atleast_nd, a, 4, pos=5)
179+
180+
128181
class TestHstack(object):
129182
def test_non_iterable(self):
130183
assert_raises(TypeError, hstack, 1)

numpy/lib/info.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,7 @@
7070
atleast_1d Force arrays to be >= 1D
7171
atleast_2d Force arrays to be >= 2D
7272
atleast_3d Force arrays to be >= 3D
73+
atleast_nd Force arrays to be >= ND, where N is specified
7374
vstack Stack arrays vertically (row on row)
7475
hstack Stack arrays horizontally (column on column)
7576
column_stack Stack 1D arrays as columns into 2D array

numpy/lib/shape_base.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -506,7 +506,7 @@ def expand_dims(a, axis):
506506
--------
507507
squeeze : The inverse operation, removing singleton dimensions
508508
reshape : Insert, remove, and combine dimensions, and resize existing ones
509-
doc.indexing, atleast_1d, atleast_2d, atleast_3d
509+
doc.indexing, atleast_1d, atleast_2d, atleast_3d, atleast_nd
510510
511511
Examples
512512
--------

0 commit comments

Comments
 (0)
0