Julia个人流程演示代码

本文为我的个人经验分享视频Julia高性能代码入门流程相应的代码。

除了官方文档的性能建议必须知道的以外，总的来说有几个套路：

1. 参数分类。
   • 物理常数直接默认全局变量，并声明const类型。Julia编译成静态语言的时候就不会害怕有别的程序会改动它，从而大幅提升性能。
   • 利用struct来打包机器参数，节省开发时间；
2. 数组预先分配其在内存中的大小；
3. 做好缓存。利用多重分派，避免循环和矢量化操作的重复计算问题。
4. 做好多线程。这个是提升最大的。
5. 一个提醒。通常来说，数组for循环第一个index放最里面——因为它改动最频繁，而数组访问第一个index最快。但实际上，谁放外面谁放里面需要具体问题具体分析。只按照性能建议中的将第一个index放最内层未必最快。

数据生成

using BSON: @save

zvec = randn(10000)*0.12
δvec = randn(10000)*1.06e-3

@save "zvec_δvec.bson" zvec δvec

MATLAB代码

data = readtable("data.csv");
zvec = table2array(data(1:end, 1));
deltavec = table2array(data(1:end, 2));

%%

restenergy = 0.511e6;
clight = 3e8;

nturns = 10001;
nparticle = length(zvec);

v1 = 3.6395053705870048e+06;
v2 = 655751.6971214055228756;
r = v2/v1;
phi_s = 2.1575815005097385;
phi_2s = 6.0620746573056303;
h1 = 756;
h = 3;
circum = 1360.4;
centerenergy = 6e9;
alpha_c = 0.0000156;

%% 10000个粒子相空间坐标，不保存历史

tic;

for i = 1:nturns
    deltavec = deltavec + v1/centerenergy * (sin(phi_s - 2 * h1 * pi / circum .* zvec) - sin(phi_s) + r * sin(phi_2s - 2 * h1 * h * pi / circum .* zvec) - r * sin(phi_2s));
    etavec = alpha_c - 1./(centerenergy/restenergy .* (1 + deltavec)).^2;
    zvec = zvec - circum .* etavec .* deltavec;
end

toc

% %%
% 
% tic;
% 
% k1 = 2 * h1 * pi / circum;
% k2 = 2 * h1 * h * pi / circum;
% gamma = centerenergy/restenergy;
% c1 = v1 / centerenergy;
% c2 = sin(phi_s) + r * sin(phi_2s);
% c3 = 1/gamma^2;
% 
% parfor j = 1:nparticle
%     for i = 1:nturns
%         deltavec(j) = deltavec(j) + c1 * (sin(phi_s - k1 * zvec(j)) + r * sin(phi_2s - k2 * zvec(j)) - c2);
%         eta = alpha_c - c3/(1 + deltavec(j))^2;
%         zvec(j) = zvec(j) - circum * eta * deltavec(j);
%     end
% end

% toc

%% 更新10000个粒子相空间坐标，同时保存历史

tic;

zmat = zeros(nparticle, nturns);
delta_mat = zeros(nparticle, nturns);
for i = 1:nturns
    deltavec = deltavec + v1/centerenergy * (sin(phi_s - 2 * h1 * pi / circum .* zvec) - sin(phi_s) + r * sin(phi_2s - 2 * h1 * h * pi / circum .* zvec) - r * sin(phi_2s));
    etavec = alpha_c - 1./(centerenergy/restenergy .* (1 + deltavec)).^2;
    zvec = zvec - circum .* etavec .* deltavec;
    zmat(1:end, i) = zvec;
    delta_mat(1:end, i) = deltavec;
end

toc

Julia代码

using BSON: @load

@load "zvec_δvec.bson" zvec δvec

# 物理常数
const restenergy = 0.511e6
const clight = 3e8

# 计算参数
nturns = 10001

