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

1 files changed, 27 insertions, 17 deletions

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));
+}