有限差分求解器中稀疏矩阵的高效生成

Question

我正在编写一个程序，用有限差分法求解三维薛定谔方程。我的代码的1D和2D版本运行得很好，但在3D版本中，我发现生成矩阵(对于那些熟悉QM的人来说，这是Hamiltonian矩阵；对于那些不知道的，它并不重要)是花费了最多的时间(典型网格间隔的分钟，而其他操作的秒，包括最小的特征值-查找器！)。

我想知道是否有人对如何更有效地编写矩阵生成有任何建议。我在下面包含了我的代码的两个版本:一个应该相对容易理解，另一个版本遵循MATLAB的文档建议，即在创建稀疏矩阵时不应该直接索引条目，而是应该生成三个向量(行和列索引及其各自的值)，并从它们生成稀疏矩阵。不幸的是，后者根本没有帮助加快速度，因为我仍然在使用一个愚蠢的三层嵌套循环，而且我想不出一个好办法来避免它。

delta = 0.1e-9;
Lx = 2e-9;
x = 0:delta:Lx;
Nx = length(x);
Ly = 2e-9;
y = 0:delta:Ly;
Ny = length(y);
Lz = 2e-9;
z = 0:delta:Lz;
Nz = length(z);

map = inline('((idx_x-1) * Ny*Nz) + ((idx_y-1) * Nz) + idx_z','idx_x','idx_y','idx_z','Ny','Nz'); % define an inline helper function for mapping (x,y,z) indices to a linear index

Tsparse = sparse([],[],[],Nx*Ny*Nz, Nx*Ny*Nz, 7*(Nx-2)*(Ny-2)*(Nz-2)); % kinetic part of Hamiltonian matrix: (d^2/dx^2 + d^2/dy^2 + d^2/dz^2); NOTE: we'll have 7*(Nx-2)*(Ny-2)*(Nz-2) non-zero entries in this matrix, so we get the sparse() function to preallocate enough memory for this

for idx_x = 2:Nx-1
    for idx_y = 2:Ny-1
        for idx_z = 2:Nz-1
            Tsparse( map(idx_x,idx_y,idx_z,Ny,Nz) , map(idx_x ,idx_y , idx_z , Ny, Nz) ) = -6/delta^2;
            Tsparse( map(idx_x,idx_y,idx_z,Ny,Nz) , map(idx_x+1,idx_y , idx_z , Ny, Nz) ) = 1/delta^2;
            Tsparse( map(idx_x,idx_y,idx_z,Ny,Nz) , map(idx_x-1,idx_y , idx_z , Ny, Nz) ) = 1/delta^2;
            Tsparse( map(idx_x,idx_y,idx_z,Ny,Nz) , map(idx_x ,idx_y+1, idx_z , Ny, Nz) ) = 1/delta^2;
            Tsparse( map(idx_x,idx_y,idx_z,Ny,Nz) , map(idx_x ,idx_y-1, idx_z , Ny, Nz) ) = 1/delta^2;
            Tsparse( map(idx_x,idx_y,idx_z,Ny,Nz) , map(idx_x ,idx_y , idx_z+1, Ny, Nz) ) = 1/delta^2;
            Tsparse( map(idx_x,idx_y,idx_z,Ny,Nz) , map(idx_x ,idx_y , idx_z-1, Ny, Nz) ) = 1/delta^2;
        end
   end
end

这段代码形成了一个矩阵，在7个对角线上只有非零项(而不是每个对角线中的所有条目都是非零的)。

下面是我尝试创建T矩阵的代码版本，它更接近MATLAB的文档建议我这样做的方式：

delta = 0.1e-9;
Lx = 2e-9;
x = 0:delta:Lx;
Nx = length(x);
Ly = 2e-9;
y = 0:delta:Ly;
Ny = length(y);
Lz = 2e-9;
z = 0:delta:Lz;
Nz = length(z);

map = inline('((idx_x-1) * Ny*Nz) + ((idx_y-1) * Nz) + idx_z','idx_x','idx_y','idx_z','Ny','Nz'); % define an inline helper function for mapping (x,y,z) indices to a linear index

