Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Evaluate parameter derivatives for NLPP #4402

Merged
merged 10 commits into from
Jan 26, 2023
Merged

Evaluate parameter derivatives for NLPP #4402

Show file tree
Hide file tree
Changes from 9 commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
6163899
Add arguments to evaluateDerivRatios
markdewing Jan 18, 2023
02e73bf
Add parameter derivative functions
markdewing Jan 19, 2023
c1968ed
Add and update tests
markdewing Jan 19, 2023
fb2e6a8
Fix section name
markdewing Jan 19, 2023
a5737cb
Fix test name
markdewing Jan 19, 2023
5b24452
Use data stored in VirtualParticleSet
markdewing Jan 23, 2023
edb11bc
Enable NLPP parameter derivatives in batched code
markdewing Jan 23, 2023
e24eb3a
Enable Hamiltonian NLPP derivatives in batched code
markdewing Jan 23, 2023
6cf307a
Merge branch 'develop' into eval_derivs
markdewing Jan 23, 2023
f8930df
Merge branch 'develop' into eval_derivs
prckent Jan 26, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 54 additions & 3 deletions src/QMCHamiltonians/tests/test_RotatedSPOs_NLPP.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ namespace qmcplusplus
// Copied and extended from QMCWaveFunctions/tests/test_RotatedSPOs.cpp


TEST_CASE("RotatedSPOs SplineR2R hcpBe values multi det", "[wavefunction]")
void test_hcpBe_rotation(bool use_single_det, bool use_nlpp_batched)
{
using RealType = QMCTraits::RealType;

Expand Down Expand Up @@ -104,7 +104,7 @@ TEST_CASE("RotatedSPOs SplineR2R hcpBe values multi det", "[wavefunction]")
WaveFunctionPool wp(pp, c);
REQUIRE(wp.empty() == true);

const char* wf_input = R"(
const char* wf_input_multi_det = R"(
<wavefunction name="psi0" target="e">
<sposet_builder type="bspline" href="hcpBe.pwscf.h5" tilematrix="1 0 0 0 1 0 0 0 1" twistnum="0" source="ion0" version="0.10" meshfactor="1.0" precision="double" truncate="no" save_coefs="no">
<sposet type="bspline" name="spo_up" size="2" spindataset="0" optimize="yes">
Expand All @@ -121,6 +121,30 @@ TEST_CASE("RotatedSPOs SplineR2R hcpBe values multi det", "[wavefunction]")
</determinantset>
</wavefunction>)";

const char* wf_input_single_det = R"(
<wavefunction name="psi0" target="e">
<sposet_builder type="bspline" href="hcpBe.pwscf.h5" tilematrix="1 0 0 0 1 0 0 0 1" twistnum="0" source="ion0" version="0.10" meshfactor="1.0" precision="double" truncate="no" save_coefs="no">
<sposet type="bspline" name="spo_up" size="2" spindataset="0" optimize="yes">
</sposet>
<sposet type="bspline" name="spo_down" size="2" spindataset="0" optimize="yes">
</sposet>
</sposet_builder>
<determinantset>
<slaterdeterminant>
<determinant id="spo_up" size="1">
<occupation mode="ground" spindataset="0"/>
</determinant>
<determinant id="spo_down" size="1">
<occupation mode="ground" spindataset="0"/>
</determinant>
</slaterdeterminant>
</determinantset>
</wavefunction>)";

const char* wf_input = wf_input_multi_det;
if (use_single_det)
wf_input = wf_input_single_det;

Libxml2Document doc;
bool okay = doc.parseFromString(wf_input);
REQUIRE(okay);
Expand All @@ -134,13 +158,23 @@ TEST_CASE("RotatedSPOs SplineR2R hcpBe values multi det", "[wavefunction]")


// Note the pbc="no" setting to turn off long-range summation.
const char* ham_input = R"(
const char* ham_input_nlpp_nonbatched = R"(
<hamiltonian name="h0" type="generic" target="e">
<pairpot type="pseudo" name="PseudoPot" source="ion0" wavefunction="psi0" format="xml" algorithm="non-batched" pbc="no">
<pseudo elementType="Be" href="Be.BFD.xml" nrule="2" disable_randomize_grid="yes"/>
</pairpot>
</hamiltonian>)";

