How to use the pyqmc.obdm.OBDMAccumulator function in pyqmc

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

lowdin = lo.orth_ao(mol, "lowdin")

    # MOs in the Lowdin basis.
    mo = solve(lowdin, mf.mo_coeff)

    # make AO to localized orbital coefficients.
    mfobdm = mf.make_rdm1(mo, mf.mo_occ)

    ### Test OBDM calculation.
    nconf = 500
    nsteps = 400
    warmup = 15
    wf = PySCFSlater(mol, mf)
    configs = initial_guess(mol, nconf)
    obdm_dict = dict(mol=mol, orb_coeff=lowdin, nsweeps=5, warmup=15)
    obdm = OBDMAccumulator(**obdm_dict)
    obdm_up = OBDMAccumulator(**obdm_dict, spin=0)
    obdm_down = OBDMAccumulator(**obdm_dict, spin=1)

    df, coords = vmc(
        wf,
        configs,
        nsteps=nsteps,
        accumulators={"obdm": obdm, "obdm_up": obdm_up, "obdm_down": obdm_down},
    )

    obdm_est = {}
    for k in ["obdm", "obdm_up", "obdm_down"]:
        avg_norm = np.mean(df[k + "norm"][warmup:], axis=0)
        avg_obdm = np.mean(df[k + "value"][warmup:], axis=0)
        obdm_est[k] = normalize_obdm(avg_obdm, avg_norm)

ovlp = mf.get_ovlp()[kinds]
    lowdin = [lo.vec_lowdin(l, o) for l, o in zip(loiao, ovlp)]
    lreps = [np.linalg.multi_dot([l.T, o, c]) for l, o, c in zip(lowdin, ovlp, coefs)]

    # make AO to localized orbital coefficients.
    mfobdm = [np.einsum("ij,j,kj->ik", l.conj(), o, l) for l, o in zip(lreps, occs)]

    ### Test OBDM calculation.
    nconf = 800
    nsteps = 50
    warmup = 6
    wf = PySCFSlater(mol, mf)
    configs = initial_guess(mol, nconf)
    obdm_dict = dict(mol=mol, orb_coeff=lowdin, kpts=kpts, nsweeps=4, warmup=10)
    obdm = OBDMAccumulator(**obdm_dict)
    obdm_up = OBDMAccumulator(**obdm_dict, spin=0)
    obdm_down = OBDMAccumulator(**obdm_dict, spin=1)

    df, coords = vmc(
        wf,
        configs,
        nsteps=nsteps,
        accumulators={"obdm": obdm, "obdm_up": obdm_up, "obdm_down": obdm_down},
        verbose=True,
    )

    obdm_est = {}
    for k in ["obdm", "obdm_up", "obdm_down"]:
        avg_norm = np.mean(df[k + "norm"][warmup:], axis=0)
        avg_obdm = np.mean(df[k + "value"][warmup:], axis=0)
        obdm_est[k] = normalize_obdm(avg_obdm, avg_norm)

# MOs in the Lowdin basis.
    mo = solve(lowdin, mf.mo_coeff)

    # make AO to localized orbital coefficients.
    mfobdm = mf.make_rdm1(mo, mf.mo_occ)

    ### Test OBDM calculation.
    nconf = 500
    nsteps = 400
    warmup = 15
    wf = PySCFSlater(mol, mf)
    configs = initial_guess(mol, nconf)
    obdm_dict = dict(mol=mol, orb_coeff=lowdin, nsweeps=5, warmup=15)
    obdm = OBDMAccumulator(**obdm_dict)
    obdm_up = OBDMAccumulator(**obdm_dict, spin=0)
    obdm_down = OBDMAccumulator(**obdm_dict, spin=1)

    df, coords = vmc(
        wf,
        configs,
        nsteps=nsteps,
        accumulators={"obdm": obdm, "obdm_up": obdm_up, "obdm_down": obdm_down},
    )

    obdm_est = {}
    for k in ["obdm", "obdm_up", "obdm_down"]:
        avg_norm = np.mean(df[k + "norm"][warmup:], axis=0)
        avg_obdm = np.mean(df[k + "value"][warmup:], axis=0)
        obdm_est[k] = normalize_obdm(avg_obdm, avg_norm)

    print("Average OBDM(orb,orb)", obdm_est["obdm"].diagonal().round(3))

