GPAW Pipek–Mezey局域轨道

from pathlib import Path
import os, json, hashlib, signal, numpy as np
from gpaw import GPAW
from gpaw.utilities.partition import AtomPartition
from gpaw.wannier.pipekmezey.pipek_mezey_wannier import PipekMezey
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))

def timeout_handler(*_):
    raise TimeoutError('Finite native Pipek-Mezey postprocessing time budget reached')
signal.signal(signal.SIGALRM, timeout_handler)
from ase import Atoms
from gpaw import Mixer
geom = np.load(R.parent / 'input-data/structure.npz')
atoms = Atoms('OH2', positions=geom['positions_A'], cell=geom['cell_A'], pbc=False)
gs = GPAW(legacy_gpaw=True, mode='lcao', basis='dzp', xc='PBE', h=0.22, nbands='nao', spinpol=False, occupations={'name': 'fixed-uniform'}, mixer=Mixer(0.2, 5, 50), convergence=dict(energy=1e-06, density=1e-05, eigenstates=1e-07), maxiter=200, txt=str(R / 'ground-native.txt'))
atoms.calc = gs
atoms.get_potential_energy()
gs.write(str(R / 'ground.gpw'), mode='all')
p['seed'] = str(R / 'ground.gpw')
p['seed_SHA'] = hashlib.sha256((R / 'ground.gpw').read_bytes()).hexdigest()
seed = Path(p['seed'])
assert hashlib.sha256(seed.read_bytes()).hexdigest() == p['seed_SHA']
for method in ['W', 'H']:
    D = R / method
    D.mkdir()
    c = GPAW(str(seed), legacy_gpaw=True, txt=str(D / 'read-only-native.txt'))
    c.calculate = lambda *a, **kw: (_ for _ in ()).throw(RuntimeError('No SCF repeat in Pipek-Mezey postprocessing'))
    C = c.wfs.kpt_u[0].C_nM.copy()
    c.wfs.set_positions(c.spos_ac, AtomPartition(c.wfs.gd.comm, c.wfs.gd.get_ranks_from_positions(c.spos_ac), name='gd'))
    c.wfs.initialize_wave_functions_from_lcao()
    psi = c.wfs.kpt_u[0].psit_nG[:4].copy()
    S = c.wfs.S_qMM[0].copy()
    S = np.tril(S) + np.tril(S, -1).T.conj()
    signal.alarm(240)
    pm = PipekMezey(calc=c, method=method, seed=42)
    before = float(pm.get_function_value())
    pm.localize(tolerance=1e-08)
    signal.alarm(0)
    history = np.array(pm.P_n)
    after = float(pm.get_function_value())
    W = pm.W_k[0].copy()
    localized = np.einsum('rn,rxyz->nxyz', W, psi)
    gram = (W.T @ C[:4]).conj() @ S @ (W.T @ C[:4]).T
    weights = np.array([pm.get_weight_function_atom(i) for i in range(3)])
    projector = (abs(psi) ** 2).sum(0)
    local_projector = (abs(localized) ** 2).sum(0)
    norm_defect = float((abs(1 - weights.sum(0)) * projector).sum() * c.wfs.gd.dv)
    Q = np.sum(abs(pm.Qadk_nn), axis=2) / pm.Nk
    independent = float(sum(((np.diagonal(Q[:, d], axis1=1, axis2=2) ** 2).sum() * wd for d, wd in enumerate(pm.wd))) / sum(pm.wd))
    relative = float(abs(history[-1] - history[-2]) / max(abs(history[-1]), 1e-15))
    ck(method + ':four_occupied_orbitals', pm.nocc == 4 and pm.Nk == 1)
    ck(method + ':localization_increases_objective', after >= before - 1e-08, value=after - before)
    ck(method + ':optimizer_original_relative_tolerance', relative <= 1e-08, value=relative)
    ck(method + ':independent_native_objective', abs(independent - after) < 1e-10)
    ck(method + ':unitary_rotation', abs(W.T.conj() @ W - np.eye(4)).max() < 1e-10)
    ck(method + ':PAW_metric_preserved', abs(gram - np.eye(4)).max() < 1e-08)
    ck(method + ':occupied_projector_fullgrid_preserved', abs(local_projector - projector).max() < 1e-10)
    ck(method + ':density_weighted_partition_defect', norm_defect < 0.001, value=norm_defect)
    ck(method + ':weights_bounds', weights.min() > -1e-12 and weights.max() < 1 + 1e-12)
    ck(method + ':finite_actual_fields', np.isfinite(localized).all() and np.isfinite(pm.Qadk_nn).all())
    ck(method + ':original_coefficients_unchanged', np.array_equal(C, c.wfs.kpt_u[0].C_nM))
    np.savez_compressed(D / 'fields.npz', rotation=W, canonical_occupied_pseudo_orbitals=psi, localized_pseudo_orbitals=localized, occupied_projector=projector, localized_projector=local_projector, PAW_Gram=gram, partition_weights=weights, native_charge_matrices=pm.Qadk_nn, charge_matrix_before_rotation=pm.Qadk_nm, direction_weights=pm.wd, objective_history=history, cell_A=c.atoms.cell.array, positions_A=c.atoms.positions, grid_spacing_Bohr=c.wfs.gd.h_cv)
    rows.append(dict(method=method, before=before, after=after, relative_last_step=relative, iterations=pm.niter, occupied_orbitals=pm.nocc, partition_defect=norm_defect))
    (R / 'PARTIAL-RESULTS.json').write_text(json.dumps(rows, indent=2) + '\n')
    print(method, before, after, pm.niter, flush=True)
ck('immutable_original_checkpoint', hashlib.sha256(seed.read_bytes()).hexdigest() == p['seed_SHA'])
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, SCF_repeats=0, H_uses_native_gaussian_stockholder_not_selfconsistent_free_atom_Hirshfeld=True, no_unique_orbital_phase_or_permutation_claim=True)
(R / 'VALIDATION.json').write_text(json.dumps(v, indent=2) + '\n')
print('414', v['passed'], v['checks'], flush=True)
assert v['passed']