GPAW meta-GGA与动能密度

from pathlib import Path
import json, numpy as np
from ase import Atoms
from ase.build import bulk
from gpaw import GPAW, PW, Mixer, Davidson
R = Path(__file__).resolve().parent / 'results'
R.mkdir()
rows = []
checks = []

def ck(n, b, **kw):
    checks.append(dict(name=n, passed=bool(b), **kw))
for name in ['H2-PBE', 'Al-PBE', 'Al-TPSS']:
    D = R / name
    D.mkdir()
    if name == 'H2-PBE':
        a = Atoms('H2', [[0, 0, 0], [0, 0, 0.75]])
        a.center(vacuum=3.5)
        a.pbc = False
        c = GPAW(legacy_gpaw=True, mode='fd', xc='PBE', h=0.25, nbands=2, spinpol=False, eigensolver=Davidson(3), mixer=Mixer(0.2, 5, 1), convergence=dict(energy=1e-06, density=1e-05, eigenstates=1e-07), maxiter=250, txt=str(D / 'native.txt'))
    elif name == 'Al-PBE':
        a = bulk('Al', 'fcc', a=4.05)
        c = GPAW(legacy_gpaw=True, mode=PW(200), xc='PBE', kpts=(4, 4, 4), nbands=8, spinpol=False, occupations={'name': 'fermi-dirac', 'width': 0.2}, mixer=Mixer(0.2, 5, 50), convergence=dict(energy=1e-06, density=1e-05, eigenstates=1e-07), maxiter=200, txt=str(D / 'native.txt'))
    else:
        c = GPAW(str(R / 'Al-PBE/state.gpw'), legacy_gpaw=True, xc='TPSS', eigensolver=Davidson(5), mixer=Mixer(0.05, 5, 50), convergence=dict(energy=1e-06, density=1e-05, eigenstates=1e-07), maxiter=240, txt=str(D / 'native.txt'))
        a = c.get_atoms()
    a.calc = c
    e = float(a.get_potential_energy())
    c.write(str(D / 'state.gpw'), mode='all')
    force = a.get_forces()
    rho = c.get_all_electron_density(gridrefinement=1)
    tau = c.wfs.calculate_kinetic_energy_density()
    occs = np.array([c.get_occupation_numbers(kpt=k) for k in range(c.wfs.kd.nibzkpts)])
    eigs = np.array([c.get_eigenvalues(kpt=k) for k in range(c.wfs.kd.nibzkpts)])
    ne = float(rho.sum() * a.get_volume() / rho.size)
    val = 2 if name.startswith('H2') else 3
    Z = 2 if name.startswith('H2') else 13
    typ = 'MGGA' if name == 'Al-TPSS' else 'GGA'
    ck(name + ':AE', abs(ne - Z) < 0.02)
    ck(name + ':valence', abs(occs.sum() - val) < 1e-08)
    ck(name + ':actual_functional', c.hamiltonian.xc.type == typ)
    ck(name + ':nonnegative_tau', np.isfinite(tau).all() and tau.min() > -1e-12 and (tau.max() > 0))
    ck(name + ':finite', np.isfinite([e, *force.ravel()]).all() and np.isfinite(rho).all())
    np.savez_compressed(D / 'fields.npz', AE_total_density_e_A3=rho, pseudo_valence_kinetic_density_Bohr_minus5=tau, weighted_occupations=occs, eigenvalues_eV=eigs, IBZ_weights=c.get_k_point_weights(), cell_A=a.cell.array, positions_A=a.positions, forces_eV_A=force)
    row = dict(name=name, energy_eV=e, AE_electrons=ne, valence=float(occs.sum()), XC_type=typ)
    if name == 'H2-PBE':
        row['TPSS_fixed_density_energy_eV'] = e + float(c.get_xc_difference('TPSS'))
        ck('H2_fixed_density_TPSS_finite', np.isfinite(row['TPSS_fixed_density_energy_eV']))
    rows.append(row)
    print(name, e, flush=True)
v = dict(passed=all((c['passed'] for c in checks)), checks=len(checks), items=checks, results=rows, H2_TPSS_only_fixed_density=True, Al_TPSS_actual_SCF=True, whole_group_complete=False)
(R / 'VALIDATION.json').write_text(json.dumps(v, indent=2) + '\n')
assert v['passed']