【位相計算:第2回】ホログラムを綺麗にしたい ― WGS法とPyTorchでどこまで改善できるか

光学

はじめに

 前回の記事で、特定の光強度パターンを作るための位相を求めるアルゴリズムとして、GS法を解説しました。実際に「YOHAKU」という文字を作ってみると、文字の形は綺麗に再現できました。
ところが、よく見ると輝度がボツボツしていて、均一ではありません。

「GS法のアルゴリズムを改善すれば、この輝度ムラも綺麗にできるのでは?」

 そこで本記事では、GS法を改良したWGS法と、PyTorchを用いた勾配最適化を使って、ホログラムの輝度ムラがどこまで改善できるのか試してみます。
果たして、アルゴリズムを変えることでホログラムはどこまで綺麗になるのでしょうか?

前回記事 : アイキャッチ画像がWS法の結果。まだらになっています

WGS(Weighted Gerchberg–Saxton)法

概要

 WGS(Weighted Gerchberg–Saxton)法は、GS法に重み付けを導入した手法です。目標画像と再構成画像の輝度差に応じて各画素の重みを更新することで、輝度ムラを抑え、より均一な強度分布を得ることを目指します。

WGS法の計算結果


WGS法のポイントはweightを導入している点です。
WGS法では、各スポットの再構成された振幅 $A_i$​ と目標振幅$A_i^{target}$ ​ の比$r_i$を用いて、スポットごとのweightを更新します。

$$ r_i(n)​=\frac{A_i^{target}}{A_i(n)​+ε} $$​​

ここで、$A_i(n)$​ は n 回目の反復で得られた再構成振幅、$A_i^{target}$​ は目標振幅、$ε$ はゼロ除算を防ぐための小さな値を足しこみます。

この振幅比を用いて、weightを

$$ w_i(n+1)​=w_i(n)(\frac{A_i^{target}}{A_i(n)​+ε})^p $$

と更新して再構成振幅が目標振幅より小さい領域ほど、大きな重みが与えられるように補正します。
$p$は0.5~1程度で調整します。1に近くなるほど補正が強くなります。

実際に計算してみます。ターゲットは前回記事と同じです。

前回記事と同じターゲット画像

以下が計算コードで、”target”変数にnumpy array形式で画像読み込まれているとします。

import numpy as np
import matplotlib.pyplot as plt

# ========================================
# Parameters
# ========================================

iterations = 200
p = 0.8

# target amplitude
target_amp = np.sqrt(target)

# ========================================
# Input amplitude
# ========================================
#ガウシアンビーム

x = np.linspace(-1, 1, N)
y = np.linspace(-1, 1, N)

X, Y = np.meshgrid(x, y )

w0 = 0.3

amplitude = np.exp(
    -(X**2 + Y**2) / w0**2
)

# ========================================
# Initial phase
# ========================================

phase = 2 * np.pi * np.random.rand(*target.shape)

# ========================================
# Initial weights
# ========================================

weights = np.ones_like(target_amp)

# ========================================
# WGS
# ========================================

for i in range(iterations):

    # -----------------------------
    # Forward propagation
    # -----------------------------

    field = amplitude * np.exp(1j * phase)

    propagated = np.fft.fftshift(
        np.fft.fft2(field, norm="ortho")
    )

    recon_amp = np.abs(propagated)

    # -----------------------------
    # Weight update
    # -----------------------------

    eps = 1e-12

    ratio = target_amp / (recon_amp + eps)

    weights *= ratio ** p

    # -----------------------------
    # Apply target amplitude
    # -----------------------------

    new_amp = target_amp * weights

    # -----------------------------
    # Back propagation
    # -----------------------------

    target_field = (
        new_amp *
        np.exp(1j * np.angle(propagated))
    )

    field_back = np.fft.ifft2(
        np.fft.ifftshift(target_field),
        norm="ortho"
    )

    # -----------------------------
    # Phase-only constraint
    # -----------------------------

    phase = np.angle(field_back)



intensity = np.abs(propagated) ** 2

display = intensity / intensity.max()


plt.figure(figsize=(5, 5))
plt.imshow(
    display,
    cmap="gray",
    vmin=0,
    vmax=1
)
plt.axis("off")
plt.show()

計算結果はこうなりました。

WGS法による結果


輝点があるため、全体的に暗くみえます。あまり改善しているように見えません。
$p$の値を振ってもあまり変化がある感じもしません。

PyTorchによる勾配最適化

GS法やWGS法とは異なるアプローチとして、近年の機械学習や数値最適化で広く利用されている勾配最適化を試してみます。

PyTorchと勾配最適化法の概要

 PyTorchは、主にディープラーニングで利用されているPythonライブラリです。
ニューラルネットワークの学習では、モデルの出力と正解との誤差(Loss)を計算し、その勾配を使ってモデルのパラメータを更新します。

このような勾配を利用した最適化は、近年の機械学習や数値最適化で広く使われているアプローチです。ニューラルネットワークに限らず、微分可能な計算であれば同じ仕組みを利用できます。

そこで今回は、ニューラルネットワークのパラメータではなく、ホログラムの位相を最適化する変数として扱います。目標画像との誤差(Loss)が小さくなるように位相の勾配を計算し、繰り返し更新することで目標とする光強度パターンを再現します。

計算結果

PyTorchの一番の特徴は自動微分機能です。
Lossを最適化したいパラメータで微分することで、値の更新を行います。自動微分のために、NumPy形式そのままでは使えず、Torch.tensorという形式で扱う必要があり少し癖があります。
今回最適化するパラメータは位相で512 x 512点もありますが、大量のパラメータを同時に最適化できる点が大きな利点です。

 損失(loss)を計算する際は、少し工夫してターゲットの明るい部分と暗い部分で別々に計算して足し合わせました。optimizerはよく使われるAdamを用いました。

