Generalized code to easily allow mixed p-spins.

1 files changed, 64 insertions, 37 deletions

diff --git a/p-spin.hpp b/p-spin.hpp
index 6111b75..a5f6c9a 100644
--- a/p-spin.hpp
+++ b/p-spin.hpp
@@ -12,66 +12,93 @@ Vector<typename Derived::Scalar> normalize(const Eigen::MatrixBase<Derived>& z)
   return z * sqrt((Real)z.size() / (typename Derived::Scalar)(z.transpose() * z));
 }
 
-template <class Scalar, int p, int ... ps>
-std::tuple<Scalar, Vector<Scalar>, Matrix<Scalar>> hamGradHess(const Tensor<Scalar, p>& J, const Tensor<Scalar, ps>& ... Js, const Vector<Scalar>& z) {
-  Matrix<Scalar> Jz = contractDown(J, z); // Contracts J into p - 2 copies of z.
+template <class Scalar>
+std::tuple<Matrix<Scalar>, Tensor<Scalar, 3>> hamGradTensorHelper(const Vector<Scalar>& z, const Tensor<Scalar, 2>& J) {
+  Tensor<Scalar, 3> J3(z.size(), z.size(), z.size());;
+  J3.setZero();
+  Matrix<Scalar> Jz = Eigen::Map<const Matrix<Scalar>>(J.data(), z.size(), z.size());
+
+  return {Jz, J3};
+}
+
+template <class Scalar, int p>
+std::tuple<Matrix<Scalar>, Tensor<Scalar, 3>> hamGradTensorHelper(const Vector<Scalar>& z, const Tensor<Scalar, p>& J) {
+  Tensor<Scalar, 3> J3 = contractDown(J, z);
+  Tensor<Scalar, 1> zT = Eigen::TensorMap<Tensor<const Scalar, 1>>(z.data(), {z.size()});
+  Tensor<Scalar, 2> J3zT = J3.contract(zT, ip00);
+  Matrix<Scalar> Jz = Eigen::Map<const Matrix<Scalar>>(J3zT.data(), z.size(), z.size());
+
+  return {Jz, J3};
+}
+
+template <class Scalar, int p, int... ps>
+std::tuple<Scalar, Vector<Scalar>, Matrix<Scalar>, Tensor<Scalar, 3>> hamGradHessHelper(const Vector<Scalar>& z, const Tensor<Scalar, p>& J, const Tensor<Scalar, ps>& ...Js) {
+  auto [Jz, J3] = hamGradTensorHelper(z, J);
+  
   Vector<Scalar> Jzz = Jz * z;
   Scalar Jzzz = Jzz.transpose() * z;
 
   Real pBang = factorial(p);
 
+  Tensor<Scalar, 3> dddH = ((p - 2) * (p - 1) * p / pBang) * J3;
   Matrix<Scalar> ddH = ((p - 1) * p / pBang) * Jz;
   Vector<Scalar> dH = (p / pBang) * Jzz;
   Scalar H = Jzzz / pBang;
 
   if constexpr (sizeof...(Js) > 0) {
-    auto [Hs, dHs, ddHs] = hamGradHess(Js..., z);
+    auto [Hs, dHs, ddHs, dddHs] = hamGradHess(z, Js...);
     H += Hs;
     dH += dHs;
     ddH += ddHs;
+    dddH += dddHs;
   }
 
-  return {H, dH, ddH};
+  return {H, dH, ddH, dddH};
 }
 
-template <class Scalar, int p>
-Scalar getHamiltonian(const Tensor<Scalar, p>& J, const Vector<Scalar>& z) {
-  Scalar H;
-  std::tie(H, std::ignore, std::ignore) = hamGradHess(J, z);
-  return H;
-}
+template <typename Scalar, int... ps>
+class pSpinModel {
+public:
+  std::tuple<Tensor<Scalar, ps>...> Js;
 
-template <class Scalar, int p>
-Vector<Scalar> getGradient(const Tensor<Scalar, p>& J, const Vector<Scalar>& z) {
-  Vector<Scalar> dH;
-  std::tie(std::ignore, dH, std::ignore) = hamGradHess(J, z);
-  return dH;
-}
+  unsigned dimension() const {
+    return std::get<0>(Js).dimension(0);
+  }
 
-template <class Scalar, int p>
-Matrix<Scalar> getHessian(const Tensor<Scalar, p>& J, const Vector<Scalar>& z) {
-  Matrix<Scalar> ddH;
-  std::tie(std::ignore, std::ignore, ddH) = hamGradHess(J, z);
-  return ddH;
-}
+  template <typename NewScalar>
+  pSpinModel<NewScalar, ps...> cast() const {
+    pSpinModel<NewScalar, ps...> M;
+    M.Js = std::apply(
+        [] (const Tensor<Scalar, ps>& ...Ks) -> std::tuple<Tensor<NewScalar, ps>...> {
+          return std::make_tuple(Ks.template cast<NewScalar>()...);
+        }, Js
+        );
 
-template <class Scalar>
-Real getThresholdEnergyDensity(unsigned p, Scalar κ, Scalar ε, Real a) {
-  Real apm2 = pow(a, p - 2);
-  Scalar δ = κ / apm2;
-  Real θ = arg(κ) + 2 * arg(ε);
+    return M;
+  }
 
-  return (p - 1) * apm2 / (2 * p) * pow(1 - norm(δ), 2) / (1 + norm(δ) - 2 * abs(δ) * cos(θ));
-}
+  std::tuple<Scalar, Vector<Scalar>, Matrix<Scalar>, Tensor<Scalar, 3>> hamGradHess(const Vector<Scalar>& z) const {
+    return std::apply([&z](const Tensor<Scalar, ps>& ...Ks) -> std::tuple<Scalar, Vector<Scalar>, Matrix<Scalar>, Tensor<Scalar, 3>> { return hamGradHessHelper(z, Ks...); }, Js);
+  }
 
-template <class Scalar, int p>
-Real getProportionOfThreshold(Scalar κ, const Tensor<Scalar, p>& J, const Vector<Scalar>& z) {
-  Real N = z.size();
-  Scalar ε = getHamiltonian(J, z) / N;
-  Real a = z.squaredNorm() / N;
+  Scalar getHamiltonian(const Vector<Scalar>& z) const {
+    Scalar H;
+    std::tie(H, std::ignore, std::ignore, std::ignore) = hamGradHess(z);
+    return H;
+  }
 
-  return norm(ε) / getThresholdEnergyDensity(p, κ, ε, a);
-}
+  Vector<Scalar> getGradient(const Vector<Scalar>& z) const {
+    Vector<Scalar> dH;
+    std::tie(std::ignore, dH, std::ignore, std::ignore) = hamGradHess(z);
+    return dH;
+  }
+
+  Matrix<Scalar> getHessian(const Vector<Scalar>& z) const {
+    Matrix<Scalar> ddH;
+    std::tie(std::ignore, std::ignore, ddH, std::ignore) = hamGradHess(z);
+    return ddH;
+  }
+};
 
 template <class Scalar>
 Vector<Scalar> zDot(const Vector<Scalar>& z, const Vector<Scalar>& dH) {