Add PySCF FCI solver plugin and pyscf integral-array model

[Copilot]

Adds a PySCF-compatible FCI solver (qmp/plugins/pyscf.py) that wraps QMP algorithms behind the standard solver.kernel(h1e, eri, norb, nelec) interface, plus a supporting model (qmp/models/pyscf.py) that builds a Hamiltonian directly from integral arrays instead of a file.

Usage

from qmp.plugins.pyscf import Solver

solver = Solver({
    "action": {"name": "haar", "params": {}},
    "network": {"name": "mlp/u1u1", "params": {}},
    "optimizer": {"name": "Adam", "params": {"lr": 1e-3}},
    "common": {
        "device": "cpu",
        "dtype": "float64",
        "max_absolute_step": 10,
    },
})
energy, ci = solver.kernel(h1e, eri, norb, nelec)

# Warm-start from previous result
energy2, ci2 = solver.kernel(h1e, eri, norb, nelec, ci0=ci)

New files

• qmp/models/pyscf.py — Model built from (h1e, eri, n_orbit, n_electron, n_spin) numpy arrays. Mirrors fcidump.py: same permute(0,2,3,1)/2 + spin-orbital expansion, same four network configs (mlp/u1u1, transformers/u1u1, mlp/u1, transformers/u1). Includes _restore_eri() to expand 4-fold symmetric compressed (nij, nij) ERI to full (norb,norb,norb,norb).
• qmp/plugins/pyscf.py — Solver class whose __init__ accepts a single config dict/DictConfig in the same format as __main__.py's YAML config (action, network, optimizer, common). common.max_absolute_step controls how many global algorithm steps to run before returning. Key internals:
  • _PyscfRuntimeContext subclasses RuntimeContext and only overrides folder() (returns a temporary directory) and create_model() (injects the pre-built model). setup() and save() are inherited unchanged and work via normal polymorphism.
  • kernel() closely mimics __main__.py: imports modules, creates the context via dacite from config.common, calls setup(), instantiates and runs the algorithm. When max_absolute_step is reached, the base save() writes the checkpoint to the temp dir and calls sys.exit(0). kernel() catches the resulting SystemExit and reads the saved data.pth checkpoint file to extract the energy and CI data.
  • Returns (energy, ci) where energy is extracted from the checkpoint (for haar: data["haar"]["excited"][N-1][0][0]) and ci is the full checkpoint minus random state. For the haar algorithm, ci["haar"]["pool"] contains a (configs, psi) tuple representing the sampled basis and Lanczos-optimised amplitudes.
  • ci0 warm-starts the algorithm by saving it to tmpdir/data.pth before setup() loads it, exactly mirroring how __main__.py handles checkpoint resume.
• qmp/plugins/__init__.py — package init.

Modified files

• pyproject.toml — include = ["qmp"] → include = ["qmp*"] so sub-packages including qmp.plugins are correctly discovered on install.

Original prompt

---

This section details on the original issue you should resolve

<issue_title>add a solver for pyscf.</issue_title>
<issue_description>Add a file qmp/plugins/pyscf.py , where would be used as something like energy, ci = solver.kernel(h1, eri, norb, nelec) , inside the implementation, it invoke algorithm in main.py . the solver init accept configurations which was passed to main.py (in context.py, max_absolute_step in considered to be set as plugin, since it need to return). besides, an new kind of models may need to be created, which is similar to openfermion.py and fcidump.py (very similar to fcidump.py).</issue_description>

Comments on the Issue (you are @copilot in this section)

• Fixes add a solver for pyscf. #207

---

🔒 GitHub Advanced Security automatically protects Copilot coding agent pull requests. You can protect all pull requests by enabling Advanced Security for your repositories. Learn more about Advanced Security.

[hzhangxyz]

修改一下, 输入参数应该是__main__.py里面那个配置文件的内容, 操作方式应该是模仿__main__.py的函数, 只是需要注意, 集成一下现有的 RuntimeContext , 把 folder 重写到一个临时目录, 跑了固定步数后, 从保存的checkpoint文件中读取结果. 通过这个方式, 参数的格式完全复用, 另外, checkpoint中要传给pyscf的数据可以参考haar.py, 其实应该最多只需要config和psi(psi要不要我不知道).

[Copilot]

Pull request overview

Adds a PySCF-compatible “FCI solver” wrapper around QMP algorithms, plus a PySCF integral-array-backed model so QMP can be driven directly from (h1e, eri, norb, nelec) without an FCIDUMP file.

Changes:

