電子衍射模擬：基于GPU加速的MATLAB/Julia實現

點擊 “AladdinEdu，同學們用得起的【H卡】算力平臺”，注冊即送-H卡級別算力，80G大顯存，按量計費，靈活彈性，頂級配置，學生更享專屬優惠。

引言：電子衍射模擬的重要性與計算挑戰

電子衍射是材料科學、凝聚態物理和結構生物學中不可或缺的表征技術。通過分析電子與物質相互作用產生的衍射圖案，研究人員可以推斷出樣品的原子結構、晶體取向和缺陷信息。傳統的電子衍射模擬往往需要大量的計算資源，特別是在處理大尺寸模型或進行動態模擬時，計算時間可能長達數小時甚至數天。

近年來，圖形處理器（GPU）的并行計算能力為科學計算帶來了革命性的變化。與傳統的中央處理器（CPU）相比，GPU擁有數千個計算核心，能夠同時處理大量相似的計算任務，特別適合電子衍射模擬中高度并行的計算需求。本文將詳細介紹如何使用MATLAB和Julia這兩種科學計算常用語言，利用GPU加速電子衍射模擬過程，實現數量級的速度提升。

1. 電子衍射理論基礎

1.1 電子衍射基本原理

電子衍射遵循布拉格定律： $2dsin?θ=nλ2d\sin\theta = n\lambda$ ，其中 $d$ 是晶面間距， $θ\theta$ 是衍射角， $λ\lambda$ 是電子波長， $n$ 是衍射級數。當高速電子束照射到晶體樣品時，會在特定方向產生相干衍射波，形成衍射圖案。

1.2 多片層傳播公式

對于較厚的樣品，我們需要使用多片層方法進行計算。電子波函數 $ψ(r)\psi(\mathbf{r})$ 的傳播可以通過以下公式描述：

$ψn+1(r)=[ψn(r)?pn(r)]?hn(r)\psi_{n+1}(\mathbf{r}) = [\psi_n(\mathbf{r}) \cdot p_n(\mathbf{r})] \otimes h_n(\mathbf{r})$

其中：

$pn(r)=exp?[iσVn(r)Δz]p_n(\mathbf{r}) = \exp[i\sigma V_n(\mathbf{r})\Delta z]$ 是樣品透射函數
$hn(r)=1iλΔzexp?[iπλΔz∣r∣2]h_n(\mathbf{r}) = \frac{1}{i\lambda\Delta z} \exp\left[\frac{i\pi}{\lambda\Delta z}|\mathbf{r}|^2\right]$ 是自由空間傳播函數
$σ\sigma$ 是相互作用常數
$Vn(r)V_n(\mathbf{r})$ 是投影勢
$Δz\Delta z$ 是切片厚度

1.3 計算復雜度分析

電子衍射模擬的計算復雜度主要來自兩個方面：

傅里葉變換：每個切片都需要進行兩次傅里葉變換（正變換和逆變換）
透射函數計算：需要計算復數指數函數，操作次數與網格點數成正比

對于 $\times N$ 的網格和 $M$ 個切片，計算復雜度為 $\cdot N^2 \log N)$ 。當 $N$ 較大時（如2048×2048），計算量非常可觀。

2. GPU加速原理與優勢

2.1 GPU并行計算架構

現代GPU采用大規模并行架構，以NVIDIA的GPU為例：

流多處理器(SM)：包含多個CUDA核心
全局內存：高帶寬但高延遲
共享內存：低延遲的片上內存
線程層次：線程→線程塊→網格

2.2 電子衍射模擬中的并行性

電子衍射模擬中存在多個層次的并行性：

像素級并行：每個像素點的計算相互獨立
切片級并行：不同切片可以并行處理（需要特殊處理）
能級并行：不同能量或不同取向的計算相互獨立

2.3 預期加速比

