solve_mol.py

#!/usr/bin/env python
#SBATCH -N 1
#SBATCH -c 10
#SBATCH -t 24:00:00
#SBATCH --mem=32G

import time
import numpy as np
from pyscf import gto, scf

BOHR = 0.52917721092

_NO_FORCE = False
_NO_DM = False
_MUST_UNRES = False
_USE_NEWTON = False

def parse_xyz(filename, basis='ccpvdz', **kwargs):
    with open(filename) as fp:
        natoms = int(fp.readline())
        comments = fp.readline()
        xyz_str = "".join(fp.readlines())
    mol = gto.Mole()
    mol.atom = xyz_str
    mol.basis = basis
    mol.set(**kwargs)
    if "spin" not in kwargs:
        mol.spin = mol.nelectron % 2
    mol.build(0,0,unit="Ang")
    return mol  


def get_method(name: str):
    lname = name.lower()
    if lname == "hf":
        return calc_hf
    if lname[:3] == "dft":
        xc = lname.split("@")[1] if "@" in lname else "pbe"
        return lambda mol, **scfargs: calc_dft(mol, xc, **scfargs)
    if lname == "mp2":
        return calc_mp2
    if lname == "ccsd":
        return calc_ccsd
    if lname.startswith(("ccsd_t", "ccsd-t", "ccsd(t)")):
        return calc_ccsd_t
    if lname == "fci":
        return calc_fci
    raise ValueError(f"Unknown calculation method: {name}")

def solve_scf(mol, **scfargs):
    HFmethod = scf.HF if not _MUST_UNRES else scf.UHF
    mf = HFmethod(mol).set(init_guess_breaksym=True)
    init_dm = mf.get_init_guess()
    # if _MUST_UNRES:
    #     init_dm[1][:2,:2] = 0
    mf.kernel(init_dm)
    if _USE_NEWTON:
        mf = scf.fast_newton(mf)
    return mf

def calc_hf(mol, **scfargs):
    mf = solve_scf(mol, **scfargs)
    if not mf.converged:
        raise RuntimeError("SCF not converged!")
    etot = mf.e_tot
    grad = mf.nuc_grad_method().kernel() if not _NO_FORCE else None
    rdm = mf.make_rdm1() if not _NO_DM else None
    return etot, grad, rdm

def calc_dft(mol, xc="pbe", **scfargs):
    from pyscf import dft
    KSmethod = dft.KS if not _MUST_UNRES else dft.UKS
    mf = KSmethod(mol, xc).run(**scfargs)
    if not mf.converged:
        raise RuntimeError("SCF not converged!")
    etot = mf.e_tot
    if _NO_FORCE or dft.libxc.xc_type(xc) in ('MGGA', 'NLC'):
        grad = None
    else:
        grad = mf.nuc_grad_method().kernel()
    rdm = mf.make_rdm1() if not _NO_DM else None
    return etot, grad, rdm

def calc_mp2(mol, **scfargs):
    import pyscf.mp
    mf = solve_scf(mol, **scfargs)
    if not mf.converged:
        raise RuntimeError("SCF not converged!")
    postmf = pyscf.mp.MP2(mf).run()
    etot = postmf.e_tot
    grad = postmf.nuc_grad_method().kernel() if not _NO_FORCE else None
    return etot, grad, None

def calc_ccsd(mol, **scfargs):
    import pyscf.cc
    mf = solve_scf(mol, **scfargs)
    if not mf.converged:
        raise RuntimeError("SCF not converged!")
    mycc = mf.CCSD().run()
    etot = mycc.e_tot
    grad = mycc.nuc_grad_method().kernel() if not _NO_FORCE else None
    ccdm = np.einsum('...pi,...ij,...qj->...pq', 
        mf.mo_coeff, mycc.make_rdm1(), mf.mo_coeff.conj()) if not _NO_DM else None
    return etot, grad, ccdm

def calc_ccsd_t(mol, **scfargs):
    import pyscf.cc
    mf = solve_scf(mol, **scfargs)
    if not mf.converged:
        raise RuntimeError("SCF not converged!")
    mycc = mf.CCSD().run()
    et_correction = mycc.ccsd_t()
    etot = mycc.e_tot + et_correction
    if _NO_FORCE:
        return etot, None, None
    import pyscf.grad.ccsd_t as ccsd_t_grad
    grad = ccsd_t_grad.Gradients(mycc).kernel()
    return etot, grad, None

def calc_fci(mol, **scfargs):
    import pyscf.fci
    mf = solve_scf(mol, **scfargs)
    if not mf.converged:
        raise RuntimeError("SCF not converged!")
    myci = pyscf.fci.FCI(mf)
    etot, fcivec = myci.kernel()
    rdm = np.einsum('...pi,...ij,...qj->...pq', 
            mf.mo_coeff, 
            myci.make_rdm1s(fcivec, mol.nao, mol.nelec), 
            mf.mo_coeff.conj()).sum(0) if not _NO_DM else None
    return etot, None, rdm


if __name__ == "__main__":
    import argparse
    import os
    parser = argparse.ArgumentParser(description="Calculate and save mp2 energy and mo_coeffs for given xyz files.")
    parser.add_argument("files", nargs="+", help="input xyz files")
    parser.add_argument("-d", "--dump-dir", help="dir of dumped files, default is same dir as xyz file")
    parser.add_argument("-v", "--verbose", default=1, type=int, help="output calculation information")
    parser.add_argument("-B", "--basis", default="ccpvdz", type=str, help="basis used to do the calculation")
    parser.add_argument("-C", "--charge", default=0, type=int, help="net charge of the molecule")
    parser.add_argument("-S", "--spin", default=0, type=int, help="net spin of the molecule")
    parser.add_argument("-M", "--method", default="ccsd", help="method used to do the calculation. support MP2, CCSD and CCSD(T)")
    parser.add_argument("-U", "--unrestrict", action="store_true", help="force using unrestricted methods")
    parser.add_argument("-NF", "--no-force", action="store_true", help="do not calculate force")
    parser.add_argument("-ND", "--no-dm", action="store_true", help="do not calculate dm")
    parser.add_argument("-SO", "--newton", action="store_true", help="allow using newton method when scf not converged")
    parser.add_argument("--scf-input", help="yaml file to specify scf arguments")
    args = parser.parse_args()
    
    if args.unrestrict: _MUST_UNRES = True
    if args.no_force: _NO_FORCE = True
    if args.no_dm: _NO_DM = True
    if args.newton: _USE_NEWTON = True

    scfargs = {}
    if args.scf_input is not None:
        import ruamel.yaml as yaml
        with open(args.scf_input, 'r') as fp:
            scfargs = yaml.safe_load(fp)        
    if args.dump_dir is not None:
        os.makedirs(args.dump_dir, exist_ok = True)
    calculator = get_method(args.method)

    for fn in args.files:
        tic = time.time()
        mol = parse_xyz(fn, args.basis, verbose=args.verbose, charge=args.charge, spin=args.spin)
        try:
            res = calculator(mol, **scfargs)
        except RuntimeError as err:
            print(fn, f"failed, {err}")
            continue
        etot, grad, rdm = res
        if args.dump_dir is None:
            dump_dir = os.path.dirname(fn)
        else:
            dump_dir = args.dump_dir
        dump = os.path.join(dump_dir, os.path.splitext(os.path.basename(fn))[0])
        np.save(dump+".energy.npy", [etot])
        if grad is not None:
            force = -grad / BOHR
            np.save(dump+".force.npy", force)
        if rdm is not None:
            np.save(dump+".dm.npy", rdm)
        if args.verbose:
            print(fn, f"done, time = {time.time()-tic}")