8000 ENH: Added atleast_nd to numpy and ma · numpy/numpy@2474c47 · GitHub
[go: up one dir, main page]
Skip to content

Commit 2474c47

Browse files
madphysicistmadphysicist
committed
ENH: Added atleast_nd to numpy and ma
1 parent f6a71fe commit 2474c47

File tree

14 files changed

+262
-73
lines changed

14 files changed

+262
-73
lines changed

‎doc/release/upcoming_changes/18386.new_function.rst

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
Added `atleast_nd` function to `numpy` and `numpy.ma`
2+
-----------------------------------------------------
3+
`atleast_nd` generalizes ``atleast_1d``, ``atleast_2d`` and ``atleast_3d`` to
4+
arbitrary numbers of dimensions.

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

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

‎doc/source/reference/routines.ma.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -127,6 +127,7 @@ Changing the number of dimensions
127127
ma.atleast_1d
128128
ma.atleast_2d
129129
ma.atleast_3d
130+
ma.atleast_nd
130131
ma.expand_dims
131132
ma.squeeze
132133

‎doc/source/user/quickstart.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -967,6 +967,7 @@ Conversions
967967
`atleast_1d`,
968968
`atleast_2d`,
969969
`atleast_3d`,
970+
`atleast_nd`,
970971
`mat`
971972
Manipulations
972973
`array_split`,

‎numpy/__init__.pyi

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -326,6 +326,7 @@ from numpy.core.shape_base import (
326326
atleast_1d as atleast_1d,
327327
atleast_2d as atleast_2d,
328328
atleast_3d as atleast_3d,
329+
atleast_nd as atleast_nd,
329330
block as block,
330331
hstack as hstack,
331332
stack as stack,

‎numpy/core/shape_base.py

Lines changed: 132 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
1-
__all__ = ['atleast_1d', 'atleast_2d', 'atleast_3d', 'block', 'hstack',
2-
'stack', 'vstack']
1+
__all__ = ['atleast_1d', 'atleast_2d', 'atleast_3d', 'atleast_nd',
2+
'block', 'hstack', 'stack', 'vstack']
33

44
import functools
55
import itertools
@@ -17,6 +17,12 @@
1717
overrides.array_function_dispatch, module='numpy')
1818

1919

20+
def _unpack(lst):
21+
if len(lst) == 1:
22+
return lst[0]
23+
return lst
24+
25+
2026
def _atleast_1d_dispatcher(*arys):
2127
return arys
2228

@@ -42,7 +48,7 @@ def atleast_1d(*arys):
4248
4349
See Also
4450
--------
45-
atleast_2d, atleast_3d
51+
atleast_2d, atleast_3d, atleast_nd
4652
4753
Examples
4854
--------
@@ -61,18 +67,7 @@ def atleast_1d(*arys):
6167
[array([1]), array([3, 4])]
6268
6369
"""
64-
res = []
65-
for ary in arys:
66-
ary = asanyarray(ary)
67-
if ary.ndim == 0:
68-
result = ary.reshape(1)
69-
else:
70-
result = ary
71-
res.append(result)
72-
if len(res) == 1:
73-
return res[0]
74-
else:
75-
return res
70+
return _unpack([atleast_nd(a, 1, 0) for a in arys])
7671

7772

7873
def _atleast_2d_dispatcher(*arys):
@@ -87,9 +82,9 @@ def atleast_2d(*arys):
8782
Parameters
8883
----------
8984
arys1, arys2, ... : array_like
90-
One or more array-like sequences. Non-array inputs are converted
91-
to arrays. Arrays that already have two or more dimensions are
92-
preserved.
85+
One or more array-like sequences. Non-array inputs are
86+
converted to arrays. Arrays that already have two or more
87+
dimensions are preserved.
9388
9489
Returns
9590
-------
@@ -100,7 +95,7 @@ def atleast_2d(*arys):
10095
10196
See Also
10297
--------
103-
atleast_1d, atleast_3d
98+
atleast_1d, atleast_3d, atleast_nd
10499
105100
Examples
106101
--------
@@ -117,20 +112,7 @@ def atleast_2d(*arys):
117112
[array([[1]]), array([[1, 2]]), array([[1, 2]])]
118113
119114
"""
120-
res = []
121-
for ary in arys:
122-
ary = asanyarray(ary)
123-
if ary.ndim == 0:
124-
result = ary.reshape(1, 1)
125-
elif ary.ndim == 1:
126-
result = ary[_nx.newaxis, :]
127-
else:
128-
result = ary
129-
res.append(result)
130-
if len(res) == 1:
131-
return res[0]
132-
else:
133-
return res
115+
return _unpack([atleast_nd(a, 2, 0) for a in arys])
134116

