How to use pyqmc - 10 common examples

from pyqmc.supercell import get_supercell

    mol = gto.Cell(atom="He 0.00 0.00 0.00", basis="ccpvdz", unit="B")
    mol.a = 5.61 * np.eye(3)
    mol.build()

    mf = scf.KRHF(mol, kpts=mol.make_kpts([2, 2, 2])).density_fit()
    ehf = mf.kernel()

    supercell = get_supercell(mol, 2 * np.eye(3))
    kinds = [0, 1]
    dm_orbs = [mf.mo_coeff[i][:, :2] for i in kinds]
    wf, to_opt = pyqmc.default_sj(mol, mf)
    accumulators = {
        "pgrad": pyqmc.gradient_generator(mol, wf, to_opt, ewald_gmax=10),
        "obdm": OBDMAccumulator(mol, dm_orbs, kpts=mf.kpts[kinds]),
        "Sq": pyqmc.accumulators.SqAccumulator(mol.lattice_vectors()),
    }
    info_functions(mol, wf, accumulators)

def runtest(mol, mf, kind=0):
    for k, occ in enumerate(mf.mo_occ):
        print(k, occ)
    kpt = mf.kpts[kind]
    twist = np.dot(kpt, mol.lattice_vectors().T / (2 * np.pi))
    print("kpt", kpt)
    print("twist", twist)
    wf0 = pyqmc.PySCFSlater(mol, mf)
    wft = pyqmc.PySCFSlater(mol, mf, twist=twist)

    #####################################
    ## compare values across boundary
    ## psi, KE, ecp,
    #####################################
    nconfig = 100
    coords = pyqmc.initial_guess(mol, nconfig, 1)
    nelec = coords.configs.shape[1]
    epos, wrap = enforce_pbc(coords.lvecs, coords.configs)
    coords = PeriodicConfigs(epos, coords.lvecs)

    shift_ = np.random.randint(10, size=coords.configs.shape) - 5
    phase = np.exp(2j * np.pi * np.einsum("ijk,k->ij", shift_, twist))

    shift = np.dot(shift_, mol.lattice_vectors())
    epos, wrap = enforce_pbc(coords.lvecs, epos + shift)

def runtest(mol, mf, kind=0):
    for k, occ in enumerate(mf.mo_occ):
        print(k, occ)
    kpt = mf.kpts[kind]
    twist = np.dot(kpt, mol.lattice_vectors().T / (2 * np.pi))
    print("kpt", kpt)
    print("twist", twist)
    wf0 = pyqmc.PySCFSlater(mol, mf)
    wft = pyqmc.PySCFSlater(mol, mf, twist=twist)

    #####################################
    ## compare values across boundary
    ## psi, KE, ecp,
    #####################################
    nconfig = 100
    coords = pyqmc.initial_guess(mol, nconfig, 1)
    nelec = coords.configs.shape[1]
    epos, wrap = enforce_pbc(coords.lvecs, coords.configs)
    coords = PeriodicConfigs(epos, coords.lvecs)

    shift_ = np.random.randint(10, size=coords.configs.shape) - 5
    phase = np.exp(2j * np.pi * np.einsum("ijk,k->ij", shift_, twist))

    shift = np.dot(shift_, mol.lattice_vectors())

def runtest(mol, mf, kind=0, do_mc=False):
    if do_mc:
        from pyscf import mcscf

        mc = mcscf.CASCI(mf, ncas=4, nelecas=(1, 1))
        mc.kernel()
        wf = pyqmc.default_msj(mol, mf, mc)[0]
        kpt = mf.kpt
        dm = mc.make_rdm1()
        if len(dm.shape) == 4:
            dm = np.sum(dm, axis=0)
    else:
        kpt = mf.kpts[kind]
        wf = pyqmc.PySCFSlater(mol, mf)
        dm = mf.make_rdm1()
        print("original dm shape", dm.shape)
        if len(dm.shape) == 4:
            dm = np.sum(dm, axis=0)
        dm = dm[kind]

    #####################################
    ## evaluate KE in PySCF
    #####################################
    ke_mat = mol.pbc_intor("int1e_kin", hermi=1, kpts=np.array(kpt))
    print("ke_mat", ke_mat.shape)
    print("dm", dm.shape)
    pyscfke = np.real(np.einsum("ij,ji->", ke_mat, dm))
    print("PySCF kinetic energy: {0}".format(pyscfke))

    #####################################

def test():
    """ Optimize a Helium atom's wave function and check that it's 
    better than Hartree-Fock"""

    mol = gto.M(atom="He 0. 0. 0.", basis="bfd_vdz", ecp="bfd", unit="bohr")
    mf = scf.RHF(mol).run()
    wf, to_opt = default_sj(mol, mf)
    print(to_opt)
    nconf = 500
    wf, dfgrad = line_minimization(
        wf, initial_guess(mol, nconf), gradient_generator(mol, wf, to_opt)
    )

    dfgrad = pd.DataFrame(dfgrad)
    print(dfgrad)
    mfen = mf.energy_tot()
    enfinal = dfgrad["energy"].values[-1]
    enfinal_err = dfgrad["energy_error"].values[-1]
    assert mfen > enfinal

