From b91b2d694a2bf25b02e0323ac6140db75beaba30 Mon Sep 17 00:00:00 2001
From: Philip Fackler <facklerpw@ornl.gov>
Date: Mon, 25 Sep 2023 15:54:54 -0400
Subject: [PATCH] Add new bits to RotatedSPOsT

---
 src/QMCWaveFunctions/RotatedSPOsT.cpp | 146 ++++++++++++++++++++++++++
 src/QMCWaveFunctions/RotatedSPOsT.h   |  70 ++++++++++++
 2 files changed, 216 insertions(+)

diff --git a/src/QMCWaveFunctions/RotatedSPOsT.cpp b/src/QMCWaveFunctions/RotatedSPOsT.cpp
index 128bca9798..dabdc282a9 100644
--- a/src/QMCWaveFunctions/RotatedSPOsT.cpp
+++ b/src/QMCWaveFunctions/RotatedSPOsT.cpp
@@ -1688,6 +1688,152 @@ RotatedSPOsT<T>::makeClone() const
     return myclone;
 }
 
+template <typename T>
+void
+RotatedSPOsT<T>::mw_evaluateDetRatios(
+    const RefVectorWithLeader<SPOSetT<T>>& spo_list,
+    const RefVectorWithLeader<const VirtualParticleSetT<T>>& vp_list,
+    const RefVector<ValueVector>& psi_list,
+    const std::vector<const ValueType*>& invRow_ptr_list,
+    std::vector<std::vector<ValueType>>& ratios_list) const
+{
+    auto phi_list = extractPhiRefList(spo_list);
+    auto& leader = phi_list.getLeader();
+    leader.mw_evaluateDetRatios(
+        phi_list, vp_list, psi_list, invRow_ptr_list, ratios_list);
+}
+
+template <typename T>
+void
+RotatedSPOsT<T>::mw_evaluateValue(
+    const RefVectorWithLeader<SPOSetT<T>>& spo_list,
+    const RefVectorWithLeader<ParticleSetT<T>>& P_list, int iat,
+    const RefVector<ValueVector>& psi_v_list) const
+{
+    auto phi_list = extractPhiRefList(spo_list);
+    auto& leader = phi_list.getLeader();
+    leader.mw_evaluateValue(phi_list, P_list, iat, psi_v_list);
+}
+
+template <typename T>
+void
+RotatedSPOsT<T>::mw_evaluateVGL(const RefVectorWithLeader<SPOSetT<T>>& spo_list,
+    const RefVectorWithLeader<ParticleSetT<T>>& P_list, int iat,
+    const RefVector<ValueVector>& psi_v_list,
+    const RefVector<GradVector>& dpsi_v_list,
+    const RefVector<ValueVector>& d2psi_v_list) const
+{
+    auto phi_list = extractPhiRefList(spo_list);
+    auto& leader = phi_list.getLeader();
+    leader.mw_evaluateVGL(
+        phi_list, P_list, iat, psi_v_list, dpsi_v_list, d2psi_v_list);
+}
+
+template <typename T>
+void
+RotatedSPOsT<T>::mw_evaluateVGLWithSpin(
+    const RefVectorWithLeader<SPOSetT<T>>& spo_list,
+    const RefVectorWithLeader<ParticleSetT<T>>& P_list, int iat,
+    const RefVector<ValueVector>& psi_v_list,
+    const RefVector<GradVector>& dpsi_v_list,
+    const RefVector<ValueVector>& d2psi_v_list,
+    OffloadMatrix<ComplexType>& mw_dspin) const
+{
+    auto phi_list = extractPhiRefList(spo_list);
+    auto& leader = phi_list.getLeader();
+    leader.mw_evaluateVGLWithSpin(
+        phi_list, P_list, iat, psi_v_list, dpsi_v_list, d2psi_v_list, mw_dspin);
+}
+
+template <typename T>
+void
+RotatedSPOsT<T>::mw_evaluateVGLandDetRatioGrads(
+    const RefVectorWithLeader<SPOSetT<T>>& spo_list,
+    const RefVectorWithLeader<ParticleSetT<T>>& P_list, int iat,
+    const std::vector<const ValueType*>& invRow_ptr_list,
+    OffloadMWVGLArray& phi_vgl_v, std::vector<ValueType>& ratios,
+    std::vector<GradType>& grads) const
+{
+    auto phi_list = extractPhiRefList(spo_list);
+    auto& leader = phi_list.getLeader();
+    leader.mw_evaluateVGLandDetRatioGrads(
+        phi_list, P_list, iat, invRow_ptr_list, phi_vgl_v, ratios, grads);
+}
+
+template <typename T>
+void
+RotatedSPOsT<T>::mw_evaluateVGLandDetRatioGradsWithSpin(
+    const RefVectorWithLeader<SPOSetT<T>>& spo_list,
+    const RefVectorWithLeader<ParticleSetT<T>>& P_list, int iat,
+    const std::vector<const ValueType*>& invRow_ptr_list,
+    OffloadMWVGLArray& phi_vgl_v, std::vector<ValueType>& ratios,
+    std::vector<GradType>& grads, std::vector<ValueType>& spingrads) const
+{
+    auto phi_list = extractPhiRefList(spo_list);
+    auto& leader = phi_list.getLeader();
+    leader.mw_evaluateVGLandDetRatioGradsWithSpin(phi_list, P_list, iat,
+        invRow_ptr_list, phi_vgl_v, ratios, grads, spingrads);
+}
+
+template <typename T>
+void
+RotatedSPOsT<T>::mw_evaluate_notranspose(
+    const RefVectorWithLeader<SPOSetT<T>>& spo_list,
+    const RefVectorWithLeader<ParticleSetT<T>>& P_list, int first, int last,
+    const RefVector<ValueMatrix>& logdet_list,
+    const RefVector<GradMatrix>& dlogdet_list,
+    const RefVector<ValueMatrix>& d2logdet_list) const
+{
+    auto phi_list = extractPhiRefList(spo_list);
+    auto& leader = phi_list.getLeader();
+    leader.mw_evaluate_notranspose(phi_list, P_list, first, last, logdet_list,
+        dlogdet_list, d2logdet_list);
+}
+
+template <typename T>
+void
+RotatedSPOsT<T>::createResource(ResourceCollection& collection) const
+{
+    Phi->createResource(collection);
+}
+
+template <typename T>
+void
+RotatedSPOsT<T>::acquireResource(ResourceCollection& collection,
+    const RefVectorWithLeader<SPOSetT<T>>& spo_list) const
+{
+    auto phi_list = extractPhiRefList(spo_list);
+    auto& leader = phi_list.getLeader();
+    leader.acquireResource(collection, phi_list);
+}
+
+template <typename T>
+void
+RotatedSPOsT<T>::releaseResource(ResourceCollection& collection,
+    const RefVectorWithLeader<SPOSetT<T>>& spo_list) const
+{
+    auto phi_list = extractPhiRefList(spo_list);
+    auto& leader = phi_list.getLeader();
+    leader.releaseResource(collection, phi_list);
+}
+
+template <typename T>
+RefVectorWithLeader<SPOSetT<T>>
+RotatedSPOsT<T>::extractPhiRefList(
+    const RefVectorWithLeader<SPOSetT<T>>& spo_list)
+{
+    auto& spo_leader = spo_list.template getCastedLeader<RotatedSPOsT>();
+    const auto nw = spo_list.size();
+    RefVectorWithLeader<SPOSetT<T>> phi_list(*spo_leader.Phi);
+    phi_list.reserve(nw);
+    for (int iw = 0; iw < nw; iw++) {
+        RotatedSPOsT& rot =
+            spo_list.template getCastedElement<RotatedSPOsT>(iw);
+        phi_list.emplace_back(*rot.Phi);
+    }
+    return phi_list;
+}
+
 // Class concrete types from ValueType
 template class RotatedSPOsT<double>;
 template class RotatedSPOsT<float>;