135117

136118
def _atleast_3d_dispatcher(*arys):
@@ -145,22 +127,31 @@ def atleast_3d(*arys):
145127
Parameters
146128
----------
147129
arys1, arys2, ... : array_like
148-
One or more array-like sequences. Non-array inputs are converted to
149-
arrays. Arrays that already have three or more dimensions are
150-
preserved.
130+
One or more array-like sequences. Non-array inputs are
131+
converted to arrays. Arrays that already have three or more
132+
dimensions are preserved.
151133
152134
Returns
153135
-------
154136
res1, res2, ... : ndarray
155-
An array, or list of arrays, each with ``a.ndim >= 3``. Copies are
156-
avoided where possible, and views with three or more dimensions are
157-
returned. For example, a 1-D array of shape ``(N,)`` becomes a view
158-
of shape ``(1, N, 1)``, and a 2-D array of shape ``(M, N)`` becomes a
159-
view of shape ``(M, N, 1)``.
137+
An array, or list of arrays, each with ``a.ndim >= 3``. Copies
138+
are avoided where possible, and views with three or more
139+
dimensions are returned. For example, a 1-D array of shape
140+
``(N,)`` becomes a view of shape ``(1, N, 1)``, and a 2-D array
141+
of shape ``(M, N)`` becomes a view of shape ``(M, N, 1)``.
160142
161143
See Also
162144
--------
163-
atleast_1d, atleast_2d
145+
atleast_1d, atleast_2d, atleast_nd
146+
147+
Notes
148+
-----
149+
As mentioned in the `Returns` section, the results of this
150+
function are not consistent with any of the other `atleast*`
151+
functions. `atleast_2d` prepends the unit dimension to a 1D array
152+
while `atleast_3d` appends it to a 2D array. The 1D array case
153+
both appends and prepends a dimension, while `atleast_nd` can only
154+
add dimensions to one end at a time.
164155
165156
Examples
166157
--------
@@ -187,22 +178,105 @@ def atleast_3d(*arys):
187178
[[[1 2]]] (1, 1, 2)
188179
189180
"""
190-
res = []
191-
for ary in arys:
192-
ary = asanyarray(ary)
193-
if ary.ndim == 0:
194-
result = ary.reshape(1, 1, 1)
195-
elif ary.ndim == 1:
196-
result = ary[_nx.newaxis, :, _nx.newaxis]
197-
elif ary.ndim == 2:
198-
result = ary[:, :, _nx.newaxis]
199-
else:
200-
result = ary
201-
res.append(result)
202-
if len(res) == 1:
203-
return res[0]
204-
else:
205-
return res
181+
return _unpack([atleast_nd(atleast_nd(a, 2, 0), 3, -1) for a in arys])
182+
183+
184+
def _atleast_nd_dispatcher(ary, ndim, pos=None):
185+
return (ary,)
186+
187+
188+
@array_function_dispatch(_atleast_nd_dispatcher)
189+
def atleast_nd(ary, ndim, pos=0):
190+
"""
191+
View input as array with at least `ndim` dimensions.
192+
193+
New unit dimensions are inserted at the index given by `pos` if
194+
necessary.
195+
196+
Parameters
197+
----------
198+
ary : array_like
199+
The input array. Non-array inputs are converted to arrays.
200+
Arrays that already have `ndim` or more dimensions are
201+
preserved.
202+
ndim : int
203+
The minimum number of dimensions required.
204+
pos : int, optional
205+
The index to insert the new dimensions. May range from
206+
``-ary.ndim - 1`` to ``+ary.ndim`` (inclusive). Non-negative
207+
indices indicate locations before the corresponding axis:
208+
``pos=0`` means to insert at the very beginning. Negative
209+
indices indicate locations after the corresponding axis:
210+
``pos=-1`` means to insert at the very end. 0 and -1 are always
211+
guaranteed to work. Any other number will depend on the
212+
dimensions of the existing array. Default is 0.
213+
214+
Returns
215+
-------
216+
res : ndarray
217+
An array with ``res.ndim >= ndim``. A view is returned for array
218+
inputs. Dimensions are prepended if `pos` is 0, so for example,
219+
a 1-D array of shape ``(N,)`` with ``ndim=4`` becomes a view of
220+
shape ``(1, 1, 1, N)``. Dimensions are appended if `pos` is -1,
221+
so for example a 2-D array of shape ``(M, N)`` becomes a view of
222+
shape ``(M, N, 1, 1)`` when ``ndim=4``.
223+
224+
See Also
225+
--------
226+
atleast_1d, atleast_2d, atleast_3d
227+
228+
Notes
229+
-----
230+
This function does not follow the convention of the other
231+
``atleast_*d`` functions in numpy in that it only accepts a single
232+
array argument. To process multiple arrays, use a comprehension or
233+
loop around the function call. See examples below.
234+
235+
Setting ``pos=0`` is equivalent to how the array would be
236+
interpreted by numpy's broadcasting rules. There is no need to call
237+
this function for simple broadcasting. This is also roughly
238+
(but not exactly) equivalent to
239+
``np.array(ary, copy=False, subok=True, ndmin=ndim)``.
240+
241+
It is easy to create functions for specific dimensions similar to
242+
the other ``atleast_*d`` functions using Python's
243+
`functools.partial` function. An example is shown below.
244+
245+
Examples
246+
--------
247+
>>> np.atleast_nd(3.0, 4)
248+
array([[[[ 3.]]]])
249+
250+
>>> x = np.arange(3.0)
251+
>>> np.atleast_nd(x, 2).shape
252+
(1, 3)
253+
254+
>>> x = np.arange(12.0).reshape(4, 3)
255+
>>> np.atleast_nd(x, 5).shape
256+
(1, 1, 1, 4, 3)
257+
>>> np.atleast_nd(x, 5).base is x.base
258+
True
259+
260+
>>> [np.atleast_nd(x, 2) for x in ((1, 2), [[3, 4]], [[[5, 6]]])]
261+
[array([[1, 2]]), array([[3, 4]]), array([[[5, 6]]])]
262+
263+
>>> np.atleast_nd((1, 2), 5, pos=0).shape
264+
(1, 1, 1, 1, 2)
265+
>>> np.atleast_nd((1, 2), 5, pos=-1).shape
266+
(2, 1, 1, 1, 1)
267+
268+
>>> from functools import partial
269+
>>> atleast_4d = partial(np.atleast_nd, ndim=4)
270+
>>> atleast_4d([1, 2, 3])
271+
[[[[1, 2, 3]]]]
272+
"""
273+
ary = array(ary, copy=False, subok=True)
274+
pos = normalize_axis_index(pos, ary.ndim + 1)
275+
extra = operator.index(ndim) - ary.ndim
276+
if extra > 0:
277+
ind = pos * (slice(None),) + extra * (None,) + (Ellipsis,)
278+
ary = ary[ind]
279+
return ary
206280

207281

208282
def _arrays_for_stack_dispatcher(arrays, stacklevel=4):

‎numpy/core/shape_base.pyi

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,8 @@ def atleast_3d(__arys: ArrayLike) -> ndarray: ...
2626
@overload
2727
def atleast_3d(*arys: ArrayLike) -> List[ndarray]: ...
2828

29+
def atleast_nd(ary: ArrayLike, ndim: int, pos: int = ...) -> ndarray: ...
30+
2931
def vstack(tup: Sequence[ArrayLike]) -> ndarray: ...
3032
def hstack(tup: Sequence[ArrayLike]) -> ndarray: ...
3133
@overload

‎numpy/core/tests/test_shape_base.py

Lines changed: 68 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
1+
from decimal import Decimal
12
import pytest
23
import numpy as np
34
from numpy.core import (
4-
array, arange, atleast_1d, atleast_2d, atleast_3d, block, vstack, hstack,
5-
newaxis, concatenate, stack
6-
)
5+
array, arange, atleast_1d, atleast_2d, atleast_3d, atleast_nd,
6+
block, vstack, hstack, newaxis, concatenate, stack)
77
from numpy.core.shape_base import (_block_dispatcher, _block_setup,
88
_block_concatenate, _block_slicing)
99
from numpy.testing import (
@@ -123,6 +123,71 @@ def test_3D_array(self):
123123
assert_array_equal(res, desired)
124124

125125

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

‎numpy/lib/shape_base.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -542,7 +542,7 @@ def expand_dims(a, axis):
542542
--------
543543
squeeze : The inverse operation, removing singleton dimensions
544544
reshape : Insert, remove, and combine dimensions, and resize existing ones
545-
doc.indexing, atleast_1d, atleast_2d, atleast_3d
545+
doc.indexing, atleast_1d, atleast_2d, atleast_3d, atleast_nd
546546
547547
Examples
548548
--------

0 commit comments

Comments
 (0)
0