CODEWITHSUNDEEP
matlabfor-loopvectormatrixvectorization
wigging
wigging

Reputation: 9180

Vectorize loop that creates tridiagonal matrix in Matlab

I have a for-loop that creates a tridiagonal matrix like so:

m = 5;
Fo = 0.35;

A(1,1) = 1+2*Fo;
A(1,2) = 1+Fo;

for i = 2:m-1
    A(i,i-1) = 1;
    A(i,i) = 2;
    A(i,i+1) = 3;
end

A(m,m-1) = 4;
A(m,m) = 5;

where the output is:

A =
    1.7000    1.3500         0         0         0
    1.0000    2.0000    3.0000         0         0
         0    1.0000    2.0000    3.0000         0
         0         0    1.0000    2.0000    3.0000
         0         0         0    4.0000    5.0000

I am trying to vectorize the creation of the tridiagonal matrix using the following to replace the for-loop:

i = 2:m-1;
A(i,i-1) = 1;
A(i,i) = 2;
A(i,i+1) = 3;

Unfortunately the the output is not correct:

A =
    1.7000    1.3500         0         0         0
    1.0000    2.0000    3.0000    3.0000    3.0000
    1.0000    2.0000    3.0000    3.0000    3.0000
    1.0000    2.0000    3.0000    3.0000    3.0000
         0         0         0    4.0000    5.0000

Is it possible to create such a matrix using vectorization instead of a for-loop? I will eventually need to create a much larger and more complicated tridiagonal matrix thus hoping to use vectorization to speed up the process.

Upvotes: 1

Views: 3032

Answers (2)

user85109
user85109

Reputation:

As Luis point out, sub2ind DOES work. I suppose you could also use diag. For example, you MIGHT do this:

m = 5;
Fo = 0.35;

d = [1+2*Fo; repmat(2,m - 2,1); 5];
ud = [1+Fo;repmat(3,m-2,1)];
sd = [ones(m-2,1);4];

A = diag(d) + diag(ud,1) + diag(sd,-1)
A =
          1.7         1.35            0            0            0
            1            2            3            0            0
            0            1            2            3            0
            0            0            1            2            3
            0            0            0            4            5

A problem is that we have done many additions in this composition, and most of the adds are zeros. Another reason to dislike this solution is because it generates a full matrix. I don't advise it in general, but it is a nice, visually intuitive solution. (By the way, I just noticed that the help for diag shows exactly this solution to create a tridiagonal matrix.)

A far better choice is to learn to use sparse matrices. Tridiagonal matrices are SPARSE. Use that capability. To do so, learn how to use spdiags, or at the very least, learn about sparse.

Let us build A as a sparse tridiagonal matrix. I'll use a larger value of m so we can see the savings.

m = 500;
Fo = 0.35;

d = [1+2*Fo; repmat(2,m - 2,1); 5];
ud = [1+Fo;repmat(3,m-2,1)];
sd = [ones(m-2,1);4];

A = diag(d) + diag(ud,1) + diag(sd,-1);
As = spdiags([[sd;0],d,[0;ud]],[-1 0 1],m,m);

whos A*
  Name        Size               Bytes  Class     Attributes
  A         500x500            2000000  double              
  As        500x500              27976  double    sparse

So the sparse form took 28k to store, while the full version took 2 megabytes.

The real benefit will come when you start to use these arrays in sparse form. For example, use backslash:

y = rand(m,1);

tic,x = A\y;toc
Elapsed time is 0.002847 seconds.

tic,xs = As\y;toc
Elapsed time is 0.000290 seconds.

I suppose I should also show how to do this with my own code, blktridiag, found on the MATLAB Central file exchange. It is really designed to generate block tridiagonal arrays, but it will work on this problem too, since we just have scalar blocks.

Ab = blktridiag(reshape(d,1,1,m),reshape(sd,1,1,m-1),reshape(ud,1,1,m-1));
As - Ab
ans =
   All zero sparse: 500-by-500

Finally, as natan points out, if your goal is a full matrix, then the function full will do that, given sparse matrix input. However, in many cases, the sparse form will be a great benefit. Learn to use sparse matrices. You will be happy that you did.

Upvotes: 2

Luis Mendo
Luis Mendo

Reputation: 112749

You can vectorize the for loop using sub2ind:

m = 5;
Fo = 0.35;

A = zeros(m); % initialize
A(sub2ind([m m],2:m, 1:m-1)) = 1;
A(sub2ind([m m],1:m, 1:m)) = 2;
A(sub2ind([m m],1:m-1, 2:m)) = 3;

A(1,1) = 1+2*Fo;
A(1,2) = 1+Fo;

A(m,m-1) = 4;
A(m,m) = 5;

The problem with your approach to replace the loop is that the two indices are considered to define a block (all combinations of the two indices), not a diagonal (each value of first index with each corresponding value in the second).

Upvotes: 1

Related Questions