用 NumPy 与 Autograd 实现可微光学逆向设计

Runnable code

完整项目见 GitHub 仓库，代码按职责组织如下：

• 完整可运行实验：autograd-inverse-design.ipynb，包含可微 ASM、Gaussian 输入与目标场、mode-overlap loss、Autograd 梯度、Adam 优化，以及最终光场和三张相位图的可视化。
• 正文中的代码片段均摘自该 notebook，并按计算链路拆分讲解。

最近在学习论文 TorchOptics: An open-source Python library for differentiable Fourier optics simulations。

一、问题场景：用三块相位调制器完成光束变换

优化目标

给定一个输入光场和目标光场，不再人工设计每一块相位板，而是把相位分布当成待优化参数，通过光场传播、loss 和梯度下降，自动寻找合适的 phase profile。

论文 Figure 4 给出了一个很小但非常完整的例子：

光学系统布局

其中三块 phase modulator 分别位于

目标平面位于

数值参数

论文使用 $250\times250$ 的采样网格、$10,\mu{\rm m}$ pixel spacing、$700,{\rm nm}$ 波长和 $150,\mu{\rm m}$ Gaussian waist，并用 Adam、$lr=0.1$ 训练 400 iterations。

初始条件与待优化参数

输入是 waist 为 $150,\mu{\rm m}$ 的单束 Gaussian，目标是在四个指定位置形成四束 Gaussian。三块调制器的 phase profile

均从零相位开始，因此初始系统只是让输入 Gaussian 在自由空间传播 $0.6,{\rm m}$。优化要寻找这三张 phase map，使输出 complex field 与目标 mode 尽可能重合。

二、整体思路：把传播模型接入梯度优化

TorchOptics 用 PyTorch 把这些步骤全部封装好了。但如果只调用：

1
2

loss.backward()
optimizer.step()

其实很难真正理解 inverse design 到底发生了什么。

所以这次我不用 TorchOptics，而是只使用：

1
2
3

NumPy
Autograd
Matplotlib

自己搭出完整链路：

这里真正想验证的并不是能不能重新造一个 TorchOptics，而是：

一个普通的 Fourier optics simulator，到底怎样变成一个可以做 inverse design 的 simulator？

从 forward simulation 到 inverse design

以前写光场传播时，我解决的问题通常是：

输入场、传播距离和系统参数都是已知的，我们只计算输出。

例如 Angular Spectrum Method：

其中 transfer function 为

这属于典型的 forward problem：

inverse design 则反过来了。

我们知道输入：

也知道希望得到的目标：

但是不知道中间的 phase modulator 应该长什么样。

于是：

变成待求参数。

问题就变成：

这已经和训练神经网络非常像了。

只是我们训练的不是 neural-network weight，而是一个真实光学系统中的 phase profile。

三、建立可优化的光学问题

三平面 phase modulator 的 forward model

每块 phase modulator 的 transmission function 为

所以通过一块 phase-only SLM 后：

三块调制器组成的完整 forward model 为：

最后

用代码写出来反而非常简单：

1
2
3
4
5
6
7
8
9
10
11

def forward(U_in, phi1, phi2, phi3, H):
    U = phase_modulate(U_in, phi1)

    U = propagate(U, H)
    U = phase_modulate(U, phi2)

    U = propagate(U, H)
    U = phase_modulate(U, phi3)

    U = propagate(U, H)
    return U

这里的 H 是传播距离为 $0.2,{\rm m}$ 的 ASM transfer function；三段自由空间传播使用同一个 H。phase_modulate() 与 propagate() 会在后文分别定义。这个函数就是整个 inverse-design forward model。

输入 Gaussian beam

输入使用 Gaussian field：

注意这里定义的是 field amplitude，而不是 intensity。

因此：

代码：

1
2
3

def gaussian_beam(w0, C, X, Y, x0=0.0, y0=0.0):
    R2 = (X - x0)**2 + (Y - y0)**2
    return C * np.exp(-R2 / w0**2)

采用论文参数：

1
2
3
4
5
6
7
8
9

wavelength = 700e-9

Nx = 250
Ny = 250