# 机器参数
struct Par
    # 这些参数要提前准备
    v1::Float64
    v2::Float64
    ϕs::Float64
    ϕ2s::Float64
    h1::Int64
    h::Int64
    circum::Float64
    centerenergy::Float64
    αc::Float64
    # 这些参数会自动计算
    r::Float64
end
function Par(v1, v2, ϕs, ϕ2s, h1, h, circum, centerenergy, αc)
    Par(v1, v2, ϕs, ϕ2s, h1, h, circum, centerenergy, αc, v2/v1)
end

function _revolution_cache(p::Par)
    c1 = p.v1/p.centerenergy
    c2 =  sin(p.ϕs) + p.r * sin(p.ϕ2s)
    c3 = 1/(p.centerenergy/restenergy)^2
    k1 = 2*p.h1*pi/p.circum
    k2 = 2*p.h1*p.h*pi/p.circum
    (; c1, c2, c3, k1, k2)
end

function revolution_with_cache(z::Number, δ::Number, p::Par, c)
    δ = δ + c.c1 * (sin(p.ϕs - c.k1 * z) + p.r * sin(p.ϕ2s - c.k2 * z) - c.c2)
    η = p.αc - c.c3/(1+δ)^2
    z = z - p.circum * η * δ
    z, δ
end

function revolution_multithreads!(zvec::Vector, δvec::Vector, nturns::Int64, p::Par)
    c = _revolution_cache(p)
    nparticles = length(zvec)
    @sync @inbounds for i=1:nparticles
        Threads.@spawn begin
            z, δ = zvec[i], δvec[i]
            @inbounds for j=1:nturns
                z, δ = revolution_with_cache(z, δ, p, c)
            end
            zvec[i], δvec[i] = z, δ
        end
    end
    zvec, δvec
end

function revolution_singlethread!(zvec::Vector, δvec::Vector, nturns::Int64, p::Par)
    c = _revolution_cache(p)
    nparticles = length(zvec)
    @inbounds for j=1:nturns
        @inbounds @simd for i=1:nparticles
            zvec[i], δvec[i] = revolution_with_cache(zvec[i], δvec[i], p, c)
        end
    end
    zvec, δvec
end

function revolution_history_1(zvec::Vector, δvec::Vector, nturns::Int64, p::Par)
    c = _revolution_cache(p)
    nparticles = length(zvec)
    zmat = zeros(eltype(zvec), nturns, nparticles)
    δmat = zeros(eltype(zvec), nturns, nparticles)
    @sync @inbounds for i=1:nparticles
        Threads.@spawn begin
            z, δ = zvec[i], δvec[i]
            @inbounds for j=1:nturns
                z, δ = revolution_with_cache(z, δ, p, c)
                zmat[j, i], δmat[j, i] = z, δ
            end
        end
    end
    zmat, δmat
end

function revolution_history_2(zvec::Vector, δvec::Vector, nturns::Int64, p::Par)
    c = _revolution_cache(p)
    nparticles = length(zvec)
    # 数组内存预分配
    zmat = zeros(eltype(zvec), nparticles, nturns)
    δmat = zeros(eltype(zvec), nparticles, nturns)
    ηvec = zeros(eltype(zvec), nparticles)
    tmpzvec, tmpδvec = zvec[:], δvec[:]
    @inbounds for i = 1:nturns
        tmpδvec .= tmpδvec .+ c.c1 .* (sin.(p.ϕs .- c.k1 .* tmpzvec) .+ p.r .* sin.(p.ϕ2s .- c.k2 .* tmpzvec) .- c.c2)
        ηvec .= p.αc .- c.c3 ./ (1 .+ tmpδvec) .^ 2
        tmpzvec .= tmpzvec .- p.circum .* ηvec .* tmpδvec
        zmat[:, i] = tmpzvec
        δmat[:, i] = tmpδvec
    end
    zmat, δmat
end


p = Par(3.6395053705870048e+06, 655751.6971214055228756, 2.1575815005097385, 6.0620746573056303, 756, 3, 1360.4, 6e9, 1.56e-5)