nan を含む配列の要素をシフトする

配列の要素をシフトさせる関数 shift() のことを前回書いたけど、この関数は配列内の要素にnanが含まれると使えない（出力する配列の要素が全てnanになる）という欠点がある。そして残念ながら今書いてるスクリプトでは配列にnanが入るので、結局自分でコードを書いた。

 前回記事はこちら。

spcx8.hatenablog.com

まずは1次元配列をシフトさせるコード。

import numpy as np

def shift_1darray(mat, ys):
    mat_shift = np.zeros_like(mat)
    mat_shift[:,:] = np.nan
    if   ys>0: mat_shift[ys:] = mat[:-ys]
    elif ys<0: mat_shift[:ys] = mat[-ys:]
    else: mat_shift = mat
    return mat_shift

 それから2次元配列をシフトさせるコード。

import numpy as np

def shift_2darray(mat, ys, xs):
    mat_yshift = np.zeros_like(mat)
    mat_yshift[:,:] = np.nan
    if   ys>0: mat_yshift[ys:  ,:] = mat[   :-ys,:]
    elif ys<0: mat_yshift[  :ys,:] = mat[-ys:   ,:]
    else: mat_yshift = mat

    mat_xshift = np.zeros_like(mat)
    mat_xshift[:,:] = np.nan
    if   ys>0: mat_xshift[:,xs:  ] = mat_yshift[:,   :-xs]
    elif ys<0: mat_xshift[:,  :xs] = mat_yshift[:,-xs:   ]
    else: mat_xshift = mat_yshift
    return mat_xshift

もっと簡単にできるよ、という方いたらぜひ教えてください。行数減らしたいです。