if文を含む自作関数をndarrayに一括適用したい

Question

したいこと
if文を含む以下の関数をndarrayにfor文を使わずに適用したいです。

def to3dim(x):
    zeros = np.zeros_like(x)
    if x[0] - x[-1] <= 0:
        y = np.dstack((x, zeros)).transpose((2, 0, 1))
    else:
        y = np.dstack((zeros, x)).transpose((2, 0, 1))
    return y

適用したいndarrayの形状は(n, 4)です。それを(2, n, 4)としたいです。

[[0, 1, 2, 3],
 [7, 6, 5, 4]]
を
[[[0, 1, 2, 3],
  [0, 0, 0, 0]],
 [[0, 0, 0, 0],
  [7, 6, 5, 4]]]
のようにする。

frompyfuncやvectorizeをただ適用しただけではエラーが出てしまいました。
どうすればよいでしょうか。お力添えお願いします。

Comments

[tiitoi]

> x[0] - x[-1] <= 0
は、(1行目 - n行目) のベクトルの"すべて"の値が0以下という意図でしょうか？

[yuru-camper]

いえ、各行の要素のつもりです。
もしかしてこれだと行単位で考えていますか？

[tiitoi]

print(x[0] - x[-1] <= 0) # [ True True True True]
となります。

[[0, 1, 2, 3],
[7, 6, 5, 4]]
からどのような規則で以下の出力が導出されたのかが質問のコードからは推察できないため、回答できません。
[[[0, 1, 2, 3],
[0, 0, 0, 0]],
[[0, 0, 0, 0],
[7, 6, 5, 4]]]

[tiitoi]

文章で入力から出力が作られる過程を書いてください。
それがわかれば、numpy でのコードを提示できます。

[yuru-camper]

x[0] と x[-1] のどちらが大きいかでdstackの引数の順番が変わる、というコードです。
例えばx = [0, 1, 2, 3] は x[0] < x[-1] なのでnp.dstack((x, zeros))、
x = [7, 6, 5, 4] は x[0] < x[-1] なのでnp.dstack((zeros, x))、という感じです。

Answer 1

1. x が (N, M) の配列だとすると、(2 * N, M) という結果を格納する配列を用意しておく。

2. x[:, 0] <= x[:, -1] を満たす行は0が下、そうでない行は0が上にくる。bool を int にキャストすると、True は1、False は0になることを利用して、0を入れる行を特定する。

import numpy as np

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

def process(x):
    # x が (N, M) の配列だとすると、(2 * N, M) という結果を格納する配列を用意しておく。
    rows, cols = x.shape
    ret = np.empty((rows * 2, cols), dtype=x.dtype)

    pos = np.arange(0, len(ret), 2)
    condition = x[:, 0] <= x[:, -1]
    print('pos', pos)  # [0 2]
    print('condition', condition)  # [ True False]

    # condition を満たす行は0が下、そうでない行は0が上にくる。
    zero_rows = pos + condition.astype(int)  # zero の行
    not_zero_rows = pos + (~condition).astype(int)  # zero でない行

    ret[zero_rows] = 0
    ret[not_zero_rows] = x
    
    return ret

print(process(x))
# [[0 1 2 3]
#  [0 0 0 0]
#  [0 0 0 0]
#  [7 6 5 4]]