主页 / coding / 问题 / 77783527
Accepted
Nick Brady
Asked: 2024-01-09 06:43:15 +0800 CST

从数组构建 Julia 块带状矩阵 - 提高速度

  • 772

我正在尝试构建并求解一个如下所示的矩阵系统：

矩阵系统求解

\begin{equation}
    \begin{bmatrix}
        B & D & 0 & \cdots &        & 0       \\
        A & B & D & \ddots &        & \vdots  \\
        0 & A & B & D      & \ddots &         \\
        \vdots & \ddots &   & \ddots & \\
        \vdots & \cdots & \cdots & A & B & D \\
        0      & \cdots & \cdots & 0 & A & B
    \end{bmatrix}
    \
    \begin{bmatrix}
       \Delta c
    \end{bmatrix}
    \
    =
    \
    \begin{bmatrix}
       G
    \end{bmatrix}
\end{equation}

该系统来自用于求解方程组 (PDE) 的一维有限体积方法。

大矩阵 ( BBM) 由数组构造而成A, B, D。以下代码可以正常运行，但我有兴趣使该代码的计算效率更高。

我很好奇堆栈溢出社区是否有任何建议。特别是有没有一种方法可以消除for函数内的循环？

using BlockBandedMatrices

function calculate_BBM(; A, B, D)
    # initialize the Block Banded Matrix (BBM)
    int_array = map(Int, N*ones(NJ));
    BBM = BlockBandedMatrix(Zeros{Float64}(N*NJ, N*NJ), int_array, int_array, (1,1))

    # Precompute indices to avoid calculating them in each iteration
    ind_begin = [(1 + N * (j - 1)) for j in 1:NJ]
    ind_end = [N * j for j in 1:NJ]

    for j in 2:NJ - 1
        # Use precomputed indices
        ib = ind_begin[j]
        ie = ind_end[j]

        BBM[ib:ie, (ib - N):(ie - N)] .= A[:, :, j - 1]
        BBM[ib:ie, ib:ie] .= B[:, :, j]
        BBM[ib:ie, (ib + N):(ie + N)] .= D[:, :, j]
    end

    j = 1
    ib = ind_begin[j]
    ie = ind_end[j]

    BBM[ib:ie, ib:ie] .= B[:, :, j]
    BBM[ib:ie, (ib + N):(ie + N)] .= D[:, :, j]


    j = NJ
    ib = ind_begin[j]
    ie = ind_end[j]

    BBM[ib:ie, (ib - N):(ie - N)] .= A[:, :, j - 1]
    BBM[ib:ie, ib:ie] .= B[:, :, j]

    return BBM

end

N = 3
NJ = 102

A = randn(N, N, NJ-1)
B = randn(N, N, NJ)
D = randn(N, N, NJ-1)

# A = -ones(N, N, NJ-1)
# B = ones(N, N, NJ) * 2
# D = ones(N, N, NJ-1)

BBM = calculate_BBM(A=A, B=B, D=D);
julia

1 个回答

  1. Best Answer

    使用Block(...)BlockArrays.jl 中解释的索引方法，该函数最初可以简化为看起来更好（并且速度加快 2 倍）：

    # initialize the Block Banded Matrix (BBM)
function calculate_BBM(NJ; A, B, D)
    N = size(A,1)
    NJ_Ns = fill(N, NJ)
    BBM = BlockBandedMatrix(Zeros{Float64}(N*NJ, N*NJ), 
      NJ_Ns, NJ_Ns, (1,1))
           
    @views BBM[Block(1,1)] .= B[:, :, 1]
    @views BBM[Block(2,1)] .= D[:, :, 1]
           
    for j in 2:NJ - 1
        @views BBM[Block(j-1,j)] .= A[:, :, j - 1]
        @views BBM[Block(j,j)] .= B[:, :, j]
        @views BBM[Block(j+1,j)] .= D[:, :, j]
    end

    @views BBM[Block(NJ-1,NJ)] .= A[:, :, NJ - 1]
    @views BBM[Block(NJ,NJ)] .= B[:, :, NJ]

    return BBM
end

    另请注意，该函数现在将其NJ作为参数并N从输入块的大小中获取。

    添加：改进版本（无 NJ 参数）：

    # initialize the Block Banded Matrix (BBM)
function calculate_BBM(; A, B, D)
    N,NJ = size(A,1), size(B,3)
    size(A) == size(D) == (N,N,NJ-1) || 
      error("input matrix sizes clash")
    size(B) == (N,N,NJ) || error("input matrix sizes clash")
    NJ_Ns = fill(N, NJ)
    BBM = BlockBandedMatrix(Zeros{Float64}(N*NJ, N*NJ), NJ_Ns, NJ_Ns, (1,1))

    for (i, M) in enumerate(eachslice(B; dims=3))
        BBM[Block(i,i)] .= M
    end
    for (i, (S, T)) in enumerate(zip(eachslice(D; dims=3), 
                                     eachslice(A; dims=3)))
        BBM[Block(i+1,i)] .= S
        BBM[Block(i,i+1)] .= T
    end
    return BBM
end

    更新：修复了要使用的第二个解决方案Zeros。

    • 1

相关问题