GPAW非共线磁矩与方向约束

from pathlib import Path
import os, json, hashlib, numpy as np
from ase import Atoms
from gpaw.new.ase_interface import GPAW
from gpaw.new.constraints import SpinDirectionConstraint
from ase.units import Hartree
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
cases = [('free-x', [1.0, 0, 0], None), ('free-z', [0, 0, 1.0], None), ('constraint-x', [1.0, 0.9, 0.9], [1.0, 0, 0]), ('constraint-z', [0.9, 0.9, 1.0], [0, 0, 1.0])]
for name, initial, target in cases:
    D = R / name
    D.mkdir()
    a = Atoms('H', positions=[[2.0, 2.0, 2.0]], cell=[4.0] * 3, pbc=False)
    m = np.array([initial])
    m /= np.linalg.norm(m)
    constraint = SpinDirectionConstraint({0: target}, penalty=2.0) if target is not None else None
    c = GPAW(mode={'name': 'pw', 'ecut': 300}, xc='LDA', nbands=2, symmetry='off', eigensolver={'niter': 8}, mixer={'backend': 'msr1'}, soc=True, magmoms=m, parallel={'domain': 1, 'band': 1}, extensions=[] if constraint is None else [constraint], 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())
    c.write(str(D / 'state.gpw'), mode='all')
    mv, ma = c.dft.density.calculate_magnetic_moments()
    rho = c.dft.densities().all_electron_densities(grid_spacing=0.2)
    data = np.array(rho.data)
    integral = data.reshape((4, -1)).sum(1) * a.get_volume() / np.prod(data.shape[1:])
    raw = c.get_occupation_numbers(raw=True).copy()
    ne = float(raw.sum() * c.dft.ibzwfs.spin_degeneracy)
    D4 = np.array(c.dft.density.D_asii[0])
    N = np.array(c.dft.density.N0_aii[0])
    local = np.einsum('vij,ij->v', D4[1:4].real, N)
    ck(name + ':noncollinear_four_components', data.shape[0] == 4 and c.dft.density.ncomponents == 4)
    ck(name + ':finite', np.isfinite([e, ne, *mv, *ma.ravel(), *integral]).all() and np.isfinite(data).all())
    ck(name + ':one_valence_electron', abs(ne - 1) < 1e-08, value=ne)
    ck(name + ':AE_grid_one_electron', abs(integral[0] - 1) < 0.02, value=float(integral[0]))
    ck(name + ':magnetic_norm1', abs(np.linalg.norm(mv) - 1) < 0.003, value=float(np.linalg.norm(mv)))
    ck(name + ':AE_magnetization_matches_native', abs(integral[1:] - mv).max() < 0.02)
    ck(name + ':PAW_local_moment_reconstruction', abs(local - ma[0]).max() < 1e-10)
    direction = np.array(initial if target is None else target, dtype=float)
    direction /= np.linalg.norm(direction)
    perp = mv - direction * (mv @ direction)
    ck(name + ':spin_axis_target', np.linalg.norm(perp) < 0.01, value=float(np.linalg.norm(perp)))
    fields = dict(AE_density_and_magnetization=data, cell_A=a.cell.array, positions_A=a.positions, native_total_magnetic_moment=mv, native_atomic_magnetic_moments=ma, atomic_spin_density_matrices=D4, PAW_inner_product=N, raw_occupations=raw, eigenvalues_eV=c.get_eigenvalues(), initial_magnetization=m, target_axis=direction)
    if constraint is not None:
        setup = c.dft.setups[0]
        M = D4[1:4].real
        penalty, V = constraint.calculate(M, 0, setup.l_j, setup.N0_q, return_energy=True)
        expected = 2 / Hartree * (local @ local - (direction @ local) ** 2)
        ck(name + ':penalty_from_actual_PAW_moment', abs(penalty - expected) < 1e-12)
        dh = 1e-05
        plus = M.copy()
        minus = M.copy()
        plus[0, 0, 0] += dh / 2
        minus[0, 0, 0] -= dh / 2
        ep = constraint.calculate(plus, 0, setup.l_j, setup.N0_q, return_energy=True)[0]
        em = constraint.calculate(minus, 0, setup.l_j, setup.N0_q, return_energy=True)[0]
        ck(name + ':actual_density_penalty_derivative', abs((ep - em) / dh - V[0, 0, 0]) < 1e-08)
        fields['constraint_Hamiltonian_Ha'] = V
        fields['penalty_energy_Ha'] = np.array(penalty)
    np.savez_compressed(D / 'fields.npz', **fields)
    rows.append(dict(name=name, energy_eV=e, magnetic_moment=mv.tolist(), atomic_magnetic_moments=ma.tolist(), AE_integrals=integral.tolist(), valence=ne, initial=m.tolist(), target=target))
    (R / 'PARTIAL-RESULTS.json').write_text(json.dumps(rows, indent=2) + '\n')
    print(name, e, mv, flush=True)
by = {x['name']: x for x in rows}
ck('H_x_z_spin_rotation_energy_invariance', abs(by['free-x']['energy_eV'] - by['free-z']['energy_eV']) < 2e-05, value=by['free-x']['energy_eV'] - by['free-z']['energy_eV'])
for f, h in p['source_SHA'].items():
    ck('source_immutable:' + f, hashlib.sha256(Path(f).read_bytes()).hexdigest() == h)
v = dict(passed=all((c['passed'] for c in checks)), checks=len(checks), items=checks, results=rows, whole_group_complete=False, penalty_selects_axis_not_sign=True, no_converged_box_or_magnetic_anisotropy_claim=True)
(R / 'VALIDATION.json').write_text(json.dumps(v, indent=2) + '\n')
print('404', v['passed'], v['checks'], flush=True)
assert v['passed']