JuliaMolSim/DFTK.jl

Blochpocalypse

Opened this issue · 7 comments

This is a writeup of a discussion with @eric-cances.

Let me prefix this by saying that this would imply a major rewrite of DFTK, add a major layer of complexity to the code, and looks like So Much Work, so it's probably not something we actually want to do. Still, fun to dream :-)

The starting question is: how do you use AD to express "q-dependent" quantities (ie, quantities that depend on perturbations that break the cell periodicity, such as phonons, dielectric properties, etc) ? Although we are able to express quite nicely periodicity-preserving response properties using AD, periodicity-breaking response computations involve basically completely new code paths, involving a lot of error-prone bookkeeping (CC @epolack ). Higher-order derivatives are even worse.

A good answer to this is to make explicit the fact that periodic DFT codes are essentially optimizations of supercell computations (by supercell, I mean the equivalent supercell implied by the kpoint sampling of the unit cell). So we could structure the code in reverse: pretend that we are actually solving the full supercell problem, and rely on clever data structures to perform the optimization in code rather than on paper. A nice way to do this is to use the Bloch representation systematically. In this way, every quantity is expressed on the supercell:

  • arrays of positions (eg atomic displacements) are represented by their Bloch components, one Vec3 per cell atom and per k point
  • potentials, densities, etc are an array per kpoint
  • orbitals are general supercell orbitals (each orbital is not associated with a kpoint)
  • the Hamiltonian is a nkpt x nkpt matrix
    Of course this is super inefficient if stored as such, but clever data structures can be used. Eg for the periodic problem (where displacements, densities, potentials and hamiltonian are cell periodic) the positions and potentials are sparse vectors with only the "DC" (periodic) component populated, and the Hamiltonian is diagonal. Operations can be implemented with these data structures. Eg a sparse vector with component k filled times a sparse vectors with component k' filled is a sparse vector with component k+k' filled. Ideally, this is more or less transparent.

Then, self_consistent_field checks if the problem is periodic, and then solves it using the standard methods. The dual of self_consistent_field, however, does not need the variation of the problem to be periodic; it only needs the primal to be. It then does rhs = δV * ψ [note: this is the product of supercell quantities; eg in the phonon case, if δV is exp(iqr) δVq(r) this is just shifting k to k+q and multiplying by δVq], passes this to the sternheimer solver, etc

Pros:

  • clearly The Right Thing to do: very elegant, represents the underlying math very well
  • one single code path
  • makes it easy to implement higher-order response with AD
  • nobody's been crazy enough to do it like this before (I think...), so it's a nice challenge
  • that kind of API can possibly be written very nicely in julia

Cons:

  • basically need to rewrite a large portion of DFTK
  • large plumbing work on getting the internal API right
  • tricky datastructures to do this efficiently
  • large abstraction layer that can be confusing to newcomers
  • possible silent performance issues (I'm not too concerned about this because if the sparsity code fails the computation is so large it'll just crash)

Interesting idea. I like it very much.

You guys probably have more details in your head, but I'm not sure it needs that much rewriting. Could we not essentially just "wrap" what we have now to get a first version working?

My suggestion would be to try this in a simpler (e.g. 1D) setting with a good master student, basically starting from scratch. That would allow us to explore the API problem and what is needed to really get the new features right.

Happy to discuss this further at some point. Definitely fun to think about.

Yes that would have to be done as you suggest in three stages: 1D prototype, parallel code that uses DFTK's routines, then integration into DFTK.

with a good master student

Know somebody?

Not right now, but I'll keep it mind when I'm searching.

Just out of curiosity: would one be able to compute phonons with q incommensurate with the k-point grid?

Not really. Nor would it handle response to homogeneous fields. I wonder if there's a way to do that... But the workflows are significantly different in these cases, no? Eg in the incommensurate case you have to compute a new set of eigen functions. Don't people compute on q values in the kgrid then interpolate?

Indeed, the workflow would be different. q interpolation is possible, but it's still useful to compute arbitrary q point. Utilizing AD also for such workflows would be nice, especially for computing dV*psi.

That might be a good compromise, to only target ADing dV*psi, and handle the rest by hand. I haven't thought through the implications, but it's probably less work than trying to shoehorn the whole problem into AD.