direct_minimization with Manopt.jl as a weak dependency

[kellertuer]

Following the discussion in #1102, here is a work in progress on moving direct_minimization to an extension and to use Manopt.jl.

Due to the discussions in the linked issue I feel a bit unsure how and whether to continue. Currently left to do

• 1 document direct_minimization in src/scf/ that it requires Manopt and Manifolds to work
• 2 (maybe) rename ManoptPreconditioner to maybe something more reasonable. It basically wraps a DFTK preconditioned
• 3 CostGradFunctor! should get a nicer name and could maybe also become the fixed energy/gradient computation but also maybe stay a keyword argument
• 4 Document CostGradFunctor!
• [ ] 5 clarify what the use of alphaguess is and what else than `nothing it can/should be.(removed it)
• 6 the preconditioner has to be passed differently to gradient_descent/cg than Quasi newton, this needs to be handled somewhat.
• 7 discuss what to do in the end with the info or in other words what to return from that function.
• 8 add tests and see where we can cover this in the docs

Points 1-4 and 8 are merely to work though as soon as I find time
On 5 and 7 I need some help and feedback
7 is maybe a bit tricky to get right in general.

For now this function (and its helpers) are a bit work in progress – I open this a bit early, since I am not sure the philosophies of the packages agree well enough that this approach is useful. in the issue above it sounds currently like the idea here is that the philosophy is too different and that it might not fit.

[kellertuer]

I had a meeting that was shorter than I thought, so here is a small snippet that somewhat works on this branch already (thanks to @elorannug for the setup)

using LinearAlgebra, DFTK, Manopt, Manifolds, RecursiveArrayTools, Random
using LineSearches
using PseudoPotentialData, AtomsBuilder, Unitful, UnitfulAtomic

#
#
# Load data and setup model
Random.seed!(1)
Ecut = 5 # Energy cut-off.   Realistic value: Ecut = 20
kg = 1   # Number of k-points Realistic value:   kg = 5
tol = 1e-1 # Tolerance for the minimization, default is 1e-6
# Set up the model using planewave basis
pseudopotentials = PseudoFamily("dojo.nc.sr.pbesol.v0_4_1.standard.upf");
model = model_DFT(bulk(:Si); functionals=PBEsol(), pseudopotentials);
basis = PlaneWaveBasis(model; Ecut=Ecut, kgrid=[kg, kg, kg]);
model_to_save_psi = basis.model
filled_occ_s = DFTK.filled_occupation(model_to_save_psi)
n_spin_s = model_to_save_psi.n_spin_components
n_bands_s = div(model_to_save_psi.n_electrons, n_spin_s * filled_occ_s, RoundUp)
random_ψ = [DFTK.random_orbitals(basis, kpt, n_bands_s) for kpt in basis.kpoints]

direct_minimization(basis;
    ψ=random_ψ,
    tol=tol,
    maxiter=1000,
    #TODO: Check that it also works with: other manifolds, retractions, vector transports, debug & record:
    debug = [DebugCost(), DebugGradientNorm(), "\n", 10, :Stop],
    # record = [:Cost, :GradientNorm],
    # manifold = Grassmann,
    # stepsize = Manopt.LineSearchesStepsize(LineSearches.BackTracking()),
)

It still requires a bit of careful default parameter tuning, since the default (Armijo) line search currently uses the QR retraction instead of the default we prescribe and such. But a very first small generic interface does work.

[kellertuer]

While I did not yet address the points above, I spent this evening debugging the code, now it is relatively easy to

• change the solver
• change the manifold
• change the linesearch
• change the retraction
• change the vector transport
• activate debug & record
• change the stopping criterion (though that “deactivates” the values passed to tol and maxiter

Every now and then it seems the gradient is not so correct because the step size collapses after a few iterations, I have to check that (or the example is too simple so it finishes after < 4 steps anyways).

6 is nearly done, just requires a bit of tolerance on the gradient descent side in manopt;
Feedback on 5 and 7 would be nice; besides a bit of debugging the rest, only documentation and testing (and maybe a bit of renaming) is missing.

[kellertuer]

I worked a bit mainly on a bit of testing (as far as my knowledge reaches), naming and documentation.

This should work for now, with the small caveat that for GradientDescentState (i.e. using gradient descent) as well as to use callbacks, we have to wait for the next version of Manopt.jl to be released.
Besides waiting for these, the PR is finished up to the following points to discuss

Questions on the code

(also written within the code in the corresponding places)

• In the stopping criterion – if I just know the manifold, can I then already compute (at least the size/type) of the density ρ? That would make a much nicer constructor only requiring the tolerance. (This is currently not possible since from M we can not infer what ρ was)
• Does it makes sense to be able to record ρ ? A user could then easily use `record = [:ρ] to record if we do an easy recorder (but maybe also no one likes that here). Provided, maybe it is useful for someone
• a debug is maybe not so useful, since it seems to be a large array, but its norm maybe? (Implemented)
• Should a user be able to provide their own cost/grad? (see ...and beyond below as well)
  Then we would have to change a few small things in the setup. (-> Implemented)
