Merge pull request #4169 from ye-luo/support-continue-opt

Support multi step WFopt
QMCPACK · Aug 19, 2022 · 13d81e2 · 13d81e2
2 parents 4081d48 + 81a1297
commit 13d81e2
Show file tree

Hide file tree

Showing 12 changed files with 297 additions and 22 deletions.
diff --git a/docs/methods.rst b/docs/methods.rst
@@ -551,6 +551,31 @@ The cost function consists of three components: energy, unreweighted variance, a
      <cost name="unreweightedvariance">     0.00 </cost>
      <cost name="reweightedvariance">       0.05 </cost>
 
+Varational parameter selection
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+The predominant way of selecting varational paramemters is via ``<wavefunction>`` input.
+``<coefficients>`` entries support ``optimize="yes"/"no"`` to enable/disable variational parameters in the wavefunction optimization.
+The secondary way of selecting varational paramemters is via ``variational_subset`` parameter in the ``<qmc>`` driver input.
+It allows controlling optimization granularity at each optimization step.
+If ``variational_subset`` is not provided or empty, all the varational paramemters are selected.
+If variational paramemters are set as not optimizable in the predominant way, the secondary way won't be able to set them optimizable even they are selected.
+
+The following example shows optimizing subsets of parameters in stages in a single QMCPACK run.
+
+::
+    <qmc method="linear">
+      ...
+      <parameter name="variational_subset"> uu ud </parameter>
+    </qmc>
+    <qmc method="linear">
+      ...
+      <parameter name="variational_subset"> uu ud eH </parameter>
+    </qmc>
+    <qmc method="linear">
+      ...
+      <parameter name="variational_subset"> uu ud eH CI </parameter>
+    </qmc>
+
 Optimizers
 ~~~~~~~~~~
 

diff --git a/docs/running_docker.rst b/docs/running_docker.rst
diff --git a/src/QMCDrivers/WFOpt/QMCCostFunctionBase.cpp b/src/QMCDrivers/WFOpt/QMCCostFunctionBase.cpp
@@ -21,7 +21,7 @@
 #include "OhmmsData/ParameterSet.h"
 #include "OhmmsData/XMLParsingString.h"
 #include "Message/CommOperators.h"
-#include <set>
+#include "Message/UniformCommunicateError.h"
 //#define QMCCOSTFUNCTION_DEBUG
 
 
@@ -414,7 +414,7 @@ bool QMCCostFunctionBase::put(xmlNodePtr q)
     if (obj.isOptimized())
       obj.checkInVariablesExclusive(OptVariablesForPsi);
   OptVariablesForPsi.resetIndex();
-  app_log() << "In total " << OptVariablesForPsi.size() << " parameters being optimzied." << std::endl;
+  app_log() << " Variational subset selects " << OptVariablesForPsi.size() << " parameters." << std::endl;
 
   //synchronize OptVariables and OptVariablesForPsi
   OptVariables  = OptVariablesForPsi;
@@ -482,6 +482,9 @@ bool QMCCostFunctionBase::put(xmlNodePtr q)
   {
     APP_ABORT("QMCCostFunctionBase::put No valid optimizable variables are found.");
   }
+  else
+    app_log() << " In total " << NumOptimizables << " parameters being optimized after applying constraints."
+              << std::endl;
   //     app_log() << "<active-optimizables> " << std::endl;
   //     OptVariables.print(app_log());
   //     app_log() << "</active-optimizables>" << std::endl;
