fields.py

import numpy as np
from typing import List, Callable
from dataclasses import dataclass, field

# Field = namedtuple("Field", ["name", "alias", "calc", "shape"])
# LabelField = namedtuple("LabelField", ["name", "alias", "calc", "shape", "required_labels"])
@dataclass
class Field:
      name: str
      alias: List[str]
      calc: Callable
      shape: str
      required_labels: List[str] = field(default_factory=list)


def select_fields(names):
    names = [n.lower() for n in names]
    scfs   = [fd for fd in SCF_FIELDS 
                  if fd.name in names 
                  or any(al in names for al in fd.alias)]
    grads  = [fd for fd in GRAD_FIELDS 
                  if fd.name in names 
                  or any(al in names for al in fd.alias)]
    return {"scf": scfs, "grad": grads}


BOHR = 0.52917721092

def isinbohr(mol):
    return mol.unit.upper().startswith(("B", "AU"))

def _Lunit(mol):
    return (1. if isinbohr(mol) else BOHR)

def atom_data(mol):
    raw_data = np.concatenate(
        [mol.atom_charges().reshape(-1,1), mol.atom_coords(unit='Bohr')], 
        axis=1)
    non_ghost = [ii for ii in range(mol.natm) 
        if not mol.elements[ii].startswith("X")]
    return raw_data[non_ghost]


SCF_FIELDS = [
    Field("atom",
          ["atoms", "mol", "molecule"],
          lambda mf: atom_data(mf.mol),
          "(nframe, natom, 4)"),
    Field("e_base", 
          ["ebase", "ene_base", "e0",
           "e_hf", "ehf", "ene_hf", 
           "e_ks", "eks", "ene_ks"], 
          lambda mf: mf.energy_tot0(),
          "(nframe, 1)"),
    Field("e_tot", 
          ["e_cf", "ecf", "ene_cf", "etot", "ene", "energy", "e"],
          lambda mf: mf.e_tot,
          "(nframe, 1)"),
    Field("rdm",
          ["dm"],
          lambda mf: mf.make_rdm1(),
          "(nframe, nao, nao)"),
    Field("proj_dm",
          ["pdm"],
          lambda mf: mf.make_pdm(flatten=True),
          "(nframe, natom, -1)"),
    Field("dm_eig",
          ["eig"],
          lambda mf: mf.make_eig(),
          "(nframe, natom, nproj)"),
    Field("hcore_eig",
          ["heig"],
          lambda mf: mf.make_eig(mf.get_hcore()),
          "(nframe, natom, nproj)"),
    Field("ovlp_eig",
          ["oeig"],
          lambda mf: mf.make_eig(mf.get_ovlp()),
          "(nframe, natom, nproj)"),
    Field("veff_eig",
          ["veig"],
          lambda mf: mf.make_eig(mf.get_veff()),
          "(nframe, natom, nproj)"),
    Field("fock_eig",
          ["feig"],
          lambda mf: mf.make_eig(mf.get_fock()),
          "(nframe, natom, nproj)"),
    Field("conv", 
          ["converged", "convergence"], 
          lambda mf: mf.converged,
          "(nframe, 1)"),
    Field("mo_coef_occ", # do not support UHF
          ["mo_coeff_occ, orbital_coeff_occ"],
          lambda mf: mf.mo_coeff[:,mf.mo_occ>0].T,
          "(nframe, -1, nao)"),
    Field("mo_ene_occ", # do not support UHF
          ["mo_energy_occ, orbital_ene_occ"],
          lambda mf: mf.mo_energy[mf.mo_occ>0],
          "(nframe, -1)"),
    # below are fields that requires labels
    Field("l_e_ref", 
          ["e_ref", "lbl_e_ref", "label_e_ref", "le_ref"],
          lambda mf, **lbl: lbl["energy"],
          "(nframe, 1)",
          ["energy"]),
    Field("l_e_delta", 
          ["le_delta", "lbl_e_delta", "label_e_delta", "lbl_ed"],
          lambda mf, **lbl: lbl["energy"] - mf.energy_tot0(),
          "(nframe, 1)",
          ["energy"]),
    Field("err_e", 
          ["e_err", "err_e_tot", "err_e_cf"],
          lambda mf, **lbl: lbl["energy"] - mf.e_tot,
          "(nframe, 1)",
          ["energy"]),
]

