RK4ODEIntegrator.h

Go to the documentation of this file.

#ifndef _utl_RK4ODEIntegrator_h_
#define _utl_RK4ODEIntegrator_h_


namespace utl {

  template<class DerivativeFunctor, class VectorType>
  class RK4Iterator {
  public:
    RK4Iterator(const double x, const VectorType& y,
                DerivativeFunctor& d)
      : fX(x), fDerivativeFunctor(d), fStatus(true)
    { Copy(d, y, fY); }

    double GetX() const { return fX; }

    VectorType& GetY() { return fY; }

    const VectorType& GetY() const { return fY; }

    bool
    Advance(const double dx, const VectorType& dYdX)
    {
      const unsigned int n = fDerivativeFunctor;
      // first step: x -> x + dx/2
      const double dx2 = 0.5 * dx;
      VectorType y1;
      for (unsigned int i = 0; i < n; ++i)
        y1[i] = fY[i] + dx2 * dYdX[i];
      // second step
      const double x1 = fX + dx2;
      VectorType dYdX1;
      if (!fDerivativeFunctor(x1, y1, dYdX1))
        return fStatus = false;
      for (unsigned int i = 0; i < n; ++i)
        y1[i] = fY[i] + dx2 * dYdX1[i];
      // third step
      VectorType dYdX2;
      if (!fDerivativeFunctor(x1, y1, dYdX2))
        return fStatus = false;
      for (unsigned int i = 0; i < n; ++i) {
        y1[i] = fY[i] + dx * dYdX2[i];
        dYdX2[i] += dYdX1[i];
      }
      // fourth step
      fX += dx;
      if (!fDerivativeFunctor(fX, y1, dYdX1))
        return fStatus = false;
      const double dx6 = (1./6) * dx;
      for (unsigned int i = 0; i < n; ++i)
        fY[i] += dx6 * (dYdX[i] + dYdX1[i] + 2.*dYdX2[i]);
      return fStatus = true;
    }

    bool
    Advance(const double dx)
    {
      VectorType dYdX;
      fDerivativeFunctor(fX, fY, dYdX);
      return Advance(dx, dYdX);
    }

    RK4Iterator&
    operator+=(const double dx)
    { Advance(dx); return *this; }

    RK4Iterator&
    operator-=(const double dx)
    { Advance(-dx); return *this; }

    explicit operator bool() const { return fStatus; }

  private:
    static
    void
    Copy(const unsigned int n, const VectorType& a, VectorType& b)
    {
      for (unsigned int i = 0; i < n; ++i)
        b[i] = a[i];
    }

    double fX;
    VectorType fY;
    DerivativeFunctor& fDerivativeFunctor;
    bool fStatus;
  };


  template<class DerivativeFunctor>
  class RK4ODEIntegrator {
  public:
    RK4ODEIntegrator(DerivativeFunctor& d)
      : fDerivativeFunctor(d) { }

    template<class VectorType>
    RK4Iterator<DerivativeFunctor, VectorType>
    Begin(const double x, const VectorType& y)
    {
      return
        RK4Iterator<DerivativeFunctor, VectorType>(x, y, fDerivativeFunctor);
    }

    DerivativeFunctor& GetDerivativeFunctor()
    { return fDerivativeFunctor; }

  private:
    DerivativeFunctor& fDerivativeFunctor;
  };

}


#endif