光子数基底での光量子計算を実装する(Numpyはえらい)

$$
\def\bra#1{\mathinner{\left\langle{#1}\right|}}
\def\ket#1{\mathinner{\left|{#1}\right\rangle}}
\def\braket#1#2{\mathinner{\left\langle{#1}\middle|#2\right\rangle}}
$$

## 前書き
[こちらの記事](https://qiita.com/ryuNagai/items/df0b69565e4dc24fba62)ではDisplacement Operatorを例に光子数基底での光量子状態の発展を実装しました。
今回は密度演算子表記での状態発展を実装し、Wigner関数を計算します。
状態ベクトル表記ではなく密度演算子表記を用いる理由ですが、率直に言うとウィグナー関数関連の文献が基本的に密度演算子表記で書かれていることです。

余談
(真面目な理由としては、個々の位相情報が不完全なアンサンブル系を考えた時に状態ベクトル表記はボロが出るので密度演算子表記のほうが良いという記述があります[[1]](https://onlinelibrary.wiley.com/doi/book/10.1002/3527602976)。
プログラムとしての実装に影響があるかはわかりません、今のところ無いと思っています。)


## 光子数密度演算子の状態発展
光子数基底の密度演算子は以下のように書けます。

$\rho = \sum_{l, m}C_{lm}\ket{l}\bra{m}$

これにDisplacement Operatorを作用させます。

$\rho' = D^{\dagger}(\alpha)\rho D(\alpha)$

Gaussian Wigner functionで確認した通りなら、位相平面($\hat{x}, \hat{p}$)上のWigner関数を並行移動させるはずです。

## 光子数密度演算子のWigner関数
$\rho = \ket{n'}\bra{n}$ に対するWigner関数は次のように定式化されています[2]。
(文献では"Moyal function for harmonic oscillator"と呼ばれます)

$n \leq n'$の場合
$W_{nn'}(q, p) = \frac{2(-1)^n}{\pi} \sqrt{\frac{2^{n'} n!}{2^{n} n'!}} (q-ip)^{n'-n}e^{-(q^2+p^2)}L_{n}^{n'-n}(2(q^2+p^2))$  
$n' \leq n$の場合
$W_{nn'}(q, p) = \frac{2(-1)^{n'}}{\pi} \sqrt{\frac{2^{n} n'!}{2^{n'} n!}} (q+ip)^{n-n'}e^{-(q^2+p^2)}L_{n'}^{n-n'}(2(q^2+p^2))$  

$L_l^{m}(x)$はラゲール陪多項式と文献では紹介されています。
ただラゲール陪多項式は文献によって異なる定義をされることがあるようで、実際に上式とScipyモジュールのラゲール陪多項式では異なります。

調べてみるとソニン多項式というものがラゲール陪多項式と呼ばれることがあるそうです[[3]](http://www.phys.shimane-u.ac.jp/tanaka_lab/lecture/math1/math1.pdf)。
上式もそのパターンでした。なんじゃそりゃ。
とりあえず実装に必要な式だけつまみ食いします。

ソニン多項式
-漸化式

```math
nS^{\alpha}_{n}(x)= (2n-1+\alpha-x)S^{\alpha}_{n-1}(x)-(n + \alpha -1)S^{\alpha}_{n-2}(x) \\
S^{\alpha}_0 (x) = 1 \\
S^{\alpha}_1 (x) = \alpha + 1 - x
```
-展開式

```math
S_{n}^{\alpha}(x) = \sum^{n}_{k=0}\frac{(-1)^k (n+\alpha)!}{(n-k)!(\alpha + k)! k!}x^k
```

## 実装

### Displacement Gate

まず[こちらの記事](https://qiita.com/ryuNagai/items/df0b69565e4dc24fba62)で作ったup(), down(), exp_annihiration(), exp_creation(), displacement()を持ってきます。

```python
import numpy as np
from scipy import special
import matplotlib.pyplot as plt
import time

def down(order, n):
    if order > n:
        return [0, 0]
    n_ = n - order
    coeff = np.prod(np.sqrt(np.arange(n, n - order, -1)))
    return [coeff, n_]

def up(order, n, cutoff = 10):
    n_ = n + order
    if n_ > cutoff:
        return [0, 0]
    coeff = np.prod(np.sqrt(np.arange(n + 1, n + 1 + order)))
    return [coeff, n_]

def exp_annihiration(fockState, alpha, cutoff):
    if np.isscalar(fockState):
        list_ = [fockState]
        fockState = np.ones(fockState + 1)
    else:
        list_ = np.nonzero(fockState)[0]
        
    state = np.zeros(cutoff + 1) + 0j
    for j in list_:
        for i in range(cutoff + 1):
            tmp =  down(i, j)
            state[np.int(tmp[1])] += tmp[0] * fockState[j] / special.factorial([i]) * alpha ** i
    return state

def exp_creation(fockState, alpha, cutoff):
    if np.isscalar(fockState):
        list_ = [fockState]
        fockState = np.ones(fockState + 1)
    else:
        list_ = np.nonzero(fockState)[0]
        
    state = np.zeros(cutoff + 1) + 0j
    for j in list_:
        for i in range(cutoff + 1):
            tmp =  up(i, j)
            state[np.int(tmp[1])] += tmp[0] * fockState[j] / special.factorial([i]) * alpha ** i
    return state
```

これを使って、密度演算子に対するDisplacement Gateを実装します。

```python
def Displacement(rho, alpha, cutoff=10):
    dim = cutoff + 1
    rho_ = np.zeros([dim, dim]) + 0j
    for row in range(dim):
        for col in range(dim):
            ket = displacement(row, alpha, cutoff)
            bra = displacement(col, alpha, cutoff)
            rho_ += rho[row, col] * np.outer(np.conj(ket), bra)
    return rho_
```

早速vacuum stateに作用させて見ます。
最後に対角成分の和を出力します(1になるはず)。

```python
 # prepare vacuum state
rho = np.zeros([11, 11])
rho[0, 0] = 1
# Displacement
rho_ = Displacement(rho, 1+1j)
print(np.sum(np.diag(rho_)))
```

出力結果

```python
(0.9999916917756311+0j)
```
ここではこれをもって良しとします。

### ソニン多項式
漸化式を使ってナイーブに実装します。

```python
def Sonin(n, alpha, x):
    S = np.empty([n + 1])
    for i in range(n + 1):
        if i == 0:
            S[0] = 1
        elif i == 1:
            S[1] = alpha + 1 - x
        else:
            res = (2 * i - 1 + alpha - x) * S[i - 1] - (i + alpha - 1) * S[i - 2]
            S[i] = res / i
    return S[n]
```

### Wigner関数 (Moyal function for harmonic oscillator)

これまたナイーブに実装します。
密度演算子の定義
$\rho = \sum_{l, m}C_{lm}\ket{l}\bra{m}$
に従い、

FockWignerElement()で$\ket{l}\bra{m}$に対するWigner関数を計算します。

FockWigner()で、すべての$l, m$について係数$C_{lm}$をかけて足し合わせ、$\rho$を算出します。

```python
def FockWignerElement(x, p, l, m):
    r = x**2 + p**2
    if l >= m:
        W = 2 * (-1)**m * np.sqrt(2**(l - m) * special.factorial(m) / special.factorial(l)) * (x + 1j * p)**(l - m) * np.exp(-r)
        S = Sonin(m, l - m, 2 * r)
        return W * S
    else:
        W = 2 * (-1)**l * np.sqrt(2**(m - l) * special.factorial(l) / special.factorial(m)) * (x - 1j * p)**(m - l) * np.exp(-r)
        S = Sonin(l, m - l, 2 * r)
        return W * S

def FockWigner(alpha, rho, tol = 1e-10):
    x = alpha[0]
    p = alpha[1]
    dim = rho.shape[0]
    W = 0 + 0j
    for row in range(dim):
        for col in range(dim):
            W += rho[row, col] * FockWignerElement(x, p, row, col)
    if np.max(np.imag(W)) < tol:
        W = np.real(W)
    else:
        raise ValueError("Wigner function has imaginary value.")
    return W
```

これらを使って、位相平面($\hat{x}, \hat{p}$)上の各点についてWigner関数の値を計算し、プロットします。

```python
x = np.arange(-5, 5, 0.1)
p = np.arange(-5, 5, 0.1)
rho = np.zeros([11, 11])
rho[0, 0] = 1
rho = Displacement(rho, 1-1j)
m = len(x)
xx, pp = np.meshgrid(x, p)
xi_array = np.dstack((pp, xx))
start = time.time()
W = np.zeros([m, m], dtype = "complex128")
for i in range(m):
    for j in range(m):
        W[i][j] = FockWigner(xi_array[j][i], rho)

h = plt.contourf(x, p, W)
t = time.time() - start
print(t)
```

出力結果

```python
47.98835206031799 # Wigner関数計算にかかった時間(s)
```
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/243619/b77615cf-cb70-d83a-73ce-7a6cc10a9565.png)

プロットはまあ良いでしょう。vacuum stateが並行移動しています。ちょっと余計な縞が見えていますが多分数値計算上の誤差で、Strawberry Fieldsもこんなもんです。
問題は実行時間で、50秒近くかかっています。

[環境]
MacBook Pro (13-inch, 2017)
2.3 GHz Intel Core i5
8 GB 2133 MHz LPDDR3

そう、これが古典計算機の限界...ではないです。Pythonのforループの限界です。上の実装だと4重forループになっています(Sonin多項式の部分だけ見れば5重)。
このままでは使い物にならないので高速化します。
Pythonの高速化は語られ尽くしていますが、今回はforループをNumpyの配列演算にひたすら押し込む方向で頑張りました。

### 高速化

まずはソニン関数をNumpy配列演算で書き直します。
結局1回だけforループを回していますが、ループ数は光子のカットオフ数なのでそこまで大きな値にならないです。
(カットオフ数増やしてくと多分先にメモリが問題になります。)

```python
def to_2d_ndarray(a):
    if isinstance(a,(np.ndarray)):
        return a
    else:
        return np.array([[a]])

def Sonin_vec(n, alpha, x):
    n = to_2d_ndarray(n)
    alpha = to_2d_ndarray(alpha)
    x = to_2d_ndarray(x)        
    N = np.max(n)
    a = special.factorial(n + alpha)
    S = np.zeros([N + 1, x.shape[0], x.shape[0], n.shape[0], n.shape[0]], dtype = "float64")
    for i in range(N + 1):
        if i == 0:
            S[0] = 1
        elif i == 1:
            I = np.where(n - i >= 0, 1, 0)
            I = I[np.newaxis, np.newaxis, :, :]
            I_bar = np.logical_not(I).astype("float64")
            S[1] = (alpha + 1 - x) * I + I_bar
        else:
            I = np.where(n - i >= 0, 1, 0)
            I = I[np.newaxis, np.newaxis, :, :].astype("float64")
            I_bar = np.logical_not(I).astype(int)
            res = (2 * i - 1 + alpha - x) * S[i - 1] - (i + alpha - 1) * S[i - 2]
            S[i] = res / i * I + S[i -1] * I_bar
    return S[N]
```

Wigner関数の計算もNumpy配列演算に押し込みます。
色々やって4次元配列演算にすべて押し込みました。

```python
def FockWignerElement_vec(xmat, pmat, l, m):
    A = np.max(np.dstack([l, m]), axis=2)
    B = np.abs(l - m)
    C = np.min(np.dstack([l, m]), axis=2)
    xmat = xmat[:, :, np.newaxis, np.newaxis]
    pmat = pmat[:, :, np.newaxis, np.newaxis]
    R = xmat**2 + pmat**2
    X = xmat + np.sign(l-m) * 1j * pmat
    W = 2 * (-1)**C * np.sqrt(2**(B) * special.factorial(C) / special.factorial(A)) * X**(B) * np.exp(-R)
    S = Sonin_vec(C, B, 2 * R)
    return W * S

def FockWigner_vec(xmat, pmat, rho, tol=1e-10):
    dim = rho.shape[0]
    grid = np.indices([dim, dim])
    W = FockWignerElement_vec(xmat, pmat, grid[0], grid[1])
    W = rho * W
    W = np.sum(np.sum(W, axis = -1), axis = -1)
    if np.max(np.imag(W)) < tol:
        W = np.real(W)
    else:
        raise ValueError("Wigner plot has imaginary value.")
    return W
```

後は上の関数を使ってプロットするだけです。

```python
x = np.arange(-5, 5, 0.1)
p = np.arange(-5, 5, 0.1)
rho = np.zeros([11, 11])
rho[0, 0] = 1
rho = Displacement(rho, 1-1j)
m = len(x)
xx, pp = np.meshgrid(x, p)
xi_array = np.dstack((pp, xx))
start = time.time()
        
W = FockWigner_vec(xx, pp, rho)

h = plt.contourf(x, p, W)
t = time.time() - start
print(t)
```

出力結果

```python
0.42597198486328125 # Wigner関数計算にかかった時間(s)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/243619/8989d62c-cf6c-0c9c-5020-0ade6c51a454.png)

頑張った甲斐あり、2桁も速くなりました。
Python forループとは何だったのか...。

他ゲートや複数モードを含めた実装の目処がついたら下のライブラリに上げると思います。
先は長い...。
https://github.com/Blueqat/Photonqat

以上、読んで頂きありがとうございました。


## 参考文献
[1] "Quantum Optics in Phase Space"(https://onlinelibrary.wiley.com/doi/book/10.1002/3527602976)
[2] 10.1103/PhysRevA.86.062117
[3] http://www.phys.shimane-u.ac.jp/tanaka_lab/lecture/math1/math1.pdf