dx = 10e-6
dy = 10e-6

w0 = 150e-6

所以整个 simulation window 为

构造 four-Gaussian target

论文 Listing 2 并没有直接生成 target，而是：

1

target_field = Field(torch.load("target.pt"), z=0.6)

也就是说，论文没有在 Listing 2 中给出四个 Gaussian 的具体位置参数。

因此这里并不是逐像素恢复论文的 target.pt，而是自己构造一个等价任务：

四个 Gaussian 分别位于：

这里取：

代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

def four_gaussian_target(X, Y, offset, waist, C):
    centers = [
        (-offset, -offset),
        (-offset, +offset),
        (+offset, -offset),
        (+offset, +offset),
    ]

    U = np.zeros_like(X, dtype=np.complex128)

    for x0, y0 in centers:
        U += gaussian_beam(
            waist,
            C,
            X,
            Y,
            x0,
            y0,
        )

    return U

因此这里的目标是：

归一化输入场与目标场

如果直接比较两个 field，它们的总功率也会影响 loss。

但现在真正关心的是：

输出 field 的 spatial mode 是否和 target 一样？

因此先定义 optical power：

离散化以后：

于是做 normalization：

1
2
3
4
5
6
7
8

def normalize(U):
    power = (
        np.sum(np.abs(U)**2)
        * dx
        * dy
    )

    return U / np.sqrt(power)

这样：

输入场和 target 都进行归一化。

Phase modulator 与零相位初始化

Phase-only modulator 不改变 amplitude，只改变 phase：

所以代码只有一行：

1
2

def phase_modulate(U, phi):
    return U * anp.exp(1j * phi)

一开始我们完全不知道 phase profile 应该是什么，因此初始化：

1
2
3

phi1 = np.zeros((Ny, Nx))
phi2 = np.zeros((Ny, Nx))
phi3 = np.zeros((Ny, Nx))

此时：

三块 SLM 什么都没有做。

因此最初的 optical system 实际上只是：

显然不可能自动变成四个 Gaussian。

接下来就需要优化。

四、构建可微的前向传播

Differentiable ASM

我这里使用 Angular Spectrum Method 作为 propagation backend。

先把输入场变换到 spatial-frequency domain：

自由空间传播会给每个 plane-wave component 加上不同 phase：

其中

以及

最后：

关键不在 ASM 本身。

关键在于现在：

依赖：

因此 FFT 也必须存在于 automatic differentiation 的 computational graph 中。

这里将固定的 frequency grid 和 transfer function 用普通 NumPy 预先计算，而依赖 phase 的传播路径使用 autograd.numpy：

1
2

import numpy as np
import autograd.numpy as anp

而不能在 differentiable path 中调用普通：

1

numpy.fft.fft2

否则 grad() 会把 phase 包装成 ArrayBox，普通 NumPy FFT 并不知道怎样处理它。

Autograd 官方文档说明它支持大部分 NumPy FFT 操作以及 complex number differentiation，因此这种包含 complex FFT 的 real-valued objective 可以进行 reverse-mode differentiation。

传播距离、波长和采样网格在优化过程中保持不变，因此先在 computational graph 外构造 $H$：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

def make_asm_transfer_function(
    Nx, Ny, dx, dy, wavelength, z
):
    fx = np.fft.fftshift(np.fft.fftfreq(Nx, d=dx))
    fy = np.fft.fftshift(np.fft.fftfreq(Ny, d=dy))
    FX, FY = np.meshgrid(fx, fy)

    k = 2 * np.pi / wavelength
    KX = 2 * np.pi * FX
    KY = 2 * np.pi * FY
    KZ = np.sqrt(
        (k**2 - KX**2 - KY**2).astype(np.complex128)
    )

    return np.exp(1j * z * KZ)


H = make_asm_transfer_function(
    Nx=Nx,
    Ny=Ny,
    dx=dx,
    dy=dy,
    wavelength=wavelength,
    z=propagation_distance,
)

这里的 Nx、Ny、dx、dy、wavelength 和 propagation_distance 就是问题场景中已经给出的数值参数。随后才定义可微的 FFT 传播：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

