NumPyのpiecewiseをCuPy向けに書き換える

NumPyのpiecewise

NumPyのpiecewiseで、区分関数を実現することができる。

f(x) = \left\{
\begin{array}{ll}
x^2 & (x \geq 0) \\
0 & (x \lt 0)
\end{array}
\right.

これをpiecewiseで書くと、次のようになる。

import numpy as np

x = np.arange(-4, 5)
np.piecewise(x, [x >= 0, x < 0], [lambda v: v ** 2, 0])

これを実行すると、

array([ 0,  0,  0,  0,  0,  1,  4,  9, 16])

となる。
x >= 0 の範囲で x の2乗、x < 0 の範囲で0になっている。

piesewiseを使用しない方法

CuPyはpiecewiseに対応していない（v7.2時点）ため、別の方法で記述することになる。
パフォーマンスが悪いif文は使わずに記述する。

書き方は簡単で、区分の条件と関数の値の掛け合わせた結果を合計すればよい。

import cupy as cp

x = cp.arange(-4, 5)
positive_condition = x >= 0
positive_value = x ** 2
negative_condition = x < 0
negative_value = 0
positive_condition * positive_value + negative_condition * negative_value

これを実行すると、

array([ 0,  0,  0,  0,  0,  1,  4,  9, 16])

と、piecewiseと同じ結果となる。

この書き方はもちろんNumPyでも有効。