AdamsMoulton.hpp

#pragma once
#include <algorithm>
#include <array>
#include <cassert>
#include <complex>
#include <tuple>
#include <type_traits>
 
namespace AdamsMoulton {
 
//==============================================================================
 
 
template <typename T = double, typename Y = double>
struct DerivativeMatrix {
  virtual Y a(T t) const = 0;
  virtual Y b(T t) const = 0;
  virtual Y c(T t) const = 0;
  virtual Y d(T t) const = 0;
  virtual Y Sf(T) const { return Y(0); };
  virtual Y Sg(T) const { return Y(0); };
  virtual ~DerivativeMatrix() = default;
};
 
//==============================================================================
 
// User-defined type-trait: Checks whether T is a std::complex type
template <typename T>
struct is_complex : std::false_type {};
// User-defined type-trait: Checks whether T is a std::complex type
template <typename T>
struct is_complex<std::complex<T>> : std::true_type {};
 
template <typename T>
constexpr bool is_complex_v = is_complex<T>::value;
 
//==============================================================================
 
 
template <typename T, typename U, std::size_t N, std::size_t M>
constexpr T inner_product(const std::array<T, N> &a,
                          const std::array<U, M> &b) {
  static_assert(std::is_convertible_v<U, T>,
                "In inner_product, type of second array (U) must be "
                "convertable to that of dirst (T)");
  constexpr std::size_t Size = std::min(N, M);
 
  if constexpr (Size == 0) {
    return T{0};
  } else if constexpr (!std::is_same_v<T, U> && is_complex_v<T>) {
    // This is to avoid float conversion warning in case that U=double,
    // T=complex<float>; want to case b to float, then to complex<float>
    T res = a[0] * static_cast<typename T::value_type>(b[0]);
    for (std::size_t i = 1; i < Size; ++i) {
      res += a[i] * static_cast<typename T::value_type>(b[i]);
    }
    return res;
  } else {
    T res = a[0] * static_cast<T>(b[0]);
    for (std::size_t i = 1; i < Size; ++i) {
      res += a[i] * static_cast<T>(b[i]);
    }
    return res;
  }
}
 
//==============================================================================
 
namespace helper {
 
// Simple struct for storing "Raw" Adams ("B") coefficients.
/*
Stored as integers, with 'denominator' factored out. \n
Converted to double  ("A" coefs) once, at compile time (see below). \n
Adams coefficients, a_k, defined such that: \n
 \f[ F_{n+K} = F_{n+K-1} + dx * \sum_{k=0}^K a_k y_{n+k} \f]
where:
 \f[ y = d(F)/dr \f]
Note: the 'order' of the coefs is reversed compared to some sources.
The final coefficient is separated, such that: \n
 \f[ a_k = b_k / denom \f]
for k = {0,1,...,K-1} \n
and
 \f[ a_K = b_K / denom \f]
*/
template <std::size_t K>
struct AdamsB {
  long denom;
  std::array<long, K> bk;
  long bK;
};
 
// List of Adams coefficients data
/*
Note: there is (of course) no 0-step Adams method.
The entry at [0] is invalid, and will produce 0.
It is included so that the index of this list matches the order of the method.
Program will not compile (static_asser) is [0] is requested.
Note: assumes that the kth element corresponds to k-order AM method.
*/
static constexpr auto ADAMS_data = std::tuple{
    AdamsB<0>{1, {}, 0}, // invalid entry, but want index to match order
    AdamsB<1>{2, {1}, 1},
    AdamsB<2>{12, {-1, 8}, 5},
    AdamsB<3>{24, {1, -5, 19}, 9},
    AdamsB<4>{720, {-19, 106, -264, 646}, 251},
    AdamsB<5>{1440, {27, -173, 482, -798, 1427}, 475},
    AdamsB<6>{60480, {-863, 6312, -20211, 37504, -46461, 65112}, 19087},
    AdamsB<7>{
        120960, {1375, -11351, 41499, -88547, 123133, -121797, 139849}, 36799},
    AdamsB<8>{3628800,
              {-33953, 312874, -1291214, 3146338, -5033120, 5595358, -4604594,
               4467094},
              1070017},
    AdamsB<9>{7257600,
              {57281, -583435, 2687864, -7394032, 13510082, -17283646, 16002320,
               -11271304, 9449717},
              2082753},
    AdamsB<10>{479001600,
               {-3250433, 36284876, -184776195, 567450984, -1170597042,
                1710774528, -1823311566, 1446205080, -890175549, 656185652},
               134211265},
    AdamsB<11>{958003200,
               {5675265, -68928781, 384709327, -1305971115, 3007739418,
                -4963166514, 6043521486, -5519460582, 3828828885, -2092490673,
                1374799219},
               262747265},
    AdamsB<12>{2615348736000,
               {-13695779093, 179842822566, -1092096992268, 4063327863170,
                -10344711794985, 19058185652796, -26204344465152,
                27345870698436, -21847538039895, 13465774256510, -6616420957428,
                3917551216986},
               703604254357}};
 
} // namespace helper
 
//==============================================================================
 
static constexpr std::size_t K_max =
    std::tuple_size_v<decltype(helper::ADAMS_data)> - 1;
 
//==============================================================================
 
 
template <std::size_t K, typename = std::enable_if_t<(K > 0)>,
          typename = std::enable_if_t<(K <= K_max)>>
struct AM_Coefs {
 
private:
  // Forms the (double) ak coefficients, from the raw (int) bk ones
  static constexpr std::array<double, K> make_ak() {
    const auto &am = std::get<K>(helper::ADAMS_data);
    static_assert(
        am.bk.size() == K,
        "Kth Entry in ADAMS_data must correspond to K-order AM method");
    std::array<double, K> tak{};
    for (std::size_t i = 0; i < K; ++i) {
      tak.at(i) = double(am.bk.at(i)) / double(am.denom);
    }
    return tak;
  }
 