@@ -1066,17 +1069,29 @@ bool QMCCostFunctionBase::isEffectiveWeightValid(EffectiveWeight effective_weigh
 UniqueOptObjRefs QMCCostFunctionBase::extractOptimizableObjects(TrialWaveFunction& psi) const
 {
   const auto& names(variational_subset_names);
-  /// survey all the optimizable objects
+  // survey all the optimizable objects
   const auto opt_obj_refs = psi.extractOptimizableObjectRefs();
+  // check if input names are valid
+  for (auto& name : names)
+    if (std::find_if(opt_obj_refs.begin(), opt_obj_refs.end(),
+                     [&name](const OptimizableObject& obj) { return name == obj.getName(); }) == opt_obj_refs.end())
+    {
+      std::ostringstream msg;
+      msg << "Variational subset entry '" << name << "' doesn't exist in the trial wavefunction which contains";
+      for (OptimizableObject& obj : opt_obj_refs)
+        msg << " '" << obj.getName() << "'";
+      msg << "." << std::endl;
+      throw UniformCommunicateError(msg.str());
+    }
+
   for (OptimizableObject& obj : opt_obj_refs)
     obj.setOptimization(names.empty() || std::find_if(names.begin(), names.end(), [&obj](const std::string& name) {
                                            return name == obj.getName();
                                          }) != names.end());
   return opt_obj_refs;
 }
 
-void QMCCostFunctionBase::resetOptimizableObjects(TrialWaveFunction& psi,
-                                                  const opt_variables_type& opt_variables) const
+void QMCCostFunctionBase::resetOptimizableObjects(TrialWaveFunction& psi, const opt_variables_type& opt_variables) const
 {
   const auto opt_obj_refs = extractOptimizableObjects(psi);
   for (OptimizableObject& obj : opt_obj_refs)

diff --git a/src/QMCWaveFunctions/Fermion/MultiSlaterDetTableMethod.cpp b/src/QMCWaveFunctions/Fermion/MultiSlaterDetTableMethod.cpp
@@ -704,7 +704,10 @@ void MultiSlaterDetTableMethod::extractOptimizableObjectRefs(UniqueOptObjRefs& o
 void MultiSlaterDetTableMethod::checkInVariablesExclusive(opt_variables_type& active)
 {
   if (CI_Optimizable && myVars->size())
+  {
+    myVars->setIndexDefault();
     active.insertFrom(*myVars);
+  }
 }
 
 void MultiSlaterDetTableMethod::checkInVariables(opt_variables_type& active)

diff --git a/src/QMCWaveFunctions/Jastrow/BsplineFunctor.h b/src/QMCWaveFunctions/Jastrow/BsplineFunctor.h
@@ -682,6 +682,7 @@ struct BsplineFunctor : public OptimizableFunctorBase
   {
     if (notOpt)
       return;
+    myVars.setIndexDefault();
     active.insertFrom(myVars);
   }
 

diff --git a/src/QMCWaveFunctions/Jastrow/PolynomialFunctor3D.h b/src/QMCWaveFunctions/Jastrow/PolynomialFunctor3D.h
@@ -1070,9 +1070,21 @@ struct PolynomialFunctor3D : public OptimizableFunctorBase
     reset_gamma();
   }
 
-  void checkInVariables(opt_variables_type& active) override { active.insertFrom(myVars); }
+  void checkInVariables(opt_variables_type& active) override
+  {
+    if (notOpt)
+      return;
+
+    myVars.setIndexDefault();
+    active.insertFrom(myVars);
+  }
 
-  void checkOutVariables(const opt_variables_type& active) override { myVars.getIndex(active); }
+  void checkOutVariables(const opt_variables_type& active) override
+  {
+    if (notOpt)
+      return;
+    myVars.getIndex(active);
+  }
 
   void print(std::ostream& os)
   {

diff --git a/src/QMCWaveFunctions/VariableSet.cpp b/src/QMCWaveFunctions/VariableSet.cpp
@@ -173,10 +173,10 @@ void VariableSet::getIndex(const VariableSet& selected)
   }
 }
 
-void VariableSet::setDefaults(bool optimize_all)
+void VariableSet::setIndexDefault()
 {
   for (int i = 0; i < Index.size(); ++i)
-    Index[i] = optimize_all ? i : -1;
+    Index[i] = i;
 }
 
 void VariableSet::print(std::ostream& os, int leftPadSpaces, bool printHeader) const

diff --git a/src/QMCWaveFunctions/VariableSet.h b/src/QMCWaveFunctions/VariableSet.h
@@ -46,11 +46,11 @@ struct VariableSet
   using value_type = qmcplusplus::QMCTraits::ValueType;
   using real_type  = qmcplusplus::QMCTraits::RealType;
 
-  using pair_type         = std::pair<std::string, value_type>;
-  using index_pair_type   = std::pair<std::string, int>;
-  using iterator          = std::vector<pair_type>::iterator;
-  using const_iterator    = std::vector<pair_type>::const_iterator;
-  using size_type         = std::vector<pair_type>::size_type;
+  using pair_type       = std::pair<std::string, value_type>;
+  using index_pair_type = std::pair<std::string, int>;
+  using iterator        = std::vector<pair_type>::iterator;
+  using const_iterator  = std::vector<pair_type>::const_iterator;
+  using size_type       = std::vector<pair_type>::size_type;
 
   ///number of active variables
   int num_active_vars;
@@ -327,10 +327,9 @@ struct VariableSet
    */
   void getIndex(const VariableSet& selected);
 
-  /** set default Indices
-   * @param optimize_all if true, all the variables are active
+  /** set default Indices, namely all the variables are active
    */
-  void setDefaults(bool optimize_all);
+  void setIndexDefault();
 
   void print(std::ostream& os, int leftPadSpaces = 0, bool printHeader = false) const;
 

diff --git a/tests/molecules/H4_ae/CMakeLists.txt b/tests/molecules/H4_ae/CMakeLists.txt
@@ -315,6 +315,19 @@ if(NOT QMC_CUDA)
       H4_OPT_SCALARS # OPT step 5
     )
 
