Minor refactoring of the FFTs

[abussy]

In this PR, I propose some refactoring of the FFTs, in order to try and help disentangle the massive PlaneWaveBasis type.

The rational behind the proposed changes is that FFTs can, in principle, take place independently of the basis. As a result, a new FFTBundle type is created, containing all data required to perform such operations. The creation of a FFTBundle can take place independently of a PlaneWaveBasis, and only requires fft_size, unit_cell_volume and architecture. All subsequent calls to the fft() and ifft() functions now require a FFTBundle as argument, rather than a PlaneWaveBasis. Instead of having a loose collection of FFT plans, the PlaneWaveBasus now has a single FFTBundle as a field.

The main changes are in these files:

• fft.jl, where the FFTBundle type is defined. All functions in this file are now independent of the PlaneWaveBasis type. An extra couple of functions (fft_matrix() and ifft_matrix()) were also moved to this file from PlaneWaveBasis.jl, as they are also basis independent.
• A Kpoint.jl file was created, containing the definition of the Kpoint type and related functions. This was also taken out of PlaneWaveBasis.jl, for the main reason that Kpoint must be defined for FFTs, which should be defined for the creation of a PlaneWaveBasis. Note that this also reduces the size and complexity of PlaneWaveBasis.jl.
• PlaneWaveBasis.jl, which was simplified.

The changes in all other files essentially consist in replacing argument in FFT function calls, e.g. fft(basis, ...) --> fft(basis.fft_bundle, ...). Note that the original fft() calls with the basis as argument could be easily retrieved by defining functions such as fft(basis, ...) = fft(basis.fft_bundle, ...) in PlaneWaveBasis.jl. I guess this is a stylistic question, and I personally like having these explicit calls (I am happy to change if the DFTK standard is different).

In the future, I plan to open a new PR where I take some of the K-point complexity out of PlaneWaveBasis.jl, by creating a new type containing all the necessary data for a set K-points.

[antoine-levitt]

Nice, thanks! I'll have more time for a proper review later, but just some comments from a quick look. Some of this I like, some not so sure. FFTBundle - > FFTPlans ? Also gvectors and rvectors do belong to the pwbasis, not just the fft thing. I kind of like just passing basis to fft, since we do this for a number of routines, and also because the normalization and interpretation do depend on the actual system, it's not just an fft, but we can discuss that.

[abussy]

Of course, I am happy to discuss!
FFTBundle is not just the FFT plans: it also contains normalization constants, G_vectors and r_vectors.

Also note that, in the current state of the PR, I left G_vectors and r_vectors fields in the PlaneWaveBasis. They are taken from the fft_bundle field of the basis for consistency.

[mfherbst]

@abussy This is a very good start.

For me what has always bothered me is that the ffts and the basis are indeed super entangled, but the relationship is not always 1:1. E.g. one may have multiple fft plans in different precisions for the same basis. That is how the idea for this PR came up. In that sense I don't think it's bad to have the G_vectors and r_vectors (which are not the cartesian things at this stage) stored in the fft object and I would even go one step further and remove them from the basis. Similarly I like the idea to enforce that fft calls need to be done with the fft object, but we can discuss the details there.

In that spirit I actually find it weird it has a unit_cell_volume. Maybe this should just be used during construction time for convenience and not stored ?

For the name. I would even go more high-level and just call it FFT or FFTsetup ?

[antoine-levitt]

OK then it's more something like a FFTGrid that handles all the operations related to that (including r and G vectors), it's more than what I had in mind but sure why not. Still not too sure about the signature of fft(), maybe if you're going to change the kpoints structure wait until then? Maybe it'll lead to some simplifications, I never really liked the fft(basis, rho) vs fft(basis, kpoint, psi) thing. Maybe if kpoint contains a reference to the fft grid we can do fft(basis.grid, rho) and fft(kpoint.grid, psi) or something?

[antoine-levitt]

Also that's a potentially pretty disrupting change, so let's merge all the non breaking internal structure change first and then do the external API change in a separate PR?

[abussy]

For the name, I like @antoine-levitt idea of FFTGrid: it's more descriptive than FFTBundle and less vague than simply FFT.

In that sense I don't think it's bad to have the G_vectors and r_vectors (which are not the cartesian things at this stage) stored in the fft object and I would even go one step further and remove them from the basis.

I personally see no problem from removing it from the basis. Then we can either define a G_vectors() function that takes a basis and returns the G_vectors of the fft_grid field, or refer to the G_vectors explicitly throughout the code (i.e. basis.fft_grid.G_vectors)

In that spirit I actually find it weird it has a unit_cell_volume. Maybe this should just be used during construction time for convenience and not stored ?

I agree with that. The unit_cell_volume is only used at initialization for the normalization constants. Removing it from the fields would further separate the FFTGrid type from a basis/model.

Still not too sure about the signature of fft(), maybe if you're going to change the kpoints structure wait until then? Maybe it'll lead to some simplifications, I never really liked the fft(basis, rho) vs fft(basis, kpoint, psi) thing.

I am not really sure how I will tackle the kpoints structure yet, but I will keep that in mind. As for now, I would leave the fft() calls as they are in fft.jl, i.e. not taking a basis as argument, and effectively separating the FFTs from the PlaneWaveBasis. In order to not disrupt the API and the user experience, I can add some short function definitions in PlaneWaveBasis.jl that forward the FFT calls: fft(basis, ...) = fft(basis.fft_grid, ...)