• What should the debug_info contain? (There is a list but it is not yet realised)

Questions in general

• does DFTK have a certain code style I should apply? (not very strict it seems)
• does DFTK want to use DocumenterCitations.jl for literature? I could set it up, either here or in a separate PR (Setup, but needs some fine tuning)
• does DFTK want to use DocumenterInterlinks.jl to link to other-peoples-package-docs? (Setup, but not yet used)

...and beyond

I had a discussion with @jonas-pueschel to also include his CG approache; that however would first mean, that in severa places, especially in Armijo and Wolfe, the inner product with the gradient has to ne replaced an evaluation of the differentia, see JuliaManifolds/Manopt.jl#482. Then we can probably either update the cost/grad here or provide a second alternative.

[mfherbst]

I left a few high-level comments. Maybe we should have a zoom to discuss a few details, that I don't get yet.

[kellertuer]

Yes a zoom is probably the best, in Text it might take a few iterations too much.

edit: one that could easily be resolved I already resolved – and added a few notes to the comments, maybe also just for me to remember to discuss this :)

[kellertuer]

@mfherbst I think I have finished everything we discussed and I was able to do. Here is a few remaining things I am unsure about

• documentation: I am currently not sure what is documented where, since I have not yet understood the structure of the docs. running them locally does at least not throw errors that docs are missing, but it currently also does not seem to load/document extensions.

• the wannier example from the ecosystem errors when running the docs, it seems that is unrelated to what I did.

• every now and then in the examples/gross_pitaevskii.jl-L89 fails and I have not yet found out why
  I checked this example more closely and the line fails with

   julia> E.total
   147.37290591541824
  
   julia> scfres.energies.total
   147.37290591541844
  

  it is just a nuance, and I think this is somewhere a rounding error since it just stems from evaluating the energy an the same point.
  edit: it is probably just rounding errors.

• I can run the tests locally and they run really long, but it seems all of them pass. Indeed after 4.5 hours, only a few tests also concerning the wannier90 fail, so that is unrelated to this PR.

• currently direct_minimization can not yet report result.converged see the explanation here

I think this is how far I can provide code here. For a next step I would probably need help. Feedback on the code is of course also welcome.

[kellertuer]

This is a careful bump after 2 months without any reaction on my last post with questions and/or a review.

If this is too much work or too much of a change to ask, we can also not do this.

[mfherbst]

Thanks for the reminder and sorry for the delay. I'll try to follow up in the next days.

[mfherbst]

every now and then in the examples/gross_pitaevskii.jl-L89 fails

If it's just a small energetic glitch as you mention it, I would not be too worried. But if fails means it crashes, then I am.

while most old examples still run, we could consider writing a separate example for this new solver, maybe using Grassmann?

We can, but not extremely important from my point of view. If you have it handy, it does not hurt, though.

the internal InsulatorEnergy currently stores both occupation and filled_occ, since as I wrote, I am not able to change that with the knowledge I have.

If it's clear now feel free to do it. If not, I'll do a cleanup pass and then I can take a look.

documentation: I am currently not sure what is documented where, since I have not yet understood the structure of the docs. running them locally does at least not throw errors that docs are missing, but it currently also does not seem to load/document extensions.

I don't know about the extensions. Basically what we do is that we re-use the examples also as documentation. Is that your confusion ? It may also be that the literate step causes trouble ?

[mfherbst]

I think it's converging. Take a look what you can fix with the additional information from my end, else I'd suggest we have a call to clean things up and then I can take over the rest whatever is needed. Things we should clarify:

