Open abussy opened 1 month ago
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.
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.
@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
?
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?
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?
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, ...)
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 aFFTBundle
can take place independently of aPlaneWaveBasis
, and only requiresfft_size
,unit_cell_volume
andarchitecture
. All subsequent calls to thefft()
andifft()
functions now require aFFTBundle
as argument, rather than aPlaneWaveBasis
. Instead of having a loose collection of FFT plans, thePlaneWaveBasus
now has a singleFFTBundle
as a field.The main changes are in these files:
fft.jl
, where theFFTBundle
type is defined. All functions in this file are now independent of thePlaneWaveBasis
type. An extra couple of functions (fft_matrix()
andifft_matrix()
) were also moved to this file fromPlaneWaveBasis.jl
, as they are also basis independent.Kpoint.jl
file was created, containing the definition of theKpoint
type and related functions. This was also taken out ofPlaneWaveBasis.jl
, for the main reason thatKpoint
must be defined for FFTs, which should be defined for the creation of aPlaneWaveBasis
. Note that this also reduces the size and complexity ofPlaneWaveBasis.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 originalfft()
calls with the basis as argument could be easily retrieved by defining functions such asfft(basis, ...) = fft(basis.fft_bundle, ...)
inPlaneWaveBasis.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.