【CFD算法】伪势LBM模拟相分离（python-numba加速）

模拟相分离（Python算法-Numba加速）

问题描述

给定随机初始状态模拟两相分离的效果，结果如下：
分离过程

代码优化说明

  • 原始Python循环较慢（200×200网格100循环步耗时166秒）
  • 使用Numba优化后降至11秒
  • MATLAB参考实现仅需5.6秒（后续计划尝试Taichi加速）
    优化结果

Numba加速版代码

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
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
# -*- coding:UTF-8 -*-
from numpy import *
import numpy as np
import numba as nb
from numba import jit
from matplotlib import cm
import matplotlib.pyplot as plt

# 参数设置
lx, ly = 200, 200
w = [4/9, 1/9,1/9,1/9,1/9, 1/36,1/36,1/36,1/36]
cx = [0, 1, 0, -1, 0, 1, -1, -1, 1]
cy = [0, 0, 1, 0, -1, 1, 1, -1, -1]
tau, mstep = 1, 5000
G, psi0, rho0 = -140, 4, 200

# 初始化随机密度场
delta_rho = 200 + (1-2.0 * np.random.random((lx, ly)))

# 分配内存
u_x = zeros((lx,ly)); u_y = zeros((lx,ly))
u = zeros((lx,ly)); v = zeros((lx,ly))
ueq = zeros((lx,ly)); veq = zeros((lx,ly))
rho = zeros((lx,ly)); psi = zeros((lx,ly))
fin = zeros((9,lx,ly)); feq = zeros((9,lx,ly))
fout = zeros((9,lx,ly))

# 初始分布函数
for k in range(9):
    fin[k, :, :] = w[k] * delta_rho

@jit(nopython=True)
def uv(fin, rho, u, v, cx, cy):
    # 计算宏观量（密度和速度）
    for i in range(lx):
        for j in range(ly):
            rho[i,j] = fin[:,i,j].sum()
            uSum = vSum = 0
            for k in range(9):
                uSum += fin[k,i,j] * cx[k]
                vSum += fin[k,i,j] * cy[k]
            u[i,j] = uSum / rho[i,j]
            v[i,j] = vSum / rho[i,j]
    return rho, u, v

@jit(nopython=True)
def fxy(psi, rho):
    # 计算势函数
    for i in range(lx):
        for j in range(ly):
            psi[i,j] = 4 * exp(-rho0 / rho[i,j])
    return psi

@jit(nopython=True)
def fff(ueq, veq, rho, fin, feq, fout, w, cx, cy):
    # 计算分布函数
    for i in range(lx):
        for j in range(ly):
            t1 = ueq[i,j]**2 + veq[i,j]**2
            for k in range(9):
                t2 = ueq[i,j]*cx[k] + veq[i,j]*cy[k]
                feq[k,i,j] = rho[i,j]*w[k]*(1 + 3*t2 + 4.5*t2**2 - 1.5*t1)
                fout[k,i,j] = fin[k,i,j] - (fin[k,i,j] - feq[k,i,j])/tau
    return feq, fout

# 主循环
for counter in range(mstep):
    rho, u, v = uv(fin, rho, u, v, nb.typed.List(cx), nb.typed.List(cy))
    psi = fxy(psi, rho)
    
    # 计算外力
    Fx = Fy = zeros((lx,ly))
    for k in range(9):
        if k != 0:
            shifted_psi = np.roll(np.roll(psi, cx[k], axis=0), cy[k], axis=1)
            Fx += G * psi * w[k] * shifted_psi * cx[k]
            Fy += G * psi * w[k] * shifted_psi * cy[k]
    
    # 更新平衡速度
    ueq = u + Fx * tau / rho
    veq = v + Fy * tau / rho
    
    # 碰撞和迁移
    feq, fout = fff(ueq, veq, rho, fin, feq, fout, 
                   nb.typed.List(w), nb.typed.List(cx), nb.typed.List(cy))
    for k in range(9):
        fin[k,:,:] = np.roll(np.roll(fout[k,:,:], cx[k], axis=0), cy[k], axis=1)
    
    # 可视化
    if counter % 100 == 0:
        plt.imshow(rho, cmap=cm.bwr)
        plt.pause(0.01)
        plt.clf()

为了方便代码的理解，再放一个没有优化的代码。

未加速版python代码

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
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
# -*- coding:UTF-8 -*- 
# Author:DA BAIMENG
# Time:2023/10/26 11:20
from numpy import *
import numpy as np
from matplotlib import cm
import matplotlib.pyplot as plt

lx=200
ly=200
w  = [4/9, 1/9,1/9,1/9,1/9, 1/36,1/36,1/36,1/36]
cx = [  0,   1,  0, -1,  0,    1,  -1,  -1,   1]
cy = [  0,   0,  1,  0, -1,    1,   1,  -1,  -1]

tau=1
mstep=5000

G=-140
psi0=4
rho0=200
# delta_rho = 200+(1-2.0 * random.rand(lx))
delta_rho = 200+(1-2.0 * np.random.random((lx, ly)))