mftbdm = singledet_tbdm(mf, mfobdm)

    ### Test TBDM calculation.
    # VMC params
    nconf = 500
    n_vmc_steps = 400
    vmc_tstep = 0.3
    vmc_warmup = 30
    # TBDM params
    tbdm_sweeps = 4
    tbdm_tstep = 0.5

    wf = PySCFSlater(mol, mf)  # Single-Slater (no jastrow) wf
    configs = initial_guess(mol, nconf)
    energy = EnergyAccumulator(mol)
    obdm_up = OBDMAccumulator(mol=mol, orb_coeff=iaos[0], nsweeps=tbdm_sweeps, spin=0)
    obdm_down = OBDMAccumulator(mol=mol, orb_coeff=iaos[1], nsweeps=tbdm_sweeps, spin=1)
    tbdm_upup = TBDMAccumulator(
        mol=mol, orb_coeff=iaos, nsweeps=tbdm_sweeps, tstep=tbdm_tstep, spin=(0, 0)
    )
    tbdm_updown = TBDMAccumulator(
        mol=mol, orb_coeff=iaos, nsweeps=tbdm_sweeps, tstep=tbdm_tstep, spin=(0, 1)
    )
    tbdm_downup = TBDMAccumulator(
        mol=mol, orb_coeff=iaos, nsweeps=tbdm_sweeps, tstep=tbdm_tstep, spin=(1, 0)
    )
    tbdm_downdown = TBDMAccumulator(
        mol=mol, orb_coeff=iaos, nsweeps=tbdm_sweeps, tstep=tbdm_tstep, spin=(1, 1)
    )

    print("VMC...")
    df, coords = vmc(

def find_basis_evaluate(mfchk, hdf_opt, hdf_vmc, hdf_final):
    """Given a wave function in hdf_opt, compute the 1-RDM (stored in hdf_vmc) , generate a minimal atomic basis and compute the energy/OBDM/TBDM and store in hdf_final """
    from pyqmc.obdm import OBDMAccumulator
    from pyqmc.tbdm import TBDMAccumulator
    from pyqmc import EnergyAccumulator

    sys = pyqmc_from_hdf(mfchk)

    mol = sys["mol"]
    a = lo.orth_ao(mol, "lowdin")
    obdm_up = OBDMAccumulator(mol=mol, orb_coeff=a, spin=0)
    obdm_down = OBDMAccumulator(mol=mol, orb_coeff=a, spin=1)
    with h5py.File(hdf_opt, "r") as hdf_in:
        if f"wf" in hdf_in.keys():
            print("reading in wave function")
            grp = hdf_in[f"wf"]
            for k in grp.keys():
                sys["wf"].parameters[k] = np.array(grp[k])

    configs = pyqmc.initial_guess(sys["mol"], 1000)
    pyqmc.vmc(
        sys["wf"],
        configs,
        nsteps=500,
        hdf_file=hdf_vmc,
        accumulators={"obdm_up": obdm_up, "obdm_down": obdm_down},
    )

configs = pyqmc.initial_guess(sys["mol"], 1000)
    pyqmc.vmc(
        sys["wf"],
        configs,
        nsteps=500,
        hdf_file=hdf_vmc,
        accumulators={"obdm_up": obdm_up, "obdm_down": obdm_down},
    )

    with h5py.File(hdf_vmc, "r") as vmc_hdf:
        obdm_up = np.mean(np.array(vmc_hdf["obdm_upvalue"]), axis=0)
        obdm_down = np.mean(np.array(vmc_hdf["obdm_downvalue"]), axis=0)
    basis_up = gen_basis(mol, sys["mf"], obdm_up)
    basis_down = gen_basis(mol, sys["mf"], obdm_down)
    obdm_up_acc = OBDMAccumulator(mol=mol, orb_coeff=basis_up, spin=0)
    obdm_down_acc = OBDMAccumulator(mol=mol, orb_coeff=basis_down, spin=1)
    tbdm = TBDMAccumulator(mol, np.array([basis_up, basis_down]), spin=(0, 1))
    acc = {
        "energy": EnergyAccumulator(mol),
        "obdm_up": obdm_up_acc,
        "obdm_down": obdm_down_acc,
        "tbdm": tbdm,
    }

    configs = pyqmc.initial_guess(sys["mol"], 1000)
    pyqmc.vmc(sys["wf"], configs, nsteps=500, hdf_file=hdf_final, accumulators=acc)