Lots of progress towards Hessian implementation.

1 files changed, 176 insertions, 0 deletions

diff --git a/stokes.hpp b/stokes.hpp
index b2694a9..948a484 100644
--- a/stokes.hpp
+++ b/stokes.hpp
@@ -333,6 +333,168 @@ public:
   }
 
   template <int p>
+  std::pair<Vector<Scalar>, Matrix<Scalar>> gsGradHess(const Tensor<Scalar, p>& J, Real t) const {
+    Vector<Scalar> z = f(t);
+    auto [H, dH, ddH] = hamGradHess(J, z);
+    Tensor<Scalar, 3> dddH = ((p - 2) * (p - 1) * p / (Real)factorial(p)) * J;
+    Vector<Scalar> ż = zDot(z, dH);
+    Tensor<Scalar, 1> żT = Eigen::TensorMap<Tensor<const Scalar, 1>>(ż.data(), {z.size()});
+    Vector<Scalar> dz = df(t);
+    Tensor<Scalar, 1> dzT = Eigen::TensorMap<Tensor<const Scalar, 1>>(dz.data(), {z.size()});
+    Matrix<Scalar> dż = dzDot(z, dH);
+    Matrix<Scalar> dżc = dzDotConjugate(z, dH, ddH);
+    Tensor<Scalar, 3> ddż = ddzDot(z, dH);
+    Tensor<Scalar, 3> ddżc = ddzDotConjugate(z, dH, ddH, dddH);
+
+    Eigen::array<Eigen::IndexPair<int>, 1> ip20 = {Eigen::IndexPair<int>(2, 0)};
+
+    Tensor<Scalar, 2> ddżcdzT = ddżc.contract(dzT, ip20);
+    Matrix<Scalar> ddżcdz = Eigen::Map<Matrix<Scalar>>(ddżcdzT.data(), z.size(), z.size());
+
+    Tensor<Scalar, 2> ddżdzT = ddż.contract(dzT.conjugate(), ip20);
+    Matrix<Scalar> ddżdz = Eigen::Map<Matrix<Scalar>>(ddżcdzT.data(), z.size(), z.size());
+
+    Tensor<Scalar, 2> ddżcżT = ddżc.contract(żT, ip20);
+    Matrix<Scalar> ddżcż = Eigen::Map<Matrix<Scalar>>(ddżcżT.data(), z.size(), z.size());
+
+    Tensor<Scalar, 2> ddżżcT = ddż.contract(żT.conjugate(), ip20);
+    Matrix<Scalar> ddżżc = Eigen::Map<Matrix<Scalar>>(ddżżcT.data(), z.size(), z.size());
+
+    Vector<Scalar> x(z.size() * gs.size());
+    Matrix<Scalar> M(x.size(), x.size());
+
+    for (unsigned i = 0; i < gs.size(); i++) {
+      Real fdg = gCoeff(i, t);
+      Real dfdg = dgCoeff(i, t);
+      Vector<Scalar> dC = - 0.5 / ż.norm() / dz.norm() * (
+          dfdg * ż.conjugate() + fdg * dżc * dz + fdg * dż * dz.conjugate()
+          - real(dz.dot(ż)) * (
+              dfdg * dz.conjugate() / dz.squaredNorm() +
+              fdg * (dżc * ż + dż * ż.conjugate()) / ż.squaredNorm()
+            )
+        );
+
+      for (unsigned j = 0; j < dC.size(); j++) {
+        x(z.size() * i + j) = dC(j);
+      }
+
+      for (unsigned j = 0; j < gs.size(); j++) {
+        Real fdgm = gCoeff(j, t);
+        Real dfdgm = dgCoeff(j, t);
+        Vector<Scalar> dCm = - 0.5 / ż.norm() / dz.norm() * (
+            dfdgm * ż.conjugate() + fdgm * dżc * dz + fdgm * dż * dz.conjugate()
+            - real(dz.dot(ż)) * (
+                dfdgm * dz.conjugate() / dz.squaredNorm() +
+                fdgm * (dżc * ż + dż * ż.conjugate()) / ż.squaredNorm()
+              )
+          );
+        Matrix<Scalar> ddC =
+          - 0.5 / ż.norm() / dz.norm() * (
+                dfdg * fdgm * dżc.transpose() + dfdgm * fdg * dżc +
+                fdg * fdgm * ddżcdz + fdg * fdgm * ddżdz
+            )
+          + 0.5 / dz.squaredNorm() * (
+              dfdg * dz.conjugate() * dCm.transpose() +
+              dfdgm * dC * dz.adjoint()
+            )
+          + 0.5 / ż.squaredNorm() * (
+              fdg * (dżc * ż + dż * ż.conjugate()) * dCm.transpose() +
+              fdgm * dC * (dżc * ż + dż * ż.conjugate()).transpose()
+            )
+          + 0.5 / ż.norm() / dz.norm() * real(dz.dot(ż)) * fdg * fdgm * (
+                  ddżżc + ddżcż + dżc * dż.transpose() + dż * dżc.transpose()
+                ) / ż.squaredNorm()
+          - 0.5 / ż.norm() / dz.norm() * real(dz.dot(ż)) * dfdg * dfdgm * dz.conjugate() * dz.adjoint() / pow(dz.squaredNorm(), 2)
+          - 0.5 / ż.norm() / dz.norm() * real(dz.dot(ż)) / pow(ż.squaredNorm(), 2)
+                * fdg * (dżc * ż + dż * ż.conjugate())
+                * fdgm * (dżc * ż + dż * ż.conjugate()).transpose()
+            ;
+
+        for (unsigned k = 0; k < z.size(); k++) {
+          for (unsigned l = 0; l < z.size(); l++) {
+            M(z.size() * i + k, z.size() * j + l) = ddC(k, l);
+          }
+        }
+      }
+    }
+
+    Cord cNew(*this);
+    Matrix<Scalar> Mf = Matrix<Scalar>::Zero(x.size(), x.size());
+    for (unsigned i = 0; i < x.size(); i++) {
+      for (unsigned j = 0; j < x.size(); j++) {
+        Scalar hi = std::complex<Real>(0.03145, 0.08452);
+        Scalar hj = std::complex<Real>(0.09483, 0.02453);
+        Scalar c0 = cost(J, t);
+        cNew.gs[i / z.size()](i % z.size()) += hi;
+        Scalar ci = cNew.cost(J, t);
+        cNew.gs[j / z.size()](j % z.size()) += hj;
+        Scalar cij = cNew.cost(J, t);
+        cNew.gs[i / z.size()](i % z.size()) -= hi;
+        Scalar cj = cNew.cost(J, t);
+        cNew.gs[j / z.size()](j % z.size()) -= hj;
+
+        Mf(i, j) = (cij - ci - cj + c0) / (hi * hj);
+      }
+    }
+
+    std::cout << M << std::endl;
+    std::cout << Mf << std::endl;
+    getchar();
+
+    return {x, M};
+  }
+
+  template <int p>
+  std::pair<Real, Real> relaxStepNewton(const Tensor<Scalar, p>& J, unsigned nt, Real δ₀) {
+    Vector<Scalar> dgsTot = Vector<Scalar>::Zero(z0.size() * gs.size());
+    Matrix<Scalar> HessTot = Matrix<Scalar>::Zero(z0.size() * gs.size(), z0.size() * gs.size());
+
+    for (unsigned i = 0; i < nt; i++) {
+      Real t = (i + 1.0) / (nt + 1.0);
+      auto [dgsi, Hessi] = gsGradHess(J, t);
+
+      dgsTot += dgsi / nt;
+      HessTot += Hessi / nt;
+    }
+
+    Real stepSize = dgsTot.norm();
+
+    Cord cNew(*this);
+
+    Real δ = δ₀;
+    Real oldCost = totalCost(J, nt);
+
+    Vector<Scalar> δg = HessTot.partialPivLu().solve(dgsTot);
+
+    for (unsigned i = 0; i < gs.size(); i++) {
+      for (unsigned j = 0; j < z0.size(); j++) {
+        cNew.gs[i](j) = gs[i](j) - δg[z0.size() * i + j];
+      }
+    }
+
+    Real newCost = cNew.totalCost(J, nt);
+    if (newCost < oldCost) {
+      std::cout << "Newton step accepted!" << std::endl;
+    }
+
+    while (newCost > oldCost) {
+      for (unsigned i = 0; i < gs.size(); i++) {
+        for (unsigned j = 0; j < z0.size(); j++) {
+          cNew.gs[i](j) = gs[i](j) - δ * conj(dgsTot[z0.size() * i + j]);
+        }
+      }
+
+      newCost = cNew.totalCost(J, nt);
+
+      δ /= 2;
+    }
+
+    gs = cNew.gs;
+
+    return {2*δ, stepSize};
+  }
+
+  template <int p>
   std::pair<Real, Real> relaxStep(const Tensor<Scalar, p>& J, unsigned nt, Real δ₀) {
     std::vector<Vector<Scalar>> dgsTot(gs.size(), Vector<Scalar>::Zero(z0.size()));
 
@@ -380,6 +542,20 @@ public:
     while (stepSize > 1e-7 && steps < maxSteps) {
       std::tie(δ, stepSize) = relaxStep(J, nt, δ);
       std::cout << totalCost(J, nt) << " " << δ << " " << stepSize << std::endl;
+      δ *= 2;
+      steps++;
+    }
+  }
+
+  template <int p>
+  void relaxNewton(const Tensor<Scalar, p>& J, unsigned nt, Real δ₀, unsigned maxSteps) {
+    Real δ = δ₀;
+    Real stepSize = std::numeric_limits<Real>::infinity();
+    unsigned steps = 0;
+    while (stepSize > 1e-7 && steps < maxSteps) {
+      std::tie(δ, stepSize) = relaxStepNewton(J, nt, δ);
+      std::cout << totalCost(J, nt) << " " << δ << " " << stepSize << std::endl;
+      δ *= 2;
       steps++;
     }
   }