理論上，GPU加速比可達：
$Speedup=TCPUTGPU≈NCPU?cores1×fCPUfGPU×η\text{Speedup} = \frac{T_{\text{CPU}}}{T_{\text{GPU}}} \approx \frac{N_{\text{CPU cores}}}{1} \times \frac{f_{\text{CPU}}}{f_{\text{GPU}}} \times \eta$

其中 $η\eta$ 是并行效率，通常可達0.6-0.8。對于高端GPU，加速比可達50-100倍。

3. MATLAB實現與GPU加速

3.1 環境配置與GPU檢測

% 檢查GPU可用性
fprintf('檢查GPU設備...\n');
if gpuDeviceCount > 0gpu = gpuDevice();fprintf('找到GPU: %s\n', gpu.Name);fprintf('計算能力: %s\n', gpu.ComputeCapability);fprintf('總內存: %.2f GB\n', gpu.TotalMemory/1e9);
elseerror('未找到可用的GPU設備');
end% 設置隨機數生成器種子
rng(42);

3.2 基本電子衍射模擬函數（CPU版本）

function diffraction_pattern = electron_diffraction_cpu(crystal, params)% 參數提取voltage = params.voltage;          % 電子電壓(kV)alpha_max = params.alpha_max;      % 最大接收角(mrad)defocus = params.defocus;          % 離焦量(nm)Cs = params.Cs;                    % 球差系數(mm)% 計算電子波長wavelength = 12.26 / sqrt(voltage * 1000 + 0.9784 * voltage^2); % ?% 創建頻率空間網格[N1, N2] = size(crystal);[kx, ky] = meshgrid(-N2/2:N2/2-1, -N1/2:N1/2-1);kx = kx / N2;  % 歸一化頻率ky = ky / N1;% 計算散射矢量k2 = kx.^2 + ky.^2;k_max = alpha_max / (wavelength * 1000);  % 最大空間頻率% 創建襯度傳遞函數(CTF)chi = pi * wavelength * defocus * 1e-10 * k2 + ...0.5 * pi * Cs * 1e7 * wavelength^3 * k2.^2;CTF = sin(chi);% 計算晶體傅里葉變換crystal_ft = fft2(crystal);crystal_ft = fftshift(crystal_ft);% 應用CTF并計算衍射圖案diffraction_ft = crystal_ft .* CTF;diffraction_pattern = abs(ifft2(ifftshift(diffraction_ft))).^2;% 應用孔徑函數aperture = k2 <= k_max^2;diffraction_pattern = diffraction_pattern .* aperture;
end

3.3 GPU加速版本

function diffraction_pattern = electron_diffraction_gpu(crystal, params)% 將數據傳輸到GPUcrystal_gpu = gpuArray(single(crystal));% 參數提取voltage = params.voltage;alpha_max = params.alpha_max;defocus = params.defocus;Cs = params.Cs;% 計算電子波長wavelength = 12.26 / sqrt(voltage * 1000 + 0.9784 * voltage^2);% 創建頻率空間網格（在GPU上）[N1, N2] = size(crystal_gpu);[kx, ky] = meshgrid(-N2/2:N2/2-1, -N1/2:N1/2-1);kx = gpuArray(single(kx / N2));ky = gpuArray(single(ky / N1));% 計算散射矢量k2 = kx.^2 + ky.^2;k_max = alpha_max / (wavelength * 1000);% 創建襯度傳遞函數(CTF)chi = pi * wavelength * defocus * 1e-10 * k2 + ...0.5 * pi * Cs * 1e7 * wavelength^3 * k2.^2;CTF = sin(chi);% 計算晶體傅里葉變換（使用GPU加速）crystal_ft = fft2(crystal_gpu);crystal_ft = fftshift(crystal_ft);% 應用CTF并計算衍射圖案diffraction_ft = crystal_ft .* CTF;diffraction_pattern = abs(ifft2(ifftshift(diffraction_ft))).^2;% 應用孔徑函數aperture = k2 <= k_max^2;diffraction_pattern = diffraction_pattern .* aperture;% 將結果傳回CPUdiffraction_pattern = gather(diffraction_pattern);
end