const char* ham_input_nlpp_batched = R"(
<hamiltonian name="h0" type="generic" target="e">
<pairpot type="pseudo" name="PseudoPot" source="ion0" wavefunction="psi0" format="xml" algorithm="batched" pbc="no">
<pseudo elementType="Be" href="Be.BFD.xml" nrule="2" disable_randomize_grid="yes"/>
</pairpot>
</hamiltonian>)";

const char* ham_input = ham_input_nlpp_nonbatched;
if (use_nlpp_batched)
ham_input = ham_input_nlpp_batched;

HamiltonianFactory hf("h0", elec, pp.getPool(), wp.getPool(), c);

Expand Down Expand Up @@ -240,4 +274,21 @@ TEST_CASE("RotatedSPOs SplineR2R hcpBe values multi det", "[wavefunction]")
CHECK(dhpsi_over_psi_list2[0][1] == Approx(dhpsioverpsi2[1]));
}

TEST_CASE("RotatedSPOs SplineR2R hcpBe values", "[wavefunction]")
{
SECTION("nlpp non-batched")
{
bool use_single_det = GENERATE(true, false);
bool use_nlpp_batched = false;
test_hcpBe_rotation(use_single_det, use_nlpp_batched);
}

SECTION("nlpp batched")
{
bool use_single_det = true;
bool use_nlpp_batched = true;
test_hcpBe_rotation(use_single_det, use_nlpp_batched);
}
}

} // namespace qmcplusplus
20 changes: 20 additions & 0 deletions src/QMCWaveFunctions/Fermion/DiracDeterminant.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -456,6 +456,18 @@ void DiracDeterminant<DU_TYPE>::mw_evaluateRatios(const RefVectorWithLeader<Wave
}
}

template<typename DU_TYPE>
void DiracDeterminant<DU_TYPE>::evaluateDerivRatios(const VirtualParticleSet& VP,
const opt_variables_type& optvars,
std::vector<ValueType>& ratios,
Matrix<ValueType>& dratios)
{
const int WorkingIndex = VP.refPtcl - FirstIndex;
assert(WorkingIndex >= 0);
std::copy_n(psiM[WorkingIndex], invRow.size(), invRow.data());
Phi->evaluateDerivRatios(VP, optvars, psiV, invRow, ratios, dratios, FirstIndex, LastIndex);
}

template<typename DU_TYPE>
void DiracDeterminant<DU_TYPE>::evaluateRatiosAlltoOne(ParticleSet& P, std::vector<ValueType>& ratios)
{
Expand Down Expand Up @@ -673,6 +685,14 @@ void DiracDeterminant<DU_TYPE>::evaluateDerivatives(ParticleSet& P,
Phi->evaluateDerivatives(P, active, dlogpsi, dhpsioverpsi, FirstIndex, LastIndex);
}

template<typename DU_TYPE>
void DiracDeterminant<DU_TYPE>::evaluateDerivativesWF(ParticleSet& P,
const opt_variables_type& active,
Vector<ValueType>& dlogpsi)
{
Phi->evaluateDerivativesWF(P, active, dlogpsi, FirstIndex, LastIndex);
}

template<typename DU_TYPE>
std::unique_ptr<DiracDeterminantBase> DiracDeterminant<DU_TYPE>::makeCopy(std::unique_ptr<SPOSet>&& spo) const
{
Expand Down
9 changes: 9 additions & 0 deletions src/QMCWaveFunctions/Fermion/DiracDeterminant.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,10 @@ class DiracDeterminant : public DiracDeterminantBase
Vector<ValueType>& dlogpsi,
Vector<ValueType>& dhpsioverpsi) override;

void evaluateDerivativesWF(ParticleSet& P,
const opt_variables_type& optvars,
Vector<ValueType>& dlogpsi) override;

void registerData(ParticleSet& P, WFBufferType& buf) override;

void updateAfterSweep(const ParticleSet& P, ParticleSet::ParticleGradient& G, ParticleSet::ParticleLaplacian& L);
Expand Down Expand Up @@ -109,6 +113,11 @@ class DiracDeterminant : public DiracDeterminantBase
const RefVectorWithLeader<const VirtualParticleSet>& vp_list,
std::vector<std::vector<ValueType>>& ratios) const override;

