Skip to content

TWFFastDerivWrapper: add support for MultiDet wfs - #5398

Merged
ye-luo merged 34 commits into
QMCPACK:developfrom
kgasperich:twffastderiv-multidet
Apr 3, 2025
Merged

TWFFastDerivWrapper: add support for MultiDet wfs#5398
ye-luo merged 34 commits into
QMCPACK:developfrom
kgasperich:twffastderiv-multidet

Conversation

@kgasperich

@kgasperich kgasperich commented Mar 25, 2025

Copy link
Copy Markdown
Contributor

Proposed changes

The plan is to extend the functionality of the TWFFastDerivWrapper to be able to handle MultiDeterminant wavefunctions (based on https://doi.org/10.1063/1.4948778).
This will probably just be a simple reference implementation at first, with some refactoring/optimization to come later if needed.

TODO:

  • we will need access to multidet excitation data (could make some container/struct to hold pointers to the required data, but maybe simpler to just have a pointer to the MultiDiracDeterminant?)
    • ndets_per_excitation_level_
    • detData
    • uniquePairs we might be able to use the specific pairs, but for first ref implementation I will probably just use the full cartesian product of holes x particles (not sure if this is available elsewhere, but we can generate it from the pairs)
    • DetSigns only needed when we're contracting/summing over dets; we can pass it all the way through and include it in the terms for each excited det, but can also keep at the level of the alpha/beta CI coefs and use it when we do the val = data_alpha[a]*signs_alpha[a] * CI_coeffs_ab[a,b] * data_beta[b]*signs_beta[b]
    • Done: add std::vector<const MultiSlaterDetTableMethod*> slaterdets_ to wrapper; access through there (also make sure SPOSets from MultiDiracDets in MultiSlaterDet are added to the wrapper)
      • I have this as a vector in case there's ever a time that we want to support multiple (not sure what that would look like), but I have checks in a few places to ensure that we don't have more than one for now (downstream functions will not handle more than one yet, and those will also check this)
  • add vector of pointers to multidet data in wrapper class
    • Done (see above)
  • overload or modify TWFFastDerivWrapper::addGroup to add MultiDiracDeterminant (or maybe a MultiSlaterDet?)
    • Done: made TWFFastDerivWrapper::addMultiSlaterDet; access MultiDiracDets through there; add necessary SPOSets at same time
  • decide how to handle MultiSlaterDet info (can pass alpha/beta MultiDiracDet individually, but will need to be able to associate the MultiSlaterDet with the MultiDiracDets; maybe better to just add the MultiSlaterDet directly and then add the constituent MultiDiracDets through the MultiSlaterDet)
    • Done: see above
  • add function similar to TWFFastDerivWrapper::computeGSDerivative computeMDDerivative
    • in progress: need to add some small matrix ops, verify dims of some inputs
    • Done (still a lot of room for optimization, especially with limiting virtuals to only the ones that we actually excite into)
  • consider adding something similar to SmallMatrixDetCalculator for small matrix inversion
    • save for later PR
  • modify QMCHamiltonian::evaluateIonDerivsFast to call multidet function
    • in progress: still need to contract alpha/beta det info with CI coefs
    • should be finished, but testing might reveal some bugs
  • uncomment multidet tests (will need some modification for new functions; see single-det examples with evaluateIonDerivsFast in same file)
    • in progress

What type(s) of changes does this code introduce?

Delete the items that do not apply

  • New feature
  • Refactoring (no functional changes, no api changes)
  • Testing changes (e.g. new unit/integration/performance tests)

Does this introduce a breaking change?

  • Yes
  • No

What systems has this change been tested on?

Checklist

Update the following with a yes where the items apply. If you're unsure about any of them, don't hesitate to ask. This is
simply a reminder of what we are going to look for before merging your code.

  • Yes/No. This PR is up to date with current the current state of 'develop'
  • Yes/No. Code added or changed in the PR has been clang-formatted
  • Yes/No. This PR adds tests to cover any new code, or to catch a bug that is being fixed
  • Yes/No. Documentation has been added (if appropriate)

@prckent

prckent commented Mar 25, 2025

Copy link
Copy Markdown
Contributor

Thanks for laying out a roadmap. Would be excellent to have this of course. Please do an absolute minimum implementation first, to get something working. Likely Ray has already thought this through , so check with him if you have not already.

In case they are needed, add accessor functions to get excitation data. Do not break encapsulation.

Forgive the reminder, but please start writing units tests early. i.e. Now. This will make reviewing and merging easier and faster. For this kind of code it is possible to achieve excellent coverage. Do not rely in integration tests.

Happy to discuss further if you want.

@kgasperich

Copy link
Copy Markdown
Contributor Author

@prckent thanks for the feedback; I have been in touch with Ray about this

In case they are needed, add accessor functions to get excitation data. Do not break encapsulation.

I'm trying to come up with a clean way to do this.
Currently in the wrapper, we construct corresponding lists of SPOSet* and particle group IDs, so that we can just iterate over those together.

For the multidet case, we will need to do something similar with at least the MultiSlaterDetTableMethod.
If we have access to that, then we can access the MultiDiracDeterminants through MultiSlaterDetTableMethod::Dets, and we can also access the SPOSets through the MultiDiracDeterminants.
I suppose that should be good enough for now -- I was getting hung up on trying to still keep pointers to all of the SPOSets directly in the wrapper, but I'm now thinking that that isn't really necessary, and we can just access them through the multidet objects.

I'm also trying to decide how to have the singledet and multidet code paths live nicely together.
If I'm not using the old lists of sposet/group ids, then we can just have the singledet functions iterate over those (so nothing will need to be done if they're empty). The multidet part of the code can iterate over the MultiSlaterDet list, and then within a MultiSlaterDet iteration, the relevant quantities for each multidiracdet will be evaluated.

