Proposal to replace `direct_minimization` and switch to Manopt solvers

[kellertuer]

This semester I have student (@elorannug) working on the direct minimization for scf and port it to Manopt.

In this issue I would like to check how to best continue. Before describing the current for of direct_minimization and a proposed new form, let's first state the main questions that I would like to discuss

Should the old method stay but be moved into an Optim.jl extension? That makes a new release at least only “mildly breaking”, a user has to load Optim and potentially Linesearches themselves.
The disadvantage is, if someone loads Optim and Manopt, both extensions have to work nicely together still.

What do you think @antoine-levitt ?

Current form

In the current scf/direct_minimization.jl the user has an easy interface based on a PlaneWaveInterface that allows to set the usual maxiter/tol stopping criteria, choose a solver from Optim and (optionally= provide a start point.

The only keyword I personally do not fully understand is alpha_guess – is that a function to have a good initial step size guess for the line search?

New form

The basis will stay the only mandatory parameter.

All existing kwargs would stay, we we just use them differently internally

Keyword arguments

• stopping_criterion= will use the existing maxiter/tol to introduce a default stopping criterion based on that, but a user could also provide their own one
  one challenge might be that we currently have one that we do define ourselves – if that is only introduced in the extension, it would be quite private, and hence not so easy to use for an end user, but I think I know a trick
• solver= provide an AbstractManoptSolverState to specify which solver to use
• peconditioner= a manopt preconditioner, especially when using LBFGS (there is also a second one its, type)
• manifold= specify whether to use Stiefel or Grassmann; one could also allow for their Generalized (i.e. weighted) versions
• kwargs... can be used to decorate both problem and state, such that debug= and such would also work in this function call.

Since this requires both Manopt anf Manifolds, I would propose to do this in a DFTKManoptManifoldsExt.jl with the above keywords. It would still only be “mildly bearking” for end users.

Let me know what you think.

[antoine-levitt]

Hi, that's awesome, I'm very much looking forward to the results! I remember benchmarking LBFGS against manopt and seeing no significant difference, but that was a long time ago and not very extensive testing. CC @gkemlin @jonas-pueschel who might be interested as well. I would say in a first stage don't worry about API, just do a direct_minimization_manopt outside of DFTK and benchmark it against the current one. If it's better, then we can just throw away the old one and depend on Manopt.

is that a function to have a good initial step size guess for the line search?

as far as I remember yes. Stepsize control in Optim is very confusing and what we do now is possibly suboptimal, that would be interesting to try.

[mfherbst]

I'm ok with the extmodule of Manopt and Manifolds and breaking changes. I'm not even sure the direct minimisation is exported at all.

update: is exported ... still breaking change is fine here.

[kellertuer]

We currently have a direct_minimization_new´ that is benchmarked and it looks quite nice, since we do not need the unsafe_unpack` any longer – the current function is not yet as generic as I described above, for example the manifold is still hard-coded. That is what I want to do either soon-ish or after the thesis is handed in.

For now it is better in most experiments if I remember correctly, might also be, because we added a small cache in cost&grad.

But I wanted to also check with the breaking ideas. Good that I also can just break stuff then ;)

[mfherbst]

Yeah. In the end where I would like to arrive is that we have one function like ground_state(basis; alg=DirectMinimisation(; manopt args or something like that), other_kwargs...) in DFTK, where the other_kwargs would be roughly standardised and behave roughly similarly across all methods to compute the ground state. Like that we could finally unify the zoo of methods (self_consistent_field, newton, potential_mixing and direct_minimization) we have for this purpose under a common interface.

But I'd suggest you first do the implementation and than we massage it in a second step to get there.

[kellertuer]

I would probably start with a direct_minimization but we can check on that PR then how to maybe already take a few steps towards your ground_state function, where I do not have a good overview about the rest of the zoo.

[mfherbst]

Oh yes, sure, I agree! This was more to point at the longer-term plan.

[kellertuer]

I have two questions from a discussion yesterday,

1. There are three “callbacks” that to see look a bit more like stopping criteria, cf
   

   DFTK.jl/src/scf/scf_callbacks.jl

   Lines 111 to 152 in 5efd87f

   	# TODO Convergence ideas:
   	# - Flag convergence only after two subsequent steps converged
   	"""
   	Flag convergence as soon as total energy change drops below a tolerance.
   	"""
   	struct ScfConvergenceEnergy
   	tolerance::Float64
   	end
   	function (conv::ScfConvergenceEnergy)(info)
   	if last(info.history_Δρ) > 10sqrt(conv.tolerance)
   	return false # The ρ change should also be small to avoid the SCF being just stuck
   	end
   	length(info.history_Etot) < 2 && return false
   	ΔE = (info.history_Etot[end-1] - info.history_Etot[end])
   	ΔE < conv.tolerance
   	end
   	"""
   	Flag convergence by using the L2Norm of the density change in one SCF step.
   	"""
   	struct ScfConvergenceDensity
   	tolerance::Float64
   	end
   	(conv::ScfConvergenceDensity)(info) = last(info.history_Δρ) < conv.tolerance
   	"""
   	Flag convergence on the change in Cartesian force between two iterations.
   	"""
   	mutable struct ScfConvergenceForce
   	tolerance
   	previous_force
   	end
   	ScfConvergenceForce(tolerance) = ScfConvergenceForce(tolerance, nothing)
   	function (conv::ScfConvergenceForce)(info)
   	# If first iteration clear a potentially cached previous force
   	info.n_iter ≤ 1 && (conv.previous_force = nothing)
   	force = compute_forces_cart(info.basis, info.ψ, info.occupation; ρ=info.ρout)
   	error = isnothing(conv.previous_force) ? NaN : norm(conv.previous_force - force)
   	conv.previous_force = force
   	error < conv.tolerance
   	end

   
   Should I provide those as stopping criteria as well? In a current draft, I only have the first – and still have to find a good way to do that as a stopping Criterion in Manopt – but a question is, whether it is worth to thing generic enough to do the other two as well.
   I would probably not depend on the info dict but on the objective/cost that has these available anyways (even cached).
   The current one does only the first as far as I can see.

2. The current direct_minimization has a callback= keyword, cf

   DFTK.jl/src/scf/direct_minimization.jl

   Line 74 in 5efd87f

   callback=ScfDefaultCallback(),

   that is called in the very end. What is that for? Should I do something like that as well?

[mfherbst]

1. Yes, that's correct they are stopping criteria. Ideally their effect should somehow be available for direct minimisation. In particular the "important" stopping criteria are ScfConvergenceEnergy and ScfConvergenceDensity. In practice people use both. Let me add that if you have good proposals that make life easier for you, we can change the inner workings of these guys. But the basic principle should be that all algorithms in DFTK can be tuned using the same set of structures / is_converged callbacks. That also applies to the info and how it's handled. Our approach is certainly not perfect there and we can change it, but it should be done consistently across the

2. That is intended for printing and "information extraction". Ideally it's not just called at the end but somehow in every algorithmic step. Again ... if you have a better mechanism in ManOpt, we can discuss and adapt our approach.

[kellertuer]

Concerning 1) my main problem is that info is no where documented and just something magic to me. I have no clue what to expect when in there and within Manopt that thingy simply does not exist. So I can not use exactly the same callback ideas.
What I can do is to accept those as parameters and “wrap” them accordingly to make something similar (I can access the tolerance then and know which tolerance is meant).

An alternative is that I write 3 Manopt-Stopping-Criteria to do the same – than it is the same functionality but different structures.

For 2) Manopt has both debug/record capabilities anyways – all of those will be available, so I would for now use those.
The callback I linked to is literally really only called at the end. I will first use the Manopt thingies (since it bridges to Manopt) and we can check how to also “wrap” a callback in there as well – I think that is doable. But functionality (i.e. the solver runs) first ;)

[antoine-levitt]

Yeah, the design we settled with is that the solver calls back its parent both to inform it (with the info field). That's very flexible, because the caller can decide by itself when it wants to stop. The reason info is not documented is that we want to keep it relatively flexible because when we add features we sometimes add info to it. DFTK currently has no great concept of "state": conceptually it's (occupation, psi) but we precompute lots of stuff from it (eg density) that has to be carried around. So there is the "basic state" which is just occupation and psi, and the "derived state" (all things computed from it).

[kellertuer]

It is very flexible, sure, but for me always unclear what it actually is. So for me it is all magic and I have to hope it works.

but then I will first work on a “internal Manopt style” and then we can check how to best map that into the DFTK scheme of things.

[mfherbst]

Sounds good. First make it work and than we can massage it into our info magic 😄.

[kellertuer]

I worked a bit the last few days with all this, and for now I think I am not able to program in the same style as the current direct_minimization

Everywhere it seems there are functions within functions that seem to use variables at least from parent scope if not parent-parent scope (sometimes a bit hard to track), and this magic info dictionary being handed around feels really strange to see, because only in rare cases things are actually stored in there but in a lot of places info.SuperMagicName is accessed (and hence assumed to just always magically exist). One example is info.history_Etot.

I think I tracked those down and understand them now a bit, I might even extend a few of the one-line-doc-strings a bit,s Ince I had to write that for myself anyways; but I will still not be able to provide that Manopt fills these during the iterations magically.
I can provide a “what-ever-you-like-recorded” history after a solver run for sure, but just a callback that magically assumes info.fnumps to exist to store something, Manopt can not offer that, and it is also not a planned feature.

But we will see. For now I am much much much much slower than I thought I would be, I will report back once I have something that maybe runs.

[mfherbst]

I see. Sounds like we should have a discussion of your gripes and see where we can meet in the middle. Having a slightly better documented info can help. Let me know if you have the need to talk.

[kellertuer]

The main problem might be to connect both approaches in the end (and how to export structs I only define in the extension for now).
I think I have a good path to follow to get my extension to work, but it will first be nearly without “info-struct-magic”.

I think I for now reverse-engineered for myself that the info provides when and where.
I can not mag that functionality into Manopt (i.e. “handing around” an info-dictionary within the elements of Manopt), so we have to check how to best construct one afterwards.

For now I might also extend a few of the “one-line-function-docs” because I needed to understand their input&output more precisely, e.g. I turned

DFTK.jl/src/terms/Hamiltonian.jl

Lines 187 to 191 in 5efd87f

	"""
	Get energies and Hamiltonian
	kwargs is additional info that might be useful for the energy terms to precompute
	(eg the density ρ)
	"""

into (±typos)

https://github.com/kellertuer/DFTK.jl/blob/36675995ec671c6e95794b187f38b05a34d86926/src/terms/Hamiltonian.jl#L187-L206