Moved some repeated code into the pSpinModel class.

1 files changed, 61 insertions, 42 deletions

diff --git a/p-spin.hpp b/p-spin.hpp
index 7a49bca..544abf1 100644
--- a/p-spin.hpp
+++ b/p-spin.hpp
@@ -12,73 +12,92 @@ 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>
-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());
+template <typename Scalar, int... ps>
+class pSpinModel {
+private:
+  std::tuple<Matrix<Scalar>, Tensor<Scalar, 3>> hamGradTensorHelper(const Vector<Scalar>& z, const Tensor<Scalar, 2>& J) const {
+    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};
-}
+    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());
+  template <int p>
+  std::tuple<Matrix<Scalar>, Tensor<Scalar, 3>> hamGradTensorHelper(const Vector<Scalar>& z, const Tensor<Scalar, p>& J) const {
+    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};
-}
+    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;
+  template <int p, int... qs>
+  std::tuple<Scalar, Vector<Scalar>, Matrix<Scalar>, Tensor<Scalar, 3>> hamGradHessHelper(const Vector<Scalar>& z, const Tensor<Scalar, p>& J, const Tensor<Scalar, qs>& ...Js) const {
+    auto [Jz, J3] = hamGradTensorHelper(z, J);
+    
+    Vector<Scalar> Jzz = Jz * z;
+    Scalar Jzzz = Jzz.transpose() * z;
 
-  Real pBang = factorial(p);
+    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;
+    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, dddHs] = hamGradHessHelper(z, Js...);
-    H += Hs;
-    dH += dHs;
-    ddH += ddHs;
-    dddH += dddHs;
-  }
+    if constexpr (sizeof...(Js) > 0) {
+      auto [Hs, dHs, ddHs, dddHs] = hamGradHessHelper(z, Js...);
+      H += Hs;
+      dH += dHs;
+      ddH += ddHs;
+      dddH += dddHs;
+    }
 
-  return {H, dH, ddH, dddH};
-}
+    return {H, dH, ddH, dddH};
+  }
 
-template <typename Scalar, int... ps>
-class pSpinModel {
 public:
   std::tuple<Tensor<Scalar, ps>...> Js;
 
+  pSpinModel() {}
+
+  pSpinModel(const std::tuple<Tensor<Scalar, ps>...>& Js) : Js(Js) {}
+
+  template <class Generator, typename... T>
+  pSpinModel<Real>(unsigned N, Generator& r, T... μs) {
+    Js = std::make_tuple(μs * generateRealPSpinCouplings<Real, ps>(N, r)...);
+  }
+
   unsigned dimension() const {
     return std::get<0>(Js).dimension(0);
   }
 
   template <typename NewScalar>
   pSpinModel<NewScalar, ps...> cast() const {
-    pSpinModel<NewScalar, ps...> M;
-    M.Js = std::apply(
+    return std::apply(
         [] (const Tensor<Scalar, ps>& ...Ks) -> std::tuple<Tensor<NewScalar, ps>...> {
           return std::make_tuple(Ks.template cast<NewScalar>()...);
         }, Js
-        );
+      );
+  }
+
+  template <typename T>
+  pSpinModel<Scalar, ps...>& operator*=(T x) {
+    std::tuple<Tensor<Scalar, ps>...> newJs = std::apply(
+        [x] (const Tensor<Scalar, ps>& ...Ks) -> std::tuple<Tensor<Scalar, ps>...> {
+          return std::make_tuple((x * Ks)...);
+        }, Js
+      );
+
+    Js = newJs;
 
-    return M;
+    return *this;
   }
 
   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);
+    return std::apply([&z, this](const Tensor<Scalar, ps>& ...Ks) -> std::tuple<Scalar, Vector<Scalar>, Matrix<Scalar>, Tensor<Scalar, 3>> { return hamGradHessHelper(z, Ks...); }, Js);
   }
 
   Scalar getHamiltonian(const Vector<Scalar>& z) const {