• Introduces qmp.plugins.pyscf.Solver implementing kernel(h1e, eri, norb, nelec, ci0=...) by running a selected QMP algorithm programmatically.
• Adds qmp.models.pyscf.Model to build QMP’s internal Hamiltonian from PySCF-style integral arrays (including restoring 4-fold-symmetric compressed ERIs).
• Updates packaging discovery to include the new qmp.plugins subpackage.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 4 comments.

File	Description
qmp/plugins/pyscf.py	Implements PySCF-style solver wrapper and Hydra-bypassing runtime context to run QMP algorithms and return (energy, state_dict)
qmp/plugins/__init__.py	Declares plugins subpackage
qmp/models/pyscf.py	Adds integral-array model, ERI restore helper, and network configs matching existing FCIDUMP model patterns
pyproject.toml	Adjusts setuptools package discovery to include qmp* (so qmp.plugins is installed)

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

When eri.ndim == 4, _restore_eri() returns the array without validating it matches (norb, norb, norb, norb). If a caller passes a mis-shaped 4D ERI, the failure will happen later (e.g., in permute) with a less clear error. Add an explicit shape check and raise a ValueError with the expected shape.

Suggested change

	if eri.ndim == 4:
	if eri.ndim == 4:
	expected_shape = (norb, norb, norb, norb)
	if eri.shape != expected_shape:
	raise ValueError(
	f"For 4D ERI, expected shape {expected_shape} for norb={norb}, got {eri.shape}"
	)

[Copilot]

prev_grad_enabled is captured here but the restoration happens only at the end of kernel(). If anything raises before the restore line (e.g., run.main(...)), the caller’s grad mode can be left incorrect. Wrap the main body in try/finally to always restore torch.set_grad_enabled(prev_grad_enabled).

Suggested change

	# Save the caller's gradient state so we can restore it on exit.
	prev_grad_enabled = torch.is_grad_enabled()
	if isinstance(nelec, int):
	nalpha = (nelec + 1) // 2
	nbeta = nelec - nalpha
	else:
	nalpha, nbeta = int(nelec[0]), int(nelec[1])
	n_electron = nalpha + nbeta
	n_spin = nalpha - nbeta
	# ------------------------------------------------------------------
	# 2. Build the PySCF model
	# ------------------------------------------------------------------
	model_config = ModelConfig(
	h1e=numpy.asarray(h1, dtype=numpy.float64),
	eri=numpy.asarray(eri, dtype=numpy.float64),
	n_orbit=norb,
	n_electron=n_electron,
	n_spin=n_spin,
	nuclear_repulsion=nuclear_repulsion,
	ref_energy=ref_energy,
	)
	# ------------------------------------------------------------------
	# Save the caller's gradient state so we can restore it on exit.
	prev_grad_enabled = torch.is_grad_enabled()
	try:
	if isinstance(nelec, int):
	nalpha = (nelec + 1) // 2
	nbeta = nelec - nalpha
	else:
	nalpha, nbeta = int(nelec[0]), int(nelec[1])
	n_electron = nalpha + nbeta
	n_spin = nalpha - nbeta
	# ------------------------------------------------------------------
	# 2. Build the PySCF model
	# ------------------------------------------------------------------
	model_config = ModelConfig(
	h1e=numpy.asarray(h1, dtype=numpy.float64),
	eri=numpy.asarray(eri, dtype=numpy.float64),
	n_orbit=norb,
	n_electron=n_electron,
	n_spin=n_spin,
	nuclear_repulsion=nuclear_repulsion,
	ref_energy=ref_energy,
	)
	finally:
	torch.set_grad_enabled(prev_grad_enabled)

[Copilot]

final_state_dict is taken from final_data.get("network", {}). If the algorithm returns normally (no _AlgorithmComplete) or doesn’t populate final_data, this will be {} and kernel() returns an empty ci even if the network was trained. Add a fallback (e.g., network.state_dict()) when final_state_dict is empty so warm-starting via ci0 remains reliable.

Suggested change

	network.load_state_dict(final_state_dict)
	network.load_state_dict(final_state_dict)
	else:
	# If the algorithm did not provide a final state dict via
	# final_data, fall back to the current network parameters so
	# that warm-starting via ci0 remains reliable.
	final_state_dict = network.state_dict()

[Copilot]

_restore_eri() expands the symmetric (nij, nij) ERI via nested Python loops over all nij**2 pair combinations, which can become a bottleneck as norb grows. Consider a vectorized NumPy approach (pair-index maps + advanced indexing) to avoid Python-level O(n^4) looping.

