#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
Solve equation A_{i,k} x_{k} = f_{i} with tridiagonal matrix elements A_{i,k} for all x_k.
All indices are 0-based, except for the lower band diagonal (in order for the lower band diagonal to match the row index of the cell in the matrix).
Therefore, the element with index 0 in the lower band diagonal is unused.

The principle used is decomposition of A into a lower left matrix L and an upper right matrix U (note: no Einstein summation convention here!):

A = L U
where A_{i,k} = L_{i,j} U_{j,k}
where L_{i,j} = δ_{i,j} + λ_i δ_{i, j + 1}
where U_{j,k} = d_j δ{j,k} + c_j δ_{j, k - 1}
where A_{i,k} = diagMain(i) δ{i,k} + diagLower(i) δ_{i, k + 1} + diagUpper(i) δ_{i, k - 1}

Therefore, in order to find the coefficients λ_i, d_j, c_j we do:

A_{i,k} = Σ_j L_{i,j} U_{j,k}
A_{i,k} = Σ_j (δ_{i,j} + λ_i δ_{i, j + 1}) (d_j δ{j,k} + c_j δ_{j, k - 1}) = diagMain(i) δ{i,k} + diagLower(i) δ_{i, k + 1} + diagUpper(i) δ_{i, k - 1}
A_{i,k} = Σ_j (δ_{i,j} d_j δ{j,k} + δ_{i,j} c_j δ_{j, k - 1}) (λ_i δ_{i, j + 1} d_j δ{j,k} + λ_i δ_{i, j + 1} c_j δ_{j, k - 1}) = diagMain(i) δ{i,k} + diagLower(i) δ_{i, k + 1} + diagUpper(i) δ_{i, k - 1}
A_{i,k} = δ_{i,k} d_i + c_i δ_{i, k - 1} + λ_i δ_{i, k + 1} d_k + λ_i δ_{i, k} c_{i - 1} = diagMain(i) δ{i,k} + diagLower(i) δ_{i, k + 1} + diagUpper(i) δ_{i, k - 1}

Comparing coefficients:

δ_{i,k} d_i + λ_i δ_{i, k} c_{i - 1}  = diagMain(i) δ{i,k}
c_i δ_{i, k - 1} = diagUpper(i) δ_{i, k - 1}
λ_i δ_{i, k + 1} d_k = diagLower(i) δ_{i, k + 1}

d_i + λ_i c_{i - 1}  = diagMain(i)
c_i = diagUpper(i)
λ_i d_{i - 1} = diagLower(i)

So we get:

d_i = diagMain(i) - λ_i diagUpper(i - 1)
c_i = diagUpper(i)
λ_i = diagLower(i)/d_{i - 1}

[ which would be d_i = diagMain(i) - (diagLower(i)/d_{i - 1}) diagUpper(i - 1), but no point in simplifying it even further.]

Seperate analysis of the corners yields:

A_{0,0} = L_{0,j} U_{j,0}
where L_{0,j} = δ_{0,j}
where U_{j,0} = d_{0} δ{j,0}
A_{0,0} = d_{0} = diagMain(0)

Combining the equations in order to find out the shape of matrix A, one gets (we're back inside):

A x = f

L U x = f

L y = f
U x = y

L_{i,j} y_j = f_i
U_{j,k} x_k = y_j
L_{i,j} = δ_{i,j} + λ_i δ_{i, j + 1}
U_{j,k} = d_j δ{j,k} + c_j δ_{j, k - 1}

Σ_j (δ_{i,j} + λ_i δ_{i, j + 1}) y_j = f_i
Σ_k (d_j δ{j,k} + c_j δ_{j, k - 1}) x_k = y_j

(Σ_j δ_{i,j} + Σ_j λ_i δ_{i, j + 1}) (Σ_k d_j δ{j,k} + Σ_k c_j δ_{j, k - 1}) x_k = f_i
Σ_j δ_{i,j} Σ_k d_j δ{j,k} x_k + Σ_j δ_{i,j} Σ_k c_j δ_{j, k - 1} x_k + Σ_j λ_i δ_{i, j + 1} Σ_k d_j δ{j,k} x_k + Σ_j λ_i δ_{i, j + 1} Σ_k c_j δ_{j, k - 1} x_k = f_i
d_i x_i + c_i x_{i + 1} + λ_i d_{i - 1} x_{i - 1} + λ_i c_{i - 1} x_i = f_i

So the matrix is tridiagonal and the composition works.

Let's return to the decompositioned form:

For a generic non-last component in x (index j) we get:

Σ_k (d_j δ{j,k} + c_j δ_{j, k - 1}) x_k = y_j
Σ_k d_j δ{j,k} x_k + Σ_k c_j δ_{j, k - 1} x_k = y_j
d_j x_j + c_j x_{j + 1} = y_j
d_j x_j = y_j - c_j x_{j + 1}
x_j = (y_j - c_j x_{j + 1})/d_j

For the last component in x (index u := N - 1) we get:

Σ_k (d_j δ{j,k} + c_j δ_{j, k - 1}) x_k = y_j
Σ_k (d_u δ{u,k} + c_u δ_{u + 1, k}) x_k = y_j
Σ_k d_u δ{u,k} x_k = y_j
d_u x_u = y_u
x_u = y_u/d_u

For y we get:

Σ_j (δ_{i,j} + λ_i δ_{i, j + 1}) y_j = f_i
Σ_j δ_{i,j} y_j + Σ_j λ_i δ_{i, j + 1} y_j = f_i
y_i + λ_i y_{i - 1} = f_i
y_i = f_i - λ_i y_{i - 1}

For the first component in y (index 0) we get:

Σ_j (δ_{0,j} + λ_0 δ_{0, j + 1}) y_j = f_0
Σ_j δ_{0,j} y_j = f_0
y_0 = f_0

So to summarize:
  x_j = (y_j - c_j x_{j + 1})/d_j    except for the last, where it is: x_u = y_u/d_u
  d_i = diagMain(i) - λ_i diagUpper(i - 1)  except for the first, where it is: d_0 = diagMain(0)
  c_i = diagUpper(i)
  λ_i = diagLower(i)/d_{i - 1}
  y_i = f_i - λ_i y_{i - 1}   except for the first, where it is: y_0 = f_0
"""

from __future__ import division
from memo import memoized

@memoized
def solveLinearTridiagonalEquation(diagUpper, diagMain, diagLower, inhomogenity):
  """
  @param diagUpper upper band diagonal.
  @param diagMain main band diagonal.
  @param diagLower lower band diagonal (and a dummy element before it).
  @param inhomogenity right-hand side of equation.
  @return result of solving "Tridiag x = inhomogenity" for x.
  """
  N = len(diagMain)
  @memoized
  def d(j):
    return diagMain[j] if j == 0 else (diagMain[j] - lam(j)*diagUpper[j - 1])
  @memoized
  def lam(j):
    return diagLower[j]/d(j - 1)
  @memoized
  def y(j):
    return inhomogenity[j] if j == 0 else (inhomogenity[j] - lam(j)*y(j - 1))
  @memoized
  def x(j):
    if (j == N - 1):
      return y(j)/d(j)
    else:
      return (y(j) - diagUpper[j]*x(j + 1))/d(j)
  return [x(j) for j in range(N)]