pwtools.num.extend_array

pwtools.num.extend_array(arr, nstep, axis=0)[source]

Repeat an array along axis by inserting a new axis (dimension) before axis. Use this to “broadcast” e.g. a 2d array (3,3) -> (3,3,nstep).

Parameters:
  • arr (ndarray)

  • nstep (int, number of times to repeat)

  • axis (axis to add)

Examples

>>> a=arange(4)
>>> extend_array(a, 3, 0)
array([[0, 1, 2, 3],
       [0, 1, 2, 3],
       [0, 1, 2, 3]])
>>> extend_array(a, 3, 0).shape
(3, 4)
>>> extend_array(a, 3, 1)
array([[0, 0, 0],
       [1, 1, 1],
       [2, 2, 2],
       [3, 3, 3]])
>>> extend_array(a, 3, 1).shape
(4, 3)
>>> a=arange(4).reshape(2,2)
>>> extend_array(a, 3, 0).shape
(3, 2, 2)
>>> extend_array(a, 3, 1).shape
(2, 3, 2)
>>> extend_array(a, 3, 2).shape
(2, 2, 3)
>>> extend_array(a, 3, 2)[...,0]
array([[0, 1],
       [2, 3]])
>>> extend_array(a, 3, 2)[...,1]
array([[0, 1],
       [2, 3]])
>>> extend_array(a, 3, 2)[...,2]
array([[0, 1],
       [2, 3]])

See also

numpy.repeat()