How to use the pyqmc.default_sj function in pyqmc

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

def test_transform():
    """ Just prints things out; 
    TODO: figure out a thing to test.
    """
    from pyscf import gto, scf

    r = 1.54 / 0.529177
    mol = gto.M(
        atom="H 0. 0. 0.; H 0. 0. %g" % r,
        ecp="bfd",
        basis="bfd_vtz",
        unit="bohr",
        verbose=1,
    )
    mf = scf.RHF(mol).run()
    wf, to_opt = pyqmc.default_sj(mol, mf)
    enacc = pyqmc.EnergyAccumulator(mol)
    print(list(wf.parameters.keys()))
    transform = LinearTransform(wf.parameters)
    x = transform.serialize_parameters(wf.parameters)

    nconfig = 10
    configs = pyqmc.initial_guess(mol, nconfig)
    wf.recompute(configs)
    pgrad = wf.pgradient()
    gradtrans = transform.serialize_gradients(pgrad)
    assert gradtrans.shape[1] == len(x)
    assert gradtrans.shape[0] == nconfig

def test_info_functions_mol():
    from pyscf import gto, scf
    from pyqmc.tbdm import TBDMAccumulator

    mol = gto.Mole()
    mol.atom = """He 0.00 0.00 0.00 """
    mol.basis = "ccpvdz"
    mol.build()

    mf = scf.RHF(mol)
    ehf = mf.kernel()

    wf, to_opt = pyqmc.default_sj(mol, mf)
    accumulators = {
        "pgrad": pyqmc.gradient_generator(mol, wf, to_opt),
        "obdm": OBDMAccumulator(mol, orb_coeff=mf.mo_coeff),
        "tbdm_updown": TBDMAccumulator(mol, np.asarray([mf.mo_coeff] * 2), (0, 1)),
    }
    info_functions(mol, wf, accumulators)

kmf = scf.KRHF(cell, exxdiv=None).density_fit()
    kmf.kpts = cell.make_kpts([nk, nk, nk])
    ehf = kmf.kernel()
    print("EHF", ehf)
    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},

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,
    )

from pyscf import gto, scf

    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},