Modernizing the GPAW code

or “Paying down technical debt” or “Future proofing the code”.

Goal:

Have a new version of the code that is

  • easier to work with for developers (and users)

  • fully backwards compatible with the old code (also gpw-files)

Agenda:

  • what is wrong with the old code

  • the plan

  • the gpaw.core module

  • layout of the new code

  • testing the new code

What is the problem?

  • The GPAW object is very big and many other objects inherit from it (LCAOTDDFT, TDDFT, AtomPAW, ExcitedState, SolvationGPAW, SJM).

    Solution: Split into smaller objects (ASECalculator + DFTCalculation + …). Use composition instead of inheritance.

  • Wave functions, electron densities and potentials are stored in numpy.ndarrays, but what about the associated k-point, the boundary conditions, the unit cell and the parallel distribution?

    Solution: Use intelligent container objects (gpaw.core module).

  • Hamiltonian too big: Split into Potential + Hamiltonian + PotentialCalculator.

  • LocalizedFunctionsCollection: Replace with AtomCenteredFunctions that has better design.

  • PWDescriptor is too big: Separate from its real-space buffer and fft-plan.

  • Sane normalization of PW-coefficients.

  • Sane normalization of occupation numbers (1.0 should mean filled).

  • Too many GPAW.get_this(...) methods: Move to specialized objects.

Transition period

Until this project is completed there will be two DFT codes in our Git repository.

Environment variable:

$ GPAW_NEW=1 python script.py

Imports:

from gpaw import GPAW  # will look at os.environ['GPAW_NEW']
from gpaw.calculator import GPAW  # old GPAW
from gpaw.new.ase_interface import GPAW  # new GPAW

The plan

  1. Make test-suite (2000+ tests) + AGTS (500+ tests) (“advanced GPAW test-suite” or “big integration tests”) pass:

    • rewrite things

    • use wrapper objects (XC-functionals, symmetry-analysis, Poisson-solvers, …)

    • add thin backwards compatibility layer (make stuff like calc.wfs.kpt_qs[7][0].P_ani[27] work)

  2. Clean up:

    • remove old code

    • move stuff from gpaw.new module to other places

    • naming of things?

  3. Refactor wrappers and compatibility layer

  4. More refactoring

  5. Think about new API’s

gpaw.core module

Homework for everybody:

The gpaw.core module implements the following objects: UGArray, PWArray, AtomArrays, Matrix and AtomCenteredFunctions.

Uses in a PAW-context: \(\tilde n_{\sigma}(\mathbf{r})\), \(\tilde v_{\sigma}(\mathbf{r})\), \(\tilde\psi_{n\mathbf{k}\sigma}(\mathbf{r})\), \(D_{\sigma,i_1,i_2}^a\), \(\Delta H_{\sigma,i_1,i_2}^a\), \(Q_{\ell m}^a\), \(P_{\sigma \mathbf{k} ni}^a\), \(\tilde{p}_{\sigma \mathbf{k} i}^a(\mathbf{r}-\mathbf{R}^a)\), \(S_{\mathbf{k}\mu\nu}\), …

Note

All the container objects have a .data attribute to expose the numpy.ndarray (or cupy.ndarray)

gpaw.core module examples

Electron density

Represented by a gpaw.core.UGArray instance (\(\tilde n_{\sigma}(\mathbf{r})\)):

>>> import numpy as np
>>> from gpaw.core import UGDesc
>>> a = 4.0  # side-length of unit cell
>>> gpts = 20  # number of grid-points
>>> grid = UGDesc(cell=[a, a, a],
...               size=(gpts, gpts, gpts)
...               # pbc, zerobc, kpt, comm, decomp, dtype
...               )
>>> nt_sR = grid.zeros(2, xp=np)  # use a numpy array
>>> nt_sR.data.shape
(2, 20, 20, 20)

Note

The container-creation methods (zeros(), empty(), new(), …) all take dims, comm and xp arguments

Wave functions

Represented by a gpaw.core.PWArray instance as sketched here:

import numpy as np
from gpaw.core import PWDesc
kpt = [0.5, 0.5, 0.25]
pw = PWDesc(gcut=8.0,  # G-vector cutoff in bohr^-1
            # ecut=...,  # plane-wave cutoff in hartree
            cell=cell,  # bohr
            kpt=kpt)
N = 5  # number of electrons
wavefunctions = pw.empty(N)

Atom-arrays

PAW atomic density-matrices for an H\(_2\)O molcule (\(D_{\sigma,i_1,i_2}^a\)):

>>> nspins = 2
>>> D_asii = AtomArraysLayout(
...     [(5, 5), (5, 5), (13, 13)],
...     # dtype=float,
...     # xp=np,
...     # atomdist=...
...     ).zeros(nspins)
>>> type(D_asii)
<class 'gpaw.core.atom_arrays.AtomArrays'>
>>> D_asii.data.shape
(2, 219)

Atom-centered functions

PAW-projector functions (\(\tilde{p}_{\sigma \mathbf{k} i}^a(\mathbf{r}-\mathbf{R}^a)\)) are represented by an AtomCenteredFunctions object:

>>> projectors = pw.atom_centered_functions(
...     splines,
...     positions,
...     # atomdist=..., xp=...
... )
>>> type(projectors)
<class 'gpaw.core.pwacf.PWAtomCenteredFunctions'>
>>> P_ani = projectors.integrate(wavefunctions)
>>> type(P_ani)
<class 'gpaw.core.atom_arrays.AtomArrays'>
>>> P_ani.data.shape
(5, 23)

Where is what?

calc = GPAW(...)

# old:
nt = calc.density.nt_sG[0]
for kpt in calc.wfs.kpt_u:
    f = kpt.f_n
    wf = kpt.psit_nG
    ...

# new:
dft = calc.dft
nt = dft.density.nt_sR.data[0]
for wfs in dft.ibzwfs:
    f = wfs.occ_n
    wf = wfs.psit_nX.data
    # or:
    f = wfs.occ_n * wfs.weight * wfs.spin_degeneracy
    wf = wfs.psit_nX.data * np.prod(nt.shape)

Full picture:

  • atoms.calc:

    • dft:

      • density: nt_sR, D_asii, …

      • potential: vt_sR, dH_asii, …

      • pot_calc:

        • xc

        • poisson_solver

      • scf_loop:

        • eigensolver

        • hamiltonian

        • occ_calc

        • mixer

      • ibzwfs: IBZWaveFunctions

        • ibz:

          • symmetries

          • bz

        • wfs_qs[q][s]:

          • psit_nX

          • occ_n

          • …

Testing

Currently, 1752 out of 2034 tests pass:

$ GPAW_NEW=1 pytest -m \
  "not (old_gpaw_only or dscf or gllb or \
        ofdft or do or legacy or lrtddft or \
        rttddft or hybrids or pipekmezey) or new_gpaw_ready"

Pytest marks:

  • old_gpaw_only

  • new_gpaw_ready

  • legacy

Pytest fixture:

gpaw_new: bool

Calculator object attribute:

calc.old: bool

Final slide

  • Try out new GPAW and let us know how/if it worked

  • To do: OFDFT, LRTDDFT(1+2), GLLBSC, core-holes, QNA, TB09, LB94, Old DSCF, SIC(1+2), …

  • Let me know if you would like to help out:

    • porting code

    • fixing tests

    • writing documentation

  • Try out some new features (only in new GPAW):