Homec4science
Diffusion yearII_reports (master) Browse

rSTICAZZI/physnum/rap7/include/tridiag.h
fe6832168f6amaster

Last Change
Raw File

tridiag.h
View Options

#pragma once
#ifndef __TRIDIAG_H__
#define __TRIDIAG_H__
template<class T>
class tridiag
{
std::vector<T> diag, lower, upper;
struct row
{
std::size_t index;
tridiag& ref;
T shared_zero = 0;
row(std::size_t i, tridiag& _ref) : index(i), ref(_ref) {}
T& operator[](std::size_t j)
{
switch(static_cast<long int>(j) - index)
{
case -1:
return ref.lower[j];
case 0:
return ref.diag[j];
case 1:
return ref.upper[index];
default:
shared_zero = 0; // assure zero
return shared_zero;
}
}
};
struct const_row
{
std::size_t index;
const tridiag& ref;
const T shared_zero = 0;
const_row(std::size_t i, const tridiag& _ref) : index(i), ref(_ref) {}
const T& operator[](std::size_t j) const
{
switch(static_cast<long int>(j) - index)
{
case -1:
return ref.lower[j];
case 0:
return ref.diag[j];
case 1:
return ref.upper[index];
default:
return shared_zero;
}
}
};
public:
tridiag(std::size_t N = 0) : diag(N), lower(N-1), upper(N-1)
{
}
void resize(std::size_t N)
{
diag.reserve(N);
lower.reserve(N-1);
upper.reserve(N-1);
}
void zeros()
{
for (std::size_t i = 0; i < size()-1; ++i)
diag[i] = lower[i] = upper[i] = 0;
diag[size()-1] = 0;
}
// set diagonal to zero
void zeros_diag()
{
for (std::size_t i = 0; i < size(); ++i)
diag[i] = 0;
}
row operator[](std::size_t i)
{
return row(i, *this);
}
const const_row operator[](std::size_t i) const
{
return const_row(i, *this);
}
std::size_t size() const
{
return diag.size();
}
std::vector<T> solve(const std::vector<T>& rhs)
{
std::vector<T> solution(diag.size());
std::vector<T> new_diag(diag);
std::vector<T> new_rhs(rhs);
for(std::size_t i(1); i < size(); ++i)
{
double pivot = lower[i-1]/new_diag[i-1];
new_diag[i] -= pivot * upper[i-1];
new_rhs[i] -= pivot * new_rhs[i-1];
}
solution[size()-1] = new_rhs[size()-1] / new_diag[size()-1];
for(int i(static_cast<int>(size()-2)); i>=0; --i)
solution[i] = (new_rhs[i] - upper[i]*solution[i+1]) / new_diag[i];
return solution;
}
};
// multiplication operator overloading
template<class T, class V>
V operator*(const tridiag<T>& M, const V& v)
{
const std::size_t N = v.size();
V u(N);
u[0] = M[0][0] * v[0] + M[0][1] * v[1];
for (int k = 1; k < N-1; ++k)
u[k] = M[k][k-1] * v[k-1] + M[k][k] * v[k] + M[k][k+1] * v[k+1];
u[N-1] = M[N-1][N-2] * v[N-2] + M[N-1][N-1] * v[N-1];
return u;
}
#endif
c4science · Help