+    list(APPEND H4_OPT_SCALARS_ONLY_JAS "totenergy" "-2.066504 0.0009") # total energy
+    qmc_run_and_check(
+      short-H4-opt-OneShiftOnly-onlyjas
+      "${qmcpack_SOURCE_DIR}/tests/molecules/H4_ae"
+      H4-OneShiftOnly
+      optm-OneShiftOnly-onlyjas.xml
+      1
+      16
+      ${SUCCESS_STATUS_MP}
+      5
+      H4_OPT_SCALARS_ONLY_JAS # OPT step 5
+    )
+
     list(APPEND H4_OPT_SCALARS_ONLY_MSD "totenergy" "-2.138635 0.0009") # total energy
     qmc_run_and_check(
       short-H4-opt-OneShiftOnly-onlymsd
@@ -328,17 +341,26 @@ if(NOT QMC_CUDA)
       H4_OPT_SCALARS_ONLY_MSD # OPT step 5
     )
 
-    list(APPEND H4_OPT_SCALARS_ONLY_JAS "totenergy" "-2.066504 0.0009") # total energy
     qmc_run_and_check(
-      short-H4-opt-OneShiftOnly-onlyjas
+      short-H4-optbatch-OneShiftOnly-multistep
       "${qmcpack_SOURCE_DIR}/tests/molecules/H4_ae"
       H4-OneShiftOnly
-      optm-OneShiftOnly-onlyjas.xml
+      optm-OneShiftOnly-multistep.xml
       1
       16
       ${SUCCESS_STATUS_MP}
-      5
-      H4_OPT_SCALARS_ONLY_JAS # OPT step 5
+      11
+      H4_OPT_SCALARS_ONLY_MSD # OPT step 011
+    )
+
+    qmc_run_and_check(
+      deterministic-input_check-invalid_variational_subset
+      "${qmcpack_SOURCE_DIR}/tests/molecules/H4_ae"
+      H4-OneShiftOnly
+      det-input-check-optm-invalid-subset.xml
+      2
+      1
+      FALSE
     )
   else()
     message(VERBOSE

diff --git a/tests/molecules/H4_ae/H4.wfs_j123.xml b/tests/molecules/H4_ae/H4.wfs_j123.xml
@@ -0,0 +1,71 @@
+<?xml version="1.0"?>
+<qmcsystem>
+  <wavefunction name="psi0" target="e">
+    <determinantset type="MolecularOrbital" name="LCAOBSet" transform="yes" source="ion0">
+      <basisset name="LCAOBSet">
+        <atomicBasisSet name="Gaussian-G2" angular="cartesian" type="Gaussian" elementType="H" normalized="no">
+          <grid type="log" ri="1.e-6" rf="1.e2" npts="1001"/>
+          <!-- Possible substitution H00 by a Slater-type orbital
+  <basisGroup rid="H00" n="0" l="0" type="Slater">
+    <radfunc exponent="1.31826" contraction="1.0"/>
+  </basisGroup>
+-->
+          <basisGroup rid="H00" n="0" l="0" type="Gaussian">
+            <radfunc exponent="3.425250900000e+00" contraction="1.543289672950e-01"/>
+            <radfunc exponent="6.239137000000e-01" contraction="5.353281422820e-01"/>
+            <radfunc exponent="1.688554000000e-01" contraction="4.446345421850e-01"/>
+          </basisGroup>
+        </atomicBasisSet>
+      </basisset>
+      <sposet basisset="LCAOBSet" name="spo-up" size="3">
+        <occupation mode="ground"/>
+        <coefficient size="4" id="updetC">
+  5.04668000000000e-01  4.50114000000000e-01  3.58423000000000e-01  1.26449000000000e-01
+ -2.40452000000000e-01 -3.20413000000000e-01  4.66777000000000e-01  7.03549000000000e-01
+  8.80080000000000e-02 -5.04842000000000e-01  8.07522000000000e-01 -7.19301000000000e-01
+  1.03323500000000e+00 -8.77213000000000e-01 -3.90492000000000e-01  2.12318000000000e-01
+</coefficient>
+      </sposet>
+      <sposet basisset="LCAOBSet" name="spo-dn" size="3">
+        <occupation mode="ground"/>
+        <coefficient size="4" id="downdetC">
+  5.04668000000000e-01  4.50114000000000e-01  3.58423000000000e-01  1.26449000000000e-01
+ -2.40452000000000e-01 -3.20413000000000e-01  4.66777000000000e-01  7.03549000000000e-01
+  8.80080000000000e-02 -5.04842000000000e-01  8.07522000000000e-01 -7.19301000000000e-01
+  1.03323500000000e+00 -8.77213000000000e-01 -3.90492000000000e-01  2.12318000000000e-01
+</coefficient>
+      </sposet>
+      <multideterminant optimize="yes" spo_up="spo-up" spo_dn="spo-dn">
+        <detlist size="2" type="CSF" nca="0" ncb="0" nea="2" neb="2" nstates="3" cutoff="0.2">
+          <csf id="CSFcoeff_0" exctLvl="0" coeff="0.605818" qchem_coeff="0.955818" occ="220">
+            <det id="csf_0-0" coeff="1" alpha="110" beta="110"/>
+          </csf>
+          <csf id="CSFcoeff_1" exctLvl="2" coeff="-5.2056644665606530e-01" qchem_coeff="-0.276685" occ="202">
+            <det id="csf_1-0" coeff="1" alpha="101" beta="101"/>
+          </csf>
+        </detlist>
+      </multideterminant>
+    </determinantset>
+    <jastrow name="J2" type="Two-Body" function="Bspline" print="yes">
+      <correlation rcut="10" size="10" speciesA="u" speciesB="u">
+        <coefficients id="uu" type="Array"> 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 </coefficients>
+      </correlation>
+      <correlation rcut="10" size="10" speciesA="u" speciesB="d">
+        <coefficients id="ud" type="Array"> 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 </coefficients>
+      </correlation>
+    </jastrow>
+    <jastrow name="J1" type="One-Body" function="Bspline" source="ion0" print="yes">
+      <correlation rcut="5" size="10" cusp="1" elementType="H">
+        <coefficients id="eH" type="Array"> 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 </coefficients>
+      </correlation>
+    </jastrow>
+    <jastrow function="polynomial" name="J3" print="yes" source="ion0" type="eeI">
+      <correlation esize="3" especies="u" isize="3" ispecies="H" rcut="10">
+        <coefficients id="uuH" type="Array"> </coefficients>
+      </correlation>
+      <correlation esize="3" especies1="u" especies2="d" isize="3" ispecies="H" rcut="10">
+        <coefficients id="udH" type="Array"> </coefficients>
+      </correlation>
+    </jastrow>
+  </wavefunction>
+</qmcsystem>
diff --git a/tests/molecules/H4_ae/det-input-check-optm-invalid-subset.xml b/tests/molecules/H4_ae/det-input-check-optm-invalid-subset.xml
@@ -0,0 +1,35 @@
+<?xml version="1.0"?>
+<simulation>
+  <project id="H4-OneShiftOnly" series="0">
+    <parameter name="driver_version">batch</parameter>
+  </project>
+  <!--random seed="605"/-->
+  <!-- input from gaussian converter -->
+  <include href="H4.ptcl.xml"/>
+  <include href="H4.wfs_j123.xml"/>
+  <hamiltonian name="h0" type="generic" target="e"> 
+    <pairpot name="ElecElec" type="coulomb" source="e" target="e"/> 
+    <pairpot name="IonElec"  type="coulomb" source="ion0"/>
+    <constant name="IonIon" type="coulomb" source="ion0" target="ion0"/> 
+  </hamiltonian> 
+
+<loop max="4">
+  <qmc method="linear" move="pbyp" checkpoint="-1" gpu="no">
+    <parameter name="walkers_per_rank">    16  </parameter>
+    <parameter name="blocks">   128  </parameter>
+    <parameter name="warmupsteps"> 5 </parameter>
+    <parameter name="steps"> 128 </parameter>
+    <parameter name="substeps"> 5 </parameter>
+    <parameter name="timestep"> 0.5 </parameter>
+    <parameter name="useDrift">  yes </parameter>
+    <parameter name="samples"> 262144 </parameter>
+         <cost name="energy">                   1.00 </cost>
+         <cost name="unreweightedvariance">     0.00 </cost>
+         <cost name="reweightedvariance">       0.00 </cost>
+    <estimator name="LocalEnergy" hdf5="no"/>
+    <parameter name="MinMethod">OneShiftOnly</parameter>
+    <parameter name="variational_subset"> uu ee </parameter>
+  </qmc>
+</loop>
+
+</simulation>
-Original file line number
+Diff line change
@@ Expand Up / @@ -682,6 +682,7 @@ struct BsplineFunctor : public OptimizableFunctorBase @@
       {
         if (notOpt)
           return;
+        myVars.setIndexDefault();
         active.insertFrom(myVars);
       }
@@ Expand Down @@