Converted more of library to templates accepting generic Scalar types and p

5 files changed, 194 insertions, 169 deletions

diff --git a/dynamics.hpp b/dynamics.hpp
new file mode 100644
index 0000000..85eef71
--- /dev/null
+++ b/dynamics.hpp
@@ -0,0 +1,117 @@
+#pragma once
+
+#include <exception>
+#include <eigen3/Eigen/LU>
+#include <random>
+
+#include "p-spin.hpp"
+#include "stereographic.hpp"
+
+class gradientDescentStallException: public std::exception {
+  virtual const char* what() const throw() {
+    return "Gradient descent stalled.";
+  }
+} gradientDescentStall;
+
+template <class Scalar, int p>
+std::tuple<double, Vector<Scalar>> gradientDescent(const Tensor<Scalar, p>& J, const Vector<Scalar>& z0, double ε, double γ0 = 1, double δγ = 2) {
+  Vector<Scalar> z = z0;
+  double γ = γ0;
+
+  auto [W, dW] = WdW(J, z);
+
+  while (W > ε) {
+    Vector<Scalar> zNewTmp = z - γ * dW.conjugate();
+    Vector<Scalar> zNew = normalize(zNewTmp);
+
+    auto [WNew, dWNew] = WdW(J, zNew);
+
+    if (WNew < W) { // If the step lowered the objective, accept it!
+      z = zNew;
+      W = WNew;
+      dW = dWNew;
+      γ = γ0;
+    } else { // Otherwise, shrink the step and try again.
+      γ /= δγ;
+    }
+
+    if (γ < 1e-50) {
+      throw gradientDescentStall;
+    }
+  }
+
+  return {W, z};
+}
+
+template <class Scalar, int p>
+Vector<Scalar> findSaddle(const Tensor<Scalar, p>& J, const Vector<Scalar>& z0, double ε, double δW = 2, double γ0 = 1, double δγ = 2) {
+  Vector<Scalar> z = z0;
+  Vector<Scalar> ζ = euclideanToStereographic(z);
+
+  double W;
+  std::tie(W, std::ignore) = WdW(J, z);
+
+  Vector<Scalar> dH;
+  Matrix<Scalar> ddH;
+  std::tie(std::ignore, dH, ddH) = stereographicHamGradHess(J, ζ, z);
+
+  while (W > ε) {
+    // ddH is complex symmetric, which is (almost always) invertible, so a
+    // partial pivot LU decomposition can be used.
+    Vector<Scalar> dζ = ddH.partialPivLu().solve(dH);
+    Vector<Scalar> ζNew = ζ - dζ;
+    Vector<Scalar> zNew = stereographicToEuclidean(ζNew);
+
+    double WNew;
+    std::tie(WNew, std::ignore) = WdW(J, zNew);
+
+    if (WNew < W) { // If Newton's step lowered the objective, accept it!
+      ζ = ζNew;
+      z = zNew;
+      W = WNew;
+    } else { // Otherwise, do gradient descent until W is a factor δW smaller.
+      std::tie(W, z) = gradientDescent(J, z, W / δW, γ0, δγ);
+      ζ = euclideanToStereographic(z);
+    }
+
+    std::tie(std::ignore, dH, ddH) = stereographicHamGradHess(J, ζ, z);
+  }
+
+  return z;
+}
+
+template <class Scalar, class Distribution, class Generator>
+Vector<Scalar> randomVector(unsigned N, Distribution d, Generator& r) {
+  Vector<Scalar> z(N);
+
+  for (unsigned i = 0; i < N; i++) {
+    z(i) = d(r);
+  }
+
+  return z;
+}
+
+template <class Scalar, int p, class Distribution, class Generator>
+std::tuple<double, Vector<Scalar>> langevin(const Tensor<Scalar, p>& J, const Vector<Scalar>& z0, double T, double γ, unsigned N, Distribution d, Generator& r) {
+  Vector<Scalar> z = z0;
+
+  double W;
+  std::tie(W, std::ignore) = WdW(J, z);
+
+  std::uniform_real_distribution<double> D(0, 1);
+
+  for (unsigned i = 0; i < N; i++) {
+    Vector<Scalar> zNewTmp = z + randomVector<Scalar>(z.size(), d, r);
+    Vector<Scalar> zNew = normalize(zNewTmp);
+
+    double WNew;
+    std::tie(WNew, std::ignore) = WdW(J, zNew);
+
+    if (exp((W - WNew) / T) > D(r)) {
+      z = zNew;
+      W = WNew;
+    }
+  }
+
+  return {W, z};
+}
diff --git a/langevin.cpp b/langevin.cpp
index a4c9cf3..870879b 100644
--- a/langevin.cpp
+++ b/langevin.cpp
@@ -1,127 +1,21 @@
-#include <exception>
 #include <getopt.h>
 
