/* trid-lu.c -- 三項方程式を Gauss の消去法で解く */

#include "trid-lu.h"

/*
 *  三項方程式とは、三重対角行列を係数行列とする連立1次方程式のこと。
 *
 *  Gauss の消去法のアルゴリズムで、係数行列を LU 分解し、それを利用し
 * て、連立1次方程式を解く。
 *
 *  (ピボットの選択はしていないので、係数行列が正定値対称行列であるな
 * ど、よい性質を持っていないと、 結果は保証されない。) */

/************************************************************************
 * trid0(), trilu0(), trisol0()
 *   添字 0 から格納されている場合
 *   trilu0()   三重対角行列を LU 分解する
 *   trisol0()  LU 分解済の三重対角行列を用いて、三項方程式を解く
 *   trid0()    三項方程式を解く (trilu0(), trisol0() を順に呼び出すだけ)
 ************************************************************************/

/*
 *  AL, AD, AU は係数行列の
 *     下三角部分 (lower part)
 *     対角部分 (diagonal part)
 *     上三角部分 (upper part)
 *  を表す。つまり係数行列は、次のような形をしている。
 *
 *  AD[0] AU[0]   0   .................. 0
 *  AL[1] AD[1] AU[1]   0  ............. 0
 *    0   AL[2] AD[2] AU[2]  0 ......... 0
 *          ....................
 *                    AL[n-2] AD[n-2] AU[n-2]
 *                      0     AL[n-1] AD[n-1]
 *
 *  ここで AL[0], AU[n-1] は意味がないことに注意。 */

/*
 *  trid0() の仕様
 *  入力
 *     n: 未知数の個数
 *     al,ad,au: 連立1次方程式の係数行列
 *       (al: 対角線の下側、 ad: 対角線、 au: 対角線の上側)
 *        al[i] = A_{i,i-1}, ad[i] = A_{i,i}, au[i] = A_{i,i+1},
 *        al[0], au[n-1] は意味がない)
 *     b: 連立1次方程式の右辺の既知ベクトル
 *       (添字は 0 から。i.e. b[0],b[1],...,b[n-1] にデータが入っている。)
 *   出力
 *     al,ad,au: 入力した係数行列を LU 分解したもの
 *     b: 連立一次方程式の解
 *   能書き
 *     一度 call すると係数行列を LU 分解したものが返されるので、
 *     以後は同じ係数行列に関する連立1次方程式を解くために、
 *     サブルーチン trisol0() が使える。
 *   注意
 *     ピボットの選択をしていないので、係数行列が正定値である
 *     などの適切な条件がない場合は結果が保証できない。
 */

void trid0(int n, double *al, double *ad, double *au, double *b)
{
  trilu0(n,al,ad,au);
  trisol0(n,al,ad,au,b);
}

/* 三重対角行列の LU 分解 (pivoting なし) */
void trilu0(int n, double *al, double *ad, double *au)
{
  int i, nm1 = n - 1;
  /* 前進消去 (forward elimination) */
  for (i = 0; i < nm1; i++) ad[i + 1] -= au[i] * al[i + 1] / ad[i];
}

/* LU 分解済みの三重対角行列を係数に持つ三項方程式を解く */
void trisol0(int n, double *al, double *ad, double *au, double *b)
{
  int i, nm1 = n - 1;
  /* 前進消去 (forward elimination) */
  for (i = 0; i < nm1; i++) b[i + 1] -= b[i] * al[i + 1] / ad[i];
  /* 後退代入 (backward substitution) */
  b[nm1] /= ad[nm1];
  for (i = n - 2; i >= 0; i--) b[i] = (b[i] - au[i] * b[i + 1]) / ad[i];
}

/************************************************************************
 * trid1(), trilu1(), trisol1()
 *   添字 1 から格納されている場合
 *   trilu1()   三重対角行列を LU 分解する
 *   trisol1()  LU 分解済の三重対角行列を用いて、三項方程式を解く
 *   trid1()    三項方程式を解く (trilu0(), trisol0() を順に呼び出すだけ)
 ************************************************************************/

/*
 *  係数行列は、次のような形をしている
 *
 *  AD[1] AU[1]   0   .................. 0
 *  AL[2] AD[2] AU[2]   0  ............. 0
 *    0   AL[3] AD[3] AU[3]  0 ......... 0
 *          ....................
 *                    AL[n-1] AD[n-1] AU[n-1]
 *                      0     AL[n]   AD[n]
 *
 *  ここで AL[1], AU[n] は意味がないことに注意。 */

void trid1(int n, double *al, double *ad, double *au, double *b)
{
  trid0(n, al+1, ad+1, au+1, b+1);
}

void trilu1(int n, double *al, double *ad, double *au)
{
  trilu0(n, al+1, ad+1, au+1);
}

void trisol1(int n, double *al, double *ad, double *au, double *b)
{
  trisol0(n, al+1, ad+1, au+1, b+1);
}