void evaluateDerivRatios(const VirtualParticleSet& VP,
const opt_variables_type& optvars,
std::vector<ValueType>& ratios,
Matrix<ValueType>& dratios) override;

PsiValueType ratioGrad(ParticleSet& P, int iat, GradType& grad_iat) override;

PsiValueType ratioGradWithSpin(ParticleSet& P, int iat, GradType& grad_iat, ComplexType& spingrad) final;
Expand Down
21 changes: 20 additions & 1 deletion src/QMCWaveFunctions/Fermion/DiracDeterminantBatched.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -815,6 +815,18 @@ void DiracDeterminantBatched<DET_ENGINE>::mw_evaluateRatios(
}
}

template<typename DET_ENGINE>
void DiracDeterminantBatched<DET_ENGINE>::evaluateDerivRatios(const VirtualParticleSet& VP,
const opt_variables_type& optvars,
std::vector<ValueType>& ratios,
Matrix<ValueType>& dratios)
{
const int WorkingIndex = VP.refPtcl - FirstIndex;
assert(WorkingIndex >= 0);
std::copy_n(det_engine_.get_psiMinv()[WorkingIndex], d2psiV.size(), d2psiV.data());
Phi->evaluateDerivRatios(VP, optvars, psiV_host_view, d2psiV_host_view, ratios, dratios, FirstIndex, LastIndex);
}

template<typename DET_ENGINE>
void DiracDeterminantBatched<DET_ENGINE>::evaluateRatiosAlltoOne(ParticleSet& P, std::vector<Value>& ratios)
{
Expand All @@ -827,7 +839,6 @@ void DiracDeterminantBatched<DET_ENGINE>::evaluateRatiosAlltoOne(ParticleSet& P,
ratios[FirstIndex + i] = simd::dot(psiMinv[i], psiV.data(), NumOrbitals);
}


template<typename DET_ENGINE>
void DiracDeterminantBatched<DET_ENGINE>::resizeScratchObjectsForIonDerivs()
{
Expand Down Expand Up @@ -1108,6 +1119,14 @@ void DiracDeterminantBatched<DET_ENGINE>::evaluateDerivatives(ParticleSet& P,
Phi->evaluateDerivatives(P, active, dlogpsi, dhpsioverpsi, FirstIndex, LastIndex);
}

template<typename DET_ENGINE>
void DiracDeterminantBatched<DET_ENGINE>::evaluateDerivativesWF(ParticleSet& P,
const opt_variables_type& active,
Vector<ValueType>& dlogpsi)
{
Phi->evaluateDerivativesWF(P, active, dlogpsi, FirstIndex, LastIndex);
}

template<typename DET_ENGINE>
void DiracDeterminantBatched<DET_ENGINE>::registerTWFFastDerivWrapper(const ParticleSet& P,
TWFFastDerivWrapper& twf) const
Expand Down
9 changes: 9 additions & 0 deletions src/QMCWaveFunctions/Fermion/DiracDeterminantBatched.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -103,6 +103,10 @@ class DiracDeterminantBatched : public DiracDeterminantBase
Vector<Value>& dlogpsi,
Vector<Value>& dhpsioverpsi) override;

void evaluateDerivativesWF(ParticleSet& P,
const opt_variables_type& optvars,
Vector<ValueType>& dlogpsi) override;

void registerData(ParticleSet& P, WFBufferType& buf) override;

LogValue updateBuffer(ParticleSet& P, WFBufferType& buf, bool fromscratch = false) override;
Expand All @@ -128,6 +132,11 @@ class DiracDeterminantBatched : public DiracDeterminantBase
const RefVectorWithLeader<const VirtualParticleSet>& vp_list,
std::vector<std::vector<Value>>& ratios) const override;

