How to use the pyqmc.gradient_generator function in pyqmc

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)

mol = gto.M(
        atom="O 0 0 0; H 0 -2.757 2.587; H 0 2.757 2.587", basis="bfd_vtz", ecp="bfd"
    )
    mf = scf.RHF(mol).run()
    return mol, mf


if __name__ == "__main__":
    cluster = LocalCluster(n_workers=ncore, threads_per_worker=1)
    client = Client(cluster)
    mol, mf = run_scf()
    from pyqmc import vmc, line_minimization, rundmc

    wf, to_opt = pyqmc.default_sj(mol, mf)
    pgrad_acc = pyqmc.gradient_generator(mol, wf, to_opt)
    configs = pyqmc.initial_guess(mol, nconfig)
    line_minimization(
        wf,
        configs,
        pgrad_acc,
        hdf_file="h2o_opt.hdf",
        client=client,
        npartitions=ncore,
        verbose=True,
    )
    df, configs = vmc(
        wf,
        configs,
        hdf_file="h2o_vmc.hdf",
        accumulators={"energy": pgrad_acc.enacc},
        client=client,

atom="O 0 0 0; H 0 -2.757 2.587; H 0 2.757 2.587", basis="bfd_vtz", ecp="bfd"
)
mf = scf.RHF(mol).run()
# clean_pyscf_objects gets rid of the TextIO objects that can't
# be sent using parsl.
mol, mf = clean_pyscf_objects(mol, mf)

# It's better to load parsl after pyscf has run. Some of the
# executors have timeouts and will kill the job while pyscf is running!
parsl.load(config)


# We make a Slater-Jastrow wave function and
# only optimize the Jastrow coefficients.
wf = pyqmc.slater_jastrow(mol, mf)
acc = pyqmc.gradient_generator(mol, wf, ["wf2acoeff", "wf2bcoeff"])


# Generate the initial configurations.
# Here we run VMC for a few steps with no accumulators to equilibrate the
# walkers.
configs = pyqmc.initial_guess(mol, nconf)
df, configs = distvmc(wf, configs, accumulators={}, nsteps=10, npartitions=ncore)

# This uses a stochastic reconfiguration step to generate parameter changes along a line,
# then minimizes the energy along that line.
wf, dfgrad, dfline = line_minimization(
    wf,
    configs,
    acc,
    npartitions=ncore,
    vmcoptions={"nsteps": 30},

def pyqmc_from_hdf(chkfile):
    """ Loads pyqmc objects from a pyscf checkfile """
    mol = lib.chkfile.load_mol(chkfile)
    mol.output = None
    mol.stdout = None

    mf = scf.RHF(mol)
    mf.__dict__.update(scf.chkfile.load(chkfile, "scf"))
    with h5py.File(chkfile, "r") as f:
        mc = mcscf.CASCI(mf, ncas=int(f["mc/ncas"][...]), nelecas=f["mc/nelecas"][...])
        mc.ci = f["mc/ci"][...]

    wf, to_opt = pyqmc.default_msj(mol, mf, mc)

    to_opt["wf1det_coeff"][...] = True
    pgrad = pyqmc.gradient_generator(mol, wf, to_opt)

    return {"mol": mol, "mf": mf, "to_opt": to_opt, "wf": wf, "pgrad": pgrad}

if __name__ == "__main__":
    import pyscf
    import pyqmc

    mol = pyscf.gto.M(atom="He 0. 0. 0.", basis="bfd_vdz", ecp="bfd", unit="bohr")

    mf = pyscf.scf.RHF(mol).run()
    wf, to_opt = pyqmc.default_sj(mol, mf)

    nconfig = 1000
    configs = pyqmc.initial_guess(mol, nconfig)

    acc = pyqmc.gradient_generator(mol, wf, to_opt)
    pyqmc.line_minimization(wf, configs, acc, hdf_file="he_opt.hdf5", verbose=True)

    pyqmc.rundmc(
        wf,
        configs,
        nsteps=5000,
        accumulators={"energy": pyqmc.EnergyAccumulator(mol)},
        tstep=0.02,
        hdf_file="he_dmc.hdf5",
        verbose=True,
    )

return cell, kmf


if __name__ == "__main__":
    # Run SCF
    cell, kmf = run_scf(nk=2)

    # Set up wf and configs
    nconfig = 100
    S = np.eye(3) * 2  # 2x2x2 supercell
    supercell = get_supercell(cell, S)
    wf, to_opt = pyqmc.default_sj(supercell, kmf)
    configs = pyqmc.initial_guess(supercell, nconfig)

    # Initialize energy accumulator (and Ewald)
    pgrad = pyqmc.gradient_generator(supercell, wf, to_opt=to_opt)

    # Optimize jastrow
    wf, lm_df = pyqmc.line_minimization(
        wf, configs, pgrad, hdf_file="pbc_he_linemin.hdf", verbose=True
    )

    # Run VMC
    df, configs = pyqmc.vmc(
        wf,
        configs,
        nblocks=100,
        accumulators={"energy": pgrad.enacc},
        hdf_file="pbc_he_vmc.hdf",
        verbose=True,
    )

from dask.distributed import Client, LocalCluster   

    r = 1.1

    ncore = 2
    sys = setuph2(r)
    cluster = LocalCluster(n_workers=ncore, threads_per_worker=1)
    client = Client(cluster)

    # Set up calculation
    nconf = 800
    configs = pyqmc.initial_guess(sys["mol"], nconf)
    wf, df = line_minimization(
        sys["wf"],
        configs,
        pyqmc.gradient_generator(sys["mol"], sys["wf"]),
        client=client,
        maxiters=5,
    )

    forcing = {}
    obj = {}
    for k in sys["descriptors"]:
        forcing[k] = 0.0
        obj[k] = 0.0

    for k in sys["descriptors_tbdm"]:
        forcing[k] = 0.0
        obj[k] = 0.0

    forcing["t"] = 0.5
    forcing["trace"] = 1.0