Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Minor refactoring of the FFTs #996

Merged
merged 6 commits into from
Oct 28, 2024
Merged

Conversation

abussy
Copy link
Contributor

@abussy abussy commented Aug 14, 2024

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
Copy link
Member

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
Copy link
Contributor Author

abussy commented Aug 14, 2024

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
Copy link
Member

@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
Copy link
Member

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
Copy link
Member

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
Copy link
Contributor Author

abussy commented Aug 15, 2024

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, ...)

src/PlaneWaveBasis.jl Outdated Show resolved Hide resolved
src/Kpoint.jl Outdated Show resolved Hide resolved
@abussy
Copy link
Contributor Author

abussy commented Oct 28, 2024

Implemented all above suggestions, should be ready for merging.

Copy link
Member

@mfherbst mfherbst left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Small nits, otherwise good to go from my end.

src/fft.jl Outdated Show resolved Hide resolved
src/fft.jl Outdated Show resolved Hide resolved
@mfherbst mfherbst enabled auto-merge (squash) October 28, 2024 18:46
@mfherbst mfherbst merged commit cc05b51 into JuliaMolSim:master Oct 28, 2024
8 checks passed
@abussy abussy mentioned this pull request Nov 25, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants