GPAW QNA空间变化泛函

from pathlib import Path
import json, os, hashlib, numpy as np
from ase import Atoms
from ase.lattice.compounds import L1_2
from gpaw import GPAW, PW, Mixer, Davidson, FermiDirac
R = Path(__file__).resolve().parent / 'results'
R.mkdir()
p = json.loads((R.parent / 'inputs.json').read_text())
rows = []
checks = []

def ck(n, b, **kw):
    checks.append(dict(name=n, passed=bool(b), **kw))
for f, h in p['source_SHA'].items():
    assert hashlib.sha256(Path(f).read_bytes()).hexdigest() == h
from types import SimpleNamespace
from ase import Atoms
from gpaw.xc.kernel import XCKernel
from gpaw.xc.gga import PurePythonGGAKernel
from gpaw.xc.qna import QNAKernel
mu = 0.2195149727645171
beta = 0.06672455060314922
n = np.logspace(-4, 2, 30)[None, :]
sigma = 0.05 * n ** (8 / 3)
kernels = {'compiledPBE': XCKernel('PBE'), 'pyPBE': PurePythonGGAKernel('pyPBE'), 'constantQNA': QNAKernel(SimpleNamespace(override_atoms=None, atoms=Atoms('H'), current_atom=0, parameters={'H': (mu, beta)}))}
probe = dict(density_Bohr_minus3=n, gradient_sigma_Bohr_minus8=sigma)
for name, kernel in kernels.items():
    e = np.zeros(n.shape[1])
    vn = np.zeros_like(n)
    vs = np.zeros_like(sigma)
    kernel.calculate(e, n.copy(), vn, sigma.copy(), vs)
    probe[name + '_energy_density_Ha_Bohr_minus3'] = e
    probe[name + '_dn'] = vn
    probe[name + '_dsigma'] = vs
for suffix in ['energy_density_Ha_Bohr_minus3', 'dn', 'dsigma']:
    ck('matched_python_kernel:' + suffix, abs(probe['pyPBE_' + suffix] - probe['constantQNA_' + suffix]).max() < 1e-12, value=float(abs(probe['pyPBE_' + suffix] - probe['constantQNA_' + suffix]).max()))
np.savez_compressed(R / 'kernel-probe.npz', **probe)
mu = 0.2195149727645171
beta = 0.06672455060314922
for name, parameters in [('Cu3Au-PBE', None), ('Cu3Au-QNA-PBE', {'Au': (mu, beta), 'Cu': (mu, beta)}), ('Cu3Au-QNA-elements', {'Au': (0.125, 0.1), 'Cu': (0.0795, 0.005)}), ('Cu3Au-pyPBE', None)]:
    D = R / name
    D.mkdir()
    a = L1_2(['Au', 'Cu'], latticeconstant=3.74)
    xc = ('pyPBE' if name == 'Cu3Au-pyPBE' else 'PBE') if parameters is None else dict(alpha=2.0, name='QNA', orbital_dependent=False, parameters=parameters, setup_name='PBE', type='qna-gga')
    c = GPAW(legacy_gpaw=True, mode=PW(200), xc=xc, nbands=30, kpts=(2, 2, 2), spinpol=False, occupations=FermiDirac(0.1), mixer=Mixer(0.15, 5, 50), convergence=dict(energy=1e-06, density=1e-05, eigenstates=1e-07), maxiter=200, txt=str(D / 'native.txt'))
    a.calc = c
    e = float(a.get_potential_energy())
    free = float(a.get_potential_energy(force_consistent=True))
    force = a.get_forces()
    rho = c.get_all_electron_density(gridrefinement=1)
    ne = float(rho.sum() * a.get_volume() / rho.size)
    occs = np.array([c.get_occupation_numbers(kpt=k) for k in range(c.wfs.kd.nibzkpts)])
    val = float(occs.sum())
    ck(name + ':AE166', abs(ne - 166) < 0.05, value=ne)
    ck(name + ':valence44', abs(val - 44) < 1e-07, value=val)
    ck(name + ':finite', np.isfinite([e, free, *force.ravel()]).all() and np.isfinite(rho).all())
    fields = dict(AE_total_density_e_A3=rho, weighted_occupations=occs, IBZ_weights=c.get_k_point_weights(), eigenvalues_eV=np.array([c.get_eigenvalues(kpt=k) for k in range(c.wfs.kd.nibzkpts)]), cell_A=a.cell.array, positions_A=a.positions, atomic_numbers=a.numbers, forces_eV_A=force)
    if parameters is not None:
        mg, bg = c.hamiltonian.xc.calculate_spatial_parameters(a)
        ck(name + ':mu_convex_bounds', mg.min() >= min((v[0] for v in parameters.values())) - 1e-12 and mg.max() <= max((v[0] for v in parameters.values())) + 1e-12)
        ck(name + ':beta_convex_bounds', bg.min() >= min((v[1] for v in parameters.values())) - 1e-12 and bg.max() <= max((v[1] for v in parameters.values())) + 1e-12)
        if name.endswith('PBE'):
            ck(name + ':constant_parameters_preserved', abs(mg - mu).max() < 1e-12 and abs(bg - beta).max() < 1e-12)
        else:
            ck(name + ':actual_spatial_parameter_variation', np.std(mg) > 0.0001 and np.std(bg) > 0.0001)
        fields['spatial_mu'] = mg
        fields['spatial_beta'] = bg
    np.savez_compressed(D / 'fields.npz', **fields)
    c.write(str(D / 'state.gpw'), mode='all')
    rows.append(dict(name=name, energy_eV=e, force_consistent_free_energy_eV=free, AE_electrons=ne, valence=val, parameters=parameters, forces_eV_A=force.tolist(), cell_A=a.cell.array.tolist(), positions_A=a.positions.tolist()))
    (R / 'PARTIAL-RESULTS.json').write_text(json.dumps(rows, indent=2) + '\n')
    print(name, e, 'valence', val, flush=True)
by = {r['name']: r for r in rows}
ck('constant_QNA_recovers_matched_python_PBE_energy', abs(by['Cu3Au-QNA-PBE']['energy_eV'] - by['Cu3Au-pyPBE']['energy_eV']) < 0.002, value=by['Cu3Au-QNA-PBE']['energy_eV'] - by['Cu3Au-pyPBE']['energy_eV'])
ck('constant_QNA_recovers_matched_python_PBE_forces', abs(np.array(by['Cu3Au-QNA-PBE']['forces_eV_A']) - by['Cu3Au-pyPBE']['forces_eV_A']).max() < 0.02)
for f, h in p['source_SHA'].items():
    ck('source_immutable:' + f, hashlib.sha256(Path(f).read_bytes()).hexdigest() == h)
compiled_difference = by['Cu3Au-QNA-PBE']['energy_eV'] - by['Cu3Au-PBE']['energy_eV']
v = dict(compiled_PBE_vs_constant_QNA_energy_difference_eV=compiled_difference, passed=all((x['passed'] for x in checks)), checks=len(checks), items=checks, results=rows, whole_group_complete=False, no_converged_lattice_binding_or_functional_accuracy_claim=True)
(R / 'VALIDATION.json').write_text(json.dumps(v, indent=2) + '\n')
print('408', v['passed'], v['checks'], flush=True)
assert v['passed']