Iidx = zeros(7*(Nx-2)*(Ny-2)*(Nz-2),1); % matrix row indices
Jidx = zeros(7*(Nx-2)*(Ny-2)*(Nz-2),1); % matrix col indices
vals = zeros(7*(Nx-2)*(Ny-2)*(Nz-2),1); % matrix non-zero values, corresponding to (row,col)
cnt = 1;
for idx_x = 2:Nx-1
    for idx_y = 2:Ny-1
        for idx_z = 2:Nz-1
            % Tsparse( map(idx_x,idx_y,idx_z,Ny,Nz) , map(idx_x ,idx_y , idx_z , Ny, Nz) ) = -6/delta^2;
            Iidx(cnt) = map(idx_x,idx_y,idx_z,Ny,Nz);
            Jidx(cnt) = map(idx_x,idx_y,idx_z,Ny,Nz);
            vals(cnt) = -6/delta^2;
            cnt = cnt + 1;
            % Tsparse( map(idx_x,idx_y,idx_z,Ny,Nz) , map(idx_x+1,idx_y , idx_z , Ny, Nz) ) = 1/delta^2;
            Iidx(cnt) = map(idx_x,idx_y,idx_z,Ny,Nz);
            Jidx(cnt) = map(idx_x+1,idx_y,idx_z,Ny,Nz);
            vals(cnt) = 1/delta^2;
            cnt = cnt + 1;
            % Tsparse( map(idx_x,idx_y,idx_z,Ny,Nz) , map(idx_x-1,idx_y , idx_z , Ny, Nz) ) = 1/delta^2;
            Iidx(cnt) = map(idx_x,idx_y,idx_z,Ny,Nz);
            Jidx(cnt) = map(idx_x-1,idx_y,idx_z,Ny,Nz);
            vals(cnt) = 1/delta^2;
            cnt = cnt + 1;
            % Tsparse( map(idx_x,idx_y,idx_z,Ny,Nz) , map(idx_x ,idx_y+1, idx_z , Ny, Nz) ) = 1/delta^2;
            Iidx(cnt) = map(idx_x,idx_y,idx_z,Ny,Nz);
            Jidx(cnt) = map(idx_x,idx_y+1,idx_z,Ny,Nz);
            vals(cnt) = 1/delta^2;
            cnt = cnt + 1;
            % Tsparse( map(idx_x,idx_y,idx_z,Ny,Nz) , map(idx_x ,idx_y-1, idx_z , Ny, Nz) ) = 1/delta^2;
            Iidx(cnt) = map(idx_x,idx_y,idx_z,Ny,Nz);
            Jidx(cnt) = map(idx_x,idx_y-1,idx_z,Ny,Nz);
            vals(cnt) = 1/delta^2;
            cnt = cnt + 1;
            % Tsparse( map(idx_x,idx_y,idx_z,Ny,Nz) , map(idx_x ,idx_y , idx_z+1, Ny, Nz) ) = 1/delta^2;
            Iidx(cnt) = map(idx_x,idx_y,idx_z,Ny,Nz);
            Jidx(cnt) = map(idx_x,idx_y,idx_z+1,Ny,Nz);
            vals(cnt) = 1/delta^2;
            cnt = cnt + 1;
            % Tsparse( map(idx_x,idx_y,idx_z,Ny,Nz) , map(idx_x ,idx_y , idx_z-1, Ny, Nz) ) = 1/delta^2;
            Iidx(cnt) = map(idx_x,idx_y,idx_z,Ny,Nz);
            Jidx(cnt) = map(idx_x,idx_y,idx_z-1,Ny,Nz);
            vals(cnt) = 1/delta^2;
            cnt = cnt + 1;

        end
    end
end
Tsparse = sparse(Iidx, Jidx, vals, Nx*Ny*Nz, Nx*Ny*Nz);

谢谢您的建议！

-- dx.dy.dz

(附带注意："map“函数用于从三维坐标系(x，y，z)到一维值。假设我的特征值问题是H=E，其中H是Hamiltonian矩阵，psi是向量，E是标量。矩阵H=T+V (V没有显示在代码样本中，只有T)是在三维psi函数离散化并从3D到1D折叠的基础上编写的。例如，假设我每个维只有两个网格点，所以x=1:1:2，y=1:1:2，z=1:1:2，那么我的哈密顿量是用{psi(1,1,1)，psi(1,1,2)，psi(1,2,1)，psi(1,2,2)，psi(2,1,1)，psi(2,1,2)，psi(2,2,2)}，即psi(2，2，2)}，即8-by-8矩阵编写的。eigs()求解器输出的特征向量psi将是一个8分量向量，如果我愿意的话，我可以将它重塑回2x2x2矩阵。

Answer 1

我想我可以给你一些建议：

• 而不是您自己的映射，您可以考虑使用sub2ind函数
• 您反复调用map(idx_x,idx_y,idx_z,Ny,Nz) --当然您可以存储它以供重用。
• 此外，邻居的相对位置也将保持不变--无需重新计算。

举个小例子，我会这样做：

siz = [4,4,4];

pos = sub2ind(siz,1,1,1)

tmp = [
    sub2ind(siz,2,1,1)-pos
    sub2ind(siz,1,2,1)-pos
    sub2ind(siz,1,1,2)-pos
    ];

neighbors = [tmp;-tmp];
%%
big_dim = prod(siz);
mat = sparse(big_dim,big_dim);
%%
for i=2:siz(1)-1
    for j=2:siz(2)-1
        for k=2:siz(3)-1
            c_pos = sub2ind(siz,i,j,k);
            mat(c_pos,c_pos)=-6;
            c_neighbors=c_pos+neighbors;
            mat(c_pos,c_neighbors)=1;
        end
    end
end