GRAD_FIELDS = [
    Field("f_base", 
          ["fbase", "force_base", "f0",
           "f_hf", "fhf", "force_hf", 
           "f_ks", "fks", "force_ks"], 
          lambda grad: - grad.get_base() / _Lunit(grad.mol),
          "(nframe, natom_raw, 3)"),
    Field("f_tot", 
          ["f_cf", "fcf", "force_cf", "ftot", "force", "f"], 
          lambda grad: - grad.de / _Lunit(grad.mol),
          "(nframe, natom_raw, 3)"),
    Field("grad_dmx",
          ["grad_dm_x", "grad_pdm_x"],
          lambda grad: grad.make_grad_pdm_x(flatten=True) / _Lunit(grad.mol),
          "(nframe, natom_raw, 3, natom, -1)"),
    Field("grad_vx",
          ["grad_eig_x", "geigx", "gvx"],
          lambda grad: grad.make_grad_eig_x()  / _Lunit(grad.mol),
          "(nframe, natom_raw, 3, natom, nproj)"),
    # below are fields that requires labels
    Field("l_f_ref", 
          ["f_ref", "lbl_f_ref", "label_f_ref", "lf_ref"],
          lambda grad, **lbl: lbl["force"],
          "(nframe, natom_raw, 3)",
          ["force"]),
    Field("l_f_delta", 
          ["lf_delta", "lbl_f_delta", "label_f_delta", "lbl_fd"],
          lambda grad, **lbl: lbl["force"] - (-grad.get_base() / _Lunit(grad.mol)),
          "(nframe, natom_raw, 3)",
          ["force"]),
    Field("err_f", 
          ["f_err", "err_f_tot", "err_f_cf"],
          lambda grad, **lbl: lbl["force"] - (-grad.de / _Lunit(grad.mol)),
          "(nframe, natom_raw, 3)",
          ["force"]),
]


# below are additional methods from addons
from deepks.scf import addons

SCF_FIELDS.extend([
    # the following two are used for regularizing the potential
    Field("grad_veg",
          ["grad_eig_egrad", "jac_eig_egrad"],
          lambda mf: addons.make_grad_eig_egrad(mf),
          "(nframe, natom, nproj, -1)"),
    Field("eg_base",
          ["ele_grad_base", "egrad0", "egrad_base"],
          lambda mf: mf.get_grad0(),
          "(nframe, -1)"),
    # the following one is used for coulomb loss optimization
    Field("grad_ldv",
          ["grad_coul_dv", "grad_coul_deig", "coulomb_grad"], 
          lambda mf, **lbl: addons.make_grad_coul_veig(mf, target_dm=lbl["dm"]),
          "(nframe, natom, nproj)",
          ["dm"]),
    Field("l_veig_raw",
          ["optim_veig_raw", "l_opt_v_raw", "l_optim_veig_raw"], 
          lambda mf, **lbl: addons.calc_optim_veig(mf, lbl["dm"], nstep=1),
          "(nframe, natom, nproj)",
          ["dm"]),
])

GRAD_FIELDS.extend([
    # the following one is used for coulomb loss optimization from grad class
    Field("l_veig",
          ["optim_veig", "l_opt_v", "l_optim_veig"], 
          lambda grad, **lbl: addons.gcalc_optim_veig(
              grad, lbl["dm"], -_Lunit(grad.mol)*lbl["force"], nstep=1),
          "(nframe, natom, nproj)",
          ["dm", "force"]),
    Field("l_veig_nof",
          ["optim_veig_nof", "l_opt_v_nof", "l_optim_veig_nof"], 
          lambda grad, **lbl: addons.gcalc_optim_veig(
              grad, lbl["dm"], grad.de, nstep=1),
          "(nframe, natom, nproj)",
          ["dm"]),
])