step.hpp

Go to the documentation of this file.

#ifndef INCLUDE_RICCATI_STEP_HPP
#define INCLUDE_RICCATI_STEP_HPP
 
#include <riccati/chebyshev.hpp>
#include <Eigen/Dense>
#include <complex>
#include <cmath>
#include <tuple>
 
namespace riccati {
 
template <typename SolverInfo, typename Scalar, typename YScalar>
inline auto nonosc_step(SolverInfo &&info, Scalar x0, Scalar h, YScalar y0,
                        YScalar dy0, Scalar epsres) {
  using complex_t = std::complex<Scalar>;
 
  Scalar maxerr = 10 * epsres;
  auto N = info.nini_;
  auto Nmax = info.nmax_;
  auto cheby = spectral_chebyshev(info, x0, h, y0, dy0, 0);
  auto yprev = std::get<0>(cheby);
  auto dyprev = std::get<1>(cheby);
  auto xprev = std::get<2>(cheby);
  int iter = 0;
  while (maxerr > epsres) {
    iter++;
    N *= 2;
    if (N > Nmax) {
      return std::make_tuple(false, complex_t(0.0, 0.0), complex_t(0.0, 0.0),
                             maxerr, yprev, dyprev, iter);
    }
    auto cheb_num = static_cast<int>(std::log2(N / info.nini_));
    auto cheby2 = spectral_chebyshev(info, x0, h, y0, dy0, cheb_num);
    auto y = std::get<0>(std::move(cheby2));
    auto dy = std::get<1>(std::move(cheby2));
    auto x = std::get<2>(std::move(cheby2));
    maxerr = std::abs((yprev(0) - y(0)) / y(0));
    if (std::isnan(maxerr)) {
      maxerr = std::numeric_limits<Scalar>::infinity();
    }
    yprev = std::move(y);
    dyprev = std::move(dy);
    xprev = std::move(x);
  }
  return std::make_tuple(true, yprev(0), dyprev(0), maxerr, yprev, dyprev,
                         iter);
}
 
template <bool DenseOut, typename SolverInfo, typename OmegaVec,
          typename GammaVec, typename Scalar, typename YScalar>
inline auto osc_step(SolverInfo &&info, OmegaVec &&omega_s, GammaVec &&gamma_s,
                     Scalar x0, Scalar h, YScalar y0, YScalar dy0,
                     Scalar epsres) {
  using complex_t = std::complex<Scalar>;
  using vectorc_t = vector_t<complex_t>;
  bool success = true;
  auto &&Dn = info.Dn();
  auto y = eval(info.alloc_, complex_t(0.0, 1.0) * omega_s);
  auto delta = [&](const auto &r, const auto &y) {
    return (-r.array() / (Scalar{2.0} * (y.array() + gamma_s.array())))
        .matrix()
        .eval();
  };
  auto R = [&](const auto &d) {
    return Scalar{2.0} / h * (Dn * d) + d.array().square().matrix();
  };
  auto Ry
      = (complex_t(0.0, 1.0) * Scalar{2.0}
         * (Scalar{1.0} / h * (Dn * omega_s) + gamma_s.cwiseProduct(omega_s)))
            .eval();
  Scalar maxerr = Ry.array().abs().maxCoeff();
 
  arena_matrix<vectorc_t> deltay(info.alloc_, Ry.size(), 1);
  Scalar prev_err = std::numeric_limits<Scalar>::infinity();
  while (maxerr > epsres) {
    deltay = delta(Ry, y);
    y += deltay;
    Ry = R(deltay);
    maxerr = Ry.array().abs().maxCoeff();
    if (maxerr >= (Scalar{2.0} * prev_err)) {
      success = false;
      break;
    }
    prev_err = maxerr;
  }
  if constexpr (DenseOut) {
    auto u1
        = eval(info.alloc_, h / Scalar{2.0} * (info.integration_matrix_ * y));
    auto f1 = eval(info.alloc_, (u1).array().exp().matrix());
    auto f2 = eval(info.alloc_, f1.conjugate());
    auto du2 = eval(info.alloc_, y.conjugate());
    auto ap_top = (dy0 - y0 * du2(du2.size() - 1));
    auto ap_bottom = (y(y.size() - 1) - du2(du2.size() - 1));
    auto ap = ap_top / ap_bottom;
    auto am = (dy0 - y0 * y(y.size() - 1))
              / (du2(du2.size() - 1) - y(y.size() - 1));
    auto y1 = eval(info.alloc_, ap * f1 + am * f2);
    auto dy1 = eval(info.alloc_,
                    ap * y.cwiseProduct(f1) + am * du2.cwiseProduct(f2));
    Scalar phase = std::imag(f1(0));
    return std::make_tuple(success, y1(0), dy1(0), maxerr, phase, u1, y,
                           std::make_pair(ap, am));
  } else {
    complex_t f1 = std::exp(h / Scalar{2.0} * (info.quadwts_.dot(y)));
    auto f2 = std::conj(f1);
    auto du2 = y.conjugate().eval();
    auto ap_top = (dy0 - y0 * du2(du2.size() - 1));
    auto ap_bottom = (y(y.size() - 1) - du2(du2.size() - 1));
    auto ap = ap_top / ap_bottom;
    auto am = (dy0 - y0 * y(y.size() - 1))
              / (du2(du2.size() - 1) - y(y.size() - 1));
    auto y1 = (ap * f1 + am * f2);
    auto dy1 = (ap * y * f1 + am * du2 * f2).eval();
    Scalar phase = std::imag(f1);
    return std::make_tuple(success, y1, dy1(0), maxerr, phase,
                           arena_matrix<vectorc_t>(info.alloc_, y.size()),
                           arena_matrix<vectorc_t>(info.alloc_, y.size()),
                           std::make_pair(ap, am));
  }
}
 
template <typename SolverInfo, typename Scalar, typename Vec>
inline auto step(SolverInfo &info, Scalar xi, Scalar xf,
                 std::complex<Scalar> yi, std::complex<Scalar> dyi, Scalar eps,
                 Scalar epsilon_h, Scalar init_stepsize, Vec &&x_eval,
                 bool hard_stop = false) {
  using vectord_t = vector_t<Scalar>;
  using complex_t = std::complex<Scalar>;
  using vectorc_t = vector_t<complex_t>;
  Scalar direction = init_stepsize > 0 ? 1 : -1;
  if (init_stepsize * (xf - xi) < 0) {
    throw std::domain_error(
        "Direction of integration does not match stepsize sign,"
        " adjusting it so that integration happens from xi to xf.");
  }
  // Check that yeval and x_eval are right size
  constexpr bool dense_output = compile_size_v<Vec> != 0;
  if constexpr (dense_output) {
    if (!x_eval.size()) {
      throw std::domain_error("Dense output requested but x_eval is size 0!");
    }
    // TODO: Better error messages
    auto x_eval_max = (direction * x_eval.maxCoeff());
    auto x_eval_min = (direction * x_eval.minCoeff());
    auto xi_intdir = direction * xi;
    auto xf_intdir = direction * xf;
    const bool high_range_err = xf_intdir < x_eval_max;
    const bool low_range_err = xi_intdir > x_eval_min;
    if (high_range_err || low_range_err) {
      if (high_range_err && low_range_err) {
        throw std::out_of_range(
            std::string{"The max and min of the output points ("}
            + std::to_string(x_eval_min) + std::string{", "}
            + std::to_string(x_eval_max)
            + ") lie outside the high and low of the integration range ("
            + std::to_string(xi_intdir) + std::string{", "}
            + std::to_string(xf_intdir) + ")!");
      }
      if (high_range_err) {
        throw std::out_of_range(
            std::string{"The max of the output points ("}
            + std::to_string(x_eval_max)
            + ") lies outside the high of the integration range ("
            + std::to_string(xf_intdir) + ")!");
      }
      if (low_range_err) {
        throw std::out_of_range(
            std::string{"The min of the output points ("}
            + std::to_string(x_eval_min)
            + ") lies outside the low of the integration range ("
            + std::to_string(xi_intdir) + ")!");
      }
    }
  }
 
  // Initialize vectors for storing results
  Scalar xs;
  complex_t ys;
  complex_t dys;
  int successes;
  int steptypes;
  Scalar phases;
  Eigen::Matrix<complex_t, -1, 1> yeval;   //(x_eval.size());
  Eigen::Matrix<complex_t, -1, 1> dyeval;  //(x_eval.size());
 
  complex_t y = yi;
  complex_t dy = dyi;
  complex_t yprev = y;
  complex_t dyprev = dy;
  auto scale_xi = scale(info.xp().array(), xi, init_stepsize).eval();
  auto omega_is = omega(info, scale_xi).eval();
  auto gamma_is = gamma(info, scale_xi).eval();
  Scalar wi = omega_is.mean();
  Scalar gi = gamma_is.mean();
  Scalar dwi = (2.0 / init_stepsize * (info.Dn() * omega_is)).mean();
  Scalar dgi = (2.0 / init_stepsize * (info.Dn() * gamma_is)).mean();
  Scalar hslo_ini = direction
                    * std::min(static_cast<Scalar>(1e8),
                               static_cast<Scalar>(std::abs(1.0 / wi)));
  Scalar hosc_ini
      = direction
        * std::min(std::min(static_cast<Scalar>(1e8),
                            static_cast<Scalar>(std::abs(wi / dwi))),
                   std::abs(gi / dgi));
 
  if (hard_stop) {
    if (direction * (xi + hosc_ini) > direction * xf) {
      hosc_ini = xf - xi;
    }
    if (direction * (xi + hslo_ini) > direction * xf) {
      hslo_ini = xf - xi;
    }
  }
  auto hslo = choose_nonosc_stepsize(info, xi, hslo_ini, Scalar(0.2));
  // o and g written here
  auto osc_step_tup = choose_osc_stepsize(info, xi, hosc_ini, epsilon_h);
  auto hosc = std::get<0>(osc_step_tup);
  // NOTE: Calling choose_osc_stepsize will update these values
  auto &&omega_n = std::get<1>(osc_step_tup);
  auto &&gamma_n = std::get<2>(osc_step_tup);
  Scalar xcurrent = xi;
  Scalar wnext = wi;
  using matrixc_t = matrix_t<complex_t>;
  matrixc_t y_eval;
  matrixc_t dy_eval;
  arena_matrix<vectorc_t> un(info.alloc_, omega_n.size(), 1);
  arena_matrix<vectorc_t> d_un(info.alloc_, omega_n.size(), 1);
  std::pair<complex_t, complex_t> a_pair;
  Scalar phase{0.0};
  bool success = false;
  bool steptype = true;
  Scalar err;
  int cheb_N = 0;
  if ((direction * hosc > direction * hslo * 5.0)
      && (direction * hosc * wnext / (2.0 * pi<Scalar>()) > 1.0)) {
    if (hard_stop) {
      if (direction * (xcurrent + hosc) > direction * xf) {
        hosc = xf - xcurrent;
        auto xp_scaled = scale(info.xp().array(), xcurrent, hosc).eval();
        omega_n = omega(info, xp_scaled);
        gamma_n = gamma(info, xp_scaled);
      }
      if (direction * (xcurrent + hslo) > direction * xf) {
        hslo = xf - xcurrent;
      }
    }
    // o and g read here
    std::tie(success, y, dy, err, phase, un, d_un, a_pair)
        = osc_step<dense_output>(info, omega_n, gamma_n, xcurrent, hosc, yprev,
                                 dyprev, eps);
    steptype = 1;
  }
  while (!success) {
    std::tie(success, y, dy, err, y_eval, dy_eval, cheb_N)
        = nonosc_step(info, xcurrent, hslo, yprev, dyprev, eps);
    steptype = 0;
    if (!success) {
      hslo *= Scalar{0.5};
    }
    if (direction * hslo < 1e-16) {
      throw std::domain_error("Stepsize became to small error");
    }
  }
  auto h = steptype ? hosc : hslo;
  if constexpr (dense_output) {
    Eigen::Index dense_size = 0;
    Eigen::Index dense_start = 0;
    // Assuming x_eval is sorted we just want start and size
    std::tie(dense_start, dense_size)
        = get_slice(x_eval, direction * xcurrent, direction * (xcurrent + h));
    yeval = Eigen::Matrix<complex_t, -1, 1>(dense_size);
    dyeval = Eigen::Matrix<complex_t, -1, 1>(dense_size);
    if (dense_size != 0) {
      auto x_eval_map
          = Eigen::Map<vectord_t>(x_eval.data() + dense_start, dense_size);
      auto y_eval_map
          = Eigen::Map<vectorc_t>(yeval.data() + dense_start, dense_size);
      auto dy_eval_map
          = Eigen::Map<vectorc_t>(dyeval.data() + dense_start, dense_size);
      if (steptype) {
        auto x_eval_scaled
            = eval(info.alloc_,
                   (2.0 / h * (x_eval_map.array() - xcurrent) - 1.0).matrix());
        auto Linterp = interpolate(info.xn(), x_eval_scaled, info.alloc_);
        auto fdense = eval(info.alloc_, (Linterp * un).array().exp().matrix());
        y_eval_map = a_pair.first * fdense + a_pair.second * fdense.conjugate();
        auto du_dense = eval(info.alloc_, (Linterp * d_un));
        dy_eval_map
            = a_pair.first * (du_dense.array() * fdense.array())
              + a_pair.second * (du_dense.array() * fdense.array()).conjugate();
      } else {
        auto xc_scaled = eval(
            info.alloc_,
            scale(std::get<2>(info.chebyshev_[cheb_N]), xcurrent, h).matrix());
        auto Linterp = interpolate(xc_scaled, x_eval_map, info.alloc_);
        y_eval_map = Linterp * y_eval;
      }
    }
  }
  // Finish appending and ending conditions
  ys = y;
  dys = dy;
  xs = xcurrent + h;
  phases = phase;
  steptypes = steptype;
  successes = success;
  Scalar dwnext;
  Scalar gnext;
  Scalar dgnext;
  if (steptype) {
    wnext = omega_n[0];
    gnext = gamma_n[0];
    dwnext = 2.0 / h * info.Dn().row(0).dot(omega_n);
    dgnext = 2.0 / h * info.Dn().row(0).dot(gamma_n);
  } else {
    wnext = omega(info, xcurrent + h);
    gnext = gamma(info, xcurrent + h);
    auto xn_scaled = scale(info.xn().array(), xcurrent, h).eval();
    dwnext = 2.0 / h * info.Dn().row(0).dot(omega(info, xn_scaled).matrix());
    dgnext = 2.0 / h * info.Dn().row(0).dot(gamma(info, (xn_scaled).matrix()));
  }
  xcurrent += h;
  if (direction * xcurrent < direction * xf) {
    hslo_ini = direction * std::min(Scalar{1e8}, std::abs(Scalar{1.0} / wnext));
    hosc_ini = direction
               * std::min(std::min(Scalar{1e8}, std::abs(wnext / dwnext)),
                          std::abs(gnext / dgnext));
    if (hard_stop) {
      if (direction * (xcurrent + hosc_ini) > direction * xf) {
        hosc_ini = xf - xcurrent;
      }
      if (direction * (xcurrent + hslo_ini) > direction * xf) {
        hslo_ini = xf - xcurrent;
      }
    }
    // o and g written here
    osc_step_tup = choose_osc_stepsize(info, xcurrent, hosc_ini, epsilon_h);
    hosc = std::get<0>(osc_step_tup);
    hslo = choose_nonosc_stepsize(info, xcurrent, hslo_ini, Scalar{0.2});
    yprev = y;
    dyprev = dy;
  }
  info.alloc_.recover_memory();
  return std::make_tuple(xs, ys, dys, hosc, hslo, successes, phases, steptypes,
                         yeval);
}
 
}  // namespace riccati
 
#endif