def fft2c(U):
    return anp.fft.fftshift(
        anp.fft.fft2(
            anp.fft.ifftshift(U)
        )
    )


def ifft2c(A):
    return anp.fft.fftshift(
        anp.fft.ifft2(
            anp.fft.ifftshift(A)
        )
    )


def propagate(U, H):
    A = fft2c(U)
    return ifft2c(A * H)

于是 computational graph 变成：

这一步完成以后，普通 Fourier-optics simulator 就变成了一个 differentiable simulator。

五、定义目标函数

Mode-overlap loss

接下来需要回答一个问题：

怎么判断 output 和 target 有多接近？

最直接的方法可能是 intensity MSE：

但论文并不是这样做的。

它优化的是两个 complex fields 的 mode overlap。论文 loss 为：

为了让实现对 field power 更稳健，我这里显式写成 normalized overlap：

然后：

其中：

当 output 完全等于 target mode：

因此：

代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

def mode_overlap(U_out, U_target, dx, dy):
    inner = (
        anp.sum(
            anp.conj(U_target)
            * U_out
        )
        * dx
        * dy
    )

    power_out = (
        anp.sum(anp.abs(U_out)**2)
        * dx
        * dy
    )

    power_target = (
        anp.sum(anp.abs(U_target)**2)
        * dx
        * dy
    )

    eta = (
        anp.abs(inner)**2
        / (power_out * power_target)
    )

    return anp.real(eta)

这里优化的不是简单的 intensity pattern。

而是 complex optical mode：

这也是 mode overlap 相比 intensity MSE 更重要的一点。

Objective function

整个 optical system 现在已经可以看成一个函数：

再把 output 送进 loss：

于是最终：

代码：

1
2
3
4

def objective(phi1, phi2, phi3, U_in, U_target, H, dx, dy):
    U_out = forward(U_in, phi1, phi2, phi3, H)
    eta = mode_overlap(U_out, U_target, dx, dy)
    return 1.0 - eta

这里把输入场、目标场、传播传输函数和积分采样间隔全部显式列出，避免 objective 依赖正文中没有交代的隐式全局状态。

六、从梯度到参数更新

Automatic differentiation

我们需要：

每个 $\phi$ 都是一个：

matrix。

也就是说单块 phase modulator 就有：

个 optimization variables。

三块一共：

个变量。

手推这样的 gradient 显然没有必要。

Autograd 可以直接根据 computational graph 使用 reverse-mode automatic differentiation。

1
2
3
4
5
6
7
8
9

grad_phi1 = grad(objective, 0)
grad_phi2 = grad(objective, 1)
grad_phi3 = grad(objective, 2)

objective_args = (U_in, U_target, H, dx, dy)

g1 = grad_phi1(phi1, phi2, phi3, *objective_args)
g2 = grad_phi2(phi1, phi2, phi3, *objective_args)
g3 = grad_phi3(phi1, phi2, phi3, *objective_args)

grad(objective, 0)、grad(objective, 1) 和 grad(objective, 2) 分别指定对前三个 phase 参数求导；其余参数只作为 forward calculation 所需的固定量传入。

至此最重要的链路已经跑通：

Adam 更新 phase

最后只剩 optimizer。

我没有再引入 PyTorch，而是简单实现了一遍 Adam。

对于 gradient $g_t$：

经过 bias correction：

最后：

代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50

class Adam:
    def __init__(
        self,
        shape,
        lr=0.1,
        beta1=0.9,
        beta2=0.999,
        eps=1e-8,
    ):
        self.lr = lr
        self.beta1 = beta1
        self.beta2 = beta2
        self.eps = eps

        self.m = np.zeros(shape)
        self.v = np.zeros(shape)
        self.t = 0

    def update(self, x, g):
        self.t += 1

        self.m = (
            self.beta1 * self.m
            + (1 - self.beta1) * g
        )

        self.v = (
            self.beta2 * self.v
            + (1 - self.beta2) * g**2
        )

        m_hat = (
            self.m
            / (1 - self.beta1**self.t)
        )

        v_hat = (
            self.v
            / (1 - self.beta2**self.t)
        )

        return (
            x
            - self.lr
            * m_hat
            / (
                np.sqrt(v_hat)
                + self.eps
            )
        )

