(※より効率的な方法を求めています)
高速フーリエ変換をする下準備として、配列の要素数を2のべき乗にしたいときがあります。
要素数を2のべき乗に切り上げ、後ろを0で埋める方法を書いておきます。
いくつか関数を定義します。また、各関数の引数に対してチェックすべきこともあるので、そちらについても適宜記し、本記事の最後に例外raiseを含めたものをまとめます。
正の整数を2のべき乗に切り上げ
まずは任意の正の整数を2のべき乗に切り上げる関数を定義します。(負数のチェックは省きます)
import numpy as np
def round_pow2(x :int) -> int:
"""xを2のべき乗に切り上げる
60 -> 64
40000 -> 65536
"""
return 1 << x.bit_length()
print(round_pow2(1000))
# => 1024
この関数はxのbit列の先頭の1までの長さを取得し、1をその分大きくなる側へシフトしています。bit列1 0 0 ... 0は2のべき乗になります。2 ** x.bit_length()でも同じです。
NumPy配列を広げある数値で埋める
ほしいのは0埋めですが、「ある数値で埋める」という抽象化が可能なのでそのように定義します。この関数もチェックすべきことはありますが省きます。
def back_padding(xs: np.ndarray, n: int, a) -> np.ndarray:
"""n個の要素になるように後ろをaで埋める"""
v = np.empty(n, dtype=type(a))
v[:xs.shape[0]] = xs
v[xs.shape[0]:] = a
return v
print(back_padding(np.arange(10), 15, 0))
#=> [0 1 2 3 4 5 6 7 8 9 0 0 0 0 0]
チェックすべきこととして、xsが1次元配列であること、およびxsの要素数がn以下である必要があります。aの型がxsの要素の型と一致していることも必要です(こちらはそれほど気にする必要はないかもしれません)。
要素数が2のべき乗になるように後ろをある数値で埋める
上記で定義した関数を組み合わせれば完成です。こちらもback_paddingと同じく一次元配列であること、配列の要素の型とaの要素の型が一致している必要があります。
def to_pow2_length(xs: np.ndarray, a) -> np.ndarray:
"""要素数が2のべき乗になるように後ろをaで埋める"""
return back_padding(xs, round_pow2(xs.shape[0]), a)
アレンジするとしたら、埋める側の前後を選択できるようにしたり、初期値としてaを0や0.0にしたりすることだと思います。私個人的にはこの定義で十分と思いましたのでこのままにしました。
まとめ
以下にdocstringおよび例外を含めた定義をまとめます。
import numpy as np
def round_pow2(x :int) -> int:
"""xを2のべき乗に切り上げる
例)
60 -> 64
40000 -> 65536
Arguments:
x (int > 0): 正の整数
Returns:
int: 2のべき乗の整数
"""
if x < 1:
raise ValueError("xは正の整数である必要があります。")
return 1 << x.bit_length()
def back_padding(xs: np.ndarray, n: int, a) -> np.ndarray:
"""n個の要素になるように後ろをaで埋める
Arguments:
xs (np.ndarray[shape=(m,)]): 配列
n (int > m): 要素数
a (xs.dtype): 埋める値
Returns:
np.ndarray[shape=(n,)]: 余分なところをaで埋めた配列
"""
if len(xs.shape) > 1:
raise ValueError("配列xsは一次元配列である必要があります。")
if xs.shape[0] > n:
raise ValueError("nは配列xsの要素数より大きい必要があります。")
if xs.dtype != type(a):
raise ValueError("aの型は配列xsの要素の型と一致している必要があります。")
v = np.empty(n, dtype=type(a))
v[:xs.shape[0]] = xs
v[xs.shape[0]:] = a
return v
def to_pow2_length(xs: np.ndarray, a) -> np.ndarray:
"""要素数が2のべき乗になるように後ろをaで埋める
Arguments:
xs (np.ndarray[shape=(n,)]): 一次元配列
a (xs.dtype): 埋める値
Returns:
np.ndarray[shape=(2**p,)]: 要素数が2のべき乗で後ろがaで埋まった配列
"""
return back_padding(xs, round_pow2(xs.shape[0]), a)