from pyscf.pbc import gto, scf
    from pyqmc.supercell import get_supercell

    mol = gto.Cell(atom="He 0.00 0.00 0.00", basis="ccpvdz", unit="B")
    mol.a = 5.61 * np.eye(3)
    mol.build()

    mf = scf.KRHF(mol, kpts=mol.make_kpts([2, 2, 2])).density_fit()
    ehf = mf.kernel()

    supercell = get_supercell(mol, 2 * np.eye(3))
    kinds = [0, 1]
    dm_orbs = [mf.mo_coeff[i][:, :2] for i in kinds]
    wf, to_opt = pyqmc.default_sj(mol, mf)
    accumulators = {
        "pgrad": pyqmc.gradient_generator(mol, wf, to_opt, ewald_gmax=10),
        "obdm": OBDMAccumulator(mol, dm_orbs, kpts=mf.kpts[kinds]),
        "Sq": pyqmc.accumulators.SqAccumulator(mol.lattice_vectors()),
    }
    info_functions(mol, wf, accumulators)

def runtest(mol, mf, kind=0, do_mc=False):
    if do_mc:
        from pyscf import mcscf

        mc = mcscf.CASCI(mf, ncas=4, nelecas=(1, 1))
        mc.kernel()
        wf = pyqmc.default_msj(mol, mf, mc)[0]
        kpt = mf.kpt
        dm = mc.make_rdm1()
        if len(dm.shape) == 4:
            dm = np.sum(dm, axis=0)
    else:
        kpt = mf.kpts[kind]
        wf = pyqmc.PySCFSlater(mol, mf)
        dm = mf.make_rdm1()
        print("original dm shape", dm.shape)
        if len(dm.shape) == 4:
            dm = np.sum(dm, axis=0)
        dm = dm[kind]

    #####################################
    ## evaluate KE in PySCF
    #####################################

def test():
    # Default multi-Slater wave function
    mol = gto.M(atom="Li 0. 0. 0.; H 0. 0. 1.5", basis="cc-pvtz", unit="bohr", spin=0)
    mf = scf.RHF(mol).run()
    mc = mcscf.CASCI(mf, ncas=2, nelecas=(1, 1))
    mc.kernel()

    wf, to_opt = default_msj(mol, mf, mc)
    old_parms = wf.parameters
    lt = LinearTransform(wf.parameters, to_opt)

    # Test serialize parameters
    x0 = lt.serialize_parameters(wf.parameters)
    x0 += np.random.normal(size=x0.shape)
    wf.parameters = lt.deserialize(x0)
    assert wf.parameters["wf1det_coeff"][0] == old_parms["wf1det_coeff"][0]
    assert np.sum(wf.parameters["wf2bcoeff"][0] - old_parms["wf2bcoeff"][0]) == 0

    # Test serialize gradients
    configs = OpenConfigs(np.random.randn(10, 4, 3))
    wf.recompute(configs)
    pgrad = wf.pgradient()
    pgrad_serial = lt.serialize_gradients(pgrad)
    assert np.sum(pgrad_serial[:, :3] - pgrad["wf1det_coeff"][:, 1:4]) == 0

def test_info_functions_pbc():
    from pyscf.pbc import gto, scf
    from pyqmc.supercell import get_supercell

    mol = gto.Cell(atom="He 0.00 0.00 0.00", basis="ccpvdz", unit="B")
    mol.a = 5.61 * np.eye(3)
    mol.build()

    mf = scf.KRHF(mol, kpts=mol.make_kpts([2, 2, 2])).density_fit()
    ehf = mf.kernel()

    supercell = get_supercell(mol, 2 * np.eye(3))
    kinds = [0, 1]
    dm_orbs = [mf.mo_coeff[i][:, :2] for i in kinds]
    wf, to_opt = pyqmc.default_sj(mol, mf)
    accumulators = {
        "pgrad": pyqmc.gradient_generator(mol, wf, to_opt, ewald_gmax=10),
        "obdm": OBDMAccumulator(mol, dm_orbs, kpts=mf.kpts[kinds]),
        "Sq": pyqmc.accumulators.SqAccumulator(mol.lattice_vectors()),
    }
    info_functions(mol, wf, accumulators)

def test():
    """ Optimize a Helium atom's wave function and check that it's 
    better than Hartree-Fock"""

    mol = gto.M(atom="He 0. 0. 0.", basis="bfd_vdz", ecp="bfd", unit="bohr")
    mf = scf.RHF(mol).run()
    wf, to_opt = default_sj(mol, mf)
    print(to_opt)
    nconf = 500
    wf, dfgrad = line_minimization(
        wf, initial_guess(mol, nconf), gradient_generator(mol, wf, to_opt)
    )

    dfgrad = pd.DataFrame(dfgrad)
    print(dfgrad)
    mfen = mf.energy_tot()
    enfinal = dfgrad["energy"].values[-1]
    enfinal_err = dfgrad["energy_error"].values[-1]
    assert mfen > enfinal