for msd_i in slaterdets:
  // allocate space for excited det derivs
  for mdd_i in msd_i.Dets:
    // compute and store excited det derivs
  // with CI coeffs from msd_i, contract over alpha/beta spindet data

@kgasperich

Copy link
Copy Markdown
Contributor Author

update on status:

I've been assuming that the fval returned here (using the notation from the article is $\frac{\partial}{\partial \mu}\frac{\hat{O}\Psi}{\Psi}$, and $\Psi$ is a multi-slaterdet wavefunction.
To build that from the individual DiracDeterminants, we need terms like $\frac{\partial}{\partial \mu}\frac{\hat{O}D}{D}$, $\frac{\partial}{\partial \mu} \log(D)$, and $\frac{\hat{O}D}{D}$.

I compute terms like these for each DiracDeterminant in computeMDDerivatives_ExcDets, and then I do the necessary contraction over the full list of SlaterDets in computeMDDerivatives_total, which returns $\frac{\partial}{\partial \mu}\frac{\hat{O}\Psi}{\Psi}$ and $\frac{\partial}{\partial \mu} \log(\Psi)$ (returned here as fval and wfcomp

If I'm reading the code correctly, then in the single-det code path, fval is $\frac{\partial}{\partial \mu}\frac{\hat{O}D}{D}$ (Eq. 15 from the paper) and wfcomp is $\frac{\partial}{\partial \mu} \log(D)$ (Eq. 4 from the paper) both seen here (where $D=D_\alpha D_\beta$ so $\frac{\partial}{\partial \mu} \log(D) = \frac{\partial}{\partial \mu} \log(D_\alpha) + \frac{\partial}{\partial \mu} \log(D_\beta)$ and $\frac{\partial}{\partial \mu}\frac{\hat{O}D}{D} = \frac{\partial}{\partial \mu}\frac{\hat{O}D_\alpha}{D_\alpha} + \frac{\partial}{\partial \mu}\frac{\hat{O}D_\beta}{D_\beta}$

I think this is mostly on the right track, but I'd welcome any feedback if any of what I've described seems incorrect.

@rcclay said that there should be some reference values already present in test_ion_derivs.cpp, so my next step is to work on checking the new functions against those values.

My first attempt was not successful, but I'm not at all confident that I'm calculating/comparing the right thing.

I added a call to the TWFFastDerivWrapper multidet function here (branched from this PR source because I don't want to make too much of a mess with the branch history if it takes a while to figure this out)
kgasperich@934e968

with that change, the output I see from the test is

dedr[0][0] = -88.3531711
dedr[0][1] = 79.53044961
dedr[0][2] = 290.585017
dedr[1][0] = 57.08550161
dedr[1][1] = -8.714815584
dedr[1][2] = 18.76148474
wf_grad[0][0] = -1.904466809
wf_grad[0][1] = 2.125775424
wf_grad[0][2] = 7.055631329
wf_grad[1][0] = 1.423336104
wf_grad[1][1] = -0.1446704418
wf_grad[1][2] = 0.1440195609
wfgradraw[0][0] = 0
wfgradraw[0][1] = 0
wfgradraw[0][2] = 0
wfgradraw[1][0] = 0
wfgradraw[1][1] = 0
wfgradraw[1][2] = 0

which doesn't seem to match any of the reference values in this test.

The wf_grad values are close enough that I suspect that these are at least representing sort of the same thing, either with some discrepancy in exactly what quantity is being compared, or some small-ish term that isn't being accounted for in the new code.

wf_grad[0][0] = -1.904466809
wf_grad[0][1] = 2.125775424
wf_grad[0][2] = 7.055631329
wf_grad[1][0] = 1.423336104
wf_grad[1][1] = -0.1446704418
wf_grad[1][2] = 0.1440195609
  CHECK(wf_grad[0][0] == Approx(-1.7045200053189544));
  CHECK(wf_grad[0][1] == Approx(2.6980932676501368));
  CHECK(wf_grad[0][2] == Approx(6.5358393587011667));
  CHECK(wf_grad[1][0] == Approx(1.6322817486980055));
  CHECK(wf_grad[1][1] == Approx(0.0091648450606385));
  CHECK(wf_grad[1][2] == Approx(0.1031883398283639));

Another thing I'm aware of is that QMCHamiltonian::evaluateIonDerivsFast seems to be evaluating all terms in the Hamiltonian ([1], [2], but some of the other reference values correspond to the kinetic and NLECP terms alone.

I could try to make a multi-det test that more closely mirrors this single-det one.
This does some of the work of QMCHamiltonian::evaluateIonDerivsFast in building the B/M/dM/etc. matrices individually for the kinetic and NLECP terms (whereas evaluateIonDerivsFast includes all terms), so if I do that for the multi-det case, I can have some more fine-grained data for debugging purposes.

(I'm not sure if I could accomplish the same thing by building a separate Hamiltonian that only includes a single term; I'd need to look at everything in more detail to figure out whether excluding the other terms would cause issues, and that would probably take more time than just adding the matrix evaluation as the single-det test does)

@rcclay

rcclay commented Mar 28, 2025

Copy link
Copy Markdown
Contributor

@kgasperich I think the path forward here is to first spot check one of the finite difference reference values just to make sure it's OK. I'm sure it is, but if there's a chance it's wrong, that concern needs to be killed now. Step 2 will be to focus entirely on making sure that wfgrad comes out right. That's the simplest/most transparent quantity, and also dodges issues related to the fact that only some of the hamiltonian quantities have ion derivatives computed in this framework. I think your plan of making a more fine grained unit test a la single det code will probably save some debugging time in the long run. Also gives you clean and independent checks of the kinetic and nonlocal force contributions.

@kgasperich

Copy link
Copy Markdown
Contributor Author

@rcclay thanks!

first spot check one of the finite difference reference values just to make sure it's OK. I'm sure it is, but if there's a chance it's wrong, that concern needs to be killed now

Is there an existing test or example somewhere that performs this? (I assume there's something somewhere in test_ion_derivs.cpp that has an example of this, but if you know exactly where it is it will save me a lot of time to dig around for it and verify that it's what I should be using)

Step 2 will be to focus entirely on making sure that wfgrad comes out right. That's the simplest/most transparent quantity, and also dodges issues related to the fact that only some of the hamiltonian quantities have ion derivatives computed in this framework. I think your plan of making a more fine grained unit test a la single det code will probably save some debugging time in the long run. Also gives you clean and independent checks of the kinetic and nonlocal force contributions.

ok, this sounds good; I will work on this part for now and update as it progresses!

@kgasperich

Copy link
Copy Markdown
Contributor Author

I fixed two small bugs, and the new code seems to be giving the correct values now.

I still get a slight discrepancy compared to numerical derivatives I computed using WFtester (I used a two-point central difference with delta of +/- 0.0001 for all dims).
The wf gradients, observables, and ion derivatives of the kinetic energy are all in agreement (to whatever tolerance Catch2 Approx uses), but five of the six nlpp ion derivs are a bit off (they're close enough that I assume it might just be an inaccurate numerical derivative rather than any bug in the new code, but I'm still planning to take a closer look)

/mnt/Work/Endymion/kgasperich/programs/qmcpack-ion-derivs/src/QMCHamiltonians/tests/test_ion_derivs.cpp:833: FAILED:
  CHECK( fnlpp[ionid][idim] == Approx(ref_nlpp[ionid][idim]) )
with expansion:
  18.9414437168 == Approx( 18.9400878321 )

/mnt/Work/Endymion/kgasperich/programs/qmcpack-ion-derivs/src/QMCHamiltonians/tests/test_ion_derivs.cpp:833: FAILED:
  CHECK( fnlpp[ionid][idim] == Approx(ref_nlpp[ionid][idim]) )
with expansion:
  -78.3304789803 == Approx( -78.327944634 )

/mnt/Work/Endymion/kgasperich/programs/qmcpack-ion-derivs/src/QMCHamiltonians/tests/test_ion_derivs.cpp:833: FAILED:
  CHECK( fnlpp[ionid][idim] == Approx(ref_nlpp[ionid][idim]) )
with expansion:
  1.2122162586 == Approx( 1.2134514881 )

/mnt/Work/Endymion/kgasperich/programs/qmcpack-ion-derivs/src/QMCHamiltonians/tests/test_ion_derivs.cpp:833: FAILED:
  CHECK( fnlpp[ionid][idim] == Approx(ref_nlpp[ionid][idim]) )
with expansion:
  -0.6163169109 == Approx( -0.6154734664 )

/mnt/Work/Endymion/kgasperich/programs/qmcpack-ion-derivs/src/QMCHamiltonians/tests/test_ion_derivs.cpp:833: FAILED:
  CHECK( fnlpp[ionid][idim] == Approx(ref_nlpp[ionid][idim]) )
with expansion:
  -3.2996553003 == Approx( -3.3012548664 )

After these fixes, I started refactoring the new multidet code for the ionic derivatives (there is/was a lot of redundant computation).
Somewhere in that process I've managed to introduce a new bug, but I still have the older code to compare to, so at least I have some reference values for debugging.

I'm working on that in another branch, and once I fix that issue I'll merge it back into this branch

@kgasperich

Copy link
Copy Markdown
Contributor Author

I see that a bunch of the complex tests are failing because Catch::Approx doesn't play nicely with complex types.
I have a bunch of checks just to compare the two different implementations that I made (the _ExcDets vs the _Obs/_dmu)

I can remove those and then use convertToReal for the rest of the checks (as is done in the other similar tests)

@kgasperich

Copy link
Copy Markdown
Contributor Author

The test should build now for QMC_COMPLEX=1

I removed the part of the test that verifies that _ExcDets agrees with _Obs/_dmu.

If there's a good reason to keep _ExcDets in the code, I can add another test for those (or add them back to the same test and be more careful about handling complex types with Catch::Approx), but I'm leaning toward just removing TWFFastDerivWrapper::computeMDDerivatives_ExcDets

@kgasperich

Copy link
Copy Markdown
Contributor Author

Just in case it's not clear from the comments, for the multidet derivatives we need terms like $\frac{\hat{O} D}{D}\text{,}\quad\partial_{\mu} \log D\text{,}\quad\partial_{\mu} \left( \frac{\hat{O} D}{D}\right)$ over all spindets $D$.

I initially made one function (TWFFastDerivWrapper::computeMDDerivatives_ExcDets) to compute all three of these for each spindet.

The $\frac{\hat{O} D}{D}$ doesn't depend on the ionic derivatives, so we can just compute those terms once (with TWFFastDerivWrapper::computeMDDerivatives_Obs) rather than 3N times.

The other two types of terms still need to be computed once for each ion dim, so those are handled by TWFFastDerivWrapper::computeMDDerivatives_dmu

@kgasperich
kgasperich requested a review from rcclay April 1, 2025 04:49
@prckent

prckent commented Apr 1, 2025

Copy link
Copy Markdown
Contributor

RE: Catch. We have our own ComplexApprox if that helps.

rcclay
rcclay previously approved these changes Apr 1, 2025

@rcclay rcclay left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

I think this looks fine. Logic is sound, design is contained to TWFFastDerivWrapper and is consistent with what is already there, and so has minimal impact to rest of code for a big improvement in capability. I indicated some examples of incomplete documentation that should be addressed--documentation of innards of computeMDDerivatives_dmu should be more verbose. Also there are a few comments in the code along the lines of "FIX ME" and "TODO" which should be revisited.

// IndexType ndet = multidiracdet_i->getNumDets();
IndexType ndet = multidiracdet_i.getNumDets();

/// FIXME: set these correctly; decide how to index into nocc if not contiguous within nOcc

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Is this still a "FIXME"?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

It's more of a possible performance improvement than a correctness issue.
Everything works now assuming that the orbital dimension of any matrix is either

  1. all of the Nparticle orbs (i.e. occupied in the reference det) in order
  2. all of the virtual orbs (i.e. OrbitalSetSize - Nparticle) in order
  3. the concatenation of [1] and [2] (i.e. all orbs in order)

For improved performance, there is the simple case where we have e.g. 10 virtuals, but we only ever excite into the first two of these. In this case, it's a simple basis truncation that could be handled very early in the construction of SPOSets, wavefunctions, etc. or even as a simple processing of input files; this would not make any difference in the current code.

There might also be a case where we have 10 virtuals, but we only ever excite into the second and sixth one. This corresponds to maybe a slightly less trivial pre-processing like the one described above where the relevant virtuals comprise a contiguous set from the start of the full set of virtuals.
With an extra level of redirection, we could handle things like this within the TWFFastDerivWrapper, and this would reduce the amount of useless work we'd do over unused orbitals.

We could do the same thing (in some places) with the occupied orbitals that we're exciting out of.
For occupied orbs, there are some intermediate quantities where we contract over the full set of ground-state orbitals, so we would need at least some matrices to be represented in that full space, but there are other quantities where we only need to consider the ones that appear as holes in the excited dets.

In the table method code for computing det ratios, this is simple, because the final quantity we need in the space of hole/particle indices is just Minv.M, so each element is a single gemv, and we can just do the ones we need rather than doing a full gemm.

For the quantities here, we could do something similar and only compute the matrix elements we need, but on my first implementation attempt that seemed like it would get a bit too messy.
Since then, I've found some more compact ways to represent some of the quantities we need (this can be seen in the difference between computeMDDerivatives_ExcDets and computeMDDerivatives_dmu). It seems more feasible now, but I haven't yet worked out exactly what it would save us in terms of scaling, and I'm not sure what the cost of the extra level of redirection would be.

Under the broader umbrella of performance considerations, there are also some things that we could avoid fully recomputing after each particle move (definitely the Minv and Minv.Mvirt, and possibly some others?). Also, I mostly implemented the _dmu and _Obs functions in a way that is easy to follow, but this isn't necessarily the most efficient; eventually it will be worth taking a closer look at this to see if we're leaving any easy performance gains on the table.

// X3 = Minv_B [occ, occ+virt]
// X4 = Minv_dB [occ, occ+virt]

// O is all occ orbs; o is occ orbs that appear as holes in exc. list; e is all elecs; v is virt orbs that appear as particles in exc. list

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

For the sake of posterity, can you document this a little better? Latex expression, or reference to equation. This is correct, but it took a while to unpack. Looks like the documentation near lines 670 could be appropriated.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

yes, I'll go through the comments again and try to make them clearer.

@kgasperich

Copy link
Copy Markdown
Contributor Author

@rcclay I'd like to remove TWFFastDerivWrapper::computeMDDerivatives_ExcDets if that's okay with you. I think it served its purpose as an initial reference implementation as I was working on this PR, but I don't think it adds anything useful to the code (there might be some comments in there that are clearer than in the other two functions that replaced it, but I can keep those alive where appropriate)

@rcclay

rcclay commented Apr 2, 2025

Copy link
Copy Markdown
Contributor

@kgasperich Sure. Slash away.

Actually, it just hit me. While we are here, can you hack something together to see if the multideterminant+Jastrow code path also works? Some reference values should already be floating around, albeit commented out. It should by construction, but if it doesn't, we need to block off that code path if it's not fixed in this PR.

@kgasperich

kgasperich commented Apr 2, 2025

Copy link
Copy Markdown
Contributor Author

@rcclay

Actually, it just hit me. While we are here, can you hack something together to see if the multideterminant+Jastrow code path also works? Some reference values should already be floating around, albeit commented out. It should by construction, but if it doesn't, we need to block off that code path if it's not fixed in this PR.

Good catch; I had forgotten to check that. I've added two new tests for md+jastrow: one is with the cn.msd-wfj.xml that was already present, and the other is with a new cn.msd-1det-wfj.xml that uses the same MOs and jastrow as the cn.wfj.xml (so we can check against the reference values that are already known), but it has a multidet block (with a single det) so that it will be treated as a multidet wavefunction in the wrapper (I wanted to make sure that it would gracefully handle a single-det wavefunction since some of the stuff involves only computing quantities for excited dets, and then adding on the ground state det separately)

I also removed the old unused TWFFastDerivWrapper::computeMDDerivatives_ExcDets, cleaned up the comments and notes about possibilities for performance improvement, and removed unused typedefs for IndexPairType and IndexPairVector.
I had been planning to use those for something like the MultiDiracDeterminant::uniquePairs, which are used to compute only the matrix elements that are involved in some hole/particle pair in the excited det list, but I think it's better to leave that for future refactoring/optimization.

@kgasperich kgasperich changed the title [WIP] TWFFastDerivWrapper: add support for MultiDet wfs TWFFastDerivWrapper: add support for MultiDet wfs Apr 2, 2025
@rcclay
rcclay self-requested a review April 2, 2025 15:30
rcclay
rcclay previously approved these changes Apr 2, 2025

@ye-luo ye-luo left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Minor change requested.

Comment thread src/QMCWaveFunctions/TWFFastDerivWrapper.h Outdated
Comment thread src/QMCWaveFunctions/TWFFastDerivWrapper.h Outdated
Comment thread src/QMCWaveFunctions/TWFFastDerivWrapper.h Outdated
Comment thread src/QMCWaveFunctions/TWFFastDerivWrapper.h Outdated
@kgasperich

Copy link
Copy Markdown
Contributor Author

@ye-luo

After making those changes, I think I actually prefer the single MSD (const WaveFunctionComponent* multislaterdet_;) for now just to make everything a bit clearer, but I can change it back if you prefer the vector (std::vector<const WaveFunctionComponent*> multislaterdets_;).

@ye-luo

ye-luo commented Apr 3, 2025

Copy link
Copy Markdown
Contributor

Test this please

@ye-luo
ye-luo enabled auto-merge April 3, 2025 01:13
@ye-luo

ye-luo commented Apr 3, 2025

Copy link
Copy Markdown
Contributor

Test this please

@ye-luo
ye-luo merged commit c302e86 into QMCPACK:develop Apr 3, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants