はじめに
これまで本ブログでは、Pythonを使って波動光学に基づく光伝播計算や位相回復計算をスクラッチで実装してきました。
今回は、波動光学計算ライブラリであるTorchOpticsを使って、これまで扱ってきた計算を再実装してみます。
TorchOpticsはPyTorchをバックエンドとしており、光の伝播計算だけでなく、自動微分を利用した勾配法による光学設計にも対応しています。
前半ではガウシアンビームや高次モードの伝播計算を行い、後半では4f光学系を構築して、目標の強度分布を得るための位相を計算してみます。
TorchOpticsってなに？
TorchOpticsはPyTorchをベースに構築された、フーリエ光学を用いて光学系をシミュレーションするためのオープンソースのPythonライブラリです。
微分可能な波動光学シミュレーションを提供していて、GPUを用いた高速化にも対応しているようです。
メリット
・シンプルなコードで伝播計算ができる
・お手軽に可視化できる
・Systemクラスで光学系をまとめた扱える
・勾配最適化を前提としていて、逆設計がやりやすい
デメリット
・少し凝ったことをやろうとすると、コードが煩雑になりやすい
興味がある方は論文を参照ください。 https://arxiv.org/abs/2411.18591
インストール方法
pipでインストールできます。
pip install torchopticsガウシアンビームの伝播計算
TorchOpticsでガウシアンビームと高次のラゲールガウシアンビームを伝播させてみます。
以前にスクラッチで伝播を計算したことがあるので、詳細に興味ある方は読んでみてください。
最初に関連ライブラリをインポートして、波長、グリッド数を指定します。
import torch
import torchoptics
from torchoptics import Field,System
from torchoptics.elements import AmplitudeModulator, Lens
from torchoptics.profiles import checkerboard, circle, gaussian
# 波長を設定
torchoptics.set_default_spacing(10e-6)
# グリッド間隔を設定
torchoptics.set_default_wavelength(532e-9)
# 1辺のピクセル数設定
shape=1000よく使われるガウシアン、エルミートガウシアン、更にレンズ、アパーチャなどは用意されています。
これらは、torch.Tensor形式で与えられる点に注意が必要です。
半径1mmのガウシアンをvisualメソッドで手軽に可視化できます。
gaussian_beam=Field(gaussian(shape,waist_radius=1e-3))
gaussian_beam.visualize(title="gaussian beam")
上図で左が強度分布、右が位相です。
伝播計算もpropagate_to_z()メソッドで計算できます。
焦点距離1mのレンズで集光した際のプロファイルを可視化してみます。
f = 1#焦点距離
lens = Lens(shape,f,z=0)
focused = lens(gaussian_beam).propagate_to_z(f)
focused.visualize(title="focused gaussian")
ビーム径が小さくなっており集光されているのが分かります。
位相は理想的にはフラットのはずですが、幾何学模様のようになっていますが、ビームが集光されている近辺では位相一定としても良さそうです。
次に、ラゲールガウシアンの伝播を計算してみます。
ラゲールガウシアンは同心円状のプロファイルであり、ファイバーレーザーのビームプロファイルの記述に使われるのをよく見ます。
$p=1,l=1$のラゲールガウシアンは以下のようなものです。
lg_gauss=Field(torchoptics.profiles.laguerre_gaussian(shape,1,1,1e-3))
lg_gauss.visualize(title="Laguerre Gaussian $LG_{p=1,l=1}$")
伝播させてみます。
形を保ったまま伝播し、焦点距離付近で一番小さく集光されていることが分かります。
1-2行程度のコードで計算出来てしまいます！
勾配を用いた位相回復計算
TorchOpticsは、torchを使用しており勾配を用いた計算ができます。
この機能を使って、ターゲットの強度分布を得るための位相を求めてみます(位相回復計算)。
問題設定
出力面で”YOHAKU”をターゲットにします。
入力面では、強度分布はガウシアンとして位相を求めます。
入力から出力面の間は4f光学系とします。
以前に本記事ではGS(Gerchberg–Saxton)法と呼ばれる古くから知られているアルゴリズムで位相を求めました。興味のある方は読んでみてください。
4f光学系とは？
4f光学系は、同じ焦点距離fをもつレンズ2枚を使った光学系です。
モノ-レンズ1-レンズ2-結像位置が順に配置された光学系です。全体で焦点距離fの4倍の距離があります。
等倍で結像され、レンズ1とレンズ2の中間に瞳面が位置する基本的な光学系です。
波動光学では瞳面が非常に重要なため、瞳面位置が明らかで扱いやすくよく使われます。
TorchOpticsではSystem クラスがあり、複数レンズからなる光学系をまとめて扱えます。
勾配法による位相回復計算
まずはターゲットを設定します。
TorchOpticsは計算にtorchをつかっているため、ターゲットはtorch.Tensor形式にします。
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt
# 白黒画像を読み込む
img = Image.open("YOHAKU.png").convert("L")
img = img.resize((shape, shape))
# numpy配列へ
target = np.array(img)
# 白文字を1、背景を0
target = target / 255.0
# 振幅
target = torch.from_numpy(np.sqrt(target))
target = Field(target,z=4*f)
target_intensity = target.data.abs().square()入射面のガウシアンビームは少し大きめに半径5mmとして規格化します。
入力面にLCOS-SLMのような位相変調装置があるとして4f光学系を設定します。
TorchOpticsではPhaseModulatorとして位相変調装置が実装されています。
import torch
from torch.nn import Parameter
import torchoptics
from torchoptics import Field, System
from torchoptics.elements import PhaseModulator,Lens
from torchoptics.profiles import gaussian
f=0.1
shape=500
gaussian_beam = gaussian(shape,waist_radius=5e-3)
gaussian_beam/= gaussian(shape,waist_radius=5e-3).data.abs().max().item()
input_field=Field(gaussian_beam)
system_4f = System(
PhaseModulator(Parameter(torch.zeros(shape,shape)),z=0),
Lens(shape,f,z=f),
Lens(shape,f,z=3*f)
)手元の環境だと少し計算が重かったため、グリッド数は半分の500に減らしました。
焦点距離は10cmとしました。
勾配計算を行います。
損失関数はMSEで設定します。
optimizer = torch.optim.Adam(system_4f.parameters(), lr=0.05)
for iteration in range(200):
optimizer.zero_grad()
output = system_4f.measure_at_z(input_field, z=4*f)
output_intensity = output.data.abs().square()
loss = torch.mean((output_intensity - target_intensity) ** 2)
loss.backward()
optimizer.step()
if (iteration + 1) % 10 == 0:
print(f"Iteration [{iteration + 1}/200], Loss: {loss.item():.4f}")Iteration [10/200], Loss: 0.0128
Iteration [20/200], Loss: 0.0083
Iteration [30/200], Loss: 0.0054
Iteration [40/200], Loss: 0.0036
Iteration [50/200], Loss: 0.0026
Iteration [60/200], Loss: 0.0020
Iteration [70/200], Loss: 0.0017
Iteration [80/200], Loss: 0.0014
Iteration [90/200], Loss: 0.0013
Iteration [100/200], Loss: 0.0012
Iteration [110/200], Loss: 0.0011
Iteration [120/200], Loss: 0.0010
Iteration [130/200], Loss: 0.0009
Iteration [140/200], Loss: 0.0009
Iteration [150/200], Loss: 0.0008
Iteration [160/200], Loss: 0.0008
Iteration [170/200], Loss: 0.0008
Iteration [180/200], Loss: 0.0007
Iteration [190/200], Loss: 0.0007
Iteration [200/200], Loss: 0.0007lossの値を見ると、収束しているようです。
可視化してみます。
with torch.no_grad():
result = system_4f.measure_at_z(input_field, z=4*f)
result.visualize(title="Optimized Output")
綺麗に計算できました。
よく見ると、中心が明るくて外側が若干暗くなっていて、元のガウシアンを反映してます。
前回記事の教訓から1辺の長さが5mmに対して、ガウシアンは半径5mmとかなり大き目にしたのですが、それでも元のガウシアンの影響は無視できませんでした。
位相の値を求めるだけでなく、光学系の結像倍率も合わせて考えないといけません。
関連記事
まとめ
Pythonの光学シミュレーションライブラリ TorchOptics を使って、ガウシアンビームの伝搬から4f光学系の構築、位相最適化による強度分布の生成までを試してみました。
TorchOpticsを使うことで、光学素子を組み合わせた光学系を比較的シンプルに記述でき、PyTorchの自動微分機能を利用して位相分布を最適化できることが確認できました。
今回は基本的な使い方の紹介が中心でしたが、今後は古典的な計算方法であるGS法との比較や、最適化条件による生成パターンの違いなども検討していきたいと思います。
参考文献
フーリエ光学を日本語で学びたい方におすすめの一冊です。 回折やフーリエ変換、ホログラムも分かりやすく解説しており、角スペクトル法の背景となる光学理論を体系的に理解できます。
PyTorchはディープラーニングに使われるライブラリであり、ニューラルネットワークの仕組みをNumPyなどを使って一から実装しながら学べる定番の入門書です。
著者の難しいことは分割して説明するスタイルが好きです。
勾配最適化をどう計算・実装しているか、なぜTorchOpticsがTorch.Tensor形式なのかを理解したい人は続編の③がおすすめです。
コメント