Apply the same cleanup to other Jastrow.

This commit is contained in:
Ye Luo 2022-01-26 09:56:30 -06:00
parent a26da86ce1
commit 8c0899b1b4
3 changed files with 76 additions and 120 deletions

View File

@ -42,6 +42,10 @@ struct J1OrbitalSoA : public WaveFunctionComponent
///use the same container
using DistRow = DistanceTable::DistRow;
using DisplRow = DistanceTable::DisplRow;
using GradDerivVec = ParticleAttrib<QTFull::GradType>;
using ValueDerivVec = ParticleAttrib<QTFull::ValueType>;
///table index
const int myTableID;
///number of ions
@ -55,8 +59,6 @@ struct J1OrbitalSoA : public WaveFunctionComponent
///reference to the sources (ions)
const ParticleSet& Ions;
///number of variables this object handles
int NumVars;
///variables handled by this orbital
opt_variables_type myVars;
@ -78,10 +80,15 @@ struct J1OrbitalSoA : public WaveFunctionComponent
std::vector<std::pair<int, int>> OffSet;
Vector<RealType> dLogPsi;
using WavefunctionFirstDerivativeType = ParticleAttrib<QTFull::GradType>;
using WavefunctionSecondDerivativeType = ParticleAttrib<QTFull::ValueType>;
std::vector<WavefunctionFirstDerivativeType*> gradLogPsi;
std::vector<WavefunctionSecondDerivativeType*> lapLogPsi;
std::vector<GradDerivVec> gradLogPsi;
std::vector<ValueDerivVec> lapLogPsi;
void resizeWFOptVectors()
{
dLogPsi.resize(myVars.size());
gradLogPsi.resize(myVars.size(), GradDerivVec(Nelec));
lapLogPsi.resize(myVars.size(), ValueDerivVec(Nelec));
}
J1OrbitalSoA(const std::string& obj_name, const ParticleSet& ions, ParticleSet& els)
: WaveFunctionComponent("J1OrbitalSoA", obj_name),
@ -89,8 +96,7 @@ struct J1OrbitalSoA : public WaveFunctionComponent
Nions(ions.getTotalNum()),
Nelec(els.getTotalNum()),
NumGroups(determineNumGroups(ions)),
Ions(ions),
NumVars(0)
Ions(ions)
{
if (myName.empty())
throw std::runtime_error("J1OrbitalSoA object name cannot be empty!");
@ -99,12 +105,6 @@ struct J1OrbitalSoA : public WaveFunctionComponent
J1OrbitalSoA(const J1OrbitalSoA& rhs) = delete;
~J1OrbitalSoA() override
{
delete_iter(gradLogPsi.begin(), gradLogPsi.end());
delete_iter(lapLogPsi.begin(), lapLogPsi.end());
}
/* determine NumGroups which controls the use of optimized code path using ion groups or not */
static int determineNumGroups(const ParticleSet& ions)
{
@ -230,7 +230,7 @@ struct J1OrbitalSoA : public WaveFunctionComponent
continue;
if (rcsingles[k])
{
dhpsioverpsi[kk] = -RealType(0.5) * ValueType(Sum(*lapLogPsi[k])) - ValueType(Dot(P.G, *gradLogPsi[k]));
dhpsioverpsi[kk] = -RealType(0.5) * ValueType(Sum(lapLogPsi[k])) - ValueType(Dot(P.G, gradLogPsi[k]));
}
}
}
@ -238,6 +238,7 @@ struct J1OrbitalSoA : public WaveFunctionComponent
void evaluateDerivativesWF(ParticleSet& P, const opt_variables_type& active, std::vector<ValueType>& dlogpsi) override
{
resizeWFOptVectors();
bool recalculate(false);
std::vector<bool> rcsingles(myVars.size(), false);
for (int k = 0; k < myVars.size(); ++k)
@ -253,10 +254,12 @@ struct J1OrbitalSoA : public WaveFunctionComponent
{
const auto& d_table = P.getDistTableAB(myTableID);
dLogPsi = 0.0;
const size_t NumVars = myVars.size();
for (int p = 0; p < NumVars; ++p)
(*gradLogPsi[p]) = 0.0;
for (int p = 0; p < NumVars; ++p)
(*lapLogPsi[p]) = 0.0;
{
gradLogPsi[p] = 0.0;
lapLogPsi[p] = 0.0;
}
std::vector<TinyVector<RealType, 3>> derivs(NumVars);
constexpr RealType cone(1);
@ -294,8 +297,8 @@ struct J1OrbitalSoA : public WaveFunctionComponent
{
dLogPsi[p] -= derivs[ip][0];
RealType dudr(rinv * derivs[ip][1]);
(*gradLogPsi[p])[j] -= dudr * dr;
(*lapLogPsi[p])[j] -= derivs[ip][2] + lapfac * dudr;
gradLogPsi[p][j] -= dudr * dr;
lapLogPsi[p][j] -= derivs[ip][2] + lapfac * dudr;
}
}
}
@ -514,23 +517,6 @@ struct J1OrbitalSoA : public WaveFunctionComponent
Lap.attachReference(buf.lendReference<valT>(Nelec), Nelec);
}
inline void setVars(const opt_variables_type& vars)
{
NumVars = vars.size();
if (NumVars == 0)
return;
myVars = vars;
dLogPsi.resize(NumVars);
gradLogPsi.resize(NumVars, 0);
lapLogPsi.resize(NumVars, 0);
for (int i = 0; i < NumVars; ++i)
{
gradLogPsi[i] = new WavefunctionFirstDerivativeType(Nelec);
lapLogPsi[i] = new WavefunctionSecondDerivativeType(Nelec);
}
}
std::unique_ptr<WaveFunctionComponent> makeClone(ParticleSet& tqp) const override
{
auto j1copy = std::make_unique<J1OrbitalSoA<FT>>(myName, Ions, tqp);
@ -543,7 +529,7 @@ struct J1OrbitalSoA : public WaveFunctionComponent
j1copy->addFunc(i, std::move(fc));
}
}
j1copy->setVars(myVars);
j1copy->myVars = myVars;
j1copy->OffSet = OffSet;
return j1copy;
}
@ -582,17 +568,9 @@ struct J1OrbitalSoA : public WaveFunctionComponent
}
}
myVars.getIndex(active);
NumVars = myVars.size();
if (NumVars && dLogPsi.size() == 0)
const size_t NumVars = myVars.size();
if (NumVars)
{
dLogPsi.resize(NumVars);
gradLogPsi.resize(NumVars, 0);
lapLogPsi.resize(NumVars, 0);
for (int i = 0; i < NumVars; ++i)
{
gradLogPsi[i] = new WavefunctionFirstDerivativeType(Nelec);
lapLogPsi[i] = new WavefunctionSecondDerivativeType(Nelec);
}
OffSet.resize(J1Functors.size());
int varoffset = myVars.Index[0];
for (int i = 0; i < J1Functors.size(); ++i)

View File

@ -42,6 +42,10 @@ struct J1Spin : public WaveFunctionComponent
///use the same container
using DistRow = DistanceTable::DistRow;
using DisplRow = DistanceTable::DisplRow;
using GradDerivVec = ParticleAttrib<QTFull::GradType>;
using ValueDerivVec = ParticleAttrib<QTFull::ValueType>;
///table index
const int myTableID;
///number of ions
@ -57,8 +61,6 @@ struct J1Spin : public WaveFunctionComponent
///reference to the sources (ions)
const ParticleSet& Ions;
///number of variables this object handles
int NumVars;
///variables handled by this orbital
opt_variables_type myVars;
@ -78,10 +80,15 @@ struct J1Spin : public WaveFunctionComponent
std::vector<std::pair<int, int>> OffSet;
Vector<RealType> dLogPsi;
using WavefunctionFirstDerivativeType = ParticleAttrib<QTFull::GradType>;
using WavefunctionSecondDerivativeType = ParticleAttrib<QTFull::ValueType>;
std::vector<WavefunctionFirstDerivativeType> gradLogPsi;
std::vector<WavefunctionSecondDerivativeType> lapLogPsi;
std::vector<GradDerivVec> gradLogPsi;
std::vector<ValueDerivVec> lapLogPsi;
void resizeWFOptVectors()
{
dLogPsi.resize(myVars.size());
gradLogPsi.resize(myVars.size(), GradDerivVec(Nelec));
lapLogPsi.resize(myVars.size(), ValueDerivVec(Nelec));
}
J1Spin(const std::string& obj_name, const ParticleSet& ions, ParticleSet& els)
: WaveFunctionComponent("J1Spin", obj_name),
@ -90,8 +97,7 @@ struct J1Spin : public WaveFunctionComponent
Nelec(els.getTotalNum()),
NumGroups(determineNumGroups(ions)),
NumTargetGroups(determineNumGroups(els)),
Ions(ions),
NumVars(0)
Ions(ions)
{
if (myName.empty())
throw std::runtime_error("J1Spin object name cannot be empty!");
@ -107,8 +113,7 @@ struct J1Spin : public WaveFunctionComponent
Nelec(rhs.Nelec),
NumGroups(rhs.NumGroups),
NumTargetGroups(rhs.NumTargetGroups),
Ions(rhs.Ions),
NumVars(0)
Ions(rhs.Ions)
{
Optimizable = rhs.Optimizable;
initialize(tqp);
@ -119,7 +124,7 @@ struct J1Spin : public WaveFunctionComponent
auto fc = std::make_unique<FT>(*rhs.J1UniqueFunctors[i * NumTargetGroups + j].get());
addFunc(i, std::move(fc), j);
}
setVars(rhs.myVars);
myVars = rhs.myVars;
OffSet = rhs.OffSet;
}
@ -274,6 +279,8 @@ struct J1Spin : public WaveFunctionComponent
void evaluateDerivativesWF(ParticleSet& P, const opt_variables_type& active, std::vector<ValueType>& dlogpsi) override
{
resizeWFOptVectors();
bool recalculate(false);
std::vector<bool> rcsingles(myVars.size(), false);
for (int k = 0; k < myVars.size(); ++k)
@ -289,10 +296,12 @@ struct J1Spin : public WaveFunctionComponent
{
const auto& d_table = P.getDistTableAB(myTableID);
dLogPsi = 0.0;
const size_t NumVars = myVars.size();
for (int p = 0; p < NumVars; ++p)
{
gradLogPsi[p] = 0.0;
for (int p = 0; p < NumVars; ++p)
lapLogPsi[p] = 0.0;
}
std::vector<TinyVector<RealType, 3>> derivs(NumVars);
constexpr RealType cone(1);
@ -566,23 +575,6 @@ struct J1Spin : public WaveFunctionComponent
Lap.attachReference(buf.lendReference<valT>(Nelec), Nelec);
}
inline void setVars(const opt_variables_type& vars)
{
NumVars = vars.size();
if (NumVars == 0)
return;
myVars = vars;
dLogPsi.resize(NumVars);
gradLogPsi.resize(NumVars);
lapLogPsi.resize(NumVars);
for (int i = 0; i < NumVars; ++i)
{
gradLogPsi[i].resize(Nelec);
lapLogPsi[i].resize(Nelec);
}
}
std::unique_ptr<WaveFunctionComponent> makeClone(ParticleSet& tqp) const override
{
auto cloned_J1Spin = std::make_unique<J1Spin<FT>>(*this, tqp);
@ -621,18 +613,10 @@ struct J1Spin : public WaveFunctionComponent
}
}
myVars.getIndex(active);
NumVars = myVars.size();
const size_t NumVars = myVars.size();
myVars.print(std::cout);
if (NumVars && dLogPsi.size() == 0)
if (NumVars)
{
dLogPsi.resize(NumVars);
gradLogPsi.resize(NumVars);
lapLogPsi.resize(NumVars);
for (int i = 0; i < NumVars; ++i)
{
gradLogPsi[i].resize(Nelec);
lapLogPsi[i].resize(Nelec);
}
OffSet.resize(J1UniqueFunctors.size());
// Find first active variable for the starting offset
int varoffset = -1;

View File

@ -94,7 +94,6 @@ class JeeIOrbitalSoA : public WaveFunctionComponent
VectorSoaContainer<valT, 9> mVGL;
// Used for evaluating derivatives with respect to the parameters
int NumVars;
Array<std::pair<int, int>, 3> VarOffset;
Vector<RealType> dLogPsi;
Array<PosType, 2> gradLogPsi;
@ -107,6 +106,29 @@ class JeeIOrbitalSoA : public WaveFunctionComponent
std::vector<std::vector<PosType>> dgrad_dalpha;
std::vector<std::vector<Tensor<RealType, 3>>> dhess_dalpha;
void resizeWFOptVectors()
{
dLogPsi.resize(myVars.size());
gradLogPsi.resize(myVars.size(), Nelec);
lapLogPsi.resize(myVars.size(), Nelec);
du_dalpha.resize(J3Unique.size());
dgrad_dalpha.resize(J3Unique.size());
dhess_dalpha.resize(J3Unique.size());
int ifunc = 0;
for (auto& j3UniquePair : J3Unique)
{
auto functorPtr = j3UniquePair.second.get();
J3UniqueIndex[functorPtr] = ifunc;
const int numParams = functorPtr->getNumParameters();
du_dalpha[ifunc].resize(numParams);
dgrad_dalpha[ifunc].resize(numParams);
dhess_dalpha[ifunc].resize(numParams);
ifunc++;
}
}
public:
///alias FuncType
using FuncType = FT;
@ -115,16 +137,13 @@ public:
: WaveFunctionComponent("JeeIOrbitalSoA", obj_name),
ee_Table_ID_(elecs.addTable(elecs, DTModes::NEED_TEMP_DATA_ON_HOST)),
ei_Table_ID_(elecs.addTable(ions, DTModes::NEED_FULL_TABLE_ANYTIME)),
Ions(ions),
NumVars(0)
Ions(ions)
{
if (myName.empty())
throw std::runtime_error("JeeIOrbitalSoA object name cannot be empty!");
init(elecs);
}
~JeeIOrbitalSoA() override {}
std::unique_ptr<WaveFunctionComponent> makeClone(ParticleSet& elecs) const override
{
auto eeIcopy = std::make_unique<JeeIOrbitalSoA<FT>>(myName, Ions, elecs, false);
@ -146,10 +165,6 @@ public:
// Ye: I don't like the following memory allocated by default.
eeIcopy->myVars.clear();
eeIcopy->myVars.insertFrom(myVars);
eeIcopy->NumVars = NumVars;
eeIcopy->dLogPsi.resize(NumVars);
eeIcopy->gradLogPsi.resize(NumVars, Nelec);
eeIcopy->lapLogPsi.resize(NumVars, Nelec);
eeIcopy->VarOffset = VarOffset;
eeIcopy->Optimizable = Optimizable;
return eeIcopy;
@ -195,25 +210,6 @@ public:
DistIndice_k.resize(Nbuffer);
}
void initUnique()
{
du_dalpha.resize(J3Unique.size());
dgrad_dalpha.resize(J3Unique.size());
dhess_dalpha.resize(J3Unique.size());
int ifunc = 0;
for (auto& j3UniquePair : J3Unique)
{
auto functorPtr = j3UniquePair.second.get();
J3UniqueIndex[functorPtr] = ifunc;
const int numParams = functorPtr->getNumParameters();
du_dalpha[ifunc].resize(numParams);
dgrad_dalpha[ifunc].resize(numParams);
dhess_dalpha[ifunc].resize(numParams);
ifunc++;
}
}
void addFunc(int iSpecies, int eSpecies1, int eSpecies2, std::unique_ptr<FT> j)
{
if (eSpecies1 == eSpecies2)
@ -246,7 +242,6 @@ public:
std::stringstream aname;
aname << iSpecies << "_" << eSpecies1 << "_" << eSpecies2;
J3Unique.emplace(aname.str(), std::move(j));
initUnique();
}
@ -331,12 +326,9 @@ public:
}
myVars.getIndex(active);
NumVars = myVars.size();
const size_t NumVars = myVars.size();
if (NumVars)
{
dLogPsi.resize(NumVars);
gradLogPsi.resize(NumVars, Nelec);
lapLogPsi.resize(NumVars, Nelec);
VarOffset.resize(iGroups, eGroups, eGroups);
int varoffset = myVars.Index[0];
for (int ig = 0; ig < iGroups; ig++)
@ -865,6 +857,8 @@ public:
std::vector<ValueType>& dlogpsi,
std::vector<ValueType>& dhpsioverpsi) override
{
resizeWFOptVectors();
bool recalculate(false);
std::vector<bool> rcsingles(myVars.size(), false);
for (int k = 0; k < myVars.size(); ++k)