import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import numpy as np

device = "cuda" if torch.cuda.is_available() else "cpu"


# -----------------------
# 1. 目標画像 YOHAKU
# -----------------------

# targetをtorchの型に変換
target_torch = torch.tensor(target, dtype=torch.float32, device=device)
target_torch = target_torch /target_torch.max()

# -----------------------
# 2. 初期位相
# -----------------------
#位相
phase = (
    2 * torch.pi *
    torch.rand((N, N), device=device)
).requires_grad_()

#ガウシアンビーム
x = torch.linspace(-1, 1, N, device=device)
y = torch.linspace(-1, 1, N, device=device)

X, Y = torch.meshgrid(x, y, indexing="ij")

w0 = 0.3

amplitude = torch.exp(
    -(X**2 + Y**2) / w0**2
)

# -----------------------
# 3. 最適化
# -----------------------

optimizer = torch.optim.Adam(
    [phase],
    lr=0.05
)
mask = target > 0.1
loss_record =[]
for i in range(500):
    #勾配の初期化
    optimizer.zero_grad()

    field = amplitude * torch.exp(1j*phase)

    #FFTによる伝搬
    propagated = torch.fft.fftshift(
        torch.fft.fft2(field,norm='ortho')
    )

    intensity = torch.abs(propagated)**2
    intensity = intensity / intensity.max()
    
    target_mean = target_torch[mask].mean()
    #ターゲットの明るい部分だけでlossを計算
    loss_shape = torch.mean(
        (intensity[mask]-target_torch[mask])**2
    )
    
    #ターゲットの暗い部分だけでlossを計算    
    loss_uniform = torch.mean(
        (intensity[mask] -target_mean ) **2
    )

    loss = loss_shape + 5*loss_uniform

    loss.backward()

    optimizer.step()

    if i%100==0:
        print(i, loss.item())


# -----------------------
# 結果
# -----------------------
plt.imshow(
    intensity.detach().cpu(),
    cmap="gray"
)
plt.axis("off")
plt.show()

>>>>0 4.649687767028809
>>>>100 3.148967742919922
>>>>200 2.396944999694824
>>>>300 2.424318790435791
>>>>400 2.5147194862365723

lossの値を表示させています。200回あたりで最適化は完了しているようです。

PyTorchを使った勾配最適化の結果。ボツボツがひどい。

輝度は一定になっておらず改善が見えません。
lossの定義やoptimizerの設定など微調整自体はまだ出来ますが、GS法と比較しても同じ感じであまり改善するような感じがしません。

計算条件を見直す

どちらの方法でもあまり改善が見られませんでした。
・位相を理論上で最適化できたとして、どこまで綺麗になるものなんだろう?
・アルゴリズムの問題でなく、もっと他の要因が効いているのかも?

WGS法に立ち戻ります。
入射はガウシアンビームとしていましたが、調査のため強度均一の光(平面波)にしてみます。

入力の定義

x = np.linspace(-1, 1, N)
X, Y = np.meshgrid(x, x)
r2 = X**2 + Y**2

beam_radius = 0.3
A_in = np.exp(-r2 / beam_radius**2)

以下で平面波に修正します。

x = np.linspace(-1, 1, N)
X, Y = np.meshgrid(x, x)

A_in = np.ones_like(X)

平面波入射+WGS法の結果です。

平面波入射の場合の結果、とても綺麗

とても綺麗です…!!ターゲットと遜色ありません。
アルゴリズムに注目していましたが、アルゴリズムの問題ではなかったようです。

考察

最初は512×512の内の30%を占めるガウシアンビームを入力として想定していました。
それに対してターゲットの”YOHAKU”は、縦幅は全体の20%程度ですが、横幅はギリギリまで広がっています。

 ガウシアンビームとすることで実効的なNAが小さくなり、回折パターンの広がりや分解能が変わったことでターゲット画像の空間的な広がりを表すことが難しくなっていたのだと思います。
アルゴリズムの問題ではなく、光学系の問題だったのです。

レーザ光で強度均一を実現するのは、難しいためガウシアンとしていました。
また、あまり詳細を詰めると話が分かりづらいのでフラウンフォーファー回折想定でFFTをつかって無限遠での回折像を検討対象としていました。

実際には入力のNAと結像の大きさの対応のチェックが必要でした。
以前結像計算の像の大きさの考え方が分からなくてまとめ記事を書いたのですが、今回の位相計算ではその点が抜け落ちていました。
当たり前のことですが、物理的に実現できないことは、どんなアルゴリズムをつかっても実現できないのです。

まとめ

今回はWGS法とPyTorchを用いた勾配最適化による位相分布計算の方法を説明しました。

・WGS法はターゲットと画像の振幅の差に基づいて重みを更新する
・PyTorchはlossを下げるように、位相を勾配を用いて更新する
・WGS法とPyTorchを用いても輝度均一は改善しなかったが、問題の光学系の設定によるものだった

参考文献

ディープラーニングの話に少し触れましたが、ニューラルネットワークの仕組みをNumPyなどを使って一から実装しながら学べる定番の入門書です。。
著者の難しいことは分割して説明するスタイルが好きです。続編も出ている人気シリーズです。

フーリエ光学を日本語で学びたい方におすすめの一冊です。 回折やフーリエ変換、ホログラムも分かりやすく解説しており、GS法の背景となる光学理論を体系的に理解できます。

光学計算や深層学習をやるにはPythonの活用が必須と言えます。最近ではレンズ設計の最適化でもバックエンドで(PyTorchのような)勾配最適化を使っている場合もあります。
体系的に学びたい方は、オンライン講座を活用するのも一つの方法です。

コメント

タイトルとURLをコピーしました