Suggested change

	eri_full = numpy.zeros((norb, norb, norb, norb), dtype=eri.dtype)
	i_idx, j_idx = numpy.tril_indices(norb)
	ij_list = list(zip(i_idx.tolist(), j_idx.tolist()))
	for ij, (i, j) in enumerate(ij_list):
	for kl, (k, m) in enumerate(ij_list):
	val = eri[ij, kl]
	eri_full[i, j, k, m] = val
	eri_full[j, i, k, m] = val
	eri_full[i, j, m, k] = val
	eri_full[j, i, m, k] = val
	eri_full[k, m, i, j] = val
	eri_full[m, k, i, j] = val
	eri_full[k, m, j, i] = val
	eri_full[m, k, j, i] = val
	# Allocate full 4D ERI tensor
	eri_full = numpy.zeros((norb, norb, norb, norb), dtype=eri.dtype)
	# Map compressed pair indices (ij) back to orbital indices (i, j) with i >= j
	i_idx, j_idx = numpy.tril_indices(norb)
	# Build broadcasted index grids for all (ij, kl) combinations
	# Shapes after broadcasting: (nij, nij) for each of I, J, K, L
	I = i_idx[:, None]
	J = j_idx[:, None]
	K = i_idx[None, :]
	L = j_idx[None, :]
	# Assign values for all 8 symmetry-related permutations in a vectorized manner
	eri_full[I, J, K, L] = eri
	eri_full[J, I, K, L] = eri
	eri_full[I, J, L, K] = eri
	eri_full[J, I, L, K] = eri
	eri_full[K, L, I, J] = eri
	eri_full[L, K, I, J] = eri
	eri_full[K, L, J, I] = eri
	eri_full[L, K, J, I] = eri

[Copilot]

修改一下, 输入参数应该是__main__.py里面那个配置文件的内容, 操作方式应该是模仿__main__.py的函数, 只是需要注意, 集成一下现有的 RuntimeContext , 把 folder ...

Done in c92e5b9. The plugin has been redesigned as requested:

• Solver.__init__ now takes a single config dict/DictConfig in exactly the same format as __main__.py's YAML config (action, network, optimizer, common). common.max_absolute_step controls when to return.

• _PyscfRuntimeContext is now minimal — only overrides folder() (to use a temp dir) and create_model() (to inject the pre-built model). setup() and save() are inherited unchanged and work via polymorphism.

• kernel() closely mimics __main__.py: imports both the algorithm and model modules, creates the context via dacite from config.common, calls setup(), instantiates the algorithm, runs main(), then catches SystemExit (raised by the base save() at max_absolute_step) and reads the saved checkpoint from tmpdir/data.pth.

• Returns (energy, ci) where energy is extracted from the checkpoint (data["haar"]["excited"][N-1][0][0] for haar) and ci is the full checkpoint minus random state — for haar it includes ci["haar"]["pool"] = (configs, psi). ci0 warm-starts by saving to tmpdir/data.pth before setup().

Example usage:

solver = Solver({
    "action": {"name": "haar", "params": {}},
    "network": {"name": "mlp/u1u1", "params": {}},
    "optimizer": {"name": "Adam", "params": {"lr": 1e-3}},
    "common": {"device": "cpu", "max_absolute_step": 10},
})
energy, ci = solver.kernel(h1e, eri, norb, nelec)

[hzhangxyz]

返回值格式不对, 给个例子:

import numpy as np
from pyscf import fci
import qmp.plugins.pyscf as qmp_pyscf

# -------------------------------
# 1. 定义系统大小
# -------------------------------
norb = 2         # 轨道数
nelec = (1, 1)   # alpha/beta 电子数

# -------------------------------
# 2. 构造简单 Hamiltonian
# -------------------------------
# one-body Hamiltonian
h1 = np.array([[1.0, 0.2],
               [0.2, 0.8]])

# two-body integrals (chemist notation)
eri = np.zeros((norb, norb, norb, norb))
# 简单填几个值，保证非零交互
eri[0,0,0,0] = 0.7
eri[1,1,1,1] = 0.6
eri[0,1,0,1] = 0.2
eri[0,1,1,0] = 0.2

# -------------------------------
# 3. 初始化 FCI solver
# -------------------------------
solver = fci.direct_spin1.FCI()
mysolver = qmp_pyscf.Solver({
    "action": {"name": "haar", "params": {}},
    "network": {"name": "mlp/u1", "params": {}},
    "optimizer": {"name": "Adam", "params": {"lr": 1e-3}},
    "common": {
        "device": "cpu",
        "dtype": "float64",
        "max_absolute_step": 10,
    },
})
# -------------------------------
# 4. 求基态
# -------------------------------
energy, ci = solver.kernel(h1, eri, norb, nelec)

print("FCI 能量:", energy)
print("CI wavefunction:", ci)

energy, ci = mysolver.kernel(h1, eri, norb, nelec)

print("FCI 能量:", energy)
#print("CI wavefunction:", ci)

这里参考的ci看起来是两个config的amplitude, 但是现在的ci是一个巨大的dict