Multicomponents

[epolack]

API question to settle

• [ ] Possibility to have components for TermKinetic.kinetic_energies, RealSpaceMultiplication.potential to prevent loops on the components; but maybe in a second time (i.e., future PR). Implementing one at a time needed generalisations? It was what was done for the preconditioner before (e.g., hash 9869eb59137997162f9660167db14177c14dbeb0)

Remaining points

∅

Future PR

• use a BlochWaves structure and BlochWave!!
• add occupation to BlochWaves
• not sure what to print by default (Base.show methods); maybe report to ψk::BlochWave.
• see for how to do in a memory friendly-way operations such that

  for (σ, n) in eachindex(ψk::BlochWave; region=[:components, :bands])
  

• Tullio.jl for when doing tricky stuff?

[epolack]

@mfherbst this is a very rough first draft, but shows that what would need to be changed in the inner working of DFTK is reasonable. (For now I have bypassed the need to change the operators in a hacky way, but it should straightforward to do in a cleaner way.)

From the user point of view, that means that by default the wave function (and after) the density will have the form of a matrix of vectors. Hence the need for the user to use a unnest or denest function that would allow the user to choose a specific set of components to view in a matrix (if single component) or a tensor. But that kinds of make sense as the user may get all the spins, or chose to plot a specific one for example.

Does it seems reasonable to you, and should I push forward to this PR, or do you find this to go in a wrong direction?

[antoine-levitt]

We discussed the following: by default things like psi (to be discussed what rho would be: do we want the full rho[alpha, beta, i, j, k] ?) would be arrays of arrays static arrays (yes...), ie psi[k][G,n][alpha]. This is the same in memory as psi[k][alpha,G,n]. There would be convenience functions nest(psi) : psi[alpha, G, n] -> psi[G, n][alpha], denest(psi) : psi[alpha, G, n] -> psi[G, n][alpha], and denest(psi, alpha) which would select only component alpha. The main objection to this is that it's pretty annoying to work with such a nested structure: eg printing by default is atrocious...

[antoine-levitt]

Ping @mfherbst does the above seem reasonable? We'd like to go ahead soon

[mfherbst]

Thanks for the reminder. I'll take a look!

[mfherbst]

Hmm. I can really see this potentially being very confusing for new users ... in particular if this is exposed to them via the scfres. The frequent constructs like eltype(eltype( ... )) or the frequent reinterpret calls make the code quite hard to follow.

I think before we do this we should try to think if we can somehow encapsulate the tricks as this would also avoid scattering them all around. Can you not use ComponentArrays for this for example? If not we should think about the effort needed to roll own array type for such things. That might also help to improve the situation for spin-related stuff ... the solution we have there is also not always super readable at places.

[mfherbst]

Ah now I get why componentarrays don't work: They encapsulate the wrong way round.

But in that sense your basic "element" becomes not T, but some StaticArray type, right? So maybe it could help to make things simpler if this is exposed more. Like right now the T of PlaneWaveBasis is used to deduce a lot of types. But now it should probably be the respective SArray{T} right?

What I really would like to avoid is to add a lot of additional complexity to the code, which then gets combinatorially multiplied with AD, GPU and other things, such that we end up writing a lot of workaround code to get things working with our highly specific data layout.

[mfherbst]

OTOH, I have to admit, it's pretty nice to see this only adds 200 lines of code so far. So our code structure in DFTK seems pretty decent in general.

[mfherbst]

BTW. I guess it might be good to schedule a call to discuss this after the long weekend.

[epolack]

ping @antoine-levitt for comments. Some stuff are still hacks, but maybe we should discuss how to do this cleanly. The setup is here and should allow to bypass DFTK's spins. It is usable, and strided FFTs work (i.e., we can use it for TBG).

Needs to investigate a bit the timings to see if some computations are slowed due to this.

[epolack]

(Comments can be done quickly tomorrow with @LaurentVidal95 if you have time, it is mostly related to spin stuff.)

[epolack]

Failing CI for documentation may be due to current memory leak that happens on lot of projects when using v1.9.3…

[antoine-levitt]

to_composite_σG
- to_composite_σG(basis, kpt, ψk) : ψk[σ,G] -> ψk[σG]
- to_composite_σG(basis, kpt, ψk) : ψk[σ,G,n] -> ψk[σG,n]
- to_composite_σG(basis, ψ) : ψ[k][σ,G] -> ψ[k][σG]
- to_composite_σG(basis, ψ) : ψ[k][σ,G,n] -> ψ[k][σG,n]

from_composite_σG

[antoine-levitt]

If I recall correctly the sizes it reported (i.e. size(qr(X)) were fine, just when you did the Matrix(qr(X).Q) then you got a full NxN array back.

OK but then this is massively expensive, and by just taking the [:, 1:n] we are hiding a massive perf issue. Do you remember how to trigger this?

[mfherbst]

It was really random. Like I had it once, but then removed the workaround later and it was ok, but had to reintroduce the workaround at some other point to make it work. It could have been to do with versions of CUDA.jl the underlying cuda library etc. Main point was to get it to work. The CUDA code is horribly slow anyway, so I have not yet bothered to investigate.

[epolack]

Antoine and ping-ponged, and we have mostly settled on this API. Happy to have your opinion @mfherbst.

[epolack]

@mfherbst Could you give me some feedback on this PR? It takes quite some time to keep it up to date with the master branch.

[mfherbst]

Could you give me some feedback on this PR? It takes quite some time to keep it up to date with the master branch.

Understood. I'm pretty busy right now, but I'll do it as soon as I can.

[epolack]

Sure, Thanks! I just wanted to make sure you had it on your radar :)