Decompose Wannier-function densities onto a Gaussian × spherical-harmonic basis

[elinscott]

This PR is a WIP. I'm very happy to discuss the big picture but there's no need to start reviewing the code details just yet

Summary

This PR adds an optional post-processing step to wannier90.x that decomposes each Wannier function's density $|w_n(r)|^2$ onto an orthonormalised basis of Gaussian radial functions × real spherical harmonics centred on the WF's own centre, writing per-WF coefficient and power-spectrum files. These SOAP-style descriptors follow Himanen et al., Adv. Sci. 7, 1902333 (2020).

The motivating use case is machine-learning workflows that need compact, rotation-covariant descriptors of individual Wannier functions. The koopmans package uses these power spectra to predict screening parameters (see Schubert et al., npj Comput. Mater. 10 299 (2024)).

Because the decomposition is very general and can be performed purely on wannier90 content, I thought it would naturally slot into wannier90. I realise it is a substantial addition and am very happy to discuss further or publish it is a separate code if you think it is genuinely beyond the scope of what you want Wannier90 to do.

Theory

Each Wannier density $\rho_n(\mathbf{r}) = |w_n(\mathbf{r})|^{2}$ is expanded about its centre $\mathbf{c}_n$ as

$$c^{(n)}_{n'lm} = \int_{|\mathbf{r}-\mathbf{c}_n| \le r_{\mathrm{cut}}} \rho_n(\mathbf{r}) g_{n'l}(|\mathbf{r}-\mathbf{c}_n|) Y_{lm}(\theta,\varphi) \mathrm{d}V,$$

evaluated in terms of radial $g_{nl}$ and angular $Y_{lm}$ basis functions over the Born–von-Kármán grid. The rotationally invariant power spectrum is then given by

$$p^{(n)}_{n_1 n_2 l} = \sum_{m=-l}^{l} c^{(n)}_{n_1 l m} c^{(n)}_{n_2 l m}, \qquad n_1 \le n_2,$$

Radial basis

The primitive radial functions are Gaussians,

$$\phi_{nl}(r) = r^{l} e^{-\alpha_{nl} r^{2}}, \qquad \alpha_{nl} = -\frac{\ln(\mathrm{thr}/r_{\mathrm{thr},n}^{l})}{r_{\mathrm{thr},n}^{2}},$$

with $\mathrm{thr} = 10^{-3}$ and the threshold radii $r_{\mathrm{thr},n}$ equally spaced on $[r_{\mathrm{min}}, r_{\mathrm{max}}]$ (decompose_r_min, decompose_r_max), i.e. each Gaussian has decayed to the threshold at its $r_{\mathrm{thr},n}$. Their overlaps are evaluated analytically,

