Lets take a look at what we can do to speed up the matrix solution, which is the slowest link in the solution of the Schrodinger equation with Crank-Nicholson.

For now, we will work without the package context to make things easier.

Map[(Unprotect[#]; Remove[#]) &, 
      Names["LinearAlgebra`Tridiagonal`*"]]

Package Code

Here is the code taken from the package.

\(TridiagonalSolve[a_List, b_List, c_List, r_List] := 
      Module[{len = Length[r], solution = Array[0, Length[r]], aux, 
            aux1 = Array[0, Length[r]], a1 = Prepend[a, 0], iter}, 
          aux = 1/b[\([1]\)]; 
solution[\([1]\)] = 
            r[\([1]\)]\ aux; 
Do[
            aux1[\([iter]\)] = 
              c[\([iter - 1]\)]\ aux; 
aux = 
              1/\((b[\([iter]\)] - 
                    a1[\([iter]\)]\ aux1[\([iter]\)])\); 
solution[\([iter]\)] = \((r[\([iter]\)] - 
                    a1[\([iter]\)]\ solution[\([iter - 1]\)])\)\ aux, {iter, 
              2, len}]; \ 			Do[
            solution[\([iter]\)] -= 
              aux1[\([iter + 1]\)]\ solution[\([iter + 1]\)], {iter, len - 1, 
              1, \(-1\)}]; 
solution] /; 
        Length[a] + 1 == Length[b] == Length[c] + 1 == 
          Length[r]\)

It turns out that I will get ranges of Timings that are very great.  A way I have found to deal with this is to define an alternative time function.

baseline  =  
      time[TridiagonalSolve[offDiag,  diag,  offDiag, χ]]

Maybe the problem is that it is not using PackedArrays (the patterns in a_List with the Condition will cause unpacking in Version 4.  In the development version, this will not be a problem.)

See about  packing

Map[
      Developer`PackedArrayQ,  {offDiag,  diag,  χ, 
        TridiagonalSolve[offDiag,  diag,  offDiag, χ]}]

Before blindly trying to pack things we should check if this is even appropriate.  For example, the precision might be bigger than $MachinePrecision.

We could use Precision to check this, but if some of the inputs are inexact zeroes, Precision will give 0, which may not reflect the effective precision of the expression.  The following is a workaround.

SymbolQ[x_]  :=  MatchQ[x,  _Symbol]

Clear[SomewhatEffectivePrecision];

SomewhatEffectivePrecision[
          x_?MachineNumberQ]  :=  $MachinePrecision;

SomewhatEffectivePrecision[x_?NumberQ]  :=  
        If[Developer`ZeroQ[x], Accuracy[x],  
          Precision[x]];

SomewhatEffectivePrecision[x_List]  :=  
        Min[Map[SomewhatEffectivePrecision, x,  
            TensorRank[x]]];

SomewhatEffectivePrecision[x_]  :=  
If[
          And[Length[x]  >  0 , 
SymbolQ[Head[x]], 
            MemberQ[Attributes[Evaluate[Head[x]]],  
              NumericFunction]], 
SomewhatEffectivePrecision[
            Apply[List, 
              x]], 
Infinity, 
];

This is a useful function for more than just how it is used below.  Since variables like x[1.] do not really affect how the expression coerces, etc. in arithmetic, this avoids looking inside non-numeric functions.  For example,

SomewhatEffectivePrecision[1`20 z[1. ]  +  0``19]

Now, using this effective precision we can make a decision about unpacking.  Basically, if we find machine numbers, try to pack all of the input.

SeeAboutPacking[input_]  :=  
      Block[{prec,  nin,  type}, 
prec  =  
          Min[Map[SomewhatEffectivePrecision,  
              input]]; 
If[
          prec  ≠  $MachinePrecision,  
          Return[False]]; 
pin  =  
          Map[Developer`ToPackedArray[#,  Union] &,  
            input]; 
If[
          Not[TrueQ[
              Apply[And,  
                Map[Developer`PackedArrayQ,  
                  pin]]]], 
Return[
            False]]; 
type  =  
          If[Count[input[[All, 1]],  _Complex]  ==  
              0, 
Real,  
            Complex]; 
pin  =  
          Map[Developer`ToPackedArray[#, type] &,  
            pin]; 
pin]

Now we redefine TridiagonalSolve with the packing step included

\(TridiagonalSolve1[ain_, bin_, cin_, rin_] := 
      Module[{len = Length[rin], solution, aux, aux1, a, b, c, r, \ iter, \ 
          zero}, 
aux\  = \ 
          SeeAboutPacking[{ain, bin, cin, rin}]; 
If[
          ListQ[aux], 
{a, b, c, r}\  = \ 
            aux; 
zero\  = \ 
            0. \ \ aux[\([1, 1]\)], 
{a, b, c, r}\  = \ 
            Developer`FromPackedArray[{ain, \ bin, \ cin, \ 
                rin}]; 
zero\  = \ 
            0]; 
a\  = \ 
          Prepend[a, \ zero]; 
solution\  = \ 
          zero\ Range[len]; 
aux1\  = \ 
          solution; 
aux = 
          1/b[\([1]\)]; 
solution[\([1]\)] = 
          r[\([1]\)]\ aux; 
Do[
          aux1[\([iter]\)] = c[\([iter - 1]\)]\ aux; 
aux = 
            1/\((b[\([iter]\)] - 
                  a[\([iter]\)]\ aux1[\([iter]\)])\); 
solution[\([iter]\)] = \((r[\([iter]\)] - 
                  a[\([iter]\)]\ solution[\([iter - 1]\)])\)\ aux, {iter, 2, 
            len}]; \ 	
Do[
          solution[\([iter]\)] -= 
            aux1[\([iter + 1]\)]\ solution[\([iter + 1]\)], {iter, len - 1, 
            1, \(-1\)}]; 
solution]\)

Lets see if this gives an improvement:

time[TridiagonalSolve1[offDiag,  diag,  offDiag, χ]]

Since the vectors giving the matrix are constant for all of the timesteps, we can do better, by packing them once rather than repeatedly:

offDiag  =  
        Developer`ToPackedArray[N[offDiag]];
diag  =  Developer`ToPackedArray[N[diag]];

time[TridiagonalSolve1[offDiag,  diag,  offDiag, χ]]

baseline/%

The packing function was set up so that if the input was packed, it all came out the same type.  This was so that replacing parts in one did not force unpacking due to type differences.  Since all are the same type, it seems like we should be able to compile it:

Try Compile

This talk is about general programming, so we will set up a real and complex version.

\(\(RealTridiagonalSolve\  = \ 
        Compile[{{a, \ _Real, \ 1}, \ {b, \ _Real, \ 1}, \ {c, \ _Real, \ 
              1}, \ {r, \ _Real, \ 1}}, 
          Module[{len = Length[r], solution = Table[0. , {Length[r]}], 
              aux\  = \ 0. , aux1 = Table[0. , {Length[r]}], 
              a1 = Prepend[a, 0. ], iter\  = \ 0}, 
            aux = 1/b[\([1]\)]; 
solution[\([1]\)] = 
              r[\([1]\)]\ aux; 
Do[
              aux1[\([iter]\)] = 
                c[\([iter - 1]\)]\ aux; 
aux = 
                1/\((b[\([iter]\)] - 
                      a1[\([iter]\)]\ aux1[\([iter]\)])\); 
solution[\([iter]\)] = \((r[\([iter]\)] - 
                      a1[\([iter]\)]\ solution[\([iter - 
                              1]\)])\)\ aux, {iter, 2, len}]; \ 			Do[
              solution[\([iter]\)] -= 
                aux1[\([iter + 1]\)]\ solution[\([iter + 1]\)], {iter, 
                len - 1, 
                1, \(-1\)}]; 
solution]];\)\)

\(\(ComplexTridiagonalSolve\  = \ 
        Compile[{{a, \ _Complex, \ 1}, \ {b, \ _Complex, \ 
              1}, \ {c, \ _Complex, \ 1}, \ {r, \ _Complex, \ 1}}, 
          Module[{len = Length[r], 
              solution = Table[0.  + 0. \ I, {Length[r]}], 
              aux\  = \ 0.  + 0. \ I, 
              aux1 = Table[0.  + 0. \ I, {Length[r]}], 
              a1 = Prepend[a, 0.  + 0. \ I], iter\  = \ 0}, 
            aux = 1/b[\([1]\)]; 
solution[\([1]\)] = 
              r[\([1]\)]\ aux; 
Do[
              aux1[\([iter]\)] = 
                c[\([iter - 1]\)]\ aux; 
aux = 
                1/\((b[\([iter]\)] - 
                      a1[\([iter]\)]\ aux1[\([iter]\)])\); 
solution[\([iter]\)] = \((r[\([iter]\)] - 
                      a1[\([iter]\)]\ solution[\([iter - 
                              1]\)])\)\ aux, {iter, 2, len}]; \ 			Do[
              solution[\([iter]\)] -= 
                aux1[\([iter + 1]\)]\ solution[\([iter + 1]\)], {iter, 
                len - 1, 
                1, \(-1\)}]; 
solution]];\)\)

Note:  in the development version, there are some mechanisms being put together to do this sort of thing automatically.  I recommend Tom-Wickham Jones talk on the compiler if you are interested.

Now we have a compiled version for real and complex types.  The Schrodinger equation will use the complex one:

speedy  =  
      time[ComplexTridiagonalSolve[offDiag,  diag,  
          offDiag, χ]]

baseline/speedy

Are we really getting the same result that much faster?

Max[
      Abs[ComplexTridiagonalSolve[offDiag,  diag,  offDiag, χ]  -  
          TridiagonalSolve[offDiag,  diag,  offDiag, χ]]]

If all we want to do is solve this problem, we are done:  we can just use the complex version.  However, since this talk is about generally efficient numerical algorithms, lets make it work for all cases:

Generalize

TridiagonalSolve2[ain_List, bin_List, cin_List, rin_List] := 
      Module[{pack  =  
            SeeAboutPacking[{ain, bin, cin, rin}]}, 
If[
          ListQ[pack], 
Apply[
            If[Head[pack[[1, 1]]]  ===  
                Complex, 
ComplexTridiagonalSolve, 
RealTridiagonalSolve],  
            pack], 
TridiagonalSolve[a, b, c, 
            r]]]

time[TridiagonalSolve2[offDiag,  diag,  offDiag, χ]]

Whoa!  Where did the speed go?

Speed trap

time[
      Map[SomewhatEffectivePrecision, {offDiag,  diag,  
          offDiag, χ}]]

Aha!  Perhaps we can rewrite this...

Clear[SomewhatMoreEffectivePrecision]; 
    SomewhatMoreEffectivePrecision[
        x_]  :=  
Which[
MachineNumberQ[
          x], 16, 
NumberQ[x], 
If[
          Developer`ZeroQ[x], Accuracy[x],  
          Precision[x]], 
ListQ[x], 
If[
          Developer`PackedArrayQ[
            x], 
SomewhatMoreEffectivePrecision[
            Part[x, Apply[Sequence,  
                Table[1, {TensorRank[x]}]]]], 
Min[
            Map[SomewhatMoreEffectivePrecision,  
              x]]], 
And[ Length[x]  >  0 , 
          SymbolQ[Head[x]], 
MemberQ[
            Attributes[Evaluate[Head[x]]],  
            NumericFunction]], 
SomewhatMoreEffectivePrecision[
          Apply[List, x]], 
True,  Infinity]

time[
      Min[Map[SomewhatMoreEffectivePrecision, {offDiag,  diag,  
            offDiag, χ}]]]

Better, but still awfully expensive:

%/speedy

time[
      SomewhatMoreEffectivePrecision[{offDiag,  diag,  
          offDiag, χ}]]

Something is fishy: you would expect these to take the same amount of time.  Now is a good time to see if the unpack message tells us anything.

The unpacking message can be a useful tool:

Developer`SetSystemOptions["UnpackMessage" -> True]

SomewhatMoreEffectivePrecision[{offDiag,  diag,  
        offDiag, χ}]

Clear[SomewhatMoreEffectivePrecision]; 
    SomewhatMoreEffectivePrecision[
        x_]  :=  
Which[
ListQ[
          x], 
If[
          Developer`PackedArrayQ[
            x], 
SomewhatMoreEffectivePrecision[
            Part[x, Apply[Sequence,  
                Table[1, {TensorRank[x]}]]]], 
Min[
            Map[SomewhatMoreEffectivePrecision,  
              x]]], 
MachineNumberQ[x], 
        16, 
NumberQ[x], 
If[
          Developer`ZeroQ[x], Accuracy[x],  
          Precision[x]], 
And[ Length[x]  >  0 , 
          SymbolQ[Head[x]], 
MemberQ[
            Attributes[Evaluate[Head[x]]],  
            NumericFunction]], 
SomewhatMoreEffectivePrecision[
          Apply[List, x]], 
True,  Infinity]

time[
      SomewhatMoreEffectivePrecision[{offDiag,  diag,  
          offDiag, χ}]]

%/speedy

Now we can rewrite the packing check using this.  Note also, that since the compiled function interpreter automatically coerces to the right type, we don't have to do this explicitly:

Clear[SeeAboutPacking]; 
    SeeAboutPacking[input_]  :=  
      Block[{prec,  nin,  type}, 
prec  =  
          Min[Map[SomewhatMoreEffectivePrecision,  
              input]]; 
If[
          prec  ≠  $MachinePrecision,  
          Return[False]]; 
pin  =  
          Map[Developer`ToPackedArray[#,  Union] &,  
            input]; 
If[
          Not[TrueQ[
              Apply[And,  
                Map[Developer`PackedArrayQ,  
                  pin]]]], 
False, 
pin]
]

time[SeeAboutPacking[{offDiag,  diag,  offDiag, χ}]]

time[TridiagonalSolve2[offDiag,  diag,  offDiag, χ]]

%/speedy

Actually, this is not quite right, because we have are only checking the type in the first vector, which may not have been coerces, so we should change to

TridiagonalSolve3[ain_List, bin_List, cin_List, rin_List] := 
      Module[{pack  =  
            SeeAboutPacking[{ain, bin, cin, rin}]}, 
If[
          ListQ[pack], 
Apply[
If[
              MemberQ[Map[Head, pack[[All, 1]]], 
                Complex], 
ComplexTridiagonalSolve, 
RealTridiagonalSolve], 
 
            pack], 
TridiagonalSolve[a, b, c, 
            r]]]

time[TridiagonalSolve3[offDiag,  diag,  offDiag, χ]]

%/speedy

Now is it fast enough?

Factorization

Mathematica 4 has a built-in general sparse solver

TridiagonalSolve4[a_List,  b_List,  c_List,  
        r_List]  :=  
Module[{n  =  Length[b], 
          spam}, 
spam  =  
          Join[Flatten[
              Table[{{i + 1, i} -> a[[i]],  {i, i} -> 
                    b[[i]],  {i, i + 1} -> c[[i]]}, {i, 
                  n - 1}], 
              1], {{n, n} -> 
                b[[n]]}]; 
Developer`SparseLinearSolve[
          spam,  r]]

time[
      TridiagonalSolve4[offDiag,  diag,  offDiag,  χ]]

What is less well know is that you can also get a sparse factorization. (Actually a sparse LU decomposition.) Since the matrix is constant, this is ideal:

TridiagonalLU[a_List,  b_List,  c_List]  :=  
      Module[{n  =  Length[b], spam}, 
spam  =  
          Join[Flatten[
              Table[{{i + 1, i} -> a[[i]],  {i, i} -> 
                    b[[i]],  {i, i + 1} -> c[[i]]}, {i, 
                  n - 1}], 
              1], {{n, n} -> 
                b[[n]]}]; 
Experimental`ExtendedLinearSolve[spam]]

This gives a special function which operates on vectors of the right length to give the system solution.

tlu  =  TridiagonalLU[offDiag,  diag,  offDiag]

time[tlu[ χ]]

%/speedy

baseline/%%