3.4 性能測試與比較

% 創建測試晶體結構
N = 1024;  % 網格大小
crystal = create_test_crystal(N, N);% 設置模擬參數
params.voltage = 200;    % 200 kV
params.alpha_max = 10;   % 10 mrad
params.defocus = -50;    % -50 nm
params.Cs = 1.0;         % 1.0 mm% CPU版本計時
tic;
pattern_cpu = electron_diffraction_cpu(crystal, params);
time_cpu = toc;
fprintf('CPU計算時間: %.2f 秒\n', time_cpu);% GPU版本計時（包含數據傳輸）
tic;
pattern_gpu = electron_diffraction_gpu(crystal, params);
time_gpu = toc;
fprintf('GPU計算時間: %.2f 秒\n', time_gpu);% 計算加速比
speedup = time_cpu / time_gpu;
fprintf('加速比: %.2f 倍\n', speedup);% 驗證結果一致性
relative_error = norm(pattern_cpu - pattern_gpu) / norm(pattern_cpu);
fprintf('相對誤差: %.2e\n', relative_error);% 可視化結果
figure;
subplot(1, 3, 1);
imagesc(log10(pattern_cpu + 1));
title('CPU計算結果');
colorbar;
axis image;subplot(1, 3, 2);
imagesc(log10(pattern_gpu + 1));
title('GPU計算結果');
colorbar;
axis image;subplot(1, 3, 3);
imagesc(log10(abs(pattern_cpu - pattern_gpu) + 1e-10));
title('差異');
colorbar;
axis image;

4. Julia實現與GPU加速

4.1 環境配置與包管理

using Pkg
Pkg.add("CUDA")
Pkg.add("FFTW")
Pkg.add("BenchmarkTools")
Pkg.add("Plots")using CUDA, FFTW, BenchmarkTools, Plots, LinearAlgebra# 檢查GPU可用性
if CUDA.functional()device = CUDA.device()println("找到GPU: ", device)println("GPU內存: ", CUDA.total_memory(device) / 1024^3, " GB")
elseerror("未找到可用的GPU設備")
end

4.2 Julia CPU版本實現

function electron_diffraction_cpu(crystal::Matrix{Float32}, params::Dict)# 參數提取voltage = params[:voltage]alpha_max = params[:alpha_max]defocus = params[:defocus]Cs = params[:Cs]# 計算電子波長wavelength = 12.26f0 / sqrt(voltage * 1000f0 + 0.9784f0 * voltage^2)# 創建頻率空間網格N1, N2 = size(crystal)kx = reshape(range(-N2÷2, N2÷2-1, length=N2) ./ N2, 1, :)ky = reshape(range(-N1÷2, N1÷2-1, length=N1) ./ N1, :, 1)# 計算散射矢量k2 = kx.^2 .+ ky.^2k_max = alpha_max / (wavelength * 1000f0)# 創建襯度傳遞函數(CTF)chi = π * wavelength * defocus * 1e-10f0 * k2 .+0.5f0 * π * Cs * 1e7f0 * wavelength^3 * k2.^2CTF = sin.(chi)# 計算晶體傅里葉變換crystal_ft = fftshift(fft(crystal))# 應用CTF并計算衍射圖案diffraction_ft = crystal_ft .* CTFdiffraction_pattern = abs.(ifft(ifftshift(diffraction_ft))).^2# 應用孔徑函數aperture = k2 .<= k_max^2diffraction_pattern .*= aperturereturn diffraction_pattern
end

4.3 Julia GPU加速版本