  // Forms the final (double) aK coefficient, from the raw (int) bK one
  static constexpr double make_aK() {
    const auto &am = std::get<K>(helper::ADAMS_data);
    return (double(am.bK) / double(am.denom));
  }
 
public:
  static constexpr std::array<double, K> ak{make_ak()};
  static constexpr double aK{make_aK()};
};
 
//==============================================================================
 
template <std::size_t K, typename T = double, typename Y = double>
class ODESolver2D {
  static_assert(K > 0, "Order (K) for Adams method must be K>0");
  static_assert(K <= K_max,
                "Order (K) requested for Adams method too "
                "large. Adams coefficients are implemented up to K_max-1 only");
  static_assert(
      is_complex_v<Y> || std::is_floating_point_v<Y>,
      "Template parameter Y (function values and dt) must be floating point "
      "or complex");
  static_assert(is_complex_v<Y> || std::is_floating_point_v<Y> ||
                    std::is_integral_v<T>,
                "Template parameter T (derivative matrix argument) must be "
                "floating point, complex, or integral");
 
private:
  // Stores the AM coefficients
  static constexpr AM_Coefs<K> am{};
  // step size:
  Y m_dt;
  // previous 't' value
  // T m_t{T(0)}; // Should be invalid (nan), but no nan for int
  // Pointer to the derivative matrix
  const DerivativeMatrix<T, Y> *m_D;
 
public:
 
  std::array<Y, K> f{}, g{};
 
  std::array<Y, K> df{}, dg{};
 
  std::array<T, K> t{};
 
  Y S_scale{1.0};
 
public:
 
  ODESolver2D(Y dt, const DerivativeMatrix<T, Y> *D) : m_dt(dt), m_D(D) {
    assert(dt != Y{0.0} && "Cannot have zero step-size in ODESolver2D");
    assert(D != nullptr && "Cannot have null Derivative Matrix in ODESolver2D");
  }
 
  constexpr std::size_t K_steps() const { return K; }
 
  Y last_f() { return f.back(); }
  Y last_g() { return g.back(); }
  T last_t() { return t.back(); }
  Y dt() { return m_dt; }
 
