利用Julia编写时域有限差分方法(FDTD)及其并行化


#1

关于FDTD

FDTD是计算电磁学领域用的最多的算法之一(其他的包括MoM,以及大名鼎鼎的FEM),FDTD也可用于求解一些量子力学中的问题,其核心思想就是对电磁场中的电场、磁场分别在时间域和空间域进行差分,然后通过时间和空间上的迭代进行时域求解,这里以一维的FDTD为例,写出一维情况下的Maxwell 方程组:

\begin{aligned} - \frac { \partial H _ { y } } { \partial z } & = \varepsilon \frac { \partial E _ { x } } { \partial t } + \sigma E _ { x } \\ \frac { \partial E _ { x } } { \partial z } & = - \mu \frac { \partial H _ { y } } { \partial t } - \sigma _ { m } H _ { y } \end{aligned}

对其分别在空间和时间上进行差分

E _ { x } ^ { n- 1 } ( k ) = \mathrm { CA } ( m ) \cdot E _ { x } ^ { n } ( k ) - \mathrm { CB } ( m ) \cdot \frac { H _ { y } ^ { n + 1 / 2 } \left( k + \frac { 1 } { 2 } \right) - H _ { y } ^ { n + 1 / 2 } \left( k - \frac { 1 } { 2 } \right) } { \Delta z }
H _ { y } ^ { n+ 1 / 2 } \left( k + \frac { 1 } { 2 } \right) = \mathrm { CP } ( m ) \cdot H _ { y } ^ { n- 1 / 2 } \left( k + \frac { 1 } { 2 } \right) - \mathrm { CQ } ( m ) \cdot \frac { E _ { x }^n ( k + 1 ) - E _ { x } ^n( k ) } { \Delta z }

其中

\begin{aligned} C A ( m ) & = \frac { \frac { \varepsilon ( m ) } { \Delta t } - \frac { \sigma ( m ) } { 2 } } { \frac { \varepsilon ( m ) } { \Delta t } + \frac { \sigma ( m ) } { 2 } } = \frac { 1 - \frac { \sigma ( m ) \Delta t } { 2 \epsilon ( m ) } } { 1 + \frac { \sigma ( m ) \Delta t } { 2 \varepsilon ( m ) } } \\ C B ( m ) & = \frac { 1 } { \frac { \varepsilon ( m ) } { \Delta t } + \frac { \sigma ( m ) } { 2 } } = \frac { \frac { \Delta t } { \varepsilon ( m ) } } { 1 + \frac { \sigma ( m ) \Delta t } { 2 \varepsilon ( m ) } } \\ C P ( m ) & = \frac { \frac { \mu ( m ) } { \Delta t } - \frac { \sigma _ { m } ( m ) } { 2 } } { \frac { \mu ( m ) } { \Delta t } + \frac { \sigma _ { m } ( m ) } { 2 } } = \frac { \frac { 1 - \sigma _ { m } ( m ) \Delta t } { 2 \mu ( m ) } } { 1 + \frac { \sigma _ { m } ( m ) \Delta t } { 2 \mu ( m ) } } \end{aligned}
\mathrm { CQ } ( m ) = \frac { 1 } { \frac { \mu ( m ) } { \Delta t } + \frac { \sigma _ { m } ( m ) } { 2 } } = \frac { \frac { \Delta t } { \mu ( m ) } } { 1 + \frac { \sigma _ { m } ( m ) \Delta t } { 2 \mu ( m ) } }

其中 kn 分别为空间步和时间步的索引,详细过程可以参考《电磁波时域有限差分方法》,葛德彪著。

利用Julia编写FDTD

之前有写过CPP版本的(其实是祖传代码),现在想把代码移植到julia,至少画图方便。话不多说,直接把我写好的代码放出来:

using Plots
using Revise
using SharedArrays
using LaTeXStrings
using Distributed
gr()

const sigma = 0.05
const ddx = 1e-10
const c = 3e8
const dt = ddx / (2 * c)
const epsz = 8.85419e-12
# eaf = dt * sigma / (2 * epsz * epsilon)
const eaf = 0 # 不考虑损耗

# 下面三个函数都是源
function Gaussian(n)
    spread = 12.0
    t0 = 40
    return exp(-0.5 * ((t0 - n) / spread) ^ 2)
end

function Modulate_Gaussian(n, freq_in)
    spread = 100000
    t0 = 0.8 * spread
    return -cos(freq_in * dt * n) * exp(-4 * pi * ((t0 - n) / spread)^2)
end

function Sin(n, freq_in)
    freq_imag = freq_in * 0.1
    #return sin(freq_in * dt *n) * exp(-freq_imag * dt * n)
    return sin(freq_in * dt *n)