void evaluateDerivRatios(const VirtualParticleSet& VP,
const opt_variables_type& optvars,
std::vector<ValueType>& ratios,
Matrix<ValueType>& dratios) override;

PsiValue ratioGrad(ParticleSet& P, int iat, Grad& grad_iat) override;

void mw_ratioGrad(const RefVectorWithLeader<WaveFunctionComponent>& wfc_list,
Expand Down
8 changes: 8 additions & 0 deletions src/QMCWaveFunctions/Fermion/SlaterDet.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,14 @@ void SlaterDet::evaluateRatiosAlltoOne(ParticleSet& P, std::vector<ValueType>& r
Dets[i]->evaluateRatiosAlltoOne(P, ratios);
}

void SlaterDet::evaluateDerivRatios(const VirtualParticleSet& VP,
const opt_variables_type& optvars,
std::vector<ValueType>& ratios,
Matrix<ValueType>& dratios)
{
return Dets[getDetID(VP.refPtcl)]->evaluateDerivRatios(VP, optvars, ratios, dratios);
}

SlaterDet::LogValueType SlaterDet::evaluateLog(const ParticleSet& P,
ParticleSet::ParticleGradient& G,
ParticleSet::ParticleLaplacian& L)
Expand Down
5 changes: 5 additions & 0 deletions src/QMCWaveFunctions/Fermion/SlaterDet.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,11 @@ class SlaterDet : public WaveFunctionComponent
return Dets[getDetID(VP.refPtcl)]->evaluateRatios(VP, ratios);
}

void evaluateDerivRatios(const VirtualParticleSet& VP,
const opt_variables_type& optvars,
std::vector<ValueType>& ratios,
Matrix<ValueType>& dratios) override;

inline void mw_evaluateRatios(const RefVectorWithLeader<WaveFunctionComponent>& wfc_list,
const RefVectorWithLeader<const VirtualParticleSet>& vp_list,
std::vector<std::vector<ValueType>>& ratios) const override
Expand Down
131 changes: 128 additions & 3 deletions src/QMCWaveFunctions/RotatedSPOs.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -306,6 +306,131 @@ void RotatedSPOs::log_antisym_matrix(ValueMatrix& mat)
}
}

void RotatedSPOs::evaluateDerivRatios(const VirtualParticleSet& VP,
const opt_variables_type& optvars,
ValueVector& psi,
const ValueVector& psiinv,
std::vector<ValueType>& ratios,
Matrix<ValueType>& dratios,
int FirstIndex,
int LastIndex)
{
Phi->evaluateDetRatios(VP, psi, psiinv, ratios);

const size_t nel = LastIndex - FirstIndex;
const size_t nmo = Phi->getOrbitalSetSize();

psiM_inv.resize(nel, nel);
psiM_all.resize(nel, nmo);
dpsiM_all.resize(nel, nmo);
d2psiM_all.resize(nel, nmo);

psiM_inv = 0;
psiM_all = 0;
dpsiM_all = 0;
d2psiM_all = 0;

const ParticleSet& P = VP.getRefPS();
int iel = VP.refPtcl;

Phi->evaluate_notranspose(P, FirstIndex, LastIndex, psiM_all, dpsiM_all, d2psiM_all);

for (int i = 0; i < nel; i++)
for (int j = 0; j < nel; j++)
psiM_inv(i, j) = psiM_all(i, j);

Invert(psiM_inv.data(), nel, nel);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Noting the full cubic cost of matrix inversion here. However, getting something working first takes precedence over possible alternatives. Plus considering that the initial electron counts will be smallish, this might not be so bad.


const ValueType* const A(psiM_all.data());
const ValueType* const Ainv(psiM_inv.data());
SPOSet::ValueMatrix T_orig;
T_orig.resize(nel, nmo);

BLAS::gemm('N', 'N', nmo, nel, nel, ValueType(1.0), A, nmo, Ainv, nel, ValueType(0.0), T_orig.data(), nmo);

SPOSet::ValueMatrix T;
T.resize(nel, nmo);

ValueVector tmp_psi;
tmp_psi.resize(nmo);

for (int iat = 0; iat < VP.getTotalNum(); iat++)
{
Phi->evaluateValue(VP, iat, tmp_psi);

for (int j = 0; j < nmo; j++)
psiM_all(iel - FirstIndex, j) = tmp_psi[j];

for (int i = 0; i < nel; i++)
for (int j = 0; j < nel; j++)
psiM_inv(i, j) = psiM_all(i, j);

Invert(psiM_inv.data(), nel, nel);
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the inverse that is likely to get expensive, as it gets call nknot times.

Maybe combining a Sherman-Morrison update with the matrix multiply below will lead to some simplifications.

Copy link
Contributor

@ye-luo ye-luo Jan 26, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I believe the math can be formulated without any update.


const ValueType* const A(psiM_all.data());
const ValueType* const Ainv(psiM_inv.data());

// The matrix A is rectangular. Ainv is the inverse of the square part of the matrix.
// The multiply of Ainv and the square part of A is just the identity.
// This multiply could be reduced to Ainv and the non-square part of A.
BLAS::gemm('N', 'N', nmo, nel, nel, ValueType(1.0), A, nmo, Ainv, nel, ValueType(0.0), T.data(), nmo);

for (int i = 0; i < m_act_rot_inds.size(); i++)
{
int kk = myVars.where(i);
const int p = m_act_rot_inds.at(i).first;
const int q = m_act_rot_inds.at(i).second;
dratios(iat, kk) = T(p, q) - T_orig(p, q); // dratio size is (nknot, num_vars)
}
}
}