  Y dfdt(Y ft, Y gt, T tt) const {
    return m_D->a(tt) * ft + m_D->b(tt) * gt + S_scale * m_D->Sf(tt);
  }
  Y dgdt(Y ft, Y gt, T tt) const {
    return m_D->c(tt) * ft + m_D->d(tt) * gt + S_scale * m_D->Sg(tt);
  }
 
 
  void drive(T t_next) {
 
    // assert (t_next = Approx[next_t(last_t())])
 
    // Use AM method to determine new values, given previous K values:
    const auto sf = f.back() + m_dt * (inner_product(df, am.ak) +
                                       am.aK * S_scale * m_D->Sf(t_next));
    const auto sg = g.back() + m_dt * (inner_product(dg, am.ak) +
                                       am.aK * S_scale * m_D->Sg(t_next));
 
    const auto a = m_D->a(t_next);
    const auto b = m_D->b(t_next);
    const auto c = m_D->c(t_next);
    const auto d = m_D->d(t_next);
 
    const auto a0 = m_dt * static_cast<Y>(am.aK);
    const auto a02 = a0 * a0;
    const auto det_inv =
        Y{1.0} / (Y{1.0} - (a02 * (b * c - a * d) + a0 * (a + d)));
    const auto fi = (sf - a0 * (d * sf - b * sg)) * det_inv;
    const auto gi = (sg - a0 * (-c * sf + a * sg)) * det_inv;
 
    // Shift values along. nb: rotate({1,2,3,4}) = {2,3,4,1}
    // We keep track of previous K values in order to determine next value
    std::rotate(f.begin(), f.begin() + 1, f.end());
    std::rotate(g.begin(), g.begin() + 1, g.end());
    std::rotate(df.begin(), df.begin() + 1, df.end());
    std::rotate(dg.begin(), dg.begin() + 1, dg.end());
    std::rotate(t.begin(), t.begin() + 1, t.end());
 
    // Sets new values:
    t.back() = t_next;
    f.back() = fi;
    g.back() = gi;
    df.back() = dfdt(fi, gi, t_next);
    dg.back() = dgdt(fi, gi, t_next);
  }
 
  void drive() { drive(next_t(last_t())); }
 
 
  void solve_initial_K(T t0, Y f0, Y g0) {
    t.at(0) = t0;
    f.at(0) = f0;
    g.at(0) = g0;
    df.at(0) = dfdt(f0, g0, t0);
    dg.at(0) = dgdt(f0, g0, t0);
    first_k_i<1>(next_t(t0));
  }
 
private:
  // only used in solve_initial_K(). Use this, because it works for double t,
  // where next value is t+dt, and for integral t, where next value is t++ is
  // driving forward (dt>0), or t-- if driving backwards (dt<0)
  T next_t(T last_t) {
    if constexpr (std::is_integral_v<T> && is_complex_v<Y>) {
      return (m_dt.real() > 0.0) ? last_t + 1 : last_t - 1;
    } else if constexpr (std::is_integral_v<T>) {
      return (m_dt > 0.0) ? last_t + 1 : last_t - 1;
    } else {
      return last_t + m_dt;
    }
  }
 
  // Used recursively by solve_initial_K() to find first K points
  template <std::size_t ik>
  void first_k_i(T t_next) {
    if constexpr (ik >= K) {
      (void)t_next; // suppress unused variable warning on old g++ versions
      return;
    } else {
      constexpr AM_Coefs<ik> ai{};
      // nb: ai.ak is smaller than df; inner_product still works
      const auto sf = f.at(ik - 1) + m_dt * (inner_product(df, ai.ak) +
                                             am.aK * S_scale * m_D->Sf(t_next));
      const auto sg = g.at(ik - 1) + m_dt * (inner_product(dg, ai.ak) +
                                             am.aK * S_scale * m_D->Sg(t_next));
      const auto a0 = m_dt * static_cast<Y>(ai.aK);
      const auto a02 = a0 * a0;
      const auto a = m_D->a(t_next);
      const auto b = m_D->b(t_next);
      const auto c = m_D->c(t_next);
      const auto d = m_D->d(t_next);
      const auto det_inv =
          Y{1.0} / (Y{1.0} - (a02 * (b * c - a * d) + a0 * (a + d)));
      const auto fi = (sf - a0 * (d * sf - b * sg)) * det_inv;
      const auto gi = (sg - a0 * (-c * sf + a * sg)) * det_inv;
      // Sets new values:
      t.at(ik) = t_next;
      f.at(ik) = fi;
      g.at(ik) = gi;
      df.at(ik) = dfdt(fi, gi, t_next);
      dg.at(ik) = dgdt(fi, gi, t_next);
      // call recursively
      first_k_i<ik + 1>(next_t(t_next));
    }
  }
};
} // namespace AdamsMoulton