end
# 迭代电场
function calcEx(ex, hy, pluse, KE, D, ex_low_m1, ex_low_m2, ex_high_m1, ex_high_m2)
    i = 1 # i用于计算AB结构一维光子晶体
    for k = 2:KE
        if k % D == 0
            i = i + 1
        end

        if k == 800
            ex[k] = pluse
        elseif i % 2 == 0 && k < i * D && k > (i - 1) * D
            ex[k] += 0.5 * (hy[k - 1] - hy[k])
        else
            ex[k] = ca * ex[k] + cb * (hy[k - 1] - hy[k])
        end
    end

    #吸收边界条件
    ex[1] = ex_low_m2;
	ex_low_m2 = ex_low_m1;
	ex_low_m1 = ex[2];

	ex[KE] = ex_high_m2;
	ex_high_m2 = ex_high_m1;
	ex_high_m1 = ex[KE - 1];
    return ex, hy, ex_low_m1, ex_low_m2, ex_high_m1, ex_high_m2
end
# 迭代磁场
function calcHy(ex, hy, KE)
    for k = 1:KE-1
        hy[k] += 0.5 * (ex[k] - ex[k + 1])
    end
    return ex, hy
end

function Time()
    ex = SharedArray{Float32}(KE)
    hy = SharedArray{Float32}(KE)
    ex_low_m1 = 0
    ex_low_m2 = 0
    ex_high_m1 = 0
    ex_high_m2 = 0
    for n = 1:TimeSteps
        pluse = Sin(n, freq_in)
        ex, hy, ex_low_m1, ex_low_m2, ex_high_m1, ex_high_m2 = calcEx(ex, hy, pluse, KE, D, ex_low_m1, ex_low_m2, ex_high_m1, ex_high_m2)
        ex, hy = calcHy(ex, hy, KE)
    end
    return ex, hy
end

epsilon = 4
ca = (1.0 - eaf) / (1.0 + eaf)
cb = 0.5 / (epsilon * (1 + eaf))
D = 1000
TimeSteps = 500000
freq_in = pi * 1e15
layer = 21 # 光子晶体层数
KE = layer * D
r_i = 100
t_i = KE - 100
@time ex, hy = Time()
x = 0.1:0.1:KE/10
plot(x, ex, xlabel = L"x(nm)", ylabel = L"E")

由于本人主要方向是光子晶体,所以这里做了一下小改动以方便计算AB结构光子晶体。

并行化

FDTD本身迭代是有多个循环的,而且既要迭代电场也要迭代磁场,两个迭代函数还一直需要互相传递数据,所以是很耗费计算资源的,能够并行化就比较好了(之前本人做过FDTD的CUDA并行,有兴趣的可以一起讨论),但是我现在在julia下试过threads和distributed效果都不太好,速度会比原来的还慢(最主要还是我菜),threads的时候这个i必须global才能计算,但是我的需求是每次循环i都初始化为1,这时候global就不太合适了。之前用传输矩阵法计算光子晶体我都是用Julia写的,而且并行效率很高,FDTD就显得麻烦一点。

其他

编程方面我实在是小菜鸡,发这个贴子只是想抛砖引玉下(没想到会是第一个计算物理的帖子)。在Python下有很多计算物理的库,我了解到的光学下的就有基于傅里叶光学的lightpipes,基于FDTD的Meep,以及用于计算光子晶体能带的MPB,Julia下光学电磁学相关的还没有了解到有,虽然这些包都可以通过pycall拿来直接用,但是也希望Julia的生态能越做越大。大家有研究光学、电磁学或者光子晶体的欢迎一起讨论鸭。


#2

我看了一下你的代码,稍微做了些基本的优化:

const c = 3e8            # speed of light
const sigma = 0.05       #
const ddx = 1e-10        # spatial discretization
const dt = ddx / (2 * c) # temporal discretization
const epsz = 8.85419e-12 # in vacuum
# eaf = dt * sigma / (2 * epsz * ϵ)
const eaf = 0.0          # no dissipation

```Gaussian source term```
function Gaussian(n)
   spread, t0 = 12.0, 40
   return exp(-0.5 * ((t0 - n) / spread)^2)
end

```Modulate Gaussian source term```
function Modulate_Gaussian(n, freq)
   spread, t0 = 100000, 0.8 * spread
   return -cos(freq * dt * n) * exp(-4*pi*((t0 - n) / spread)^2)
end

```Sin wave source term```
function Sin(n, freq)
   freq_imag = freq * 0.1
   #return sin(freq * dt *n) * exp(-freq_imag * dt * n)
   return sin(freq*dt*n)
end

