如何高效实现多元函数求值与Jacobian矩阵计算？

我正在迭代数值求解器中使用ForwardDiff.jl的对偶数实现函数求值与Jacobian矩阵计算，但Jacobian计算是当前的性能瓶颈。现寻求针对flux_results、flux_real、flux_dc和flux_d∇c环节的优化方案，当前实现代码如下：

定义部分代码 #

using ForwardDiff

function create_interface_dual_vectors(N::Int)
    #Preallocate the dual_vector with the correct type and size
    dual_vector = Vector{ForwardDiff.Dual{Nothing, Float64, 2*N}}(undef, N)
    ∇_dual_vector = Vector{ForwardDiff.Dual{Nothing, Float64, 2*N}}(undef, N)
    
    for i in 1:N
        derivative_part = zeros(N*2)
        derivative_part[i] = 1.0
        dual_vector[i] = ForwardDiff.Dual(0.0, derivative_part...)

        derivative_part = zeros(N*2)
        derivative_part[N+i] = 1.0
        ∇_dual_vector[i] = ForwardDiff.Dual(0.0, derivative_part...)
    end

    return dual_vector, ∇_dual_vector
end


function FLUX(; c_vars, ∇c_vars) # (; c_vars, ΔV)
    c1, c2, c3 = c_vars
    ∇c1, ∇c2, ∇c3 = ∇c_vars

    Diff_coef = 1.0e-4

    # Initialize 'result' based on the type of 'c_vars'
    result = fill(typeof(c_vars[1])(0), length(c_vars))

    result[1] = -Diff_coef * ∇c1
    result[2] = -Diff_coef * ∇c2
    result[3] = -Diff_coef * ∇c3
    
    return result
end

主计算代码 #

N = 3
NJ = 102


# interface variables
c_interface = randn(N, NJ-1)
∇c_prev = randn(N, NJ-1)

# create the dual arrays
interface_dual_vector, ∇_dual_vector = create_interface_dual_vectors(N)

# make the interface variables dual valued
c_interface_dual = c_interface .+ interface_dual_vector
∇c_dual = ∇c_prev .+ ∇_dual_vector

# iteratively evaluate the function FLUX
flux_results = [FLUX(c_vars = c_interface_dual[:, i], ∇c_vars = ∇c_dual[:, i]) for i in 1:NJ-1]

# partition the results into the real and dual values
flux_real = [res.value for j in 1:NJ-1 for res in flux_results[j]]
flux_dc   = [flux_results[j][i].partials[k] for j in 1:NJ-1 for i in 1:N, k in 1:N]
flux_d∇c  = [flux_results[j][i].partials[k] for j in 1:NJ-1 for i in 1:N, k in (N+1):(2*N)]

flux_dc  = reshape(flux_dc,  N, N, NJ-1)
flux_d∇c = reshape(flux_d∇c, N, N, NJ-1);

优化建议 #

1. 简化对偶向量构造，减少内存分配 #

当前create_interface_dual_vectors中循环创建zeros数组的方式效率低下，可直接用元组构造对偶向量，避免临时数组分配：

function create_interface_dual_vectors(N::Int)
    dual_vector = Vector{ForwardDiff.Dual{Nothing, Float64, 2N}}(undef, N)
    ∇_dual_vector = similar(dual_vector)
    for i in 1:N
        # 直接生成导数元组，无需zeros数组
        dual_vector[i] = ForwardDiff.Dual(0.0, ntuple(k -> k == i ? 1.0 : 0.0, 2N)...)
        ∇_dual_vector[i] = ForwardDiff.Dual(0.0, ntuple(k -> k == N+i ? 1.0 : 0.0, 2N)...)
    end
    return dual_vector, ∇_dual_vector
end

2. 向量化FLUX计算，避免循环调用 #

你的FLUX函数仅依赖∇c_vars，可直接改为支持矩阵输入的向量化版本，一次性处理所有接口数据：

function FLUX(c_vars::AbstractMatrix, ∇c_vars::AbstractMatrix)
    Diff_coef = 1.0e-4
    return -Diff_coef .* ∇c_vars
end

# 直接批量计算，替代原列表推导式
flux_results = FLUX(c_interface_dual, ∇c_dual)

3. 直接提取导数，避免嵌套列表推导 #

原代码用嵌套列表推导提取导数会产生大量临时数组，可通过预分配数组+循环填充，减少GC压力：

# 预分配结果数组
flux_real = Vector{Float64}(undef, N*(NJ-1))
flux_dc = Array{Float64}(undef, N, N, NJ-1)
flux_d∇c = similar(flux_dc)

# 循环填充结果
for j in 1:NJ-1
    flux_col = flux_results[:, j]
    for i in 1:N
        # 填充实部
        flux_real[(j-1)*N + i] = flux_col[i].value
        # 填充对c的偏导
        for k in 1:N
            flux_dc[i, k, j] = flux_col[i].partials[k]
        end
        # 填充对∇c的偏导
        for k in 1:N
            flux_d∇c[i, k, j] = flux_col[i].partials[N + k]
        end
    end
end

4. 改用ForwardDiff.jacobian直接计算 #

手动构造对偶向量的方式繁琐，可直接用ForwardDiff.jacobian计算每个接口的Jacobian矩阵，代码更简洁且内置优化：

# 包装函数：将c和∇c拼接为单个输入向量
function flux_wrapper(x)
    N = length(x) ÷ 2
    c_vars = x[1:N]
    ∇c_vars = x[N+1:end]
    return FLUX(c_vars=c_vars, ∇c_vars=∇c_vars)
end

# 批量计算每个接口的Jacobian
jacobians = [ForwardDiff.jacobian(flux_wrapper, [c_interface[:,j]; ∇c_prev[:,j]]) for j in 1:NJ-1]

# 拆分得到所需的偏导矩阵
flux_dc = cat([jac[:,1:N] for jac in jacobians]..., dims=3)
flux_d∇c = cat([jac[:,N+1:2N] for jac in jacobians]..., dims=3)

# 计算实部
flux_real = vec([FLUX(c_vars=c_interface[:,j], ∇c_vars=∇c_prev[:,j]) for j in 1:NJ-1])

5. 利用静态数组优化小尺寸向量运算 #

如果N固定为3，使用StaticArrays.jl处理小向量可大幅减少动态内存分配开销：

using StaticArrays

# 适配静态数组的FLUX函数
function FLUX(c_vars::SVector{3,T}, ∇c_vars::SVector{3,T}) where T
    Diff_coef = 1.0e-4
    return -Diff_coef .* ∇c_vars
end

# 直接构造静态对偶向量（无需循环）
function create_interface_dual_vectors()
    N = 3
    dual_vector = SVector{3, ForwardDiff.Dual{Nothing, Float64, 6}}(
        ForwardDiff.Dual(0.0, 1.0,0.0,0.0,0.0,0.0,0.0),
        ForwardDiff.Dual(0.0, 0.0,1.0,0.0,0.0,0.0,0.0),
        ForwardDiff.Dual(0.0, 0.0,0.0,1.0,0.0,0.0,0.0)
    )
    ∇_dual_vector = SVector{3, ForwardDiff.Dual{Nothing, Float64, 6}}(
        ForwardDiff.Dual(0.0, 0.0,0.0,0.0,1.0,0.0,0.0),
        ForwardDiff.Dual(0.0, 0.0,0.0,0.0,0.0,1.0,0.0),
        ForwardDiff.Dual(0.0, 0.0,0.0,0.0,0.0,0.0,1.0)
    )
    return dual_vector, ∇_dual_vector
end

内容的提问来源于stack exchange，提问作者Nick Brady