Pythonで複素数を含む配列の各要素の2乗を返す方法（numpy.square）

Pythonで配列の各要素（要素ごと）の2乗を返すには、NumPyのnumpy.square()メソッドを使用します。このメソッドは要素ごとの x*x を計算し、入力 x と同じ形状・同じdtypeを持つ結果を返します。xがスカラーの場合は、スカラー値が返されます。

numpy.square()の主なパラメータ

第1引数：x
計算対象となる入力データです。複素数を含む配列も指定できます。

第2引数：out
結果を格納する場所を指定します。指定する場合は、入力がブロードキャストされる形状と一致している必要があります。省略した場合（None）や未指定の場合は、新しく確保された配列が返されます。タプルとして渡すことも可能ですが、その場合はキーワード引数としてのみ指定でき、長さは出力の数と一致していなければなりません。

第3引数：where
この条件は入力に対してブロードキャストされます。条件がTrueの場所では、out配列にufuncの計算結果が設定されます。それ以外の場所では、out配列は元の値を保持します。なお、デフォルトの out=None によって初期化されていない配列が作成された場合、条件がFalseとなった箇所は初期化されないまま残る点に注意してください。

実行手順

まず、必要なライブラリをインポートします。

import numpy as np

array()メソッドを使ってNumPy配列を作成します。ここでは、複素数型の要素を含む2次元配列を用意しました。

arr = np.array([[3 + 4j, 5 + 7j], [2 + 6j, -2j]])

配列の内容を表示します。

print("Our Array...\n",arr)

配列の次元数を確認します。

print("\nDimensions of our Array...\n",arr.ndim)

配列のデータ型（dtype）を取得します。

print("\nDatatype of our Array object...\n",arr.dtype)

最後に、numpy.square()メソッドを使って、配列入力の各要素の2乗を返します。結果は入力と同じ形状・dtypeになります。

print("\nResult...\n",np.square(arr))

サンプルコード全体

import numpy as np

# array()メソッドでNumPy配列を作成
# 複素数型の要素を含む配列
arr = np.array([[3 + 4j, 5 + 7j], [2 + 6j, -2j]])

# 配列を表示
print("Our Array...\n",arr)

# 次元数を確認
print("\nDimensions of our Array...\n",arr.ndim)

# データ型を取得
print("\nDatatype of our Array object...\n",arr.dtype)

# numpy.square()で各要素の2乗を計算
print("\nResult...\n",np.square(arr))

出力結果

Our Array...
[[ 3.+4.j 5.+7.j]
[ 2.+6.j -0.-2.j]]

Dimensions of our Array...
2

Datatype of our Array object...
complex128

Result...
[[ -7.+24.j -24.+70.j]
[-32.+24.j -4. +0.j]]

このように、numpy.square()を使えば、複素数を含む配列でも各要素の2乗を簡単に計算できます。例えば (3+4j) の2乗は -7+24j となり、複素数の乗算規則に従った正しい結果が得られます。