python numpy roll with padding

Question

I'd like to roll a 2D numpy in python, except that I'd like pad the ends with zeros rather than roll the data as if its periodic.

Specifically, the following code

import numpy as np

x = np.array([[1, 2, 3], [4, 5, 6]])

np.roll(x, 1, axis=1)

returns

array([[3, 1, 2],[6, 4, 5]])

but what I would prefer is

array([[0, 1, 2], [0, 4, 5]])

I could do this with a few awkward touchups, but I'm hoping that there's a way to do it with fast built-in commands.

Thanks

Answer 1

There is a new numpy function in version 1.7.0 numpy.pad that can do this in one-line. Pad seems to be quite powerful and can do much more than a simple "roll". The tuple ((0,0),(1,0)) used in this answer indicates the "side" of the matrix which to pad.

import numpy as np
x = np.array([[1, 2, 3],[4, 5, 6]])

print np.pad(x,((0,0),(1,0)), mode='constant')[:, :-1]

Giving

[[0 1 2]
 [0 4 5]]

Answer 2

I just wrote the following. It could be more optimized by avoiding zeros_like and just computing the shape for zeros directly.

import numpy as np
def roll_zeropad(a, shift, axis=None):
    """
    Roll array elements along a given axis.

    Elements off the end of the array are treated as zeros.

    Parameters
    ----------
    a : array_like
        Input array.
    shift : int
        The number of places by which elements are shifted.
    axis : int, optional
        The axis along which elements are shifted.  By default, the array
        is flattened before shifting, after which the original
        shape is restored.

    Returns
    -------
    res : ndarray
        Output array, with the same shape as `a`.

    See Also
    --------
    roll     : Elements that roll off one end come back on the other.
    rollaxis : Roll the specified axis backwards, until it lies in a
               given position.

    Examples
    --------
    >>> x = np.arange(10)
    >>> roll_zeropad(x, 2)
    array([0, 0, 0, 1, 2, 3, 4, 5, 6, 7])
    >>> roll_zeropad(x, -2)
    array([2, 3, 4, 5, 6, 7, 8, 9, 0, 0])

    >>> x2 = np.reshape(x, (2,5))
    >>> x2
    array([[0, 1, 2, 3, 4],
           [5, 6, 7, 8, 9]])
    >>> roll_zeropad(x2, 1)
    array([[0, 0, 1, 2, 3],
           [4, 5, 6, 7, 8]])
    >>> roll_zeropad(x2, -2)
    array([[2, 3, 4, 5, 6],
           [7, 8, 9, 0, 0]])
    >>> roll_zeropad(x2, 1, axis=0)
    array([[0, 0, 0, 0, 0],
           [0, 1, 2, 3, 4]])
    >>> roll_zeropad(x2, -1, axis=0)
    array([[5, 6, 7, 8, 9],
           [0, 0, 0, 0, 0]])
    >>> roll_zeropad(x2, 1, axis=1)
    array([[0, 0, 1, 2, 3],
           [0, 5, 6, 7, 8]])
    >>> roll_zeropad(x2, -2, axis=1)
    array([[2, 3, 4, 0, 0],
           [7, 8, 9, 0, 0]])

    >>> roll_zeropad(x2, 50)
    array([[0, 0, 0, 0, 0],
           [0, 0, 0, 0, 0]])
    >>> roll_zeropad(x2, -50)
    array([[0, 0, 0, 0, 0],
           [0, 0, 0, 0, 0]])
    >>> roll_zeropad(x2, 0)
    array([[0, 1, 2, 3, 4],
           [5, 6, 7, 8, 9]])

    """
    a = np.asanyarray(a)
    if shift == 0: return a
    if axis is None:
        n = a.size
        reshape = True
    else:
        n = a.shape[axis]
        reshape = False
    if np.abs(shift) > n:
        res = np.zeros_like(a)
    elif shift < 0:
        shift += n
        zeros = np.zeros_like(a.take(np.arange(n-shift), axis))
        res = np.concatenate((a.take(np.arange(n-shift,n), axis), zeros), axis)
    else:
        zeros = np.zeros_like(a.take(np.arange(n-shift,n), axis))
        res = np.concatenate((zeros, a.take(np.arange(n-shift), axis)), axis)
    if reshape:
        return res.reshape(a.shape)
    else:
        return res

Answer 3

import numpy as np

