普段は量子コンピュータのシミュレータをPythonで書いています。
その都合上、複素数は大変よく使います。
Python, numpyは複素数もサポートされており、あんまり困ることはないのですが、いくつか罠もあるので、そこらへんをまとめておきます。
Pythonでの複素数の基本
このへんは、普通にリファレンスに載ってるやつです。軽くおさらい。
虚数単位は1j、型はcomplex
数学では多くの場合、虚数単位にはiが使われますが、Pythonでは1jが使われます。
(電気などの分野で、iは電流に使っているので虚数単位をjと置く、というのを見たことがありますが、それ以外の分野では他に知りません)
2jとか、1.23jのような書き方もできます。(普通に1jの2倍、1.23倍、という意味)
1と1.で型に違いがあるのに対し、1jと1.jはどちらも同じです。
a, bをfloatとして、a + b * 1.jはcomplex(a, b)と書くこともできます。
.real, .imagで実部、虚部を取り出す conjugate()で複素共役をとる
(1 + 2j).real は 1, (1 + 2j).imag は 2 です。メソッドではないので()を付けないことに注意してください。
複素共役は(1 + 2j).conjugate()のように、.conjugate()を使います。こちらはメソッドなので()が必要です。
ちなみに.real, .imag, conjugate() は、実はfloatやintに対しても使えます。
numpyも普通に使える
numpyでもnp.array([1 + 2j, 3, 4j, 5+6.7j])のように、普通に複素数が使えます。このときdtypeは、デフォルトでは'complex128'になります。
.real, .imag, .conjugate()もnumpyのarrayに対して使うことができます。
また、numpyでは、.conjugate()のかわりに.conj()でも構いません。
なお、行列のエルミート共役取るのが.conj().Tとする必要があり、やや面倒です。
np.matrixだと.Hでできたんですが。np.arrayにはないようです。
mathのかわりにcmath
import mathのかわりにimport cmathを使います。
mathの関数が定義域や値域を実数にとっているのに対し、cmathは定義域や値域に複素数も含まれます。
詳細は公式ドキュメント参照。
複素数の偏角を求めるphase()や、半径と偏角のタプルを返すpolar()や、半径と偏角から複素数を作るrect()といった便利な関数も用意されてます。
Pythonで複素数を扱うときの注意
いくつか罠やその対処を書いていきます。
複素数になりうるなら、numpyのarrayは予め複素数で作る
Pythonは自分で型を書かなくても暗黙のうちに型変換してくれますが、numpyでは、arrayは型を持っていて、勝手にarrayの型が変わることはありません。(intのarrayがfloatにならないのも同様です)
a = np.array([1., 2.]) # dtypeがfloat64
a *= 1j
# UFuncTypeError: Cannot cast ufunc 'multiply' output from dtype('complex128') to dtype('float64') with casting rule 'same_kind'
こうならないために
np.array([1., 2.], dtype=np.complex128)のように最初から型を指定するか、
np.array([1+0j, 2.], dtype=np.complex128)のように最初から複素数を含めるかが必要です。
複素ベクトルの内積はnp.dotではなくnp.vdot
これは本当に罠です。
複素ベクトルの内積は$u\cdot v = \sum_{i=1}^n u_i^* v_i$のように、$u$の複素共役と$v$との各成分の積の和を取るのが普通ですが、np.dotでは、複素共役をとらず、そのまま各成分の積の和を取ります。どうしてもnp.dotを使う必要がある場合は、自分で複素共役を取ってつかいましょう。
一方、np.vdotを使うと、複素ベクトルであっても適切に内積を取れます。
当たり前だけど、複素数同士は大小の比較ができない
他の言語で複素数を扱っている人には当たり前ですが。複素数同士は比較ができません。
より正確には、「複素数型」の大小比較ができないので、プログラマが計算結果が実数になっていると知っていても、コンピュータにとって複素数型だったら、大小比較はできません。.realなどで実数にする必要があります。
複素数配列の「絶対値の2乗」は効率がいい方法がないらしい
np.absは、内部的にはsqrt(real**2 + imag**2)のような計算をしているので、これを2乗するのは無駄が多いです。
一方、arr.real ** 2 + arr.imag **2のような計算をPython側で書くと、インプレースで計算されず、無駄なメモリ確保が生じます。
以前調べたことがあるんですが、StackOverflowには、numba使え、と書いてありました。Python, numpyだけでインプレースに無駄な開平計算をせずにやる方法は意外とないみたいです。