function electron_diffraction_gpu(crystal::Matrix{Float32}, params::Dict)# 將數據傳輸到GPUcrystal_gpu = CuArray(crystal)# 參數提取voltage = params[:voltage]alpha_max = params[:alpha_max]defocus = params[:defocus]Cs = params[:Cs]# 計算電子波長wavelength = 12.26f0 / sqrt(voltage * 1000f0 + 0.9784f0 * voltage^2)# 創建頻率空間網格（在GPU上）N1, N2 = size(crystal_gpu)kx = CuArray(reshape(range(-N2÷2, N2÷2-1, length=N2) ./ N2, 1, :))ky = CuArray(reshape(range(-N1÷2, N1÷2-1, length=N1) ./ N1, :, 1))# 計算散射矢量k2 = kx.^2 .+ ky.^2k_max = alpha_max / (wavelength * 1000f0)# 創建襯度傳遞函數(CTF)chi = π * wavelength * defocus * 1e-10f0 * k2 .+0.5f0 * π * Cs * 1e7f0 * wavelength^3 * k2.^2CTF = sin.(chi)# 計算晶體傅里葉變換（使用GPU加速）crystal_ft = fftshift(fft(crystal_gpu))# 應用CTF并計算衍射圖案diffraction_ft = crystal_ft .* CTFdiffraction_pattern = abs.(ifft(ifftshift(diffraction_ft))).^2# 應用孔徑函數aperture = k2 .<= k_max^2diffraction_pattern .*= aperture# 將結果傳回CPUreturn Array(diffraction_pattern)
end

4.4 性能測試與優化

# 創建測試晶體結構
function create_test_crystal(N1, N2)# 創建簡單的晶體測試圖案x = range(-1, 1, length=N2)y = range(-1, 1, length=N1)X = reshape(x, 1, :)Y = reshape(y, :, 1)# 多個正弦波疊加模擬晶體結構crystal = zeros(Float32, N1, N2)for k in 1:5frequency = 10 * kcrystal .+= sin.(2π * frequency * X) .* sin.(2π * frequency * Y)endreturn crystal
end# 創建測試數據
N = 1024
crystal = create_test_crystal(N, N)# 設置模擬參數
params = Dict(:voltage => 200f0,:alpha_max => 10f0,:defocus => -50f0,:Cs => 1.0f0
)# CPU版本性能測試
println("CPU版本性能測試:")
@time pattern_cpu = electron_diffraction_cpu(crystal, params)# GPU版本性能測試（預熱）
println("GPU版本性能測試（預熱）:")
@time pattern_gpu = electron_diffraction_gpu(crystal, params)# 正式性能測試
println("GPU版本性能測試（正式）:")
@time pattern_gpu = electron_diffraction_gpu(crystal, params)# 驗證結果一致性
relative_error = norm(pattern_cpu - pattern_gpu) / norm(pattern_cpu)
println("相對誤差: ", relative_error)# 可視化結果
p1 = heatmap(log10.(pattern_cpu .+ 1f0), title="CPU計算結果")
p2 = heatmap(log10.(pattern_gpu .+ 1f0), title="GPU計算結果")
p3 = heatmap(log10.(abs.(pattern_cpu - pattern_gpu) .+ 1e-10f0), title="差異")plot(p1, p2, p3, layout=(1,3), size=(1200,400))

5. 高級優化技巧

5.1 內存訪問優化