diff --git a/src/QMCWaveFunctions/RotatedSPOsT.h b/src/QMCWaveFunctions/RotatedSPOsT.h
index 971d2528b3..fa4778a6f4 100644
--- a/src/QMCWaveFunctions/RotatedSPOsT.h
+++ b/src/QMCWaveFunctions/RotatedSPOsT.h
@@ -40,6 +40,8 @@ class RotatedSPOsT : public SPOSetT<T>, public OptimizableObjectT<T>
     using IndexType = typename SPOSetT<T>::IndexType;
     using RealType = typename SPOSetT<T>::RealType;
     using ValueType = typename SPOSetT<T>::ValueType;
+    using GradType = typename SPOSetT<T>::GradType;
+    using ComplexType = typename SPOSetT<T>::ComplexType;
     using FullRealType = typename SPOSetT<T>::FullRealType;
     using ValueVector = typename SPOSetT<T>::ValueVector;
     using ValueMatrix = typename SPOSetT<T>::ValueMatrix;
@@ -49,6 +51,9 @@ class RotatedSPOsT : public SPOSetT<T>, public OptimizableObjectT<T>
     using HessMatrix = typename SPOSetT<T>::HessMatrix;
     using GGGVector = typename SPOSetT<T>::GGGVector;
     using GGGMatrix = typename SPOSetT<T>::GGGMatrix;