-#include <eigen3/Eigen/LU>
-#include <utility>
-
 #include "complex_normal.hpp"
 #include "p-spin.hpp"
-#include "stereographic.hpp"
+#include "dynamics.hpp"
 
 #include "pcg-cpp/include/pcg_random.hpp"
 #include "randutils/randutils.hpp"
 #include "unsupported/Eigen/CXX11/src/Tensor/TensorFFT.h"
 
-using Rng = randutils::random_generator<pcg32>;
-
-Vector normalize(const Vector& z) {
-  return z * sqrt((double)z.size() / (Scalar)(z.transpose() * z));
-}
-
-template <class Distribution, class Generator>
-Vector randomVector(unsigned N, Distribution d, Generator& r) {
-  Vector z(N);
-
-  for (unsigned i = 0; i < N; i++) {
-    z(i) = d(r);
-  }
-
-  return z;
-}
-
-class gradientDescentStallException: public std::exception {
-  virtual const char* what() const throw() {
-    return "Gradient descent stalled.";
-  }
-} gradientDescentStall;
-
-std::tuple<double, Vector> gradientDescent(const Tensor& J, const Vector& z0, double ε, double γ0 = 1, double δγ = 2) {
-  Vector z = z0;
-  double γ = γ0;
-
-  auto [W, dW] = WdW(J, z);
-
-  while (W > ε) {
-    Vector zNew = normalize(z - γ * dW.conjugate());
-
-    auto [WNew, dWNew] = WdW(J, zNew);
-
-    if (WNew < W) { // If the step lowered the objective, accept it!
-      z = zNew;
-      W = WNew;
-      dW = dWNew;
-      γ = γ0;
-    } else { // Otherwise, shrink the step and try again.
-      γ /= δγ;
-    }
-
-    if (γ < 1e-50) {
-      throw gradientDescentStall;
-    }
-  }
+#define PSPIN_P 3
+const int p = PSPIN_P; // polynomial degree of Hamiltonian
+using Complex = std::complex<double>;
+using ComplexVector = Vector<Complex>;
+using ComplexMatrix = Matrix<Complex>;
+using ComplexTensor = Tensor<Complex, p>;
 
-  return {W, z};
-}
-
-Vector findSaddle(const Tensor& J, const Vector& z0, double ε, double δW = 2, double γ0 = 1, double δγ = 2) {
-  Vector z = z0;
-  Vector ζ = euclideanToStereographic(z);
-
-  double W;
-  std::tie(W, std::ignore) = WdW(J, z);
-
-  Vector dH;
-  Matrix ddH;
-  std::tie(std::ignore, dH, ddH) = stereographicHamGradHess(J, ζ, z);
-
-  while (W > ε) {
-    // ddH is complex symmetric, which is (almost always) invertible, so a
-    // partial pivot LU decomposition can be used.
-    Vector dζ = ddH.partialPivLu().solve(dH);
-    Vector ζNew = ζ - dζ;
-    Vector zNew = stereographicToEuclidean(ζNew);
-
-    double WNew;
-    std::tie(WNew, std::ignore) = WdW(J, zNew);
-
-    if (WNew < W) { // If Newton's step lowered the objective, accept it!
-      ζ = ζNew;
-      z = zNew;
-      W = WNew;
-    } else { // Otherwise, do gradient descent until W is a factor δW smaller.
-      std::tie(W, z) = gradientDescent(J, z, W / δW, γ0, δγ);
-      ζ = euclideanToStereographic(z);
-    }
-
-    std::tie(std::ignore, dH, ddH) = stereographicHamGradHess(J, ζ, z);
-  }
-
-  return z;
-}
-
-std::tuple<double, Vector> langevin(const Tensor& J, const Vector& z0, double T, double γ, unsigned N, Rng& r) {
-  Vector z = z0;
-
-  double W;
-  std::tie(W, std::ignore) = WdW(J, z);
-
-  complex_normal_distribution<> d(0, γ, 0);
-
-  for (unsigned i = 0; i < N; i++) {
-    Vector dz = randomVector(z.size(), d, r.engine());
-    Vector zNew = normalize(z + dz);
-
-    double WNew;
-    std::tie(WNew, std::ignore) = WdW(J, zNew);
-
-    if (exp((W - WNew) / T) > r.uniform(0.0, 1.0)) {
-      z = zNew;
-      W = WNew;
-    }
-  }
-
-  return {W, z};
-}
+using Rng = randutils::random_generator<pcg32>;
 
 int main(int argc, char* argv[]) {
   // model parameters
@@ -178,22 +72,24 @@ int main(int argc, char* argv[]) {
     }
   }
 
