DFTK.jl icon indicating copy to clipboard operation
DFTK.jl copied to clipboard

Autodiff

Open antoine-levitt opened this issue 3 years ago • 2 comments
trafficstars

Lots of boxes! Gotta check them all!

Forward mode

  • [ ] libxc interfacing
  • [ ] upstream FFT stuff
  • [ ] figure out the T = promote_type pattern and use it consistently
  • [ ] make sure everything works with multiple partials

Reverse mode

  • [ ] make it work
    • [ ] multiple kpoints
    • [ ] symmetries
    • [ ] MPI
  • [ ] make it clean-ish
    • [ ] more efficient compute_density rrule
  • [ ] make it fast-ish

Solvers

  • [ ] make the non-split omega_plus_k work for finite temperature
  • [ ] use the Schur automatically, get rid of the distinction between psi and psi_extra

Properties

  • [ ] figure out how to make other types of properties (eg phonons) fit into the AD framework

antoine-levitt avatar Mar 04 '22 10:03 antoine-levitt

Hey I would like to contribute to the issues here. (one at a time, at least :D). Could you elaborate on "interfacing libxc" ?

dynamic-queries avatar May 09 '22 09:05 dynamic-queries

Hi @dynamic-queries. Thanks for your interest and sorry it took us so long to get back to you. Pretty busy times at the moment ;) Indeed any help on this is very much welcome!

To answer your question upfront, "interfacing libxc" refers to writing custom rules for forward-mode and adjoint mode for Libxc.jl such that differentiation through this foreign c++ code can be performed. That is actually pretty tricky and would thus not be what I would recommend you to start with if you want to get into the topic. To summarise in one sentence it involves digging into low-level C code and massaging it to work with AD, for which you need a pretty solid idea how XC functionals are usually implemented.

Of course feel free to give that a go, but I just wanted to also give you a few pointers to things that can be tackled more easily and are probably better to get started:

  • In src/workarounds/forwarddiff_rules.jl we have a bunch of workarounds related to applying forward-diff to FFTW that really should not live in DFTK. Much rather they should be upstreamed to an appropriate place in ForwardDiff. I think there is already a draft-PR on this that could be completed. @niklasschmitz can add details here if this interests you.
  • Similarly there are a bunch of libraries we commonly use (such as Optim.jl for running geometry optimisation or NLSolve.jl), which currently do not out of the box play nicely with AD. Again contributing some upstream rules to compute gradients for the optimisation procedure (in this case just a Newton step) would be appreciated.
  • A lot of the things we did for forward-mode AD have only been tested for a single partial derivative and some more testing / bug fixing is very likely needed to get it to work for multiple partials.

In case any of this interests you, just let us know. I'm happy to provide more details.

mfherbst avatar May 19 '22 18:05 mfherbst