From 7567420d7daca4d85b75ab8d7ab07352e14c16cd Mon Sep 17 00:00:00 2001
From: Austeja <austeja.jurgaityte23@gmail.com>
Date: Mon, 16 Sep 2024 23:39:22 +0300
Subject: [PATCH] Add support for 'std::atan2' and 'std::acos'

---
 .../clad/Differentiator/BuiltinDerivatives.h  | 22 ++++++++
 test/FirstDerivative/BuiltinDerivatives.C     | 53 +++++++++++++++++++
 2 files changed, 75 insertions(+)

diff --git a/include/clad/Differentiator/BuiltinDerivatives.h b/include/clad/Differentiator/BuiltinDerivatives.h
index 3f9dcae67..557274a56 100644
--- a/include/clad/Differentiator/BuiltinDerivatives.h
+++ b/include/clad/Differentiator/BuiltinDerivatives.h
@@ -183,6 +183,25 @@ CUDA_HOST_DEVICE ValueAndPushforward<T, T> floor_pushforward(T x, T /*d_x*/) {
   return {::std::floor(x), (T)0};
 }
 
+template <typename T>
+CUDA_HOST_DEVICE ValueAndPushforward<T, T> atan2_pushforward(T y, T x, T d_y,
+                                                             T d_x) {
+  return {::std::atan2(y, x),
+          -(y / ((x * x) + (y * y))) * d_x + x / ((x * x) + (y * y)) * d_y};
+}
+
+template <typename T, typename U>
+CUDA_HOST_DEVICE void atan2_pullback(T y, T x, U d_z, T* d_y, T* d_x) {
+  *d_y += x / ((x * x) + (y * y)) * d_z;
+
+  *d_x += -(y / ((x * x) + (y * y))) * d_z;
+}
+
+template <typename T>
+CUDA_HOST_DEVICE ValueAndPushforward<T, T> acos_pushforward(T x, T d_x) {
+  return {::std::acos(x), ((-1) / (::std::sqrt(1 - x * x))) * d_x};
+}
+
 template <typename T>
 CUDA_HOST_DEVICE ValueAndPushforward<T, T> ceil_pushforward(T x, T /*d_x*/) {
   return {::std::ceil(x), (T)0};
@@ -321,6 +340,9 @@ inline void free_pushforward(void* ptr, void* d_ptr) {
 // These are required because C variants of mathematical functions are
 // defined in global namespace.
 using std::abs_pushforward;
+using std::acos_pushforward;
+using std::atan2_pullback;
+using std::atan2_pushforward;
 using std::ceil_pushforward;
 using std::cos_pushforward;
 using std::exp_pushforward;
diff --git a/test/FirstDerivative/BuiltinDerivatives.C b/test/FirstDerivative/BuiltinDerivatives.C
index e608473bc..c0db1fc8f 100644
--- a/test/FirstDerivative/BuiltinDerivatives.C
+++ b/test/FirstDerivative/BuiltinDerivatives.C
@@ -248,6 +248,45 @@ double f14(double x) {
   return __builtin_pow(x, 3);
 }
 
+double f15(double y, double x) {
+  return std::atan2(y, x);
+}
+
+//CHECK:  float f15_darg0(float y, float x) {
+//CHECK-NEXT:      float _d_y = 1;
+//CHECK-NEXT:      float _d_x = 0;
+//CHECK-NEXT:      ValueAndPushforward<float, float> _t0 = clad::custom_derivatives::std::atan2_pushforward(y, x, _d_y, _d_x);
+//CHECK-NEXT:      return _t0.pushforward;
+//CHECK-NEXT: }
+
+//CHECK: float f15_darg1(float y, float x) {
+//CHECK-NEXT:     float _d_y = 0;
+//CHECK-NEXT:     float _d_x = 1;
+//CHECK-NEXT:     ValueAndPushforward<float, float> _t0 = clad::custom_derivatives::std::atan2_pushforward(y, x, _d_y, _d_x);
+//CHECK-NEXT:     return _t0.pushforward;
+//CHECK-NEXT: }
+
+void f15_grad(double y, double x, double *_d_y, double *_d_x);
+//CHECK: void f15_grad(double y, double x, double *_d_y, double *_d_x) {
+//CHECK-NEXT:     double _t0 = std::atan2(y, x);
+//CHECK-NEXT:     {
+//CHECK-NEXT:         double _r0 = 0;
+//CHECK-NEXT:         double _r1 = 0;
+//CHECK-NEXT:         clad::custom_derivatives::atan2_pullback(y, x, 2 * 1, &_r0, &_r1);
+//CHECK-NEXT:         *_d_y += _r0;
+//CHECK-NEXT:         *_d_x += _r1;
+//CHECK-NEXT:     }
+//CHECK-NEXT: }
+
+float f16(float x) {
+  return std::acos(x);
+}
+// CHECK: float f16_darg0(float x) {
+//CHECK-NEXT:    float _d_x = 1;
+//CHECK-NEXT:    ValueAndPushforward<float, float> _t0 = clad::custom_derivatives::std::acos_pushforward(x, _d_x);
+//CHECK-NEXT:    return _t0.pushforward;
+//CHECK-NEXT: }
+
 int main () { //expected-no-diagnostics
   float f_result[2];
   double d_result[2];
@@ -326,5 +365,19 @@ int main () { //expected-no-diagnostics
   auto f14_ddarg0 = clad::differentiate<2>(f14, 0);
   printf("Result is = %f\n", f14_ddarg0.execute(1)); //CHECK-EXEC: Result is = 6.000000
 
+  auto f15_darg0 = clad::differentiate(f15, 0);
+  printf("Result is = %f\n", f15_darg0.execute(4, 3)); //CHECK-EXEC: Result is = 0.120000
+
+  auto f15_darg1 = clad::differentiate(f15, 1);
+  printf("Result is = %f\n", f15_darg1.execute(4, 3)); //CHECK-EXEC: Result is = -0.160000
+
+  d_result[0] = d_result[1] = 0;
+  clad::gradient(f15);
+  f15_grad(4, 3, &d_result[0], &d_result[1]);
+  printf("Result is = {%f, %f}\n", d_result[0], d_result[1]); //CHECK-EXEC: Result is = {0.240000, -0.320000}
+
+  auto f16_darg0 = clad::differentiate(f16, 0);
+  printf("Result is = %f\n", f16_darg0.execute(0.9)); //CHECK-EXEC: Result is = -2.294157
+  
   return 0;
 }