+    using OffloadMWVGLArray = typename SPOSetT<T>::OffloadMWVGLArray;
+    template <typename DT>
+    using OffloadMatrix = Matrix<DT, OffloadPinnedAllocator<DT>>;
 
     // constructor
     RotatedSPOsT(
@@ -399,6 +404,68 @@ class RotatedSPOsT : public SPOSetT<T>, public OptimizableObjectT<T>
         use_global_rot_ = use_global_rotation;
     }
 
+    void
+    mw_evaluateDetRatios(const RefVectorWithLeader<SPOSetT<T>>& spo_list,
+        const RefVectorWithLeader<const VirtualParticleSetT<T>>& vp_list,
+        const RefVector<ValueVector>& psi_list,
+        const std::vector<const ValueType*>& invRow_ptr_list,
+        std::vector<std::vector<ValueType>>& ratios_list) const override;
+
+    void
+    mw_evaluateValue(const RefVectorWithLeader<SPOSetT<T>>& spo_list,
+        const RefVectorWithLeader<ParticleSetT<T>>& P_list, int iat,
+        const RefVector<ValueVector>& psi_v_list) const override;
+
+    void
+    mw_evaluateVGL(const RefVectorWithLeader<SPOSetT<T>>& spo_list,
+        const RefVectorWithLeader<ParticleSetT<T>>& P_list, int iat,
+        const RefVector<ValueVector>& psi_v_list,
+        const RefVector<GradVector>& dpsi_v_list,
+        const RefVector<ValueVector>& d2psi_v_list) const override;
+
+    void
+    mw_evaluateVGLWithSpin(const RefVectorWithLeader<SPOSetT<T>>& spo_list,
+        const RefVectorWithLeader<ParticleSetT<T>>& P_list, int iat,
+        const RefVector<ValueVector>& psi_v_list,
+        const RefVector<GradVector>& dpsi_v_list,
+        const RefVector<ValueVector>& d2psi_v_list,
+        OffloadMatrix<ComplexType>& mw_dspin) const override;
+
+    void
+    mw_evaluateVGLandDetRatioGrads(
+        const RefVectorWithLeader<SPOSetT<T>>& spo_list,
+        const RefVectorWithLeader<ParticleSetT<T>>& P_list, int iat,
+        const std::vector<const ValueType*>& invRow_ptr_list,
+        OffloadMWVGLArray& phi_vgl_v, std::vector<ValueType>& ratios,
+        std::vector<GradType>& grads) const override;
+
+    void
+    mw_evaluateVGLandDetRatioGradsWithSpin(
+        const RefVectorWithLeader<SPOSetT<T>>& spo_list,
+        const RefVectorWithLeader<ParticleSetT<T>>& P_list, int iat,
+        const std::vector<const ValueType*>& invRow_ptr_list,
+        OffloadMWVGLArray& phi_vgl_v, std::vector<ValueType>& ratios,
+        std::vector<GradType>& grads,
+        std::vector<ValueType>& spingrads) const override;
+
+    void
+    mw_evaluate_notranspose(const RefVectorWithLeader<SPOSetT<T>>& spo_list,
+        const RefVectorWithLeader<ParticleSetT<T>>& P_list, int first, int last,
+        const RefVector<ValueMatrix>& logdet_list,
+        const RefVector<GradMatrix>& dlogdet_list,
+        const RefVector<ValueMatrix>& d2logdet_list) const override;
+
+    void
+    createResource(ResourceCollection& collection) const override;
+
+    void
+    acquireResource(ResourceCollection& collection,
+        const RefVectorWithLeader<SPOSetT<T>>& spo_list) const override;
+
+    void
+    releaseResource(ResourceCollection& collection,
+        const RefVectorWithLeader<SPOSetT<T>>& spo_list) const override;
+
 private:
     /// true if SPO parameters (orbital rotation parameters) have been supplied
     /// by input
@@ -415,6 +482,9 @@ class RotatedSPOsT : public SPOSetT<T>, public OptimizableObjectT<T>
     /// Use global rotation or history list
     bool use_global_rot_ = true;
 
+    static RefVectorWithLeader<SPOSetT<T>>
+    extractPhiRefList(const RefVectorWithLeader<SPOSetT<T>>& spo_list);
+
     friend OptVariablesType<double>&
     testing::getMyVarsFull(RotatedSPOsT<double>& rot);
     friend OptVariablesType<float>&