tuple_tensor_dual_functions.hpp

#pragma once
 
#include "serac/numerics/functional/tuple.hpp"
#include "serac/numerics/functional/tensor.hpp"
#include "serac/numerics/functional/dual.hpp"
 
#include "mfem.hpp"
 
namespace serac {
 
template <typename T>
struct is_tensor_of_dual_number {
  static constexpr bool value = false;  
};
 
template <typename T, int... n>
struct is_tensor_of_dual_number<tensor<dual<T>, n...>> {
  static constexpr bool value = true;  
};
 
template <typename S, typename T, int m, int... n,
          typename = std::enable_if_t<std::is_arithmetic_v<S> || is_dual_number<S>::value>>
SERAC_HOST_DEVICE constexpr auto operator*(S scale, const tensor<T, m, n...>& A)
{
  tensor<decltype(S{} * T{}), m, n...> C{};
  for (int i = 0; i < m; i++) {
    C[i] = scale * A[i];
  }
  return C;
}
 
template <typename S, typename T, int m, int... n,
          typename = std::enable_if_t<std::is_arithmetic_v<S> || is_dual_number<S>::value>>
SERAC_HOST_DEVICE constexpr auto operator*(const tensor<T, m, n...>& A, S scale)
{
  tensor<decltype(T{} * S{}), m, n...> C{};
  for (int i = 0; i < m; i++) {
    C[i] = A[i] * scale;
  }
  return C;
}
 
template <typename S, typename T, int m, int... n,
          typename = std::enable_if_t<std::is_arithmetic_v<S> || is_dual_number<S>::value>>
SERAC_HOST_DEVICE constexpr auto operator/(S scale, const tensor<T, m, n...>& A)
{
  tensor<decltype(S{} * T{}), n...> C{};
  for (int i = 0; i < m; i++) {
    C[i] = scale / A[i];
  }
  return C;
}
 
template <typename S, typename T, int m, int... n,
          typename = std::enable_if_t<std::is_arithmetic_v<S> || is_dual_number<S>::value>>
SERAC_HOST_DEVICE constexpr auto operator/(const tensor<T, m, n...>& A, S scale)
{
  tensor<decltype(T{} * S{}), m, n...> C{};
  for (int i = 0; i < m; i++) {
    C[i] = A[i] / scale;
  }
  return C;
}
 
template <int i, typename S, typename T>
struct one_hot_helper;
 
template <int i, int... I, typename T>
struct one_hot_helper<i, std::integer_sequence<int, I...>, T> {
  using type = tuple<std::conditional_t<i == I, T, zero>...>;
};
 
template <int i, int n, typename T>
struct one_hot : public one_hot_helper<i, std::make_integer_sequence<int, n>, T> {
};
 
template <int i, int n, typename T>
using one_hot_t = typename one_hot<i, n, T>::type;
 
template <int i, int N>
SERAC_HOST_DEVICE constexpr auto make_dual_helper(zero /*arg*/)
{
  return zero{};
}
 
template <int i, int N>
SERAC_HOST_DEVICE constexpr auto make_dual_helper(double arg)
{
  using gradient_t = one_hot_t<i, N, double>;
  dual<gradient_t> arg_dual{};
  arg_dual.value                   = arg;
  serac::get<i>(arg_dual.gradient) = 1.0;
  return arg_dual;
}
 
template <int i, int N, typename T, int... n>
SERAC_HOST_DEVICE constexpr auto make_dual_helper(const tensor<T, n...>& arg)
{
  using gradient_t = one_hot_t<i, N, tensor<T, n...>>;
  tensor<dual<gradient_t>, n...> arg_dual{};
  for_constexpr<n...>([&](auto... j) {
    arg_dual(j...).value                         = arg(j...);
    serac::get<i>(arg_dual(j...).gradient)(j...) = 1.0;
  });
  return arg_dual;
}
 
template <typename T0, typename T1>
SERAC_HOST_DEVICE constexpr auto make_dual(const tuple<T0, T1>& args)
{
  return tuple{make_dual_helper<0, 2>(get<0>(args)), make_dual_helper<1, 2>(get<1>(args))};
}
 
template <typename T0, typename T1, typename T2>
SERAC_HOST_DEVICE constexpr auto make_dual(const tuple<T0, T1, T2>& args)
{
  return tuple{make_dual_helper<0, 3>(get<0>(args)), make_dual_helper<1, 3>(get<1>(args)),
               make_dual_helper<2, 3>(get<2>(args))};
}
 
template <bool dualify, typename T>
SERAC_HOST_DEVICE auto promote_to_dual_when(const T& x)
{
  if constexpr (dualify) {
    return make_dual(x);
  }
  if constexpr (!dualify) {
    return x;
  }
}
 
template <bool dualify, typename T, int n>
SERAC_HOST_DEVICE auto promote_each_to_dual_when(const tensor<T, n>& x)
{
  if constexpr (dualify) {
    using return_type = decltype(make_dual(T{}));
    tensor<return_type, n> output;
    for (int i = 0; i < n; i++) {
      output[i] = make_dual(x[i]);
    }
    return output;
  }
  if constexpr (!dualify) {
    return x;
  }
}
 
template <int n, typename... T, int... i>
SERAC_HOST_DEVICE constexpr auto make_dual_helper(const serac::tuple<T...>& args, std::integer_sequence<int, i...>)
{
  // Sam: it took me longer than I'd like to admit to find this issue, so here's an explanation
  //
  // note: we use serac::make_tuple(...) instead of serac::tuple{...} here because if
  // the first argument passed in is of type `serac::tuple < serac::tuple < T ... > >`
  // then doing something like
  //
  // serac::tuple{serac::get<i>(args)...};
  //
  // will be expand to something like
  //
  // serac::tuple{serac::tuple< T ... >{}};
  //
  // which invokes the copy ctor, returning a `serac::tuple< T ... >`
  // instead of `serac::tuple< serac::tuple < T ... > >`
  //
  // but serac::make_tuple(serac::get<i>(args)...) will never accidentally trigger the copy ctor
  return serac::make_tuple(promote_to_dual_when<i == n>(serac::get<i>(args))...);
}
 
template <int n, typename... T>
constexpr auto make_dual_wrt(const serac::tuple<T...>& args)
{
  return make_dual_helper<n>(args, std::make_integer_sequence<int, static_cast<int>(sizeof...(T))>{});
}
 
template <typename T1, typename T2, int n>
SERAC_HOST_DEVICE auto get_value(const tensor<tuple<T1, T2>, n>& input)
{
  tensor<decltype(get_value(tuple<T1, T2>{})), n> output{};
  for (int i = 0; i < n; i++) {
    output[i] = get_value(input[i]);
  }
  return output;
}
 
template <typename... T>
SERAC_HOST_DEVICE auto get_value(const serac::tuple<T...>& tuple_of_values)
{
  return serac::apply([](const auto&... each_value) { return serac::tuple{get_value(each_value)...}; },
                      tuple_of_values);
}
 
template <typename... T>
SERAC_HOST_DEVICE auto get_gradient(dual<serac::tuple<T...>> arg)
{
  return serac::apply([](auto... each_value) { return serac::tuple{each_value...}; }, arg.gradient);
}
 
template <typename... T, int... n>
SERAC_HOST_DEVICE auto get_gradient(const tensor<dual<serac::tuple<T...>>, n...>& arg)
{
  serac::tuple<outer_product_t<tensor<double, n...>, T>...> g{};
  for_constexpr<n...>([&](auto... i) {
    for_constexpr<sizeof...(T)>([&](auto j) { serac::get<j>(g)(i...) = serac::get<j>(arg(i...).gradient); });
  });
  return g;
}
 
template <typename... T>
SERAC_HOST_DEVICE auto get_gradient(serac::tuple<T...> tuple_of_values)
{
  return serac::apply([](auto... each_value) { return serac::tuple{get_gradient(each_value)...}; }, tuple_of_values);
}
 
template <int... n>
SERAC_HOST_DEVICE constexpr auto make_dual(const tensor<double, n...>& A)
{
  tensor<dual<tensor<double, n...>>, n...> A_dual{};
  for_constexpr<n...>([&](auto... i) {
    A_dual(i...).value          = A(i...);
    A_dual(i...).gradient(i...) = 1.0;
  });
  return A_dual;
}
 
template <typename T, int n>
SERAC_HOST_DEVICE constexpr LuFactorization<T, n> factorize_lu(const tensor<T, n, n>& A)
{
  constexpr auto abs  = [](double x) { return (x < 0) ? -x : x; };
  constexpr auto swap = [](auto& x, auto& y) {
    auto tmp = x;
    x        = y;
    y        = tmp;
  };
 
  auto U = A;
  // initialize L to Identity
  auto L = tensor<T, n, n>{};
  // This handles the case if T is a dual number
  // TODO - BT: make a dense identity that is templated on type
  for (int i = 0; i < n; i++) {
    if constexpr (is_dual_number<T>::value) {
      L[i][i].value = 1.0;
    } else {
      L[i][i] = 1.0;
    }
  }
  tensor<int, n> P(make_tensor<n>([](auto i) { return i; }));
 
  for (int i = 0; i < n; i++) {
    // Search for maximum in this column
    double max_val = abs(get_value(U[i][i]));
 
    int max_row = i;
    for (int j = i + 1; j < n; j++) {
      auto U_ji = get_value(U[j][i]);
      if (abs(U_ji) > max_val) {
        max_val = abs(U_ji);
        max_row = j;
      }
    }
 
    swap(P[max_row], P[i]);
    swap(U[max_row], U[i]);
  }
 
  for (int i = 0; i < n; i++) {
    // zero entries below in this column in U
    // and fill in L entries
    for (int j = i + 1; j < n; j++) {
      auto c  = U[j][i] / U[i][i];
      L[j][i] = c;
      U[j] -= c * U[i];
      U[j][i] = T{};
    }
  }
 
  return {P, L, U};
}
 
template <typename S, typename T, int n, int... m>
SERAC_HOST_DEVICE constexpr auto linear_solve(const tensor<S, n, n>& A, const tensor<T, n, m...>& b)
{
  // We want to avoid accumulating the derivative through the
  // LU factorization, because it is computationally expensive.
  // Instead, we perform the LU factorization on the values of
  // A, and then two backsolves: one to compute the primal (x),
  // and another to compute its derivative (dx).
  // If A is not dual, the second solve is a no-op.
 
  // Strip off derivatives, if any, and compute only x (ie no derivative)
  auto lu_factors = factorize_lu(get_value(A));
  auto x          = linear_solve(lu_factors, get_value(b));
 
  // Compute directional derivative of x.
  // If both b and A are not dual, the zero type
  // makes these no-ops.
  auto r  = get_gradient(b) - dot(get_gradient(A), x);
  auto dx = linear_solve(lu_factors, r);
 
  if constexpr (is_zero<decltype(dx)>{}) {
    return x;
  } else {
    return make_dual(x, dx);
  }
}
 
template <typename T, int n>
SERAC_HOST_DEVICE constexpr auto make_dual(const tensor<T, n>& x, const tensor<T, n>& dx)
{
  return make_tensor<n>([&](int i) { return dual<T>{x[i], dx[i]}; });
}
 
template <typename T, int m, int n>
SERAC_HOST_DEVICE constexpr auto make_dual(const tensor<T, m, n>& x, const tensor<T, m, n>& dx)
{
  return make_tensor<m, n>([&](int i, int j) { return dual<T>{x[i][j], dx[i][j]}; });
}
 
template <typename gradient_type, int n>
SERAC_HOST_DEVICE constexpr auto inv(tensor<dual<gradient_type>, n, n> A)
{
  auto invA = inv(get_value(A));
  return make_tensor<n, n>([&](int i, int j) {
    auto          value = invA[i][j];
    gradient_type gradient{};
    for (int k = 0; k < n; k++) {
      for (int l = 0; l < n; l++) {
        gradient -= invA[i][k] * A[k][l].gradient * invA[l][j];
      }
    }
    return dual<gradient_type>{value, gradient};
  });
}
 
template <typename T, int... n>
SERAC_HOST_DEVICE auto get_value(const tensor<dual<T>, n...>& arg)
{
  tensor<double, n...> value{};
  for_constexpr<n...>([&](auto... i) { value(i...) = arg(i...).value; });
  return value;
}
 
template <int... n>
SERAC_HOST_DEVICE constexpr auto get_gradient(const tensor<dual<double>, n...>& arg)
{
  tensor<double, n...> g{};
  for_constexpr<n...>([&](auto... i) { g(i...) = arg(i...).gradient; });
  return g;
}
 
template <int... n, int... m>
SERAC_HOST_DEVICE constexpr auto get_gradient(const tensor<dual<tensor<double, m...>>, n...>& arg)
{
  tensor<double, n..., m...> g{};
  for_constexpr<n...>([&](auto... i) { g(i...) = arg(i...).gradient; });
  return g;
}
 
struct SolverStatus {
  bool         converged;   
  unsigned int iterations;  
  double       residual;    
};
 
struct ScalarSolverOptions {
  double       xtol;      
  double       rtol;      
  unsigned int max_iter;  
};
 
const ScalarSolverOptions default_solver_options{.xtol = 1e-8, .rtol = 0, .max_iter = 25};
 
template <typename function, typename... ParamTypes>
auto solve_scalar_equation(function&& f, double x0, double lower_bound, double upper_bound, ScalarSolverOptions options,
                           ParamTypes... params)
{
  double x, df_dx;
  double fl = f(lower_bound, get_value(params)...);
  double fh = f(upper_bound, get_value(params)...);
 
  SLIC_ERROR_ROOT_IF(fl * fh > 0, "solve_scalar_equation: root not bracketed by input bounds.");
 
  unsigned int iterations = 0;
  bool         converged  = false;
 
  // handle corner cases where one of the brackets is the root
  if (fl == 0) {
    x         = lower_bound;
    converged = true;
  } else if (fh == 0) {
    x         = upper_bound;
    converged = true;
  }
 
  if (converged) {
    df_dx = get_gradient(f(make_dual(x), get_value(params)...));
 
  } else {
    // orient search so that f(xl) < 0
    double xl = lower_bound;
    double xh = upper_bound;
    if (fl > 0) {
      xl = upper_bound;
      xh = lower_bound;
    }
 
    // move initial guess if it is not between brackets
    if (x0 < lower_bound || x0 > upper_bound) {
      x0 = 0.5 * (lower_bound + upper_bound);
    }
 
    x                  = x0;
    double delta_x_old = std::abs(upper_bound - lower_bound);
    double delta_x     = delta_x_old;
    auto   R           = f(make_dual(x), get_value(params)...);
    auto   fval        = get_value(R);
    df_dx              = get_gradient(R);
 
    while (!converged) {
      if (iterations == options.max_iter) {
        SLIC_WARNING("solve_scalar_equation failed to converge in allotted iterations.");
        break;
      }
 
      // use bisection if Newton oversteps brackets or is not decreasing sufficiently
      if ((x - xh) * df_dx - fval > 0 || (x - xl) * df_dx - fval < 0 ||
          std::abs(2. * fval) > std::abs(delta_x_old * df_dx)) {
        delta_x_old = delta_x;
        delta_x     = 0.5 * (xh - xl);
        x           = xl + delta_x;
        converged   = (x == xl);
      } else {  // use Newton step
        delta_x_old = delta_x;
        delta_x     = fval / df_dx;
        auto temp   = x;
        x -= delta_x;
        converged = (x == temp);
      }
 
      // function and jacobian evaluation
      R     = f(make_dual(x), get_value(params)...);
      fval  = get_value(R);
      df_dx = get_gradient(R);
 
      // convergence check
      converged = converged || (std::abs(delta_x) < options.xtol) || (std::abs(fval) < options.rtol);
 
      // maintain bracket on root
      if (fval < 0) {
        xl = x;
      } else {
        xh = x;
      }
 
      ++iterations;
    }
  }
 
  // Accumulate derivatives so that the user can get derivatives
  // with respect to parameters, subject to constraing that f(x, p) = 0 for all p
  // Conceptually, we're doing the following:
  // [fval, df_dp] = f(get_value(x), p)
  // df = 0
  // for p in params:
  //   df += inner(df_dp, dp)
  // dx = -df / df_dx
  constexpr bool contains_duals =
      (is_dual_number<ParamTypes>::value || ...) || (is_tensor_of_dual_number<ParamTypes>::value || ...);
  if constexpr (contains_duals) {
    auto [fval, df] = f(x, params...);
    auto         dx = -df / df_dx;
    SolverStatus status{.converged = converged, .iterations = iterations, .residual = fval};
    return tuple{dual{x, dx}, status};
  }
  if constexpr (!contains_duals) {
    auto         fval = f(x, params...);
    SolverStatus status{.converged = converged, .iterations = iterations, .residual = fval};
    return tuple{x, status};
  }
}
 
template <typename function, int n>
auto find_root(function&& f, tensor<double, n> x0)
{
  static_assert(std::is_same_v<decltype(f(x0)), tensor<double, n>>,
                "error: f(x) must have the same number of equations as unknowns");
 
  double epsilon        = 1.0e-8;
  int    max_iterations = 10;
 
  auto x = x0;
 
  for (int k = 0; k < max_iterations; k++) {
    auto output = f(make_dual(x));
    auto r      = get_value(output);
    if (norm(r) < epsilon) break;
    auto J = get_gradient(output);
    x -= linear_solve(J, r);
  }
 
  return x;
};
 
template <typename T, int size>
auto eigenvalues(const serac::tensor<T, size, size>& A)
{
  // put tensor values in an mfem::DenseMatrix
  mfem::DenseMatrix matA(size, size);
  for (int i = 0; i < size; i++) {
    for (int j = 0; j < size; j++) {
      if constexpr (is_dual_number<T>::value) {
        matA(i, j) = A[i][j].value;
      } else {
        matA(i, j) = A[i][j];
      }
    }
  }
 
  // compute eigendecomposition
  mfem::DenseMatrixEigensystem eig_sys(matA);
  eig_sys.Eval();
 
  serac::tensor<T, size> output;
 
  for (int k = 0; k < size; k++) {
    // extract eigenvalues
    output[k] = eig_sys.Eigenvalue(k);
 
    // and calculate their derivatives, when appropriate
    if constexpr (is_dual_number<T>::value) {
      tensor<double, size> phi = make_tensor<size>([&](int i) { return eig_sys.Eigenvector(k)[i]; });
      auto                 dA  = make_tensor<size, size>([&](int i, int j) { return A(i, j).gradient; });
      output[k].gradient       = dot(phi, dA, phi);
    }
  }
 
  return output;
}
 
}  // namespace serac