Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Appearance settings

Commit 2474c47

Browse filesBrowse files
madphysicistmadphysicist
committed
ENH: Added atleast_nd to numpy and ma
1 parent f6a71fe commit 2474c47
Copy full SHA for 2474c47

File tree

14 files changed

+262
-73
lines changed
Filter options

14 files changed

+262
-73
lines changed

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

Copy file name to clipboard
+4Lines 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

Copy file name to clipboardExpand all lines: doc/source/reference/routines.array-manipulation.rst
+1Lines 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

Copy file name to clipboardExpand all lines: doc/source/reference/routines.ma.rst
+1Lines 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

Copy file name to clipboardExpand all lines: doc/source/user/quickstart.rst
+1Lines 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

Copy file name to clipboardExpand all lines: numpy/__init__.pyi
+1Lines 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

Copy file name to clipboardExpand all lines: numpy/core/shape_base.py
+132-58Lines 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

Copy file name to clipboardExpand all lines: numpy/core/shape_base.pyi
+2Lines 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

Copy file name to clipboardExpand all lines: numpy/core/tests/test_shape_base.py
+68-3Lines 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

Copy file name to clipboardExpand all lines: numpy/lib/shape_base.py
+1-1Lines 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)
Morty Proxy This is a proxified and sanitized view of the page, visit original site.