-  Scalar κ(Rκ, Iκ);
+  Complex κ(Rκ, Iκ);
   double σ = sqrt(factorial(p) / (2.0 * pow(N, p - 1)));
 
   Rng r;
 
-  Tensor J = generateCouplings<Scalar, PSPIN_P>(N, complex_normal_distribution<>(0, σ, κ), r.engine());
-  Vector z0 = normalize(randomVector(N, complex_normal_distribution<>(0, 1, 0), r.engine()));
+  complex_normal_distribution<> d(0, 1, 0);
+
+  ComplexTensor J = generateCouplings<Complex, PSPIN_P>(N, complex_normal_distribution<>(0, σ, κ), r.engine());
+  ComplexVector z0 = normalize(randomVector<Complex>(N, d, r.engine()));
 
-  Vector zSaddle = findSaddle(J, z0, ε);
-  Vector zSaddlePrev = Vector::Zero(N);
-  Vector z = zSaddle;
+  ComplexVector zSaddle = findSaddle(J, z0, ε);
+  ComplexVector zSaddlePrev = ComplexVector::Zero(N);
+  ComplexVector z = zSaddle;
 
   while (δ < (zSaddle - zSaddlePrev).norm()) { // Until we find two saddles sufficiently close...
-    std::tie(std::ignore, z) = langevin(J, z, T, γ, M, r);
+    std::tie(std::ignore, z) = langevin(J, z, T, γ, M, d, r.engine());
     try {
-      Vector zSaddleNext = findSaddle(J, z, ε);
+      ComplexVector zSaddleNext = findSaddle(J, z, ε);
       if (Δ < (zSaddleNext - zSaddle).norm()) { // Ensure we are finding distinct saddles.
         zSaddlePrev = zSaddle;
         zSaddle = zSaddleNext;
@@ -209,20 +105,20 @@ int main(int argc, char* argv[]) {
 
   complex_normal_distribution<> dJ(0, εJ * σ, 0);
 
-  std::function<void(Tensor&, std::array<unsigned, p>)> perturbJ =
-    [&dJ, &r] (Tensor& JJ, std::array<unsigned, p> ind) {
-      Scalar Ji = getJ<Scalar, p>(JJ, ind);
-      setJ<Scalar, p>(JJ, ind, Ji + dJ(r.engine()));
+  std::function<void(ComplexTensor&, std::array<unsigned, p>)> perturbJ =
+    [&dJ, &r] (ComplexTensor& JJ, std::array<unsigned, p> ind) {
+      Complex Ji = getJ<Complex, p>(JJ, ind);
+      setJ<Complex, p>(JJ, ind, Ji + dJ(r.engine()));
     };
 
   for (unsigned i = 0; i < n; i++) {
-    Tensor Jp = J;
+    ComplexTensor Jp = J;
 
-    iterateOver<Scalar, p>(Jp, perturbJ);
+    iterateOver<Complex, p>(Jp, perturbJ);
 
     try {
-      Vector zSaddleNew = findSaddle(Jp, zSaddle, ε);
-      Vector zSaddlePrevNew = findSaddle(Jp, zSaddlePrev, ε);
+      ComplexVector zSaddleNew = findSaddle(Jp, zSaddle, ε);
+      ComplexVector zSaddlePrevNew = findSaddle(Jp, zSaddlePrev, ε);
 
       std::cout << zSaddleNew.transpose() << " " << zSaddlePrevNew.transpose() << std::endl;
     } catch (std::exception& e) {
diff --git a/p-spin.hpp b/p-spin.hpp
index 3ef10e7..4621db6 100644
--- a/p-spin.hpp
+++ b/p-spin.hpp
@@ -3,49 +3,59 @@
 #include <eigen3/Eigen/Dense>
 
 #include "tensor.hpp"
+#include "factorial.hpp"
 
-#define PSPIN_P 3
-const unsigned p = PSPIN_P; // polynomial degree of Hamiltonian
+template <class Scalar>
+using Vector = Eigen::Matrix<Scalar, Eigen::Dynamic, 1>;
 
-using Scalar = std::complex<double>;
-using Vector = Eigen::VectorXcd;
-using Matrix = Eigen::MatrixXcd;
-using Tensor = Eigen::Tensor<Scalar, PSPIN_P>;
+template <class Scalar>
+using Matrix = Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic>;
 
-std::tuple<Scalar, Vector, Matrix> hamGradHess(const Tensor& J, const Vector& z) {
-  Matrix Jz = contractDown(J, z); // Contracts J into p - 2 copies of z.
-  Vector Jzz = Jz * z;
+template <class Scalar, int p>
+using Tensor = Eigen::Tensor<Scalar, p>;
+
+template <class Scalar, int p>
+std::tuple<Scalar, Vector<Scalar>, Matrix<Scalar>> hamGradHess(const Tensor<Scalar, p>& J, const Vector<Scalar>& z) {
+  Matrix<Scalar> Jz = contractDown(J, z); // Contracts J into p - 2 copies of z.
+  Vector<Scalar> Jzz = Jz * z;
   Scalar Jzzz = Jzz.transpose() * z;
 
   double pBang = factorial(p);
 
-  Matrix hessian = ((p - 1) * p / pBang) * Jz;
-  Vector gradient = (p / pBang) * Jzz;
+  Matrix<Scalar> hessian = ((p - 1) * p / pBang) * Jz;
+  Vector<Scalar> gradient = (p / pBang) * Jzz;
   Scalar hamiltonian = Jzzz / pBang;
 
   return {hamiltonian, gradient, hessian};
 }
 
-Vector project(const Vector& z, const Vector& x) {
+template <class Scalar>
+Vector<Scalar> project(const Vector<Scalar>& z, const Vector<Scalar>& x) {
   Scalar xz = x.transpose() * z;
   return x - (xz / z.squaredNorm()) * z.conjugate();
 }
 
-std::tuple<double, Vector> WdW(const Tensor& J, const Vector& z) {
-  Vector dH;
-  Matrix ddH;
+template <class Scalar, int p>
+std::tuple<double, Vector<Scalar>> WdW(const Tensor<Scalar, p>& J, const Vector<Scalar>& z) {
+  Vector<Scalar> dH;
+  Matrix<Scalar> ddH;
   std::tie(std::ignore, dH, ddH) = hamGradHess(J, z);
 
-  Vector dzdt = project(z, dH.conjugate());
+  Vector<Scalar> dzdt = project(z, dH.conjugate().eval());
 
   double a = z.squaredNorm();
   Scalar A = (Scalar)(z.transpose() * dzdt) / a;
   Scalar B = dH.dot(z) / a;
 
   double W = dzdt.squaredNorm();
-  Vector dW = ddH * (dzdt - A * z.conjugate())
+  Vector<Scalar> dW = ddH * (dzdt - A * z.conjugate())
     + 2 * (conj(A) * B * z).real()
     - conj(B) * dzdt - conj(A) * dH.conjugate();
 
   return {W, dW};
 }
+
+template <class Scalar>
+Vector<Scalar> normalize(const Vector<Scalar>& z) {
+  return z * sqrt((double)z.size() / (Scalar)(z.transpose() * z));
+}
diff --git a/stereographic.hpp b/stereographic.hpp
index 638923f..2dfa735 100644
--- a/stereographic.hpp
+++ b/stereographic.hpp
@@ -4,9 +4,10 @@
 
 #include "p-spin.hpp"
 
-Vector stereographicToEuclidean(const Vector& ζ) {
+template <class Scalar>
+Vector<Scalar> stereographicToEuclidean(const Vector<Scalar>& ζ) {
   unsigned N = ζ.size() + 1;
-  Vector z(N);
+  Vector<Scalar> z(N);
 
   Scalar a = ζ.transpose() * ζ;
   Scalar b = 2 * sqrt(N) / (1.0 + a);
@@ -20,9 +21,10 @@ Vector stereographicToEuclidean(const Vector& ζ) {
   return b * z;
 }
 
-Vector euclideanToStereographic(const Vector& z) {
+template <class Scalar>
+Vector<Scalar> euclideanToStereographic(const Vector<Scalar>& z) {
   unsigned N = z.size();
-  Vector ζ(N - 1);
+  Vector<Scalar> ζ(N - 1);
 
   for (unsigned i = 0; i < N - 1; i++) {
     ζ(i) = z(i);
@@ -31,9 +33,10 @@ Vector euclideanToStereographic(const Vector& z) {
   return ζ / (sqrt(N) - z(N - 1));
 }
 
-Matrix stereographicJacobian(const Vector& ζ) {
+template <class Scalar>
+Matrix<Scalar> stereographicJacobian(const Vector<Scalar>& ζ) {
   unsigned N = ζ.size();
-  Matrix J(N, N + 1);
+  Matrix<Scalar> J(N, N + 1);
 
   Scalar b = 1.0 + (Scalar)(ζ.transpose() * ζ);
 
@@ -52,15 +55,16 @@ Matrix stereographicJacobian(const Vector& ζ) {
   return 4 * sqrt(N + 1) * J / pow(b, 2);
 }
 
-std::tuple<Scalar, Vector, Matrix> stereographicHamGradHess(const Tensor& J, const Vector& ζ, const Vector& z) {
+template <class Scalar, int p>
+std::tuple<Scalar, Vector<Scalar>, Matrix<Scalar>> stereographicHamGradHess(const Tensor<Scalar, p>& J, const Vector<Scalar>& ζ, const Vector<Scalar>& z) {
   auto [hamiltonian, gradZ, hessZ] = hamGradHess(J, z);
-  Matrix jacobian = stereographicJacobian(ζ);
+  Matrix<Scalar> jacobian = stereographicJacobian(ζ);
 
-  Matrix metric = jacobian * jacobian.adjoint();
+  Matrix<Scalar> metric = jacobian * jacobian.adjoint();
 
   // The metric is Hermitian and positive definite, so a Cholesky decomposition can be used.
-  Vector grad = metric.llt().solve(jacobian) * gradZ;
-  Matrix hess = jacobian * hessZ * jacobian.transpose();
+  Vector<Scalar> grad = metric.llt().solve(jacobian) * gradZ;
+  Matrix<Scalar> hess = jacobian * hessZ * jacobian.transpose();
 
   return {hamiltonian, grad, hess};
 }
diff --git a/tensor.hpp b/tensor.hpp
index 66f6d59..21f2a89 100644
--- a/tensor.hpp
+++ b/tensor.hpp
@@ -5,43 +5,41 @@
 
 #include <eigen3/unsupported/Eigen/CXX11/Tensor>
 
-#include "factorial.hpp"
-
-template <class Scalar, unsigned p, std::size_t... Indices>
+template <class Scalar, int p, std::size_t... Indices>
 Eigen::Tensor<Scalar, p> initializeJHelper(unsigned N, std::index_sequence<Indices...>) {
   std::array<unsigned, p> Ns;
   std::fill_n(Ns.begin(), p, N);
   return Eigen::Tensor<Scalar, p>(std::get<Indices>(Ns)...);
 }
 
-template <class Scalar, unsigned p>
+template <class Scalar, int p>
 Eigen::Tensor<Scalar, p> initializeJ(unsigned N) {
   return initializeJHelper<Scalar, p>(N, std::make_index_sequence<p>());
 }
 
-template <class Scalar, unsigned p, std::size_t... Indices>
+template <class Scalar, int p, std::size_t... Indices>
 void setJHelper(Eigen::Tensor<Scalar, p>& J, const std::array<unsigned, p>& ind, Scalar val, std::index_sequence<Indices...>) {
   J(std::get<Indices>(ind)...) = val;
 }
 
-template <class Scalar, unsigned p>
+template <class Scalar, int p>
 void setJ(Eigen::Tensor<Scalar, p>& J, std::array<unsigned, p> ind, Scalar val) {
   do {
     setJHelper<Scalar, p>(J, ind, val, std::make_index_sequence<p>());
   } while (std::next_permutation(ind.begin(), ind.end()));
 }
 
-template <class Scalar, unsigned p, std::size_t... Indices>
+template <class Scalar, int p, std::size_t... Indices>
 Scalar getJHelper(const Eigen::Tensor<Scalar, p>& J, const std::array<unsigned, p>& ind, std::index_sequence<Indices...>) {
   return J(std::get<Indices>(ind)...);
 }
 
-template <class Scalar, unsigned p>
+template <class Scalar, int p>
 Scalar getJ(const Eigen::Tensor<Scalar, p>& J, const std::array<unsigned, p>& ind) {
   return getJHelper<Scalar, p>(J, ind, std::make_index_sequence<p>());
 }
 
-template <class Scalar, unsigned p>
+template <class Scalar, int p>
 void iterateOverHelper(Eigen::Tensor<Scalar, p>& J,
     std::function<void(Eigen::Tensor<Scalar, p>&, std::array<unsigned, p>)>& f,
     unsigned l, std::array<unsigned, p> is) {
@@ -62,13 +60,13 @@ void iterateOverHelper(Eigen::Tensor<Scalar, p>& J,
   }
 }
 
-template <class Scalar, unsigned p>
+template <class Scalar, int p>
 void iterateOver(Eigen::Tensor<Scalar, p>& J, std::function<void(Eigen::Tensor<Scalar, p>&, std::array<unsigned, p>)>& f) {
   std::array<unsigned, p> is;
   iterateOverHelper<Scalar, p>(J, f, p, is);
 }
 
-template <class Scalar, unsigned p, class Distribution, class Generator>
+template <class Scalar, int p, class Distribution, class Generator>
 Eigen::Tensor<Scalar, p> generateCouplings(unsigned N, Distribution d, Generator& r) {
   Eigen::Tensor<Scalar, p> J = initializeJ<Scalar, p>(N);