GPAW电子收敛与自定义判据

from pathlib import Path
import json, os, hashlib, numpy as np
from ase import Atoms
from ase.build import molecule
from ase.units import Bohr
from gpaw import GPAW, Mixer, Davidson
from gpaw.convergence_criteria import Criterion, Energy, Forces
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))

class DipoleWindow(Criterion):
    name = 'dipole window'
    tablename = 'dplw'
    calc_last = False

    def __init__(self, tol=1e-05, n_old=4):
        self.tol = tol
        self.n_old = n_old
        self.description = f'Dipole vector variation over {n_old} iterations below {tol} eAngstrom'
        self.reset()

    def reset(self):
        self._trace = []

    def __call__(self, context):
        mu = np.array(context.dens.calculate_dipole_moment()) * Bohr
        window = [r['dipole_eA'] for r in self._trace[-(self.n_old - 1):]] + [mu.tolist()]
        error = float(np.linalg.norm(np.array(window)[:, None] - np.array(window)[None, :], axis=-1).max()) if len(window) == self.n_old else np.inf
        self._trace.append(dict(iteration=context.niter, dipole_eA=mu.tolist(), window_error=error))
        return (error < self.tol, '%5.1f' % np.log10(max(error, 1e-30)) if np.isfinite(error) else '')
for f, h in p['source_SHA'].items():
    assert hashlib.sha256(Path(f).read_bytes()).hexdigest() == h
for name in ['H2-forces', 'water-dipole']:
    D = R / name
    D.mkdir()
    a = Atoms('H2', [[0, 0, 0], [0, 0, 0.75]]) if name == 'H2-forces' else molecule('H2O')
    a.center(vacuum=3.5)
    a.pbc = False
    conv = dict(energy=Energy(1e-06, n_old=4, relative=False), density=1e-05, eigenstates=1e-07, **{'minimum iterations': 25})
    if name == 'H2-forces':
        conv['forces'] = Forces(atol=0.001)
    else:
        conv['custom'] = [DipoleWindow(1e-05, 4)]
    c = GPAW(legacy_gpaw=True, mode='fd', h=0.25, xc='PBE', nbands=2 if name == 'H2-forces' else 8, spinpol=False, eigensolver=Davidson(3), mixer=Mixer(0.15, 5, 1), convergence=conv, maxiter=200, txt=str(D / 'native.txt'))
    a.calc = c
    force_trace = []

    def observe():
        criterion = c.scf.criteria.get('forces')
        if criterion is not None and criterion.old_F_av is not None:
            n = c.scf.niter
            if not force_trace or force_trace[-1]['iteration'] != n:
                f = criterion.old_F_av.copy()
                error = np.inf if not force_trace else float(np.linalg.norm(f - np.array(force_trace[-1]['forces_eV_A']), axis=1).max())
                force_trace.append(dict(iteration=n, forces_eV_A=f.tolist(), error=error))
    c.attach(observe, 1)
    e = float(a.get_potential_energy())
    force = a.get_forces()
    rho = c.get_all_electron_density(gridrefinement=1)
    ne = float(rho.sum() * a.get_volume() / rho.size)
    iterations = c.get_number_of_iterations()
    np.savez_compressed(D / 'SCF-extracted-before-analysis.npz', AE_total_density_e_A3=rho, cell_A=a.cell.array, positions_A=a.positions, forces_eV_A=force, eigenvalues_eV=c.get_eigenvalues(), weighted_occupations=c.get_occupation_numbers(), energy_eV=np.array(e))
    if name == 'H2-forces':
        c.write(str(D / 'state.gpw'), mode='all')
    ck(name + ':minimum25', iterations >= 25, value=iterations)
    ck(name + ':AE', abs(ne - (2 if name == 'H2-forces' else 10)) < 0.03, value=ne)
    ck(name + ':finite', np.isfinite([e, ne, *force.ravel()]).all() and np.isfinite(rho).all())
    ck(name + ':all_original_density_eigenstate_gates_active', c.scf.converged and c.scf.criteria['density'].tol == 1e-05 and (c.scf.criteria['eigenstates'].tol == 1e-07))
    energy = c.scf.criteria['energy']
    ck(name + ':absolute_total_energy_fourstep_window', energy.n_old == 4 and (not energy.relative) and (energy.tol == 1e-06) and (np.ptp(energy._old) < 1e-06))
    if name == 'H2-forces':
        trace = force_trace
        assert len(trace) >= 2
        f0 = np.array(trace[-2]['forces_eV_A'])
        f1 = np.array(trace[-1]['forces_eV_A'])
        error = float(np.linalg.norm(f1 - f0, axis=1).max())
        ck('native_force_change_reconstructed', error < 0.001 and abs(error - trace[-1]['error']) < 1e-12)
        ck('force_change_is_not_structure_optimization', len(a) == 2 and np.isclose(a.get_distance(0, 1), 0.75))
        ck('native_force_trace_last_equals_result', abs(f1 - force).max() < 1e-07)
    else:
        trace = c.scf.criteria['dipole window']._trace
        mu = np.array([r['dipole_eA'] for r in trace[-4:]])
        error = float(np.linalg.norm(mu[:, None] - mu[None, :], axis=-1).max())
        ck('custom_dipole_window_independently_reconstructed', error < 1e-05 and abs(error - trace[-1]['window_error']) < 1e-12)
        ck('custom_checked_each_iteration', len(trace) == iterations and trace[-1]['iteration'] == iterations)
        ck('nonzero_physical_water_dipole', np.linalg.norm(mu[-1]) > 0.1, value=float(np.linalg.norm(mu[-1])))
        ck('dipole_trace_matches_result', abs(mu[-1] - c.get_dipole_moment()).max() < 1e-05)
    np.savez_compressed(D / 'fields.npz', AE_total_density_e_A3=rho, cell_A=a.cell.array, positions_A=a.positions, forces_eV_A=force, eigenvalues_eV=c.get_eigenvalues(), weighted_occupations=c.get_occupation_numbers(), energy_window_eV=np.array(energy._old))
    (D / 'SCF-TRACE.json').write_text(json.dumps([{k: None if isinstance(v, float) and (not np.isfinite(v)) else v for k, v in row.items()} for row in trace], indent=2) + '\n')
    rows.append(dict(name=name, energy_eV=e, iterations=iterations, AE_electrons=ne, final_additional_error=error, forces_eV_A=force.tolist()))
    (R / 'PARTIAL-RESULTS.json').write_text(json.dumps(rows, indent=2) + '\n')
    print(name, e, iterations, error, flush=True)
for f, h in p['source_SHA'].items():
    ck('source_immutable:' + f, hashlib.sha256(Path(f).read_bytes()).hexdigest() == h)
v = dict(passed=all((x['passed'] for x in checks)), checks=len(checks), items=checks, results=rows, whole_group_complete=False, no_disabled_default_criteria=True, no_fixed_iteration_unconverged_state_claim=True)
(R / 'VALIDATION.json').write_text(json.dumps(v, indent=2) + '\n')
print('409', v['passed'], v['checks'], flush=True)
assert v['passed']