```Update electric field```
function calcEx!(Ex, Hy, pulse, KE, D, ExLow, ExHigh, ca, cb)
   @inbounds Threads.@threads for k = 2:KE
      #if k % D == 0; i += 1 end
      i = floor(Int32, k/D) + 1

      if k == 800
         Ex[k] = pulse
      elseif i % 2 == 0 && i*D < k < (i-1)*D
         Ex[k] += 0.5 * (Hy[k-1] - Hy[k])
      else
         Ex[k] = ca*Ex[k] + cb*(Hy[k-1] - Hy[k])
      end
   end

   # Absorbing BC
   Ex[1] = ExLow[2]
   ExLow[2] = ExLow[1]
   ExLow[1] = Ex[2]

   Ex[KE] = ExHigh[2]
   ExHigh[2] = ExHigh[1]
   ExHigh[1] = Ex[KE-1]
end

```Update magnetic field```
function calcHy!(Ex, Hy, KE)
   @inbounds Threads.@threads for k = 1:KE-1
      Hy[k] += 0.5*(Ex[k] - Ex[k+1])
   end
end

```计算AB结构一维光子晶体```
function main()
   ϵ = 4.0
   ca = (1.0 - eaf) / (1.0 + eaf)
   cb = 0.5 / (ϵ * (1 + eaf))
   D = 1000
   nSteps = 10000
   freq_in = pi * 1e15
   layer = 21 # 光子晶体层数
   KE = layer * D

   # Maybe distributed array is more appropriate?
   #Ex = SharedArray{Float32}(KE)
   #Hy = SharedArray{Float32}(KE)

   Ex = zeros(Float32, KE)
   Hy = zeros(Float32, KE)

   ExLow = [0.0f0, 0.0f0]
   ExHigh= [0.0f0, 0.0f0]

   @time for n = 1:nSteps
      pulse = Sin(n, freq_in)

      calcEx!(Ex, Hy, pulse, KE, D, ExLow, ExHigh, ca, cb)

      calcHy!(Ex, Hy, KE)
   end
   return Ex, Hy, KE
end

@time Ex, Hy, KE = main()

在我的Mac 上跑Julia v1.3,取步数10000,原始的code在不用sharedArray的情况下

13.969151 seconds (655.65 M allocations: 9.777 GiB, 0.81% gc time)

改进的code(去掉多线程的指令)

1.311910 seconds (185.54 k allocations: 10.281 MiB)

保留多线程,JULIA_NUM_THREADS=1

1.983146 seconds (416.44 k allocations: 28.062 MiB, 3.17% gc time)

JULIA_NUM_THREADS=2

1.196766 seconds (556.64 k allocations: 41.490 MiB, 5.42% gc time)

JULIA_NUM_THREADS=4

0.862614 seconds (837.76 k allocations: 68.670 MiB, 0.49% gc time)

并行的部分里,你原来计算下一步的电场用到了i,但根据你写的代码,这个i可以对每个k直接算出来,所以循环里其实不存在前后依赖性。
根据我的使用体验,sharedArray只是方便,实际的并行表现并不好;@threads在核数少的情况下会比@distributed快不少。其实我也不太明白为什么你这里加入@threads指令但只用单线程也会比什么都不加慢很多。


#3

感谢你的回复,受教了。在实际使用中可以用另一种方法来加速,就是对Ex的赋值不是通过循环而是直接数组的索引取值,比如Ex[1000:2000] += 0.5 * (Hy[1000-1:2000-1] - Hy[1000:2000]),会提速很多,如果对Ex初始为二维数组,Ex[t,x],一个维度用来存储时间另一个空间也可以用上面的方法赋值,不过不知道能不能并行操作。
此外我原始的程序很粗糙,看了一些别人写的FDTD,一般是把epsilon初始为一个数组,这样算复杂点的体系就会简洁而高效。
关于我threads并行的问题,我测试的时候没有加@inbounds导语,应该是这个问题导致的我程序的问题,是我对Julia不够熟悉,受教了十分感谢!


#4

在Julia中,这种vectorization的方法不可能会比显格式的循环更快,只是写起来更加方便而已。如果你以前经常使用Matlab,就会跟我一样犯很多这样的错误。而且,在Julia中为了优化内存的使用,这么写的标准方法是:

Ex[1000:2000] .+= 0.5 .* (Hy[1000-1:2000-1] .- Hy[1000:2000])

或者 (我不太记得@.的具体位置了,可以在Ex之前也可以在之后,意思有所区别)

@. Ex[1000:2000] += 0.5 * (Hy[1000-1:2000-1] - Hy[1000:2000])

离散化方法的时间维度是有先后依赖关系的,一般不能和空间一起并行。

另外,@inbounds@threads没啥关系,仅仅是为了取消边界检查(类似C的-O2


#5

再次感谢,很奇怪我使用中没有用到.*和.+也得到了结果,我要学习的还好多,感谢感谢