Code covered by the BSD License

(Block) tri-diagonal matrices

John D'Errico (view profile)

01 Apr 2006 (Updated )

Generate (block) tridiagonal matrices

blktridiag(Amd,Asub,Asup,n)
function A = blktridiag(Amd,Asub,Asup,n)
% BLKTRIDIAG: computes a sparse (block) tridiagonal matrix with n blocks
% usage: A = BLKTRIDIAG(Amd,Asub,Asup,n)  % identical blocks
% usage: A = BLKTRIDIAG(Amd,Asub,Asup)    % a list of distinct blocks
%
% BLKTRIDIAG runs in two distinct modes. The first mode
% supplies three blocks, one for the main diagonal, and
% the super and subdiagonal blocks. These blocks will be
% replicated n times down the main diagonals of the matrix,
% and n-1 times down the sub and super diagonals.
%
% The second mode is to supply a list of distinct blocks
% for each diagonal, as planes of 3d arrays. No replication
% factor is needed in this mode.
%
% arguments: (input mode 1)
%  Amd  - pxq array, forming the main diagonal blocks
%
%  Asub - pxq array, sub diagonal block
%         Asub must be the same size and shape as Amd
%
%  Asup - pxq array, super diagonal block
%         Asup must be the same size and shape as Amd
%
%  n    - scalar integer, defines the number of blocks
%         When n == 1, only a single block will be formed, A == Amd
%
% arguments: (input mode 2)
%  Amd  - pxqxn array, a list of n distinct pxq arrays
%         Each plane of Amd corresponds to a single block
%         on the main diagonal.
%
%  Asub - pxqx(n-1) array, a list of n-1 distinct pxq arrays,
%         Each plane of Asub corresponds to a single block
%         on the sub-diagonal.
%
%  Asup - pxqx(n-1) array, a list of n-1 distinct pxq arrays,
%         Each plane of Asup corresponds to a single block
%         on the super-diagonal.
%
% Note: the sizes of Amd, Asub, and Asup must be consistent
% with each other, or an error will be generated.
%
% arguments: (output)
%  A    - (n*p by n*q) SPARSE block tridiagonal array
%         If you prefer that A be full, use A=full(A) afterwards.
%
%
% Example 1:
%  Compute the simple 10x10 tridiagonal matrix, with 2 on the
%  diagonal, -1 on the off diagonal.
%
%  A = blktridiag(2,-1,-1,10);
%
%
% Example 2:
%  Compute the 5x5 lower bi-diagonal matrix, with blocks of
%  [1 1;1 1] on the main diagonal, [2 2;2 2] on the sub-diagonal,
%  and blocks of zeros above.
%
%  A = blktridiag(ones(2),2*ones(2),zeros(2),5);
%
% Example 3:
%  Compute the 3x6 tridiagonal matrix, with non-square blocks
%  that vary along the main diagonal, [2 2] on the sub-diagonal,
%  and [1 1] on the super-diagonal. Note that all blocks must have
%  the same shape.
%
%  A = blktridiag(rand(1,2,3),2*ones(1,2,2),ones(1,2,2));
%
%
%
% Author: John D'Errico
% e-mail address: woodchips@rochester.rr.com
% Release: 4.0
% Original release date: 4/01/06
% Current release date: 12/14/07

% Which mode of operation are we in?
if nargin==4
% replicated block mode

% verify the inputs in this mode are 2-d arrays.
if (length(size(Amd))~=2) || ...
(length(size(Asub))~=2) || ...
(length(size(Asup))~=2)
error 'Inputs must be 2d arrays if a replication factor is provided'
end

% get block sizes, check for consistency
[p,q] = size(Amd);
if isempty(Amd)
error 'Blocks must be non-empty arrays or scalars'
end
if any(size(Amd)~=size(Asub)) || any(size(Amd)~=size(Asup))
error 'Amd, Asub, Asup are not identical in size'
end

if isempty(n) || (length(n)>1) || (n<1) || (n~=floor(n))
error 'n must be a positive scalar integer'
end

% scalar inputs?
% since p and q are integers...
if (p*q)==1
if n==1
A = Amd;
else
% faster as Jos points out
A = spdiags(repmat([Asub Amd Asup],n,1),-1:1,n,n);
end
% no need to go any farther
return
end

% use sparse. the main diagonal elements of each array are...
v = repmat(Amd(:),n,1);
% then the sub and super diagonal blocks.
if n>1
% sub-diagonal
v=[v;repmat(Asub(:),n-1,1)];

% super-diagonal
v=[v;repmat(Asup(:),n-1,1)];
end

elseif nargin==3
% non-replicated blocks, supplied as planes of a 3-d array

% get block sizes, check for consistency
[p,q,n] = size(Amd);
if isempty(Amd)
error 'Blocks must be (non-empty) arrays or scalars'
end

if (p~=size(Asub,1)) || (q~=size(Asub,2)) || (p~=size(Asup,1)) || (q~=size(Asup,2))
error 'Amd, Asub, Asup do not have the same size blocks'
end

if (n>1) && (((n-1) ~= size(Asub,3)) || ((n-1) ~= size(Asup,3)))
error 'Asub and Asup must each have one less block than Amd'
end

% scalar inputs?
if (p*q)==1
if n==1
A = Amd(1);
else
% best to just use spdiags
A = spdiags([[Asub(:);0], Amd(:), [0;Asup(:)]],-1:1,n,n);
end
% no need to go any farther
return
end

% The main diagonal elements
v = Amd(:);
% then the sub and super diagonal blocks.
if n>1
% sub-diagonal
v=[v;Asub(:)];

% super-diagonal
v=[v;Asup(:)];
end
else
% must have 3 or 4 arguments
error 'Must have either 3 or 4 arguments to BLKTRIDIAG'
end

% now generate the index arrays. first the main diagonal
[ind1,ind2,ind3]=ndgrid(0:p-1,0:q-1,0:n-1);
rind = 1+ind1(:)+p*ind3(:);
cind = 1+ind2(:)+q*ind3(:);
% then the sub and super diagonal blocks.
if n>1
% sub-diagonal
[ind1,ind2,ind3]=ndgrid(0:p-1,0:q-1,0:n-2);
rind = [rind;1+p+ind1(:)+p*ind3(:)];
cind = [cind;1+ind2(:)+q*ind3(:)];

% super-diagonal
rind = [rind;1+ind1(:)+p*ind3(:)];
cind = [cind;1+q+ind2(:)+q*ind3(:)];
end

% build the final array all in one call to sparse
A = sparse(rind,cind,v,n*p,n*q);

% ====================================
% end mainline
% ====================================