tridiagonalOperator.hpp

#ifndef MARIAN_TRIDIAGONALOPERATOR_HPP
#define MARIAN_TRIDIAGONALOPERATOR_HPP

#include <vector>
#include <iostream>

namespace marian {

  class TridiagonalOperator {   
  public:
    TridiagonalOperator(){};
    explicit TridiagonalOperator(unsigned int size);
    TridiagonalOperator(unsigned int size, double low, double mid, double upp);
    
    TridiagonalOperator(const std::vector<double>& low, const std::vector<double>& mid, const std::vector<double>& upp):
      size_(mid.size()), low_(low), mid_(mid), upp_(upp) {};

    
    static TridiagonalOperator DPlus(int n, double h);
    static TridiagonalOperator DMinus(int n, double h);  
    static TridiagonalOperator DZero(int n, double h);
    static TridiagonalOperator DZero(const std::vector<double>& grid);
    static TridiagonalOperator DPlusMinus(int n, double h);
    static TridiagonalOperator DPlusMinus(const std::vector<double>& grid);
    static TridiagonalOperator I(int n);
    static TridiagonalOperator I(const std::vector<double>& grid);

    int size() const { return size_; }
    double low(int) const;
    double mid(int) const;
    double upp(int) const;

    void setFirstRow(double, double);
    void setMidRow(int, double, double, double);
    void setMidRows(double, double, double);
    void setLastRow(double, double);

    
    virtual ~TridiagonalOperator(){};
    
    friend std::ostream & operator<<(std::ostream &s, const TridiagonalOperator& A);
    friend TridiagonalOperator operator+(const TridiagonalOperator&, const TridiagonalOperator&);
    friend TridiagonalOperator operator-(const TridiagonalOperator&, const TridiagonalOperator&);
    friend TridiagonalOperator operator*(double, const TridiagonalOperator&);
    friend TridiagonalOperator operator*(const TridiagonalOperator&, double);
    friend TridiagonalOperator operator/(const TridiagonalOperator&, double);
    friend std::vector<double> operator*(const TridiagonalOperator&, std::vector<double>);
  private:
    unsigned int size_;    
    std::vector<double> low_;  
    std::vector<double> mid_;  
    std::vector<double> upp_;  
  };

  inline TridiagonalOperator TridiagonalOperator::DPlus(int n, double h) {
    TridiagonalOperator to(n);
    double inv = 1.0 / h;
    to.setFirstRow(1.0, 0.0);          
    to.setMidRows(0.0, -inv, inv);
    to.setLastRow(0.0, 1.0);
    return to;
  }

  
  inline TridiagonalOperator TridiagonalOperator::DMinus(int n, double h) {
    TridiagonalOperator to(n);
    double inv = 1.0 / h;
    to.setFirstRow(1.0, 0.0);          
    to.setMidRows(-inv, inv, 0.0);
    to.setLastRow(0.0, 1.0);
    return to;
  }

  inline TridiagonalOperator TridiagonalOperator::DZero(int n, double h) {
    TridiagonalOperator to(n);
    double inv = 1.0 / (2.0*h);
    to.setFirstRow(1.0, 0.0);          
    to.setMidRows(-inv, 0.0, inv);
    to.setLastRow(0.0, 1.0);
    return to;
  }

  inline TridiagonalOperator TridiagonalOperator::DZero(const std::vector<double>& grid) {
    TridiagonalOperator to(grid.size());
    to.setFirstRow(1.0, 0.0);          
    for (unsigned int i = 2; i < grid.size(); ++i) {
      double hm  = grid.at(i-1) - grid.at(i-2);
      double hp  = grid.at(i)   - grid.at(i-1);
      double num = hm*hp*(hp+hm);
      to.setMidRow(i, -hp*hp / num, (hp*hp - hm*hm) / num  ,hm * hm / num);
    }
    to.setLastRow(0.0, 1.0);
    return to;
  }
  inline TridiagonalOperator TridiagonalOperator::DPlusMinus(int n, double h) {
    TridiagonalOperator to(n);
    double inv = 1.0 / (h*h);
    to.setFirstRow(1.0, 0.0);   
    to.setMidRows(inv, -2.0*inv, inv);
    to.setLastRow(0.0, 1.0);   
    return to;
  }

  inline TridiagonalOperator TridiagonalOperator::DPlusMinus(const std::vector<double>& grid) {
    TridiagonalOperator to(grid.size());
    to.setFirstRow(1.0, 0.0);          
    for (unsigned int i = 2; i < grid.size()-1; ++i) {
      double hm  = grid.at(i-1) - grid.at(i-2);
      double hp  = grid.at(i)   - grid.at(i-1);
      double num = hm*hp*(hp+hm);
      to.setMidRow(i, 2.0*hp / num, -2.0*(hp + hm) / num  ,2.0 * hm / num);
    }
    to.setLastRow(0.0, 1.0); 
    return to;
  }

  
  inline TridiagonalOperator TridiagonalOperator::I(int n) {
    TridiagonalOperator to(n);
    to.setFirstRow(1.0, 0.0);   
    to.setMidRows(0.0, 1.0, 0.0);
    to.setLastRow(0.0, 1.0);   
    return to;
  }

  inline TridiagonalOperator TridiagonalOperator::I(const std::vector<double>& grid) {
    TridiagonalOperator to(grid.size());
    to.setFirstRow(1.0, 0.0);   
    to.setMidRows(0.0, 1.0, 0.0);
    to.setLastRow(0.0, 1.0);   
    return to;
  }
} //namespace marian

#endif /* MARIAN_TRIDIAGONALOPERATOR_HPP */