dual.hpp

Go to the documentation of this file.

// Copyright (c) 2019-2024, Lawrence Livermore National Security, LLC and
// other Serac Project Developers. See the top-level LICENSE file for
// details.
//
// SPDX-License-Identifier: (BSD-3-Clause)
 
#pragma once
 
#include <iostream>
 
#include <cmath>
 
#include "serac/infrastructure/accelerator.hpp"
 
namespace serac {
 
template <typename gradient_type>
struct dual {
  double        value;     
  gradient_type gradient;  
 
  SERAC_HOST_DEVICE constexpr auto& operator=(double b)
  {
    value    = b;
    gradient = {};
    return *this;
  }
};
 
template <typename T>
dual(double, T) -> dual<T>;
 
template <typename gradient_type>
SERAC_HOST_DEVICE constexpr auto operator+(dual<gradient_type> a, double b)
{
  return dual{a.value + b, a.gradient};
}
 
template <typename gradient_type>
SERAC_HOST_DEVICE constexpr auto operator+(double a, dual<gradient_type> b)
{
  return dual{a + b.value, b.gradient};
}
 
template <typename gradient_type_a, typename gradient_type_b>
SERAC_HOST_DEVICE constexpr auto operator+(dual<gradient_type_a> a, dual<gradient_type_b> b)
{
  return dual{a.value + b.value, a.gradient + b.gradient};
}
 
template <typename gradient_type>
constexpr auto operator-(dual<gradient_type> x)
{
  return dual{-x.value, -x.gradient};
}
 
template <typename gradient_type>
SERAC_HOST_DEVICE constexpr auto operator-(dual<gradient_type> a, double b)
{
  return dual{a.value - b, a.gradient};
}
 
template <typename gradient_type>
SERAC_HOST_DEVICE constexpr auto operator-(double a, dual<gradient_type> b)
{
  return dual{a - b.value, -b.gradient};
}
 
template <typename gradient_type_a, typename gradient_type_b>
SERAC_HOST_DEVICE constexpr auto operator-(dual<gradient_type_a> a, dual<gradient_type_b> b)
{
  return dual{a.value - b.value, a.gradient - b.gradient};
}
 
template <typename gradient_type>
SERAC_HOST_DEVICE constexpr auto operator*(const dual<gradient_type>& a, double b)
{
  return dual{a.value * b, a.gradient * b};
}
 
template <typename gradient_type>
SERAC_HOST_DEVICE constexpr auto operator*(double a, const dual<gradient_type>& b)
{
  return dual{a * b.value, a * b.gradient};
}
 
template <typename gradient_type_a, typename gradient_type_b>
SERAC_HOST_DEVICE constexpr auto operator*(dual<gradient_type_a> a, dual<gradient_type_b> b)
{
  return dual{a.value * b.value, b.value * a.gradient + a.value * b.gradient};
}
 
template <typename gradient_type>
SERAC_HOST_DEVICE constexpr auto operator/(const dual<gradient_type>& a, double b)
{
  return dual{a.value / b, a.gradient / b};
}
 
template <typename gradient_type>
SERAC_HOST_DEVICE constexpr auto operator/(double a, const dual<gradient_type>& b)
{
  return dual{a / b.value, -(a / (b.value * b.value)) * b.gradient};
}
 
template <typename gradient_type_a, typename gradient_type_b>
SERAC_HOST_DEVICE constexpr auto operator/(dual<gradient_type_a> a, dual<gradient_type_b> b)
{
  return dual{a.value / b.value, (a.gradient / b.value) - (a.value * b.gradient) / (b.value * b.value)};
}
 
#define binary_comparator_overload(x)                                             \
  template <typename T>                                                           \
  SERAC_HOST_DEVICE constexpr bool operator x(const dual<T>& a, double b)         \
  {                                                                               \
    return a.value x b;                                                           \
  }                                                                               \
                                                                                  \
  template <typename T>                                                           \
  SERAC_HOST_DEVICE constexpr bool operator x(double a, const dual<T>& b)         \
  {                                                                               \
    return a x b.value;                                                           \
  };                                                                              \
                                                                                  \
  template <typename T, typename U>                                               \
  SERAC_HOST_DEVICE constexpr bool operator x(const dual<T>& a, const dual<U>& b) \
  {                                                                               \
    return a.value x b.value;                                                     \
  };
 
binary_comparator_overload(<);   
binary_comparator_overload(<=);  
binary_comparator_overload(==);  
binary_comparator_overload(>=);  
binary_comparator_overload(>);   
 
#undef binary_comparator_overload
 
template <typename gradient_type>
SERAC_HOST_DEVICE constexpr auto& operator+=(dual<gradient_type>& a, const dual<gradient_type>& b)
{
  a.value += b.value;
  a.gradient += b.gradient;
  return a;
}
 
template <typename gradient_type>
SERAC_HOST_DEVICE constexpr auto& operator-=(dual<gradient_type>& a, const dual<gradient_type>& b)
{
  a.value -= b.value;
  a.gradient -= b.gradient;
  return a;
}
 
template <typename gradient_type>
SERAC_HOST_DEVICE constexpr auto& operator+=(dual<gradient_type>& a, double b)
{
  a.value += b;
  return a;
}
 
template <typename gradient_type>
SERAC_HOST_DEVICE constexpr auto& operator-=(dual<gradient_type>& a, double b)
{
  a.value -= b;
  return a;
}
 
template <typename gradient_type>
SERAC_HOST_DEVICE auto abs(dual<gradient_type> x)
{
  return (x.value >= 0) ? x : -x;
}
 
template <typename gradient_type>
SERAC_HOST_DEVICE auto max(dual<gradient_type> a, double b)
{
  dual<gradient_type> b_dual{b, 0.0 * a.gradient};
  return (a > b_dual) ? a : b_dual;
}
 
template <typename gradient_type>
SERAC_HOST_DEVICE auto max(double a, dual<gradient_type> b)
{
  dual<gradient_type> a_dual{a, 0.0 * b.gradient};
  return (a_dual > b) ? a_dual : b;
}
 
template <typename gradient_type>
SERAC_HOST_DEVICE auto max(dual<gradient_type> a, dual<gradient_type> b)
{
  return (a > b) ? a : b;
}
 
template <typename gradient_type>
SERAC_HOST_DEVICE auto min(dual<gradient_type> a, double b)
{
  dual<gradient_type> b_dual{b, 0.0 * a.gradient};
  return (a < b_dual) ? a : b_dual;
}
 
template <typename gradient_type>
SERAC_HOST_DEVICE auto min(double a, dual<gradient_type> b)
{
  dual<gradient_type> a_dual{a, 0.0 * b.gradient};
  return (a_dual < b) ? a_dual : b;
}
 
template <typename gradient_type>
SERAC_HOST_DEVICE auto min(dual<gradient_type> a, dual<gradient_type> b)
{
  return (a < b) ? a : b;
}
 
template <typename gradient_type>
SERAC_HOST_DEVICE auto sqrt(dual<gradient_type> x)
{
  using std::sqrt;
  return dual<gradient_type>{sqrt(x.value), x.gradient / (2.0 * sqrt(x.value))};
}
 
template <typename gradient_type>
SERAC_HOST_DEVICE auto cos(dual<gradient_type> a)
{
  using std::cos, std::sin;
  return dual<gradient_type>{cos(a.value), -a.gradient * sin(a.value)};
}
 
template <typename gradient_type>
SERAC_HOST_DEVICE auto sin(dual<gradient_type> a)
{
  using std::cos, std::sin;
  return dual<gradient_type>{sin(a.value), a.gradient * cos(a.value)};
}
 
template <typename gradient_type>
SERAC_HOST_DEVICE auto atan(dual<gradient_type> a)
{
  using std::atan, std::pow;
  return dual<gradient_type>{atan(a.value), a.gradient / (1.0 + pow(a.value, 2))};
}
 
template <typename gradient_type>
SERAC_HOST_DEVICE auto atan2(dual<gradient_type> y, dual<gradient_type> x)
{
  using std::atan2, std::pow;
  return dual<gradient_type>{atan2(y.value, x.value), y.gradient * x.value / (pow(x.value, 2) + pow(y.value, 2)) -
                                                          x.gradient * y.value / (pow(x.value, 2) + pow(y.value, 2))};
}
 
template <typename gradient_type>
SERAC_HOST_DEVICE auto atan2(double y, dual<gradient_type> x)
{
  using std::atan2, std::pow;
  return dual<gradient_type>{atan2(y, x.value), -x.gradient * y / (pow(x.value, 2) + pow(y, 2))};
}
 
template <typename gradient_type>
SERAC_HOST_DEVICE auto atan2(dual<gradient_type> y, double x)
{
  using std::atan2, std::pow;
  return dual<gradient_type>{atan2(y.value, x), y.gradient * x / (pow(x, 2) + pow(y.value, 2))};
}
 
template <typename gradient_type>
SERAC_HOST_DEVICE auto asin(dual<gradient_type> a)
{
  using std::asin, std::pow, std::sqrt;
  return dual<gradient_type>{asin(a.value), a.gradient / sqrt(1.0 - pow(a.value, 2))};
}
 
template <typename gradient_type>
SERAC_HOST_DEVICE auto acos(dual<gradient_type> a)
{
  using std::acos, std::pow, std::sqrt;
  return dual<gradient_type>{acos(a.value), -a.gradient / sqrt(1.0 - pow(a.value, 2))};
}
 
template <typename gradient_type>
SERAC_HOST_DEVICE auto exp(dual<gradient_type> a)
{
  using std::exp;
  return dual<gradient_type>{exp(a.value), exp(a.value) * a.gradient};
}
 
template <typename gradient_type>
SERAC_HOST_DEVICE auto log(dual<gradient_type> a)
{
  using std::log;
  return dual<gradient_type>{log(a.value), a.gradient / a.value};
}
 
template <typename gradient_type>
SERAC_HOST_DEVICE auto log1p(dual<gradient_type> a)
{
  using std::log1p;
  return dual<gradient_type>{log1p(a.value), a.gradient / (1.0 + a.value)};
}
 
template <typename gradient_type>
SERAC_HOST_DEVICE auto pow(dual<gradient_type> a, dual<gradient_type> b)
{
  using std::pow, std::log;
  double        value = pow(a.value, b.value);
  gradient_type grad  = pow(a.value, b.value - 1) * (a.gradient * b.value + b.gradient * a.value * log(a.value));
  return dual<gradient_type>{value, grad};
}
 
template <typename gradient_type>
SERAC_HOST_DEVICE auto pow(double a, dual<gradient_type> b)
{
  using std::pow, std::log;
  double value = pow(a, b.value);
  return dual<gradient_type>{value, value * b.gradient * log(a)};
}
 
template <typename gradient_type>
SERAC_HOST_DEVICE auto pow(dual<gradient_type> a, double b)
{
  using std::pow;
  double        value = pow(a.value, b);
  gradient_type grad  = b * pow(a.value, b - 1) * a.gradient;
  return dual<gradient_type>{value, grad};
}
 
template <typename T, int... n>
auto& operator<<(std::ostream& out, dual<T> A)
{
  out << '(' << A.value << ' ' << A.gradient << ')';
  return out;
}
 
SERAC_HOST_DEVICE constexpr auto make_dual(double x) { return dual{x, 1.0}; }
 
template <typename T>
SERAC_HOST_DEVICE constexpr auto get_value(const T& arg)
{
  return arg;
}
 
template <typename T>
SERAC_HOST_DEVICE constexpr auto get_value(dual<T> arg)
{
  return arg.value;
}
 
template <typename gradient_type>
SERAC_HOST_DEVICE constexpr auto get_gradient(dual<gradient_type> arg)
{
  return arg.gradient;
}
 
template <typename T>
struct is_dual_number {
  static constexpr bool value = false;  
};
 
template <typename T>
struct is_dual_number<dual<T> > {
  static constexpr bool value = true;  
};
 
}  // namespace serac
 
#include "serac/numerics/functional/tuple_tensor_dual_functions.hpp"