void RotatedSPOs::evaluateDerivativesWF(ParticleSet& P,
const opt_variables_type& optvars,
Vector<ValueType>& dlogpsi,
int FirstIndex,
int LastIndex)
{
const size_t nel = LastIndex - FirstIndex;
const size_t nmo = Phi->getOrbitalSetSize();

//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~PART1

psiM_inv.resize(nel, nel);
psiM_all.resize(nel, nmo);
dpsiM_all.resize(nel, nmo);
d2psiM_all.resize(nel, nmo);

psiM_inv = 0;
psiM_all = 0;
dpsiM_all = 0;
d2psiM_all = 0;

Phi->evaluate_notranspose(P, FirstIndex, LastIndex, psiM_all, dpsiM_all, d2psiM_all);

for (int i = 0; i < nel; i++)
for (int j = 0; j < nel; j++)
psiM_inv(i, j) = psiM_all(i, j);

Invert(psiM_inv.data(), nel, nel);

//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~PART2
const ValueType* const A(psiM_all.data());
const ValueType* const Ainv(psiM_inv.data());
SPOSet::ValueMatrix T;
T.resize(nel, nmo);

BLAS::gemm('N', 'N', nmo, nel, nel, ValueType(1.0), A, nmo, Ainv, nel, ValueType(0.0), T.data(), nmo);

for (int i = 0; i < m_act_rot_inds.size(); i++)
{
int kk = myVars.where(i);
const int p = m_act_rot_inds.at(i).first;
const int q = m_act_rot_inds.at(i).second;
dlogpsi[kk] = T(p, q);
}
}

void RotatedSPOs::evaluateDerivatives(ParticleSet& P,
const opt_variables_type& optvars,
Vector<ValueType>& dlogpsi,
Expand Down Expand Up @@ -401,9 +526,9 @@ void RotatedSPOs::evaluateDerivatives(ParticleSet& P,

for (int i = 0; i < m_act_rot_inds.size(); i++)
{
int kk = myVars.where(i);
const int p = m_act_rot_inds.at(i).first;
const int q = m_act_rot_inds.at(i).second;
int kk = myVars.where(i);
const int p = m_act_rot_inds.at(i).first;
const int q = m_act_rot_inds.at(i).second;
dlogpsi[kk] = T(p, q);
dhpsioverpsi[kk] = ValueType(-0.5) * Y4(p, q);
}
Expand Down
Loading