$$S_{nn'l} = \int_0^{\infty} \phi_{nl}(r) \phi_{n'l}(r) r^{2} \mathrm{d}r = \frac{1}{2} \Gamma(l + 3/2) (\alpha_{nl}+\alpha_{n'l})^{-(l+3/2)},$$

and the orthonormal radial basis is obtained by Löwdin symmetric orthogonalisation,

$$g_{nl}(r) = \sum_{n'} \beta_{n'nl} \phi_{n'l}(r), \qquad \beta^{(l)} = V \mathrm{diag}(e^{-1/2}) V^{T},$$

where $V$, $e$ are the eigenvectors/eigenvalues of $S^{(l)}$.

Angular basis

The angular part uses real spherical harmonics $Y_{lm}(\theta,\varphi)$ in the scipy real convention (Condon–Shortley phase, $\sqrt{2}$-normalised $m \neq 0$ components).

What the code does

With wannier_decompose = .true. (runs during the plotting phase, also under restart = plot):

1. Builds each WF on the real-space grid of the Born–von-Kármán supercell (the mp_grid multiple of the unit cell).
2. Normalises each |w_n|² to unit integral.
3. Projects each selected WF density onto the basis functions within a sphere of radius decompose_r_cut about the WF centre, using the minimum-image convention
4. Writes, per WF n:
   • <seed>_NNNNN.coeff — expansion coefficients c_nlm
   • <seed>_NNNNN.power — the invariant power spectrum p(n₁,n₂,l) built from those coefficients.
5. Optionally, if decompose_centres_file is given: also decomposes the group density (sum of all this run's normalised WF densities) about every centre listed in that file, writing <seed>_gc_NNNNN.coeff.

New keywords

Keyword	Type	Default	Meaning
wannier_decompose	logical	.false.	Enable the decomposition
decompose_r_cut	real (Å)	— (required)	Cutoff radius of the projection sphere
decompose_n_max	integer	6	Number of radial basis functions per l
decompose_l_max	integer	6	Maximum angular momentum
decompose_r_min	real (Å)	0.5	Smallest Gaussian decay radius
decompose_r_max	real (Å)	4.0	Largest Gaussian decay radius
decompose_list	list	all WFs	Which WFs get own-centre output (group density always sums all WFs)
decompose_centres_file	string	—	File of extra centres (Å, Cartesian) for group-density decomposition

Implementation

• src/decompose.F90 (new) — module w90_decompose: basis math (decompose_radial_params, decompose_real_ylm), projection (decompose_project), power spectra (decompose_power_orb), and the root-only driver + I/O (decompose_main)
• src/plot.F90 — the WF real-space grid construction inside plot_wannier is extracted into a shared helper plot_build_wannier_grid
• src/wannier90_readwrite.F90 / src/readwrite.F90 / src/wannier90_types.F90 — support for new keywords.

Validation and tests

• Basis functions, projection coefficients and power spectra were validated against the legacy Python implementation in koopmans (coefficients agree to ~2·10⁻⁷ relative — limited by the reference's numerical-quadrature overlaps vs. our analytic ones; power spectra to ~7·10⁻¹⁵ on identical inputs).
• New test-suite cases (serial, group plot):
  • testw90_decompose, testw90_decompose_power, testw90_decompose_gc — GaAs (inputs shared with testw90_example01), covering own-centre coefficients, power spectra and the group-density/centres-file channel, with new parsers in test-suite/tools/parsers/;
  • testw90_decompose_sphere — a synthetic spherical density in a triclinic cell)
  • testw90_decompose_l2 — a synthetic pure-l=2 density in a cubic cell

Documentation

• User guide (docs/docs/user_guide/wannier90/parameters.md), parameters.xml, plot-parameters CSV, and a CHANGELOG entry.

Limitations / notes

• Root-only implementation: the decomposition itself runs on the root rank after the (already parallel) grid accumulation. Should be fine for the intended grid sizes; parallelisation is possible later if needed.
• Spinor WFs are not supported for the moment.
• The total-density-across-runs story assumes integer occupations (each run's group density is the sum of its normalised WF densities)
• The UNK-file error-handling fix from Fix MPI hang/abort on UNK read failure in plot_wannier #647 is merged into this branch, applied inside the extracted plot_build_wannier_grid helper (where the k-point read loop now lives)

AI Declaration

I have used Claude extensively to port the legacy code (originally written in python by a student of mine) over to fortran, document the code, etc. I am in the process of reviewing all the code. It has been validated against an independent reference implementation and the new test-suite cases described above.

[jaemolihm]

Hello @elinscott , this looks very interesting! I have some questions.

1. In the decomposition, could you also decompose w_i(r) directly instead of decomposing rho_i(r) = |w_i(r)|^2? I think this would directly connect the atomic-orbital-like (s,p,d,...) nature of WFs.

2. How do you deal with the real-space grid? In plane-wave code like QE, wavefunctions are represented in plane waves, then taking overlap with spherical harmonics is much easier. In your Wannier implementation, do you also FFT to plane waves and compute there? Or do you take overlaps in real space? In the latter case, I guess there could be some aliasing error when the Wannier centers are not exactly on the FFT grid.

3. Practically, I often find storing the full UNK in a real-space grid is quite demanding for disk usage. So I prefer to postprocess using the plane-wave representation of the DFT code directly rather than converting to real-space UNK files, although this means the postprocessing can no longer be DFT-code-agnostic. I wonder what your experiences are.

[elinscott]

Hi @jaemolihm thanks for the nice questions! Let me answer them below

1. In the decomposition, could you also decompose w_i(r) directly instead of decomposing rho_i(r) = |w_i(r)|^2? I think this would directly connect the atomic-orbital-like (s,p,d,...) nature of WFs.

Sure, the machinery is agnostic. We went with density for a few reasons:

• the motivating application comes from the world of orbital-density-dependent functionals, where the orbital density is our fundamental object
• dealing with $|w_n|^2$ removes the arbitrary global sign/gauge of WFs (which would cause uniqueness issues when used as a descriptor for ML, though I'm sure there would be some way of constructing a clever descriptor that would digest it...)

2. How do you deal with the real-space grid? In plane-wave code like QE, wavefunctions are represented in plane waves, then taking overlap with spherical harmonics is much easier. In your Wannier implementation, do you also FFT to plane waves and compute there? Or do you take overlaps in real space? In the latter case, I guess there could be some aliasing error when the Wannier centers are not exactly on the FFT grid.

I do it naively: overlaps are taken in real space. WF centres don't lie on the grid, so yes, there is a discretisation error for off-grid centres. In practice it's small: the test suite includes a synthetic spherical density about a deliberately off-grid centre in a triclinic cell, and the spurious $l \ge 1$ content is $\sim 10^{-8}$ of the $l=0$ signal. (That said, that benchmark is deliberately smooth on the grid scale; for realistic densities with sharper radial basis functions the discretisation error would be larger... I will investigate more!)

I would absolutely consider a reciprocal-space grid implementation if it turned out there were situations where the aliasing causes problems.

3. Practically, I often find storing the full UNK in a real-space grid is quite demanding for disk usage. So I prefer to postprocess using the plane-wave representation of the DFT code directly rather than converting to real-space UNK files, although this means the postprocessing can no longer be DFT-code-agnostic. I wonder what your experiences are.

You're absolutely right, full real-space UNKs are heavy. This was a deliberate tradeoff to get the code sitting entirely within wannier90 rather than writing something DFT-code-specific. I'm testing it at the moment and thus far it hasn't been prohibitively cumbersome... but that experience may well change!

[jaemolihm]

Thanks a lot for the detailed reply! Looking forward to seeing how Gaussian the WFs are 😃