GPAW线性展开ΔSCF激发态

from pathlib import Path
import os, json, hashlib, numpy as np
from ase.build import molecule
from ase.units import Hartree
from gpaw import GPAW, Mixer, FermiDirac
from gpaw import dscf
R = Path(__file__).resolve().parent / 'results'
R.mkdir()
p = json.loads((R.parent / 'inputs.json').read_text())
checks = []
rows = []

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
a = molecule('CO')
a.center(vacuum=3.5)
a.pbc = False

def calc(name):
    return GPAW(legacy_gpaw=True, mode='lcao', basis='dzp', h=0.25, xc='PBE', nbands=8, spinpol=True, symmetry='off', occupations=FermiDirac(0.01), mixer={'backend': 'pulay', 'nmaxold': 6, 'beta': 0.1, 'weight': 1} if name.startswith('excited') else Mixer(0.05, 6, 50), convergence=dict(energy=1e-08, density=1e-07, eigenstates=1e-09, bands='all') if name == 'excited-AE' else dict(energy=1e-06, density=1e-05, eigenstates=1e-07, bands='all'), maxiter=400, txt=str(R / (name + '-native.txt')))

def extract(name, c, energy):
    spin = np.array([c.get_all_electron_density(spin=s, gridrefinement=1) for s in [0, 1]])
    total = c.get_all_electron_density(gridrefinement=1)
    integral = float(total.sum() * a.get_volume() / total.size)
    ks = float(sum((k.f_n.sum() for k in c.wfs.kpt_u)))
    ne = []
    coefs = []
    grams = []
    eps = []
    norms = []
    delta = 0.0
    manual = 0.0
    for k in c.wfs.kpt_u:
        C = k.C_nM
        S = c.wfs.S_qMM[k.q].copy()
        S = np.tril(S) + np.tril(S, -1).T.conj()
        gram = C.conj() @ S @ C.T
        grams.append(gram)
        eps.append(k.eps_n * Hartree)
        ck(name + ':spin' + str(k.s) + ':PAW_metric', abs(gram - np.eye(8)).max() < 1e-06, value=float(abs(gram - np.eye(8)).max()))
        if hasattr(k, 'ne_o'):
            ne.append(k.ne_o.copy())
            coefs.append(k.c_on.copy())
            cn = np.sum(abs(k.c_on) ** 2, axis=1)
            norms.append(cn)
            delta += float(np.dot(k.ne_o, cn))
            manual += float(sum((n * np.dot(abs(v) ** 2, k.eps_n) for n, v in zip(k.ne_o, k.c_on))))
            ck(name + ':spin' + str(k.s) + ':linear_expansion_normalized', abs(cn - 1).max() < 1e-08)
    ck(name + ':AE14', abs(integral - 14) < 0.03, value=integral)
    ck(name + ':KS_plus_constraint_valence10', abs(ks + delta - 10) < 1e-07, value=ks + delta)
    ck(name + ':two_spin_density_sum', abs(spin.sum(0) - total).max() < 1e-12)
    ck(name + ':finite', np.isfinite([energy, integral, ks, delta]).all() and np.isfinite(total).all())
    mag = float((spin[0] - spin[1]).sum() * a.get_volume() / total.size)
    ck(name + ':density_spin_matches_reported', abs(mag - c.get_magnetic_moment()) < 0.03, value=mag)
    if name.startswith('excited'):
        ck('native_DSCF_band_energy_identity', abs(c.wfs.occupations.calculate_band_energy(c.wfs) - manual) < 1e-10)
        ck('constrained_electron1', abs(delta - 1) < 1e-08)
        ck('ordinary_KS_electrons9', abs(ks - 9) < 1e-07)
    np.savez_compressed(R / (name + '-fields.npz'), AE_spin_density_e_A3=spin, AE_total_density_e_A3=total, cell_A=a.cell.array, positions_A=a.positions, eigenvalues_eV=np.array(eps), KS_weighted_occupations=np.array([k.f_n for k in c.wfs.kpt_u]), PAW_Gram=np.array(grams), constraint_electrons=np.array(ne), linear_expansion_coefficients=np.array(coefs), linear_expansion_norms=np.array(norms))
    row = dict(name=name, energy_eV=energy, AE_electrons=integral, ordinary_KS_valence=ks, constrained_valence=delta, total_valence=ks + delta, magnetic_moment=mag, iterations=c.get_number_of_iterations())
    rows.append(row)
    (R / 'PARTIAL-RESULTS.json').write_text(json.dumps(rows, indent=2) + '\n')
    print(row, flush=True)
c = calc('ground')
a.calc = c
eg = float(a.get_potential_energy())
c.write(str(R / 'ground.gpw'), mode='all')
extract('ground', c, eg)
target = np.load(R.parent / 'input-data/reference-orbital.npz')
assert np.allclose(a.positions, target['positions_A'], atol=1e-12, rtol=0) and np.allclose(a.cell.array, target['cell_A'], atol=1e-12, rtol=0)
wf = list(target['pseudo_target_by_spin'])
proj = [{atom: target[f'projector_spin{spin}_atom{atom}'].copy() for atom in range(2)} for spin in range(2)]
for name, kind in [('excited-AE', 'AE'), ('excited-Molecular', 'Molecular')]:
    ce = calc(name)
    ce.initialize(a)
    if kind == 'AE':
        orb = dscf.AEOrbital(ce, wf, proj, txt=str(R / 'expansion-overlap.log'))
    else:
        orb = dscf.MolecularOrbital(ce, weights={0: [0.0, 0.0, 0.0, 1.0], 1: [0.0, 0.0, 0.0, -1.0]})
    dscf.dscf_calculation(ce, [[1.0, orb, 1]], a)
    a.calc = ce
    ee = float(a.get_potential_energy())
    extract(name, ce, ee)
    ck(name + ':positive_vertical_excitation', 0 < ee - eg < 20, value=ee - eg)
for f, h in p['source_SHA'].items():
    ck('source_immutable:' + f, hashlib.sha256(Path(f).read_bytes()).hexdigest() == h)
ck('AE_and_Molecular_excitation_consistency', abs(rows[1]['energy_eV'] - rows[2]['energy_eV']) < 0.02)
v = dict(passed=all((x['passed'] for x in checks)), checks=len(checks), items=checks, results=rows, summary=dict(vertical_excitation_eV={r['name']: r['energy_eV'] - eg for r in rows[1:]}), whole_group_complete=False, no_pure_singlet_or_triplet_claim=True, not_MOM=True)
(R / 'VALIDATION.json').write_text(json.dumps(v, indent=2) + '\n')
print('402', v['passed'], v['checks'], v['summary'], flush=True)
assert v['passed']