YAKL_tridiagonal.h

Go to the documentation of this file.

 
#pragma once
 
#include "YAKL.h"
 
 
__YAKL_NAMESPACE_WRAPPER_BEGIN__
namespace yakl {
 
 
  template <class real, unsigned int n>
  YAKL_INLINE void tridiagonal_periodic(SArray<real,1,n> const &a, SArray<real,1,n> &b, SArray<real,1,n> &c, SArray<real,1,n> &d) {
    SArray<real,1,n> tmp;
    // Save the original "b0" because it's needed later on to compute ( (v^T*y) / (1 + v^T*q) )
    real b0 = b(0);
 
    // This is the vector "u"
    tmp(0  ) = -b0;
    tmp(n-1) =  c(n-1);
 
    // The new tridiagonal system "B" alters the entries of the main diagonal
    b(n-1) = b(n-1) + a(0)*c(n-1)/b(0);
    b(0  ) = b(0  ) + b(0);
 
    // Thomas algorithm for  B*y=d  and  B*q=u simultaneously to save cost
    real div = static_cast<real>(1) / b(0);
    c  (0) *= div;
    d  (0) *= div;
    tmp(0) *= div;
    for (int i = 1; i < n-1; i++) {
      div = static_cast<real>(1) / (b(i) - a(i)*c(i-1));
      c  (i) =  c(i)                  * div;
      d  (i) = (d(i) - a(i)*d  (i-1)) * div;
      tmp(i) = (     - a(i)*tmp(i-1)) * div;
    }
    div = static_cast<real>(1) / (b(n-1) - a(n-1)*c(n-2));
    d  (n-1) = (d  (n-1) - a(n-1)*d  (n-2)) * div;
    tmp(n-1) = (tmp(n-1) - a(n-1)*tmp(n-2)) * div;
    for (int i = n-2; i >= 0; i--) {
      d  (i) -= c(i)*d  (i+1);
      tmp(i) -= c(i)*tmp(i+1);
    }
 
    // Compute factor = ( (v^T*y) / (1 + v^T*q) )
    real factor;
    if ( (tmp(0) - a(0)*tmp(n-1)/b0 + static_cast<real>(1)) == 0 ) {
      factor = 0;
    } else {
      factor = (d(0) - a(0)*d(n-1)/b0)/(tmp(0) - a(0)*tmp(n-1)/b0 + static_cast<real>(1));
    }
 
    // Subtract factor * q from the previous solution to get the final solution
    for (int i = 0; i < n; i++) {
      d(i) -= factor*tmp(i);
    }
  }
 
 
 
  template <class real, unsigned int n>
  YAKL_INLINE void tridiagonal(SArray<real,1,n> const &a, SArray<real,1,n> const &b, SArray<real,1,n> &c, SArray<real,1,n> &d) {
    real tmp = static_cast<real>(1) / b(0);
    c(0) *= tmp;
    d(0) *= tmp;
    for (int i = 1; i < n-1; i++) {
      real tmp = static_cast<real>(1) / (b(i) - a(i)*c(i-1));
      c(i) *= tmp;
      d(i) = (d(i) - a(i)*d(i-1)) * tmp;
    }
    d(n-1) = (d(n-1) - a(n-1)*d(n-2)) / (b(n-1) - a(n-1)*c(n-2));
    for (int i = n-2; i >= 0; i--) {
      d(i) -= c(i)*d(i+1);
    }
  }
 
}
__YAKL_NAMESPACE_WRAPPER_END__