def shift_2d_replace(data, dx, dy, constant=False):
    """
    Shifts the array in two dimensions while setting rolled values to constant
    :param data: The 2d numpy array to be shifted
    :param dx: The shift in x
    :param dy: The shift in y
    :param constant: The constant to replace rolled values with
    :return: The shifted array with "constant" where roll occurs
    """
    shifted_data = np.roll(data, dx, axis=1)
    if dx < 0:
        shifted_data[:, dx:] = constant
    elif dx > 0:
        shifted_data[:, 0:dx] = constant

    shifted_data = np.roll(shifted_data, dy, axis=0)
    if dy < 0:
        shifted_data[dy:, :] = constant
    elif dy > 0:
        shifted_data[0:dy, :] = constant
    return shifted_data

This function would work on 2D arrays and replace rolled values with a constant of your choosing.

Answer 4

A bit late, but feels like a quick way to do what you want in one line. Perhaps would work best if wrapped inside a smart function (example below provided just for horizontal axis):

import numpy

a = numpy.arange(1,10).reshape(3,3)  # an example 2D array

print a

[[1 2 3]
 [4 5 6]
 [7 8 9]]

shift = 1
a = numpy.hstack((numpy.zeros((a.shape[0], shift)), a[:,:-shift]))

print a

[[0 1 2]
 [0 4 5]
 [0 7 8]]

Answer 5

Elaborating on the answer by Hooked (since it took me a few minutes to understand it)

The code below first pads a certain amount of zeros in the up, down, left and right margins and then selects the original matrix inside the padded one. A perfectly useless code, but good for understanding np.pad.

import numpy as np
x = np.array([[1, 2, 3],[4, 5, 6]])
y = np.pad(x,((1,3),(2,4)), mode='constant')[1:-3,2:-4]

print np.all(x==y)

now to make an upwards shift of 2 combined with a rightwards shift of 1 position one should do

print np.pad(x,((0,2),(1,0)), mode='constant')[2:0,0:-1]

Answer 6

You could also use numpy's triu and scipy.linalg's circulant. Make a circulant version of your matrix. Then, select the upper triangular part starting at the first diagonal, (the default option in triu). The row index will correspond to the number of padded zeros you want.

If you don't have scipy you can generate a nXn circulant matrix by making an (n-1) X (n-1) identity matrix and stacking a row [0 0 ... 1] on top of it and the column [1 0 ... 0] to the right of it.

Answer 7

I faced a similar problem with shifting a 2-d array in both directions

def shift_frame(img,move_dir,fill=np.inf):
    frame = np.full_like(img,fill)
    x,y = move_dir
    size_x,size_y = np.array(img.shape) - np.abs(move_dir)
    frame_x = slice(0,size_x) if x>=0 else slice(-x,size_x-x)
    frame_y = slice(0,size_y) if y>=0 else slice(-y,size_y-y)
    img_x = slice(x,None) if x>=0 else slice(0,size_x)
    img_y = slice(y,None) if y>=0 else slice(0,size_y)
    frame[frame_x,frame_y] = img[img_x,img_y]
    return frame

test = np.arange(25).reshape((5,5))
shift_frame(test,(1,1))
'''
returns:
array([[ 6,  7,  8,  9, -1],
       [11, 12, 13, 14, -1],
       [16, 17, 18, 19, -1],
       [21, 22, 23, 24, -1],
       [-1, -1, -1, -1, -1]])
'''

I haven't measured the runtime of this, but it seems to work well enough for my use, although a built-in one liner would be nice

Answer 8

You can also use ndimage.shift:

>>> from scipy import ndimage
>>> arr = np.array([[1, 2, 3], [4, 5, 6]])
>>> ndimage.shift(arr, (0,1))
array([[0, 1, 2],
       [0, 4, 5]])

Answer 9

import numpy as np

def roll_zeropad(a, dyx):
    h, w = a.shape[:2]
    dy, dx = dyx
    pad_x, start_x, end_x = ((dx,0), 0, w) if dx > 0 else ((0,-dx), -dx, w-dx)
    pad_y, start_y, end_y = ((dy,0), 0, h) if dy > 0 else ((0,-dy), -dy, h-dy)
    return np.pad(a, (pad_y, pad_x))[start_y:end_y,start_x:end_x]

test = np.arange(25).reshape((5,5))
out = roll_zeropad(test,(1,1))
print(out)
"""
returns:

[[ 0  0  0  0  0]
 [ 0  0  1  2  3]
 [ 0  5  6  7  8]
 [ 0 10 11 12 13]
 [ 0 15 16 17 18]]

"""