#初始化
u_x=zeros((lx,ly))
u_y=zeros((lx,ly))
u=zeros((lx,ly))
v=zeros((lx,ly))
ueq=zeros((lx,ly))
veq=zeros((lx,ly))
rho=zeros((lx,ly))

fin=zeros((9,lx,ly))
feq=zeros((9,lx,ly))
fout=zeros((9,lx,ly))
psi=zeros((lx,ly))


for k in range(9):
    fin[k, :, :] = w[k] * delta_rho

for counter in range(mstep):

    for i in range(lx):
        for j in range(ly):
            rho[i, j] = sum(fin[:, i, j], axis=0)
            uSum = 0
            vSum = 0
            for k in range(9):
                uSum = uSum + fin[k, i, j] * cx[k]
                vSum = vSum + fin[k, i, j] * cy[k]
            u[i, j] = uSum / rho[i, j]
            v[i, j] = vSum / rho[i, j]

    for i in range(lx):
        for j in range(ly):
            psi[i, j] = 4 * exp(-rho0 / rho[i, j])

    Fx = zeros((lx, ly))
    Fy = zeros((lx, ly))

    for k in range(9):
        if k != 0:
            Fx = Fx + G * psi * w[k] * roll(roll(psi, cx[k], axis=0), cy[k], axis=1) * cx[k]
            Fy = Fy + G * psi * w[k] * roll(roll(psi, cx[k], axis=0), cy[k], axis=1) * cy[k]

    ueq = u + Fx * tau / rho
    veq = v + Fy * tau / rho

    for i in range(lx):
        for j in range(ly):
            t1 = ueq[i, j] * ueq[i, j] + veq[i, j] * veq[i, j]
            for k in range(9):
                t2 = ueq[i, j] * cx[k] + veq[i, j] * cy[k]
                feq[k, i, j] = rho[i, j] * w[k] * (1 + 3 * t2 + 4.5 * t2 * t2 - 1.5 * t1)
                fout[k, i, j] = fin[k, i, j] - 1 / tau * (fin[k, i, j] - feq[k, i, j])

    for k in range(9):
        fin[k, :, :] = roll(roll(fout[k, :, :], cx[k], axis=0), cy[k], axis=1)
    
    print('%s / %s' %(counter,mstep))

    plt.imshow(rho, cmap=cm.bwr)  # 绘制密度场
    # 不显示图形
    plt.ioff()
    # 更新图形
    plt.pause(0.01)
    # plt.savefig(r"./计算结果/" + str(counter) + ".png")
    plt.clf()

matlab版代码

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
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
clear
clc

lx=200;
ly=200;
w = [4/9, 1/9,1/9,1/9,1/9, 1/36,1/36,1/36,1/36];
cx = [ 0, 1, 0, -1, 0, 1, -1, -1, 1];
cy = [ 0, 0, 1, 0, -1, 1, 1, -1, -1];

tau=1;
mstep=5000;

G=-140;
psi0=4;
rho0=200;
delta_rho = 200+(1-2.0*rand(lx));
%初始化

u_x(lx,ly)=0;
u_y(lx,ly)=0;
u(lx,ly)=0;
v(lx,ly)=0;
ueq(lx,ly)=0;
veq(lx,ly)=0;
for k=1:9
fin(:,:,k)=w(k).*delta_rho;
end

for kk=1:mstep
tic

%计算宏观量

rho=sum(fin,3);%计算每一点的密度

for i=1:lx
for j=1:ly
uSum=0;
vSum=0;
for k=1:9
uSum=uSum+fin(i,j,k)*cx(k);
vSum=vSum+fin(i,j,k)*cy(k);
end
u(i,j)=uSum/rho(i,j);%计算速度场
v(i,j)=vSum/rho(i,j);
end
end

for i=1:lx
for j=1:ly
psi(i,j)=4*exp(-rho0./rho(i,j));
end
end

Fx=zeros(lx,ly);
Fy=zeros(lx,ly);

for k=2:9
Fx=Fx+G.*psi.*w(k).*circshift(psi,[cx(k),cy(k),0]).*cx(k);
Fy=Fy+G.*psi.*w(k).*circshift(psi,[cx(k),cy(k),0]).*cy(k);
end

ueq=u+Fx*tau./rho;
veq=v+Fy*tau./rho;


%碰撞
for i=1:lx
for j=1:ly
t1=ueq(i,j)^2+veq(i,j)^2;
for k=1:9
t2=ueq(i,j)*cx(k)+veq(i,j)*cy(k);
feq(i,j,k)=rho(i,j)*w(k)*(1+3*t2+4.5*t2^2-1.5*t1);
fout(i,j,k) = fin(i,j,k) - 1/tau .* (fin(i,j,k)-feq(i,j,k));
end
end
end

for k=1:9
fin(:,:,k ) = circshift(fout(:,:,k ), [cx(k),cy(k),0]);
end

if mod(kk,1)==0
imagesc(fin(:,:,1 ));
axis off
axis equal
drawnow
end

toc
end
流体算法
#格子玻尔兹曼
【CFD算法】伪势LBM模拟相分离（python-numba加速）
http://example.com/2025/03/26/伪势LBM模型模拟相分离（python-numba加速）/
作者
DABAI MENG
发布于
2025年3月26日
许可协议