• How does the printing look like
• Anything else to be discussed on docs building ?
• Source code formatting: We try to keep line length around 92 (see the scarce style guide https://docs.dftk.org/stable/developer/style_guide/)

[mfherbst]

It's a little weird that one needs to using all there three packages explicitly. Can we somehow avoid that ?

[kellertuer]

Not really.

• Manopt is needed to load the solver(s) we use
• Manifolds is needed for Grassmann and the Product manifold
• The product manifold is basically a dummy until RecursiveArrayTools are loaded end the extension (in Manifolds) provides the efficient product manifold functions (baes on ArrayPartition)

so all 3 are needed.

[mfherbst]

I see, but that's quite inconvenient for a potential user. Could we avoid that by making a metapackage that depends on all three of them (like DFTKManopt.jl) ?

[kellertuer]

Uff. I would have no clue how to do that?
I also do not agree that that is inconvenient, it is what the extensions now bring – more modular packages without direct dependencies:

One can use Manopt with an own manifold (nand then not load Manifolds)
One can use Manifolds with an own implementation for the product structure (and then not load RecursiveArrayTools)

What would that package look like? Something with the Reexport all package? I am not sure what to do here.

[mfherbst]

Is there a way to avoid duplicating info between here and the low-level function below?
Perhaps we only provide docs here and the in the low-level version we don't document it, referring to this version instead ?

[kellertuer]

Hm, these two variants allow for a slightly different level of settings.

This is the version that is meant for ease of use and follow a bit closer the DFTK scheme. That limits a bit what one can set up.
The other (maybe considered a bit more internal) already requires to pass a Manopt state, but then provides more access to “Manopt-stuff”. Sure we can leave that one undocumented, but then I will have forgotten half of what I wrote in about a year and would not be able to help (or help easily) if someone comes along with a problem.
That would not be my preferred way to go.

[kellertuer]

To answer a few already:

every now and then in the examples/gross_pitaevskii.jl-L89 fails

If it's just a small energetic glitch as you mention it, I would not be too worried. But if fails means it crashes, then I am.

to me it looks like small glitches.

while most old examples still run, we could consider writing a separate example for this new solver, maybe using Grassmann?

We can, but not extremely important from my point of view. If you have it handy, it does not hurt, though.

Mainly a matter of time, and I fear it would take me quite some time.

the internal InsulatorEnergy currently stores both occupation and filled_occ, since as I wrote, I am not able to change that with the knowledge I have.

If it's clear now feel free to do it. If not, I'll do a cleanup pass and then I can take a look.

to me this is to 0% clear.

documentation: I am currently not sure what is documented where, since I have not yet understood the structure of the docs. running them locally does at least not throw errors that docs are missing, but it currently also does not seem to load/document extensions.

I don't know about the extensions. Basically what we do is that we re-use the examples also as documentation. Is that your confusion ? It may also be that the literate step causes trouble ?

Well. To be precise: The new doc strings are defined in the extension.
They are not included in the documentation at all. But the documentation also does not error, like it usually does when doc strings are not included in the docs. The simple plain reason is, that the documentation is not aware the extension exists.
See https://github.com/JuliaManifolds/Manopt.jl/blob/e4a5872c097ec839b799c5d4cacd314d9a5837f6/docs/make.jl#L162-L169 how you have to add them (and of course also have to have loaded them beforehand).
I think this is unrelated to literate.

[kellertuer]

I think it's converging. Take a look what you can fix with the additional information from my end, else I'd suggest we have a call to clean things up and then I can take over the rest whatever is needed. Things we should clarify:

• How does the printing look like

If you could make more precise what you mean here?

• Anything else to be discussed on docs building ?

As written in the post above, currently 0 documentation of extensions seems to be in the documentation.

• Source code formatting: We try to keep line length around 92 (see the scarce style guide https://docs.dftk.org/stable/developer/style_guide/)

Been there, done that.

[kellertuer]

Since we released has_converged yesterday, I readied the converged= field to the named tuple as well.

[kellertuer]

I tried at least illustrating the documentation of things in the extension

1. load all packages to trigger loading of the extension
2. adding the extension to the modules in makedocs
3. collect all those doc strings on a separate page

When trying to do 3 I usually check which signatures fail to add them on that page – otherwise it is a bit hard to specify that by hand.

However, I can not render the docs locally, the Error I get is (starting snippet)

┌ Error: failed to run `@example` block in src/ecosystem/wannier.md:177-198
│ ```@example wannier
│ using wannier90_jll  # Needed to make run_wannier90 available
│ run_wannier90(scfres;
        [... skipped rest of code block that is run ... ]
│              );
│ nothing #hide
│ ```
│   exception =
│    UndefVarError: `wannier90` not defined in `wannier90_jll`
│    Suggestion: check for spelling errors or missing imports.
│    Stacktrace:
│      [1] (::DFTKWannier90Ext.var"#2#3"{String, String, String})()

I would assume that is unlrelated to this PR? However, the last commit is hence only a first step towards documenting the extension.

I checked the other extensions, and only one has a one-line (more-like-dummy) doc string. So maybe for other extensions this is not necessary. For the one here, I am unsure how to continue – also not for the other open points.

Let me know what best to do then.

[mfherbst]

to me it looks like small glitches.

ok, good then.

Mainly a matter of time, and I fear it would take me quite some time.

then let's not

to me this is to 0% clear.

Ok then I do it. Is it ok if I push to your branch directly ?

If you could make more precise what you mean here?

How do the output look like when running the direct minimisation with ManOpt.

As written in the post above, currently 0 documentation of extensions seems to be in the documentation.

Ok, but then all it takes is to add the Base.get_extension calls, right ?

I would assume that is unlrelated to this PR?

I would think so. Are you on a windows machine ? Because Wannier90 does not compile on windows (hence no binary in Yggdrasil).

[kellertuer]

Mainly a matter of time, and I fear it would take me quite some time.

then let's not

Thanks for your understanding.

to me this is to 0% clear.

Ok then I do it. Is it ok if I push to your branch directly ?

Sure, feel free to change on this branch what is necessary or what ever you like :)

If you could make more precise what you mean here?

How do the output look like when running the direct minimisation with ManOpt.

Its Manopt (no capital o) – but sure I can check that again and post that here, though maybe only tomorrow

As written in the post above, currently 0 documentation of extensions seems to be in the documentation.

Ok, but then all it takes is to add the Base.get_extension calls, right ?

And loading the packages before, the ? get_extension` only returns extensions that are loaded already.
I already extended the docs environment (to inlucde Manopt, Manifolds and RecursiveArrayTools) and then the extension is loaded and added.
Just that I could not yet render the docs then.

I would assume that is unlrelated to this PR?

I would think so. Are you on a windows machine ? Because Wannier90 does not compile on windows (hence no binary in Yggdrasil).

No I am on a Mac (macOS Tahoe 26.0, just updated yesterday)

[kellertuer]

It's been a while again. Anything I can help with?

I think the one thing missing is still the occupation vs filled_occ that I still can not fix because I do still not yet understand how one can avoid the duplication.

Besides that

• How does the printing look like

If you could make more precise what you mean here?

This was not yet answered. I am happy to provide show methods, but I can not just guess what is the usual thing one want the printing to look like?

• Anything else to be discussed on docs building ?

As written in the post above, currently 0 documentation of extensions seems to be in the documentation.

Same here: Currently 0% of extension documentation is rendered, since extensions are not included in the modules list of the documentation. If that is how it should be, then all things for the docs have been discussed.

and finally (and fourth) the issue that now 3 packages have to be loaded. I do not have a good solution for that. If that is too annoying, then we should maybe discontinue this and stick to the old scheme of using optim, that is also fine with me, then just close this PR.

[mfherbst]

I think the one thing missing is still the occupation vs filled_occ that I still can not fix because I do still not yet understand how one can avoid the duplication.

As I said, I'll look into it, but my coding time is limited, unfortunately. But it's still on my list.

If you could make more precise what you mean here?

I meant quite naively, what does the printed output look like if I run the code. But I can check when I look at the first point .

Same here: Currently 0% of extension documentation is rendered, since extensions are not included in the modules list of the documentation. If that is how it should be, then all things for the docs have been discussed.

I thought you suggested this can be changed, which sounded very reasonable to me. You also mentioned cross-referencing documentation. These are things which would be nice to have and I just wondered if we are all set on this or if more things need to be discussed.

and finally (and fourth) the issue that now 3 packages have to be loaded. I do not have a good solution for that. If that is too annoying, then we should maybe discontinue this and stick to the old scheme of using optim, that is also fine with me, then just close this PR.

Yeah, I think that will be too annoying. But I think we can work around that by making a small DFTKManopt package, which then hard-depends on all there dependencies. If this exposes (@rexports maybe) Manopt, Manifolds and the third package then you use that in your extension. That would mean only one extra package needs to be loaded by the user.

[kellertuer]

Same here: Currently 0% of extension documentation is rendered, since extensions are not included in the modules list of the documentation. If that is how it should be, then all things for the docs have been discussed.

I thought you suggested this can be changed, which sounded very reasonable to me. You also mentioned cross-referencing documentation. These are things which would be nice to have and I just wondered if we are all set on this or if more things need to be discussed.

Yes, but to me it was unclear whether that is wanted - and currently I can not do that since I can not get the docs to render locally, see the problem with wannier.
And sure cross-refs could be introduced then as well, but I would prefer to also do that not “blind” (without being able to render the docs myself). So probably someone else has to take that over.

and finally (and fourth) the issue that now 3 packages have to be loaded. I do not have a good solution for that. If that is too annoying, then we should maybe discontinue this and stick to the old scheme of using optim, that is also fine with me, then just close this PR.

Yeah, I think that will be too annoying. But I think we can work around that by making a small DFTKManopt package, which then hard-depends on all there dependencies. If this exposes (@rexports maybe) Manopt, Manifolds and the third package then you use that in your extension. That would mean only one extra package needs to be loaded by the user.

RecursiveArrayTools is the third one. I have no experience with @reexport, but I will leave that then also do you, so I am unsure how to do this.

But ok, then I know now, that I currently can not help any further, but I am of course available if there are things that need to be clarified further, so I will keep myself subscribed here, but put it aside (no more reminders) now then.

[mfherbst]

But ok, then I know now, that I currently can not help any further, but I am of course available if there are things that need to be clarified further, so I will keep myself subscribed here, but put it aside (no more reminders) now then.

Sounds good. I'll @mention you when I'm done or need help.