1 files changed, 32 insertions, 16 deletions

diff --git a/langevin.cpp b/langevin.cpp
index 05b436e..acea609 100644
--- a/langevin.cpp
+++ b/langevin.cpp
@@ -26,16 +26,16 @@ Vector randomVector(unsigned N, Distribution d, Generator& r) {
   return z;
 }
 
-Vector findSaddle(const Tensor& J, const Vector& z0, double γ0, double δ, double ε, Rng& r) {
+std::tuple<double, Vector> gradientDescent(const Tensor& J, const Vector& z0, double γ0, double ε) {
   Vector z = z0;
 
   double W;
   Vector dW;
   std::tie(W, dW) = WdW(J, z);
 
-  while (W > δ) {
-    double γ = pow(r.variate<double, std::normal_distribution>(0, γ0), 2);
+  double γ = γ0;
 
+  while (W > ε) {
     Vector zNew = z - γ * dW.conjugate();
     zNew *= sqrt(z.size()) / sqrt((Scalar)(zNew.transpose() * zNew));
 
@@ -47,32 +47,48 @@ Vector findSaddle(const Tensor& J, const Vector& z0, double γ0, double δ, doub
       z = zNew;
       W = WNew;
       dW = dWNew;
-      std::cout << W << std::endl;
-    } 
+      γ = γ0;
+    } else {
+      γ *= 0.5;
+    }
+
+    if (γ < 1e-15) {
+      std::cout << "Gradient descent stalled." << std::endl;
+      exit(1);
+    }
   }
 
-  Vector ζ = euclideanToStereographic(z);
+  return {W, z};
+}
+
+Vector findSaddle(const Tensor& J, const Vector& z0, double ε, double δW, double γ0) {
+  double W;
+  std::tie(W, std::ignore) = WdW(J, z0);
+
+  Vector ζ = euclideanToStereographic(z0);
   Vector dH;
   Matrix ddH;
   std::tie(std::ignore, dH, ddH) = stereographicHamGradHess(J, ζ);
 
   while (W > ε) {
-    double γ = pow(r.variate<double, std::normal_distribution>(0, γ0), 2);
-
+    // ddH is complex symmetric, which is (almost always) invertible
     Vector dζ = ddH.partialPivLu().solve(dH);
     Vector ζNew = ζ - dζ;
 
     double WNew;
-    Vector dWNew;
-    std::tie(WNew, dWNew) = WdW(J, stereographicToEuclidean(ζNew));
+    std::tie(WNew, std::ignore) = WdW(J, stereographicToEuclidean(ζNew));
 
-    if (WNew < W) {
+    if (WNew < W) { // If Newton's step lowered the objective, accept it!
       ζ = ζNew;
       W = WNew;
-      dW = dWNew;
-      std::tie(std::ignore, dH, ddH) = stereographicHamGradHess(J, ζ);
-      std::cout << WNew << std::endl;
-    } 
+    } else { // Otherwise, do gradient descent until W is a factor δW smaller.
+      Vector z;
+      std::tie(W, z) = gradientDescent(J, stereographicToEuclidean(ζ), γ0, W / δW);
+      ζ = euclideanToStereographic(z);
+    }
+
+    std::tie(std::ignore, dH, ddH) = stereographicHamGradHess(J, ζ);
+    std::cout << W << std::endl;
   }
 
   return stereographicToEuclidean(ζ);
@@ -162,7 +178,7 @@ int main(int argc, char* argv[]) {
   };
 
   //Vector zm = langevin(J, z, T, γ, f, r);
-  Vector zm = findSaddle(J, z, γ, δ, ε, r);
+  Vector zm = findSaddle(J, z, ε, δ, γ);
 
   Scalar H;
   Vector dH;