using Metal using KernelAbstractions using ParallelStencil using ParallelStencil.FiniteDifferences2D @init_parallel_stencil(Metal, Float32, 2) using GLMakie # 1. 计算通量 qx 和 qy (带阻尼项的 PT 格式) @parallel function compute_flux!(qx, qy, T, kx, ky, θ, Δx, Δy) # 使用 θ (阻尼) 来限制通量的变化量 @all(qx) = (@all(qx) * θ + (-@all(kx) * @d_xi(T) / Δx)) / (1.0f0 + θ) @all(qy) = (@all(qy) * θ + (-@all(ky) * @d_yi(T) / Δy)) / (1.0f0 + θ) return nothing end # 2. 更新温度 T (基于通量散度) @parallel function compute_T!(T, qx, qy, dτ_T, Δx, Δy) @inn(T) = @inn(T) - dτ_T * (@d_xa(qx)/Δx + @d_ya(qy)/Δy) return nothing end # 3. 设置边界条件 @parallel_indices (ix, iy) function set_BCs!(T) nx, ny = size(T) if ix <= nx && iy == ny; T[ix, ny] = 0.0f0; end # Top = 0 if ix <= nx && iy == 1; T[ix, 1] = 1000.0f0; end # Bottom = 1000 # 左右边界默认绝热 (ParallelStencil 的 @inn 只更新内部，边界不变) return end function PT_solver_2D() # --- 物理与网格参数 --- L, H = 10000.0f0, 10000.0f0 nx, ny = 128, 128 x = range(0.0f0, L, length=nx) y = range(-H, 0.0f0, length=ny) Δx, Δy = L/(nx-1), H/(ny-1) # --- PT 算法核心参数 --- Vp = 1.0f0 # 数值波速 re_pt = 0.7f0 * Float32(π) # 数值雷诺数 (2D 建议在 0.5π 到 π 之间) CFL = 0.4f0 # 2D 安全库朗数 dτ_T = CFL * min(Δx, Δy) / Vp # 温度伪时间步长 θ = (re_pt * min(Δx, Δy)) / L # 阻尼系数 # ---------------------- # 数组初始化 T = @zeros(nx, ny) qx = @zeros(nx-1, ny) qy = @zeros(nx, ny-1) kx = @fill(3.0f0, nx-1, ny) ky = @fill(3.0f0, nx, ny-1) # 初始猜测值：线性初始化 (Bottom 1000 -> Top 0) 极大提高稳定性 T_cpu = [1000.0f0 * (1.0f0 + iy/ny) for ix in 1:nx, iy in 1:ny] # 简易斜坡 copyto!(T, Float32.(T_cpu)) @parallel set_BCs!(T) # 可视化设置 T_obs = Observable(Array(T)) fig = Figure(resolution=(800, 700)) ax = Axis(fig[1, 1], title="PT Steady State Diffusion", xlabel="x [km]", ylabel="y [km]") hm = heatmap!(ax, x/1e3, y/1e3, T_obs, colormap=:inferno) Colorbar(fig[1, 2], hm) display(fig) # 迭代循环 max_iter = 20000 n_check = 500 println("2D PT 求解器启动 (网格: $nx x $ny)") for iter = 1:max_iter @parallel compute_flux!(qx, qy, T, kx, ky, θ, Δx, Δy) @parallel compute_T!(T, qx, qy, dτ_T, Δx, Δy) @parallel set_BCs!(T) if iter % n_check == 0 T_host = Array(T) if any(isnan, T_host) error("发散了！请减小 CFL 或 re_pt") end T_obs[] = T_host yield() println("迭代次数: $iter / $max_iter") end end return Array(T) end T_final = PT_solver_2D()