三块 phase modulator 分别维护自己的 Adam state：

1
2
3

adam1 = Adam(phi1.shape, lr=0.1)
adam2 = Adam(phi2.shape, lr=0.1)
adam3 = Adam(phi3.shape, lr=0.1)

Optimization loop

最终训练循环其实已经非常接近 neural-network training：

1
2
3
4
5
6
7
8
9
10
11
12

objective_args = (U_in, U_target, H, dx, dy)

for i in range(num_iterations):
    loss = objective(phi1, phi2, phi3, *objective_args)

    g1 = grad_phi1(phi1, phi2, phi3, *objective_args)
    g2 = grad_phi2(phi1, phi2, phi3, *objective_args)
    g3 = grad_phi3(phi1, phi2, phi3, *objective_args)

    phi1 = adam1.update(phi1, g1)
    phi2 = adam2.update(phi2, g2)
    phi3 = adam3.update(phi3, g3)

其中 num_iterations=400，三个 Adam 实例分别保存三张 phase map 的一阶矩和二阶矩状态。

经过当前三块 phase masks，

得到：

计算它与：

之间的 mode mismatch，

再沿着：

方向调整所有 SLM pixels。

然后重新传播。

不断重复。

七、理解优化结果

怎样理解训练后的三张 phase map

最终得到：

它们通常看起来并不像某种简单 lens 或 grating。

这是正常的。

因为 optimizer 并没有被告知：

第一块负责分束，第二块负责聚焦，第三块负责修正。

它只知道最终 objective：

因此三块 phase mask 会联合工作。

可以理解为：

共同构成了一个 spatial-mode transformation。

第一块产生新的 spatial-frequency components，传播过程中这些频率分量积累不同 phase，第二和第三块继续重新组织 amplitude 和 phase，最终让目标 plane 上的 complex field 尽可能接近指定 target。

这就是 multi-plane light conversion 背后的基本思想之一。

八、边界、差异与实现注意点

与原论文的差异：这里固定使用 ASM

需要特别说明一点。

论文给出的 propagation-method selection criterion 为：

当前：

因此：

而每次 propagation distance 为：

也就是说：

按照论文的规则，TorchOptics 在这种情况下默认更倾向使用 Direct Integration，而不是 ASM。论文 Listing 2 没有显式指定 propagation method，而是让 System 自己处理。

所以这篇文章更准确的表述是：

用 NumPy + Autograd 重写 TorchOptics Figure 4 的 inverse-design architecture，并使用自己实现的 ASM 作为 differentiable propagation backend。

它不是 Figure 4 的逐像素数值复刻。

如果之后把 propagate() 换成 differentiable DI，其余 optimization pipeline 完全不需要改变。

这一点实际上也很好地体现了系统设计：

是两个可以独立替换的模块。

为什么优化中的 phase 可以超过 $2\pi$

训练过程中我没有执行：

1

phi = phi % (2 * np.pi)

因为真正进入 forward model 的是：

而：

所以数学上：

和：

完全等价。

因此 optimization 时允许 phase 自由变化没有问题。

最终显示或映射到实际 SLM 时，再使用：

九、总结：从 forward optics 到 inverse design

以前学习 Fourier optics 时，重点通常是：

也就是：

光经过一个系统以后会变成什么？

但 inverse design 把问题变成了：

也就是：

我想得到这个结果，中间的光学系统应该长什么样？

真正连接这两个问题的并不是某个新的 diffraction formula。

Forward physics 仍然是：

新增的其实只有三个东西：

也就是说：

forward model 决定：

loss 定义：

gradient 告诉我们：

optimizer 则不断执行这个修改。

这次真正从头写完之后，TorchOptics 里面：

1
2

loss.backward()
optimizer.step()

这两行代码背后的物理和计算过程也就不再是黑盒了。

许可协议

本文采用 署名-非商业性使用-相同方式共享 4.0 国际 许可协议，转载请注明出处。

分享文章