# 優化內存訪問的GPU版本
function electron_diffraction_gpu_optimized(crystal::Matrix{Float32}, params::Dict)crystal_gpu = CuArray(crystal)N1, N2 = size(crystal_gpu)# 使用共享內存優化網格創建function create_k_grid!(k2, N1, N2)i = (blockIdx().x - 1) * blockDim().x + threadIdx().xj = (blockIdx().y - 1) * blockDim().y + threadIdx().yif i <= N1 && j <= N2kx = (j - 1 - N2÷2) / N2ky = (i - 1 - N1÷2) / N1k2[i, j] = kx^2 + ky^2endreturnend# 分配GPU內存k2 = CUDA.zeros(Float32, N1, N2)# 配置線程塊和網格threads = (16, 16)blocks = (cld(N1, threads[1]), cld(N2, threads[2]))# 啟動內核計算k網格@cuda threads=threads blocks=blocks create_k_grid!(k2, N1, N2)# 其余計算部分保持不變voltage = params[:voltage]alpha_max = params[:alpha_max]defocus = params[:defocus]Cs = params[:Cs]wavelength = 12.26f0 / sqrt(voltage * 1000f0 + 0.9784f0 * voltage^2)k_max = alpha_max / (wavelength * 1000f0)chi = π * wavelength * defocus * 1e-10f0 * k2 .+0.5f0 * π * Cs * 1e7f0 * wavelength^3 * k2.^2CTF = sin.(chi)crystal_ft = fftshift(fft(crystal_gpu))diffraction_ft = crystal_ft .* CTFdiffraction_pattern = abs.(ifft(ifftshift(diffraction_ft))).^2aperture = k2 .<= k_max^2diffraction_pattern .*= aperturereturn Array(diffraction_pattern)
end

5.2 多GPU并行計算

# 多GPU并行計算框架
function multi_gpu_diffraction(crystals::Vector{Matrix{Float32}}, params::Dict)results = similar(crystals)n_gpus = length(CUDA.devices())# 為每個GPU分配任務tasks = Vector{Task}()for (i, crystal) in enumerate(crystals)gpu_id = (i - 1) % n_gpus + 1push!(tasks, @spawn beginCUDA.device!(gpu_id - 1)electron_diffraction_gpu(crystal, params)end)end# 收集結果for (i, task) in enumerate(tasks)results[i] = fetch(task)endreturn results
end

6. 實際應用案例

6.1 納米顆粒衍射模擬

# 模擬金納米顆粒的電子衍射
function simulate_gold_nanoparticle()# 創建金納米顆粒模型（FCC結構）N = 512nanoparticle = zeros(Float32, N, N)# 設置晶格參數（金：面心立方，a=4.08 ?）a = 4.08f0d_spacings = [a/√(h^2 + k^2 + l^2) for h in 1:3, k in 1:3, l in 1:3]# 創建顆粒形狀（球形）center = N ÷ 2radius = N ÷ 4for i in 1:N, j in 1:Nif √((i-center)^2 + (j-center)^2) ≤ radius# 簡化的晶體結構模型nanoparticle[i, j] = sin(2π*i/10) * sin(2π*j/10)endend# 設置衍射參數params = Dict(:voltage => 200f0,:alpha_max => 20f0,:defocus => 0f0,:Cs => 0.5f0)# 計算衍射圖案pattern = electron_diffraction_gpu(nanoparticle, params)# 可視化p1 = heatmap(nanoparticle, title="金納米顆粒模型")p2 = heatmap(log10.(pattern .+ 1f0), title="衍射圖案")plot(p1, p2, layout=(1,2), size=(800,400))
end

6.2 動態衍射模擬

# 動態電子衍射模擬（隨時間變化）
function dynamic_diffraction_simulation()# 創建初始晶體結構N = 256crystal = create_test_crystal(N, N)# 設置模擬參數params = Dict(:voltage => 200f0,:alpha_max => 15f0,:defocus => -100f0,:Cs => 1.0f0)# 時間演化參數n_frames = 60patterns = Vector{Matrix{Float32}}(undef, n_frames)# 動態模擬循環for t in 1:n_frames# 模擬結構隨時間變化（例如：熱振動）deformation = 0.1f0 * sin.(2π * t/30 * reshape(1:N, :, 1)) .*cos.(2π * t/30 * reshape(1:N, 1, :))deformed_crystal = crystal .+ deformation# 計算衍射圖案patterns[t] = electron_diffraction_gpu(deformed_crystal, params)println("完成幀: $t/$n_frames")endreturn patterns
end