Trigonometric Trick for imporved per-step complexity of Gaussian HMC

GeomScale · May 27, 2024 · 07b98bc · 07b98bc
1 parent cca44a0
commit 07b98bc
Showing 1 changed file with 16 additions and 15 deletions.
diff --git a/include/random_walks/gaussian_hamiltonian_monte_carlo_exact_walk.hpp b/include/random_walks/gaussian_hamiltonian_monte_carlo_exact_walk.hpp
@@ -224,22 +224,23 @@ private :
         }
     }
 
-    inline void update_position(Point &p, Point &v, NT const& T, NT const& omega)
-    {
-        NT C, Phi;
-        for (size_t i = 0; i < p.dimension(); i++)
-        {
-            C = std::sqrt(p[i] * p[i] + (v[i] * v[i]) / (omega * omega));
-            Phi = std::atan((-v[i]) / (p[i] * omega));
-            if (v[i] < 0.0 && Phi < 0.0) {
-                Phi += M_PI;
-            } else if (v[i] > 0.0 && Phi > 0.0) {
-                Phi -= M_PI;
-            }
-            p.set_coord(i, C * std::cos(omega * T + Phi));
-            v.set_coord(i, -C * omega * std::sin(omega * T + Phi));
-        }
+    inline void update_position(Point& p, Point& v, const NT& T, const NT& omega) {
 
+        NT next_p, next_v;
+
+        NT sinVal = std::sin(omega * T);
+        NT cosVal = std::cos(omega * T);
+
+        NT factor1 = sinVal / omega;
+        NT factor2 = -omega * sinVal;
+
+        for (size_t i = 0; i < p.dimension(); i++) {
+            next_p = cosVal * p[i] + v[i] * factor1;
+            next_v = factor2 * p[i] + cosVal * v[i];
+
+            p.set_coord(i, next_p);
+            v.set_coord(i, next_v);
+        }
     }
 
     inline double get_max_distance(std::vector<Point> &pointset, Point const& q, double &rad)