1 files changed, 15 insertions, 11 deletions

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