From e7b2874565c995bc372df2876fe3364e12f374d5 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 5 Mar 2026 21:21:02 +0000
Subject: [PATCH 1/6] Initial plan


From b8dca25939d1af159559da785cf4979a0483ff96 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 5 Mar 2026 21:43:19 +0000
Subject: [PATCH 2/6] Add pyscf model and solver plugin

- qmp/models/pyscf.py: Model class taking h1e/eri arrays directly (similar to
  fcidump.py). Includes _restore_eri helper for compressed ERI format. Registers
  all four network architectures (mlp/u1u1, transformers/u1u1, mlp/u1,
  transformers/u1).

- qmp/plugins/__init__.py: New plugins subpackage init.

- qmp/plugins/pyscf.py: PySCF-compatible FCI solver. Solver.kernel(h1, eri,
  norb, nelec) builds the model, runs a QMP algorithm (bypassing Hydra and
  sys.exit), then returns (energy, network_state_dict). ci0 warm-starts the
  network. _PyscfRuntimeContext overrides folder(), setup(), create_model(),
  create_network() and save() for programmatic use.

- pyproject.toml: Updated include to 'qmp*' so all sub-packages are properly
  discovered on installation.

Co-authored-by: hzhangxyz <11623447+hzhangxyz@users.noreply.github.com>
---
 pyproject.toml          |   2 +-
 qmp/models/pyscf.py     | 407 ++++++++++++++++++++++++++++++++++++++++
 qmp/plugins/__init__.py |   3 +
 qmp/plugins/pyscf.py    | 390 ++++++++++++++++++++++++++++++++++++++
 4 files changed, 801 insertions(+), 1 deletion(-)
 create mode 100644 qmp/models/pyscf.py
 create mode 100644 qmp/plugins/__init__.py
 create mode 100644 qmp/plugins/pyscf.py

diff --git a/pyproject.toml b/pyproject.toml
index b39d7f3..909d771 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -50,7 +50,7 @@ qmp = "qmp.__main__:main"
 qmp-kit = "qmp.__main__:main"
 
 [tool.setuptools.packages.find]
-include = ["qmp"]
+include = ["qmp*"]
 
 [tool.setuptools.package-data]
 qmp = [
diff --git a/qmp/models/pyscf.py b/qmp/models/pyscf.py
new file mode 100644
index 0000000..96368fd
--- /dev/null
+++ b/qmp/models/pyscf.py
@@ -0,0 +1,407 @@
+"""
+This file provides an interface to work with PySCF-style integral arrays.
+
+The model accepts 1-electron (h1e) and 2-electron (eri) integrals in the same
+convention used by PySCF's FCI solver, converting them to the internal Hamiltonian
+representation used by QMP.
+"""
+
+import logging
+import dataclasses
+import numpy
+import torch
+import openfermion
+from ..networks.mlp import WaveFunctionElectronUpDown as MlpWaveFunction
+from ..networks.mlp import WaveFunctionElectron as MlpWaveFunctionElectron
+from ..networks.transformers import WaveFunctionElectronUpDown as TransformersWaveFunction
+from ..networks.transformers import WaveFunctionElectron as TransformersWaveFunctionElectron
+from ..hamiltonian import Hamiltonian
+from ..utility.model_dict import model_dict, ModelProto, NetworkProto, NetworkConfigProto
+
+
+@dataclasses.dataclass
+class ModelConfig:
+    """
+    The configuration of the PySCF model.
+
+    Integrals should be in chemist's notation (ij|kl), matching PySCF's convention.
+    """
+
+    # 1-electron integrals in MO basis, shape (norb, norb)
+    h1e: numpy.ndarray
+    # 2-electron repulsion integrals in MO basis, shape (norb, norb, norb, norb)
+    # or compressed 4-fold symmetric shape (nij, nij) where nij = norb*(norb+1)//2
+    eri: numpy.ndarray
+    # Number of spatial orbitals
+    n_orbit: int
+    # Total number of electrons
+    n_electron: int
+    # Spin (n_alpha - n_beta, i.e. 2*S)
+    n_spin: int = 0
+    # Constant energy offset (e.g. nuclear repulsion energy)
+    nuclear_repulsion: float = 0.0
+    # Reference (FCI) energy for logging; defaults to zero
+    ref_energy: float = 0.0
+
+
+def _restore_eri(eri: numpy.ndarray, norb: int) -> numpy.ndarray:
+    """
+    Restore a compressed ERI array to the full (norb, norb, norb, norb) form.
+
+    Supports:
+    - 4D array of shape (norb, norb, norb, norb): returned as-is.
+    - 2D array of shape (nij, nij) in 4-fold symmetric format
+      (nij = norb*(norb+1)//2), where the lower-triangle index is
+      ij = i*(i+1)//2 + j for i >= j.
+
+    Parameters
+    ----------
+    eri : numpy.ndarray
+        The electron-repulsion integral array.
+    norb : int
+        Number of spatial orbitals.
+
+    Returns
+    -------
+    numpy.ndarray
+        Full (norb, norb, norb, norb) ERI array.
+    """
+    if eri.ndim == 4:
+        return eri
+    if eri.ndim == 2:
+        nij = norb * (norb + 1) // 2
+        if eri.shape != (nij, nij):
+            raise ValueError(
+                f"For 2D ERI, expected shape ({nij}, {nij}) for norb={norb}, got {eri.shape}"
+            )
+        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
+        return eri_full
+    raise ValueError(f"Unsupported ERI dimensionality: {eri.ndim} (expected 2 or 4)")
+
+
+class Model(ModelProto[ModelConfig]):
+    """
+    This class handles models built from PySCF-style integral arrays (h1e and eri).
+
+    The integrals are in chemist's (ij|kl) notation and are converted to the
+    internal spin-orbital Hamiltonian representation used by QMP.
+    """
+
+    network_dict: dict[str, type[NetworkConfigProto["Model"]]] = {}
+
+    config_t = ModelConfig
+
+    def __init__(self, args: ModelConfig) -> None:
+        n_orbit = args.n_orbit
+        n_electron = args.n_electron
+        n_spin = args.n_spin
+
+        logging.info(
+            "Building PySCF model: %d orbitals, %d electrons, spin=%d",
+            n_orbit,
+            n_electron,
+            n_spin,
+        )
+
+        # -- 1-electron integrals ------------------------------------------------
+        h1e = numpy.asarray(args.h1e, dtype=numpy.float64)
+        if h1e.shape != (n_orbit, n_orbit):
+            raise ValueError(f"Expected h1e shape ({n_orbit}, {n_orbit}), got {h1e.shape}")
+        energy_1: torch.Tensor = torch.as_tensor(h1e, dtype=torch.float64)
+
+        # -- 2-electron integrals ------------------------------------------------
+        eri_full = _restore_eri(numpy.asarray(args.eri, dtype=numpy.float64), n_orbit)
+        energy_2: torch.Tensor = torch.as_tensor(eri_full, dtype=torch.float64)
+
+        # Apply the same permutation used in fcidump.py to convert from
+        # chemist's (ij|kl) notation to the form expected by OpenFermion.
+        energy_2 = energy_2.permute(0, 2, 3, 1).contiguous() / 2
+
+        # -- Build spin-orbital integrals -----------------------------------------
+        energy_1_b: torch.Tensor = torch.zeros([n_orbit * 2, n_orbit * 2], dtype=torch.float64)
+        energy_2_b: torch.Tensor = torch.zeros(
+            [n_orbit * 2, n_orbit * 2, n_orbit * 2, n_orbit * 2], dtype=torch.float64
+        )
+        energy_1_b[0::2, 0::2] = energy_1
+        energy_1_b[1::2, 1::2] = energy_1
+        energy_2_b[0::2, 0::2, 0::2, 0::2] = energy_2
+        energy_2_b[0::2, 1::2, 1::2, 0::2] = energy_2
+        energy_2_b[1::2, 0::2, 0::2, 1::2] = energy_2
+        energy_2_b[1::2, 1::2, 1::2, 1::2] = energy_2
+
+        # -- Convert to FermionOperator via OpenFermion ---------------------------
+        logging.info("Converting integrals to internal Hamiltonian representation")
+        interaction_operator: openfermion.InteractionOperator = openfermion.InteractionOperator(
+            args.nuclear_repulsion, energy_1_b.numpy(), energy_2_b.numpy()
+        )  # type: ignore[no-untyped-call]
+        fermion_operator: openfermion.FermionOperator = openfermion.get_fermion_operator(
+            interaction_operator
+        )  # type: ignore[no-untyped-call]
+        openfermion_hamiltonian_dict = {
+            k: complex(v)
+            for k, v in openfermion.normal_ordered(fermion_operator).terms.items()  # type: ignore[no-untyped-call]
+        }
+        self.hamiltonian = Hamiltonian(openfermion_hamiltonian_dict, kind="fermi")
+        logging.info("Internal Hamiltonian representation successfully created")
+
+        self.n_qubits: int = n_orbit * 2
+        self.n_electrons: int = n_electron
+        self.n_spins: int = n_spin
+        self.ref_energy: float = args.ref_energy
+        logging.info(
+            "Identified %d qubits, %d electrons, spin=%d, ref_energy=%.10f",
+            self.n_qubits,
+            self.n_electrons,
+            self.n_spins,
+            self.ref_energy,
+        )
+
+    def apply_within(self, configs_i: torch.Tensor, psi_i: torch.Tensor, configs_j: torch.Tensor) -> torch.Tensor:
+        return self.hamiltonian.apply_within(configs_i, psi_i, configs_j)
+
+    def find_relative(
+        self,
+        configs_i: torch.Tensor,
+        psi_i: torch.Tensor,
+        count_selected: int,
+        configs_exclude: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        return self.hamiltonian.find_relative(configs_i, psi_i, count_selected, configs_exclude)
+
+    def list_relative(
+        self,
+        configs_i: torch.Tensor,
+        psi_i: torch.Tensor,
+        configs_exclude: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return self.hamiltonian.list_relative(configs_i, psi_i, configs_exclude)
+
+    def diagonal_term(self, configs: torch.Tensor) -> torch.Tensor:
+        return self.hamiltonian.diagonal_term(configs)
+
+    def show_config(self, config: torch.Tensor) -> str:
+        string = "".join(f"{i:08b}"[::-1] for i in config.cpu().numpy())
+        return (
+            "["
+            + "".join(self._show_config_site(string[index : index + 2]) for index in range(0, self.n_qubits, 2))
+            + "]"
+        )
+
+    def _show_config_site(self, string: str) -> str:
+        match string:
+            case "00":
+                return " "
+            case "10":
+                return "↑"
+            case "01":
+                return "↓"
+            case "11":
+                return "↕"
+            case _:
+                raise ValueError(f"Invalid string: {string}")
+
+
+model_dict["pyscf"] = Model
+
+
+@dataclasses.dataclass
+class MlpConfig:
+    """
+    The configuration of the MLP network.
+    """
+
+    # The hidden widths of the network
+    hidden: tuple[int, ...] = (512,)
+
+    def create(self, model: Model) -> NetworkProto:
+        """
+        Create a MLP network for the model.
+        """
+        logging.info("Hidden layer widths: %a", self.hidden)
+
+        network = MlpWaveFunction(
+            double_sites=model.n_qubits,
+            physical_dim=2,
+            is_complex=True,
+            spin_up=(model.n_electrons + model.n_spins) // 2,
+            spin_down=(model.n_electrons - model.n_spins) // 2,
+            hidden_size=self.hidden,
+            ordering=+1,
+        )
+
+        return network
+
+
+Model.network_dict["mlp/u1u1"] = MlpConfig
+
+
+@dataclasses.dataclass
+class TransformersConfig:
+    """
+    The configuration of the transformers network.
+    """
+
+    # Embedding dimension
+    embedding_dim: int = 512
+    # Heads number
+    heads_num: int = 8
+    # Feedforward dimension
+    feed_forward_dim: int = 2048
+    # Shared expert number
+    shared_expert_num: int = 1
+    # Routed expert number
+    routed_expert_num: int = 0
+    # Selected expert number
+    selected_expert_num: int = 0
+    # Network depth
+    depth: int = 6
+
+    def create(self, model: Model) -> NetworkProto:
+        """
+        Create a transformers network for the model.
+        """
+        logging.info(
+            "Transformers network configuration: "
+            "embedding dimension: %d, "
+            "number of heads: %d, "
+            "feed-forward dimension: %d, "
+            "shared expert number: %d, "
+            "routed expert number: %d, "
+            "selected expert number: %d, "
+            "depth: %d",
+            self.embedding_dim,
+            self.heads_num,
+            self.feed_forward_dim,
+            self.shared_expert_num,
+            self.routed_expert_num,
+            self.selected_expert_num,
+            self.depth,
+        )
+
+        network = TransformersWaveFunction(
+            double_sites=model.n_qubits,
+            physical_dim=2,
+            is_complex=True,
+            spin_up=(model.n_electrons + model.n_spins) // 2,
+            spin_down=(model.n_electrons - model.n_spins) // 2,
+            embedding_dim=self.embedding_dim,
+            heads_num=self.heads_num,
+            feed_forward_dim=self.feed_forward_dim,
+            shared_num=self.shared_expert_num,
+            routed_num=self.routed_expert_num,
+            selected_num=self.selected_expert_num,
+            depth=self.depth,
+            ordering=+1,
+        )
+
+        return network
+
+
+Model.network_dict["transformers/u1u1"] = TransformersConfig
+
+
+@dataclasses.dataclass
+class MlpElectronConfig:
+    """
+    The configuration of the MLP network with total electron number conservation.
+    """
+
+    # The hidden widths of the network
+    hidden: tuple[int, ...] = (512,)
+
+    def create(self, model: Model) -> NetworkProto:
+        """
+        Create a MLP network for the model.
+        """
+        logging.info("Hidden layer widths: %a", self.hidden)
+
+        network = MlpWaveFunctionElectron(
+            sites=model.n_qubits,
+            physical_dim=2,
+            is_complex=True,
+            electrons=model.n_electrons,
+            hidden_size=self.hidden,
+            ordering=+1,
+        )
+
+        return network
+
+
+Model.network_dict["mlp/u1"] = MlpElectronConfig
+
+
+@dataclasses.dataclass
+class TransformersElectronConfig:
+    """
+    The configuration of the transformers network with total electron number conservation.
+    """
+
+    # Embedding dimension
+    embedding_dim: int = 512
+    # Heads number
+    heads_num: int = 8
+    # Feedforward dimension
+    feed_forward_dim: int = 2048
+    # Shared expert number
+    shared_expert_num: int = 1
+    # Routed expert number
+    routed_expert_num: int = 0
+    # Selected expert number
+    selected_expert_num: int = 0
+    # Network depth
+    depth: int = 6
+
+    def create(self, model: Model) -> NetworkProto:
+        """
+        Create a transformers network for the model.
+        """
+        logging.info(
+            "Transformers network configuration: "
+            "embedding dimension: %d, "
+            "number of heads: %d, "
+            "feed-forward dimension: %d, "
+            "shared expert number: %d, "
+            "routed expert number: %d, "
+            "selected expert number: %d, "
+            "depth: %d",
+            self.embedding_dim,
+            self.heads_num,
+            self.feed_forward_dim,
+            self.shared_expert_num,
+            self.routed_expert_num,
+            self.selected_expert_num,
+            self.depth,
+        )
+
+        network = TransformersWaveFunctionElectron(
+            sites=model.n_qubits,
+            physical_dim=2,
+            is_complex=True,
+            electrons=model.n_electrons,
+            embedding_dim=self.embedding_dim,
+            heads_num=self.heads_num,
+            feed_forward_dim=self.feed_forward_dim,
+            shared_num=self.shared_expert_num,
+            routed_num=self.routed_expert_num,
+            selected_num=self.selected_expert_num,
+            depth=self.depth,
+            ordering=+1,
+        )
+
+        return network
+
+
+Model.network_dict["transformers/u1"] = TransformersElectronConfig
diff --git a/qmp/plugins/__init__.py b/qmp/plugins/__init__.py
new file mode 100644
index 0000000..45157d2
--- /dev/null
+++ b/qmp/plugins/__init__.py
@@ -0,0 +1,3 @@
+"""
+This subpackage contains plugin interfaces for third-party quantum-chemistry frameworks.
+"""
diff --git a/qmp/plugins/pyscf.py b/qmp/plugins/pyscf.py
new file mode 100644
index 0000000..476ccf6
--- /dev/null
+++ b/qmp/plugins/pyscf.py
@@ -0,0 +1,390 @@
+"""
+PySCF-compatible FCI solver plugin for QMP.
+
+This module exposes a :class:`Solver` that conforms to the PySCF FCI-solver
+interface::
+
+    solver = Solver(action_name="haar", max_absolute_step=10, ...)
+    energy, ci = solver.kernel(h1e, eri, norb, nelec)
+
+Internally the solver:
+
+1. Builds a :class:`~qmp.models.pyscf.Model` from the supplied integral arrays.
+2. Creates a temporary output directory (for TensorBoard logs).
+3. Constructs a :class:`_PyscfRuntimeContext` that bypasses Hydra and replaces
+   ``sys.exit`` with a Python exception so that the kernel can *return* the result
+   instead of terminating the process.
+4. Imports and instantiates the requested QMP algorithm (e.g. ``haar``, ``vmc``)
+   and calls its ``main`` method.
+5. After the algorithm reaches ``max_absolute_step``, computes the variational
+   energy from the trained network and returns ``(energy, network_state_dict)``.
+   The returned ``network_state_dict`` may be passed back as ``ci0`` to warm-start
+   subsequent calls.
+"""
+
+import importlib
+import logging
+import pathlib
+import tempfile
+import typing
+
+import dacite
+import numpy
+import omegaconf
+import torch
+
+from ..models.pyscf import Model, ModelConfig
+from ..utility.action_dict import action_dict
+from ..utility.context import DACITE_CAST, RuntimeContext
+from ..utility.model_dict import ModelProto, NetworkProto
+
+
+# ---------------------------------------------------------------------------
+# Internal helpers
+# ---------------------------------------------------------------------------
+
+
+class _AlgorithmComplete(Exception):
+    """
+    Raised by :class:`_PyscfRuntimeContext` when ``max_absolute_step`` is
+    reached, so that the solver can recover the final checkpoint data and
+    return it to the caller instead of terminating the process.
+    """
+
+    def __init__(self, data: dict[str, typing.Any]) -> None:
+        self.data = data
+
+
+class _PyscfRuntimeContext(RuntimeContext):
+    """
+    A :class:`~qmp.utility.context.RuntimeContext` subclass tailored for
+    programmatic (non-Hydra) use from the PySCF plugin.
+
+    Key differences from the base class:
+
+    * ``folder()`` returns an explicit temporary directory instead of reading
+      the Hydra output path.
+    * ``setup()`` creates the folder and configures random state without
+      requiring Hydra to be initialised.
+    * ``create_model()`` ignores the config argument and returns the
+      pre-built :class:`~qmp.models.pyscf.Model` that was injected at
+      construction time.
+    * ``create_network()`` delegates to the base implementation and also
+      saves a reference to the created network so the solver can evaluate
+      the final energy.
+    * ``save()`` raises :class:`_AlgorithmComplete` instead of calling
+      ``sys.exit(0)`` when ``max_absolute_step`` is reached, enabling the
+      caller to retrieve the results and return normally.
+    """
+
+    def __init__(
+        self,
+        model_instance: ModelProto,
+        folder_path: pathlib.Path,
+        **kwargs: typing.Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self._model_instance: ModelProto = model_instance
+        self._folder_path: pathlib.Path = folder_path
+        self._network: NetworkProto | None = None
+
+    # -- Hydra bypass -----------------------------------------------------------
+
+    def folder(self) -> pathlib.Path:
+        return self._folder_path
+
+    def setup(self) -> dict[str, typing.Any]:
+        self._folder_path.mkdir(parents=True, exist_ok=True)
+        logging.info("Log directory: %s", self._folder_path)
+
+        logging.info("Disabling PyTorch's default gradient computation")
+        torch.set_grad_enabled(False)
+
+        if self.random_seed is not None:
+            logging.info("Setting random seed to: %d", self.random_seed)
+            torch.manual_seed(self.random_seed)
+        else:
+            current_seed = torch.seed()
+            logging.info("Random seed not specified, using current seed: %d", current_seed)
+
+        return {}
+
+    # -- Model injection --------------------------------------------------------
+
+    def create_model(self, model_config: omegaconf.DictConfig) -> ModelProto:
+        return self._model_instance
+
+    # -- Network reference capture ----------------------------------------------
+
+    def create_network(
+        self,
+        network_config: omegaconf.DictConfig,
+        model: ModelProto,
+        state_dict: dict[str, typing.Any] | None = None,
+    ) -> NetworkProto:
+        network = super().create_network(network_config, model, state_dict)
+        self._network = network
+        return network
+
+    # -- Stop via exception instead of sys.exit ---------------------------------
+
+    def save(self, data: dict[str, typing.Any], step: int) -> None:
+        if self.max_relative_step is not None:
+            self.max_absolute_step = step + self.max_relative_step - 1
+            self.max_relative_step = None
+        if step == self.max_absolute_step:
+            logging.info("Reached the maximum step, returning from solver")
+            raise _AlgorithmComplete(data)
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+
+class Solver:
+    """
+    PySCF-compatible FCI solver backed by QMP algorithms.
+
+    The interface mirrors that of ``pyscf.fci.direct_spin1.FCISolver`` so that
+    this solver can be used as a drop-in replacement inside PySCF workflows::
+
+        from qmp.plugins.pyscf import Solver
+        solver = Solver(action_name="haar", max_absolute_step=10)
+        energy, ci = solver.kernel(h1e, eri, norb, nelec)
+
+    Parameters
+    ----------
+    action_name : str
+        Name of the QMP algorithm to run (e.g. ``"haar"`` or ``"vmc"``).
+        The corresponding module ``qmp.algorithms.<action_name>`` is imported
+        automatically.
+    action_params : dict, optional
+        Keyword arguments forwarded to the algorithm's configuration dataclass.
+    network_name : str
+        Name of the network architecture registered in the model's
+        ``network_dict`` (e.g. ``"mlp/u1u1"``).
+    network_params : dict, optional
+        Keyword arguments forwarded to the network configuration dataclass.
+    optimizer_name : str
+        Name of a ``torch.optim`` optimizer class (e.g. ``"Adam"``).
+    optimizer_params : dict, optional
+        Keyword arguments forwarded to the optimizer constructor
+        (e.g. ``{"lr": 1e-3}``).
+    max_absolute_step : int
+        Number of global algorithm steps to execute before returning.
+        This parameter is required when using the solver programmatically
+        (unlike the CLI which can run indefinitely).
+    device : str or torch.device
+        Compute device (default: ``"cpu"``).
+    dtype : str or torch.dtype, optional
+        Floating-point dtype for the network parameters (e.g. ``"float64"``).
+    random_seed : int, optional
+        Random seed for reproducibility.
+    sampling_count : int
+        Number of configurations to sample when computing the final
+        variational energy after the algorithm finishes (default: 4096).
+    """
+
+    def __init__(
+        self,
+        *,
+        action_name: str = "haar",
+        action_params: dict[str, typing.Any] | None = None,
+        network_name: str = "mlp/u1u1",
+        network_params: dict[str, typing.Any] | None = None,
+        optimizer_name: str = "Adam",
+        optimizer_params: dict[str, typing.Any] | None = None,
+        max_absolute_step: int = 10,
+        device: str | torch.device = "cpu",
+        dtype: str | torch.dtype | None = None,
+        random_seed: int | None = None,
+        sampling_count: int = 4096,
+    ) -> None:
+        self.action_name = action_name
+        self.action_params: dict[str, typing.Any] = action_params or {}
+        self.network_name = network_name
+        self.network_params: dict[str, typing.Any] = network_params or {}
+        self.optimizer_name = optimizer_name
+        self.optimizer_params: dict[str, typing.Any] = optimizer_params or {"lr": 1e-3}
+        self.max_absolute_step = max_absolute_step
+        self.device = torch.device(device) if isinstance(device, str) else device
+        self.dtype = dtype
+        self.random_seed = random_seed
+        self.sampling_count = sampling_count
+
+    def kernel(
+        self,
+        h1: numpy.ndarray,
+        eri: numpy.ndarray,
+        norb: int,
+        nelec: int | tuple[int, int],
+        ci0: dict[str, torch.Tensor] | None = None,
+        nuclear_repulsion: float = 0.0,
+        ref_energy: float = 0.0,
+        **kwargs: typing.Any,
+    ) -> tuple[float, dict[str, torch.Tensor]]:
+        """
+        Solve the FCI problem for the given integrals.
+
+        Parameters
+        ----------
+        h1 : numpy.ndarray
+            1-electron integrals in MO basis, shape ``(norb, norb)``, in
+            chemist's notation (same convention as PySCF).
+        eri : numpy.ndarray
+            2-electron repulsion integrals in MO basis in chemist's (ij|kl)
+            notation.  Accepted shapes:
+
+            * ``(norb, norb, norb, norb)`` – full 4-index array.
+            * ``(nij, nij)`` – 4-fold symmetric compressed array
+              (``nij = norb*(norb+1)//2``).
+        norb : int
+            Number of spatial orbitals.
+        nelec : int or tuple[int, int]
+            Number of electrons.  If a plain ``int``, the electrons are split
+            as evenly as possible between alpha and beta spin (rounding alpha
+            up).  If a 2-tuple ``(nalpha, nbeta)``, each spin count is used
+            directly.
+        ci0 : dict[str, torch.Tensor], optional
+            Network state dict from a previous :meth:`kernel` call.  When
+            provided the network is warm-started from this state, analogous
+            to passing an initial CI vector in a conventional FCI solver.
+        nuclear_repulsion : float
+            Constant energy offset added to the Hamiltonian (e.g. nuclear
+            repulsion energy).  Default is 0.
+        ref_energy : float
+            Reference energy used for logging only (e.g. a known FCI
+            result).  Default is 0.
+        **kwargs
+            Ignored; present for drop-in compatibility with PySCF solvers.
+
+        Returns
+        -------
+        energy : float
+            Variational energy estimate after ``max_absolute_step`` global
+            algorithm steps.
+        ci : dict[str, torch.Tensor]
+            Final network state dict.  This can be passed back as ``ci0`` in
+            a subsequent call to warm-start the optimisation.
+        """
+        # ------------------------------------------------------------------
+        # 1. Resolve nelec
+        # ------------------------------------------------------------------
+        # 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,
+        )
+        model = Model(model_config)
+
+        # ------------------------------------------------------------------
+        # 3. Import the algorithm module (registers action in action_dict)
+        # ------------------------------------------------------------------
+        importlib.import_module(f"qmp.algorithms.{self.action_name}")
+
+        # ------------------------------------------------------------------
+        # 4. Build the runtime config (OmegaConf DictConfig)
+        # ------------------------------------------------------------------
+        runtime_config = omegaconf.OmegaConf.create(
+            {
+                # "model" entry is required by the algorithm's main() signature
+                # but is ignored because _PyscfRuntimeContext.create_model()
+                # returns the pre-built model.
+                "model": {"name": "pyscf", "params": {}},
+                "network": {"name": self.network_name, "params": self.network_params},
+                "optimizer": {"name": self.optimizer_name, "params": self.optimizer_params},
+            }
+        )
+
+        # ------------------------------------------------------------------
+        # 5. Create the runtime context
+        # ------------------------------------------------------------------
+        with tempfile.TemporaryDirectory() as tmpdir:
+            context = _PyscfRuntimeContext(
+                model_instance=model,
+                folder_path=pathlib.Path(tmpdir),
+                device=self.device,
+                dtype=self.dtype,
+                random_seed=self.random_seed,
+                max_absolute_step=self.max_absolute_step,
+            )
+
+            # setup() creates the folder and sets the grad mode; it returns
+            # the initial checkpoint dict (empty when no ci0 is provided).
+            checkpoint_data: dict[str, typing.Any] = context.setup()
+
+            # If ci0 is supplied, warm-start the network from the saved state.
+            if ci0 is not None:
+                checkpoint_data["network"] = ci0
+
+            # ------------------------------------------------------------------
+            # 6. Instantiate the algorithm
+            # ------------------------------------------------------------------
+            run = dacite.from_dict(
+                data_class=action_dict[self.action_name],
+                data=self.action_params,
+                config=dacite.Config(cast=DACITE_CAST),
+            )
+
+            # ------------------------------------------------------------------
+            # 7. Run the algorithm; catch the completion signal
+            # ------------------------------------------------------------------
+            final_data: dict[str, typing.Any] = {}
+            try:
+                run.main(
+                    context=context,
+                    runtime_config=runtime_config,
+                    checkpoint_data=checkpoint_data,
+                )
+                # If main() returns normally (no max_absolute_step set, or the
+                # algorithm exits for another reason), use whatever data is
+                # available.
+            except _AlgorithmComplete as exc:
+                final_data = exc.data
+
+            # ------------------------------------------------------------------
+            # 8. Compute variational energy from the trained network
+            # ------------------------------------------------------------------
+            network = context._network
+            if network is None:
+                raise RuntimeError(
+                    "The algorithm did not create a network. "
+                    "Ensure that the selected action calls context.create_network()."
+                )
+
+            # Load the final network weights (the algorithm may have updated
+            # them after context._network was first set).
+            final_state_dict: dict[str, torch.Tensor] = final_data.get("network", {})
+            if final_state_dict:
+                network.load_state_dict(final_state_dict)
+
+            with torch.no_grad():
+                configs, psi, _, _ = network.generate_unique(self.sampling_count)
+                h_psi = model.apply_within(configs, psi, configs)
+                energy_tensor = (psi.conj() @ h_psi) / (psi.conj() @ psi)
+                energy = energy_tensor.real.item()
+
+        # Restore the caller's gradient state.
+        torch.set_grad_enabled(prev_grad_enabled)
+
+        return energy, final_state_dict

From c92e5b90dca8a197bcd4b473b9952fc234e5f10a Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 5 Mar 2026 22:13:36 +0000
Subject: [PATCH 3/6] Redesign pyscf solver plugin per review feedback

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

- _PyscfRuntimeContext is now minimal: only overrides folder() and
  create_model(). setup() and save() are inherited unchanged from
  RuntimeContext and work via polymorphism (both call self.folder()).

- kernel() closely mimics __main__.py: imports modules, builds model,
  creates context from config.common via dacite, calls setup(), instantiates
  algorithm, catches SystemExit raised by base save() at max_absolute_step,
  reads saved checkpoint from tmpdir/data.pth, returns (energy, ci).

- ci is the full checkpoint minus random state; for haar it contains
  ci['haar']['pool'] = (configs, psi). ci0 is saved to tmpdir/data.pth
  before setup() for warm-starting.

- _extract_energy extracts ground-state energy from haar checkpoint:
  data['haar']['excited'][global-1][0][0].

- Additional cleanups: rename h1 -> h1e in signature, move tmpdir
  creation inside try block, add .item() guard in _extract_energy,
  rename global_step -> haar_global_step for clarity.

Co-authored-by: hzhangxyz <11623447+hzhangxyz@users.noreply.github.com>
---
 qmp/plugins/pyscf.py | 438 +++++++++++++++++++++----------------------
 1 file changed, 209 insertions(+), 229 deletions(-)

diff --git a/qmp/plugins/pyscf.py b/qmp/plugins/pyscf.py
index 476ccf6..abd0a09 100644
--- a/qmp/plugins/pyscf.py
+++ b/qmp/plugins/pyscf.py
@@ -4,27 +4,45 @@
 This module exposes a :class:`Solver` that conforms to the PySCF FCI-solver
 interface::
 
-    solver = Solver(action_name="haar", max_absolute_step=10, ...)
+    solver = Solver(config)
     energy, ci = solver.kernel(h1e, eri, norb, nelec)
 
-Internally the solver:
-
-1. Builds a :class:`~qmp.models.pyscf.Model` from the supplied integral arrays.
-2. Creates a temporary output directory (for TensorBoard logs).
-3. Constructs a :class:`_PyscfRuntimeContext` that bypasses Hydra and replaces
-   ``sys.exit`` with a Python exception so that the kernel can *return* the result
-   instead of terminating the process.
-4. Imports and instantiates the requested QMP algorithm (e.g. ``haar``, ``vmc``)
-   and calls its ``main`` method.
-5. After the algorithm reaches ``max_absolute_step``, computes the variational
-   energy from the trained network and returns ``(energy, network_state_dict)``.
-   The returned ``network_state_dict`` may be passed back as ``ci0`` to warm-start
-   subsequent calls.
+The ``config`` argument mirrors the YAML config accepted by ``qmp.__main__``::
+
+    config = {
+        "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,
+        },
+    }
+
+Internally the solver closely mimics :func:`qmp.__main__.main`:
+
+1. Imports the requested algorithm and model modules.
+2. Builds a :class:`~qmp.models.pyscf.Model` from the supplied integral arrays.
+3. Constructs a :class:`_PyscfRuntimeContext` that overrides
+   :meth:`~qmp.utility.context.RuntimeContext.folder` to use a temporary
+   directory and :meth:`~qmp.utility.context.RuntimeContext.create_model` to
+   inject the pre-built model.  All other methods — including ``setup()`` and
+   ``save()`` — are inherited unchanged from :class:`RuntimeContext` and work
+   via normal polymorphism.
+4. Runs the algorithm via its ``main()`` method.
+5. When ``max_absolute_step`` is reached, the base ``save()`` writes the
+   checkpoint to the temporary directory and calls ``sys.exit(0)``.  The
+   solver catches the resulting :class:`SystemExit` and reads the saved
+   checkpoint file to extract the energy and CI data.
+
+The returned ``ci`` dict may be passed back as ``ci0`` to warm-start
+subsequent calls.
 """
 
 import importlib
-import logging
 import pathlib
+import shutil
 import tempfile
 import typing
 
@@ -36,7 +54,7 @@
 from ..models.pyscf import Model, ModelConfig
 from ..utility.action_dict import action_dict
 from ..utility.context import DACITE_CAST, RuntimeContext
-from ..utility.model_dict import ModelProto, NetworkProto
+from ..utility.model_dict import ModelProto
 
 
 # ---------------------------------------------------------------------------
@@ -44,37 +62,22 @@
 # ---------------------------------------------------------------------------
 
 
-class _AlgorithmComplete(Exception):
-    """
-    Raised by :class:`_PyscfRuntimeContext` when ``max_absolute_step`` is
-    reached, so that the solver can recover the final checkpoint data and
-    return it to the caller instead of terminating the process.
+class _PyscfRuntimeContext(RuntimeContext):
     """
+    A :class:`~qmp.utility.context.RuntimeContext` subclass for programmatic
+    (non-Hydra) use from the PySCF solver plugin.
 
-    def __init__(self, data: dict[str, typing.Any]) -> None:
-        self.data = data
+    The only behavioural differences from the base class are:
 
+    * :meth:`folder` returns an explicit path supplied at construction time
+      instead of reading the Hydra output directory.
+    * :meth:`create_model` returns a pre-built :class:`~qmp.models.pyscf.Model`
+      instead of instantiating one from the OmegaConf config.
 
-class _PyscfRuntimeContext(RuntimeContext):
-    """
-    A :class:`~qmp.utility.context.RuntimeContext` subclass tailored for
-    programmatic (non-Hydra) use from the PySCF plugin.
-
-    Key differences from the base class:
-
-    * ``folder()`` returns an explicit temporary directory instead of reading
-      the Hydra output path.
-    * ``setup()`` creates the folder and configures random state without
-      requiring Hydra to be initialised.
-    * ``create_model()`` ignores the config argument and returns the
-      pre-built :class:`~qmp.models.pyscf.Model` that was injected at
-      construction time.
-    * ``create_network()`` delegates to the base implementation and also
-      saves a reference to the created network so the solver can evaluate
-      the final energy.
-    * ``save()`` raises :class:`_AlgorithmComplete` instead of calling
-      ``sys.exit(0)`` when ``max_absolute_step`` is reached, enabling the
-      caller to retrieve the results and return normally.
+    All other methods, including :meth:`~RuntimeContext.setup` and
+    :meth:`~RuntimeContext.save`, are inherited unchanged.  Because they call
+    ``self.folder()`` internally, they automatically use the temporary
+    directory via normal polymorphism.
     """
 
     def __init__(
@@ -86,56 +89,13 @@ def __init__(
         super().__init__(**kwargs)
         self._model_instance: ModelProto = model_instance
         self._folder_path: pathlib.Path = folder_path
-        self._network: NetworkProto | None = None
-
-    # -- Hydra bypass -----------------------------------------------------------
 
     def folder(self) -> pathlib.Path:
         return self._folder_path
 
-    def setup(self) -> dict[str, typing.Any]:
-        self._folder_path.mkdir(parents=True, exist_ok=True)
-        logging.info("Log directory: %s", self._folder_path)
-
-        logging.info("Disabling PyTorch's default gradient computation")
-        torch.set_grad_enabled(False)
-
-        if self.random_seed is not None:
-            logging.info("Setting random seed to: %d", self.random_seed)
-            torch.manual_seed(self.random_seed)
-        else:
-            current_seed = torch.seed()
-            logging.info("Random seed not specified, using current seed: %d", current_seed)
-
-        return {}
-
-    # -- Model injection --------------------------------------------------------
-
     def create_model(self, model_config: omegaconf.DictConfig) -> ModelProto:
         return self._model_instance
 
-    # -- Network reference capture ----------------------------------------------
-
-    def create_network(
-        self,
-        network_config: omegaconf.DictConfig,
-        model: ModelProto,
-        state_dict: dict[str, typing.Any] | None = None,
-    ) -> NetworkProto:
-        network = super().create_network(network_config, model, state_dict)
-        self._network = network
-        return network
-
-    # -- Stop via exception instead of sys.exit ---------------------------------
-
-    def save(self, data: dict[str, typing.Any], step: int) -> None:
-        if self.max_relative_step is not None:
-            self.max_absolute_step = step + self.max_relative_step - 1
-            self.max_relative_step = None
-        if step == self.max_absolute_step:
-            logging.info("Reached the maximum step, returning from solver")
-            raise _AlgorithmComplete(data)
-
 
 # ---------------------------------------------------------------------------
 # Public API
@@ -146,90 +106,51 @@ class Solver:
     """
     PySCF-compatible FCI solver backed by QMP algorithms.
 
-    The interface mirrors that of ``pyscf.fci.direct_spin1.FCISolver`` so that
-    this solver can be used as a drop-in replacement inside PySCF workflows::
+    The interface mirrors that of ``pyscf.fci.direct_spin1.FCISolver``::
 
         from qmp.plugins.pyscf import Solver
-        solver = Solver(action_name="haar", max_absolute_step=10)
+        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)
 
     Parameters
     ----------
-    action_name : str
-        Name of the QMP algorithm to run (e.g. ``"haar"`` or ``"vmc"``).
-        The corresponding module ``qmp.algorithms.<action_name>`` is imported
-        automatically.
-    action_params : dict, optional
-        Keyword arguments forwarded to the algorithm's configuration dataclass.
-    network_name : str
-        Name of the network architecture registered in the model's
-        ``network_dict`` (e.g. ``"mlp/u1u1"``).
-    network_params : dict, optional
-        Keyword arguments forwarded to the network configuration dataclass.
-    optimizer_name : str
-        Name of a ``torch.optim`` optimizer class (e.g. ``"Adam"``).
-    optimizer_params : dict, optional
-        Keyword arguments forwarded to the optimizer constructor
-        (e.g. ``{"lr": 1e-3}``).
-    max_absolute_step : int
-        Number of global algorithm steps to execute before returning.
-        This parameter is required when using the solver programmatically
-        (unlike the CLI which can run indefinitely).
-    device : str or torch.device
-        Compute device (default: ``"cpu"``).
-    dtype : str or torch.dtype, optional
-        Floating-point dtype for the network parameters (e.g. ``"float64"``).
-    random_seed : int, optional
-        Random seed for reproducibility.
-    sampling_count : int
-        Number of configurations to sample when computing the final
-        variational energy after the algorithm finishes (default: 4096).
+    config : dict or omegaconf.DictConfig
+        Configuration in the same format as the ``qmp.__main__`` YAML config.
+        Required top-level keys: ``action``, ``network``, ``optimizer``,
+        ``common``.  The ``model`` key is ignored (the model is built
+        automatically from the integral arrays passed to :meth:`kernel`).
+        ``common.max_absolute_step`` controls how many global algorithm steps
+        to run before returning.
     """
 
-    def __init__(
-        self,
-        *,
-        action_name: str = "haar",
-        action_params: dict[str, typing.Any] | None = None,
-        network_name: str = "mlp/u1u1",
-        network_params: dict[str, typing.Any] | None = None,
-        optimizer_name: str = "Adam",
-        optimizer_params: dict[str, typing.Any] | None = None,
-        max_absolute_step: int = 10,
-        device: str | torch.device = "cpu",
-        dtype: str | torch.dtype | None = None,
-        random_seed: int | None = None,
-        sampling_count: int = 4096,
-    ) -> None:
-        self.action_name = action_name
-        self.action_params: dict[str, typing.Any] = action_params or {}
-        self.network_name = network_name
-        self.network_params: dict[str, typing.Any] = network_params or {}
-        self.optimizer_name = optimizer_name
-        self.optimizer_params: dict[str, typing.Any] = optimizer_params or {"lr": 1e-3}
-        self.max_absolute_step = max_absolute_step
-        self.device = torch.device(device) if isinstance(device, str) else device
-        self.dtype = dtype
-        self.random_seed = random_seed
-        self.sampling_count = sampling_count
+    def __init__(self, config: dict[str, typing.Any] | omegaconf.DictConfig) -> None:
+        if isinstance(config, dict):
+            self.config: omegaconf.DictConfig = omegaconf.OmegaConf.create(config)
+        else:
+            self.config = config
 
     def kernel(
         self,
-        h1: numpy.ndarray,
+        h1e: numpy.ndarray,
         eri: numpy.ndarray,
         norb: int,
         nelec: int | tuple[int, int],
-        ci0: dict[str, torch.Tensor] | None = None,
+        ci0: dict[str, typing.Any] | None = None,
         nuclear_repulsion: float = 0.0,
         ref_energy: float = 0.0,
         **kwargs: typing.Any,
-    ) -> tuple[float, dict[str, torch.Tensor]]:
+    ) -> tuple[float, dict[str, typing.Any]]:
         """
-        Solve the FCI problem for the given integrals.
+        Solve the electronic structure problem for the given integrals.
 
         Parameters
         ----------
-        h1 : numpy.ndarray
+        h1e : numpy.ndarray
             1-electron integrals in MO basis, shape ``(norb, norb)``, in
             chemist's notation (same convention as PySCF).
         eri : numpy.ndarray
@@ -246,34 +167,34 @@ def kernel(
             as evenly as possible between alpha and beta spin (rounding alpha
             up).  If a 2-tuple ``(nalpha, nbeta)``, each spin count is used
             directly.
-        ci0 : dict[str, torch.Tensor], optional
-            Network state dict from a previous :meth:`kernel` call.  When
-            provided the network is warm-started from this state, analogous
-            to passing an initial CI vector in a conventional FCI solver.
+        ci0 : dict, optional
+            Checkpoint data from a previous :meth:`kernel` call (the ``ci``
+            value returned by that call).  When provided it is written to the
+            temporary checkpoint directory so that the algorithm resumes from
+            the saved state, analogous to passing an initial CI vector in a
+            conventional FCI solver.
         nuclear_repulsion : float
-            Constant energy offset added to the Hamiltonian (e.g. nuclear
-            repulsion energy).  Default is 0.
+            Constant nuclear repulsion energy added to the Hamiltonian.
         ref_energy : float
-            Reference energy used for logging only (e.g. a known FCI
-            result).  Default is 0.
+            Reference energy used for logging only.
         **kwargs
             Ignored; present for drop-in compatibility with PySCF solvers.
 
         Returns
         -------
         energy : float
-            Variational energy estimate after ``max_absolute_step`` global
-            algorithm steps.
-        ci : dict[str, torch.Tensor]
-            Final network state dict.  This can be passed back as ``ci0`` in
-            a subsequent call to warm-start the optimisation.
+            Ground-state energy after ``common.max_absolute_step`` global
+            algorithm steps, extracted from the saved checkpoint.
+        ci : dict
+            Checkpoint data (without random engine state) that can be passed
+            back as ``ci0`` in a subsequent call to continue optimisation.
+            For the ``haar`` algorithm this contains, among other things,
+            ``ci["haar"]["pool"]`` — a ``(configs, psi)`` tuple representing
+            the sampled basis and its Lanczos-optimised amplitudes.
         """
         # ------------------------------------------------------------------
-        # 1. Resolve nelec
+        # 1. Resolve electron counts
         # ------------------------------------------------------------------
-        # 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
@@ -284,107 +205,166 @@ def kernel(
         n_spin = nalpha - nbeta
 
         # ------------------------------------------------------------------
-        # 2. Build the PySCF model
+        # 2. Dynamic imports — mirroring __main__.py
         # ------------------------------------------------------------------
-        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,
-        )
-        model = Model(model_config)
+        importlib.import_module(f"qmp.algorithms.{self.config.action.name}")
+        importlib.import_module("qmp.models.pyscf")
 
         # ------------------------------------------------------------------
-        # 3. Import the algorithm module (registers action in action_dict)
+        # 3. Build PySCF model from integral arrays
         # ------------------------------------------------------------------
-        importlib.import_module(f"qmp.algorithms.{self.action_name}")
+        model = Model(
+            ModelConfig(
+                h1e=numpy.asarray(h1e, 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,
+            )
+        )
 
         # ------------------------------------------------------------------
-        # 4. Build the runtime config (OmegaConf DictConfig)
+        # 4. Build OmegaConf runtime_config — same structure as the YAML
+        #    config, but with model.name = "pyscf" so that the algorithm
+        #    can log the model type.  create_model() is overridden so the
+        #    params field is never read.
         # ------------------------------------------------------------------
-        runtime_config = omegaconf.OmegaConf.create(
+        runtime_config: omegaconf.DictConfig = omegaconf.OmegaConf.create(
             {
-                # "model" entry is required by the algorithm's main() signature
-                # but is ignored because _PyscfRuntimeContext.create_model()
-                # returns the pre-built model.
+                "action": omegaconf.OmegaConf.to_container(self.config.action, resolve=True),
                 "model": {"name": "pyscf", "params": {}},
-                "network": {"name": self.network_name, "params": self.network_params},
-                "optimizer": {"name": self.optimizer_name, "params": self.optimizer_params},
+                "network": omegaconf.OmegaConf.to_container(self.config.network, resolve=True),
+                "optimizer": omegaconf.OmegaConf.to_container(self.config.optimizer, resolve=True),
             }
         )
 
         # ------------------------------------------------------------------
-        # 5. Create the runtime context
+        # 5. Construct _PyscfRuntimeContext from config.common
+        # ------------------------------------------------------------------
+        common_data = omegaconf.OmegaConf.to_container(self.config.common, resolve=True)
+        assert isinstance(common_data, dict)
+        # Use dacite for type conversion (pathlib.Path, torch.device, etc.),
+        # matching the same approach used in __main__.py.
+        rt = dacite.from_dict(
+            data_class=RuntimeContext,
+            data=common_data,
+            config=dacite.Config(cast=DACITE_CAST),
+        )
+
+        # ------------------------------------------------------------------
+        # 6. Save caller gradient state; create temp directory
         # ------------------------------------------------------------------
-        with tempfile.TemporaryDirectory() as tmpdir:
+        prev_grad_enabled = torch.is_grad_enabled()
+        tmpdir: pathlib.Path | None = None
+        try:
+            tmpdir = pathlib.Path(tempfile.mkdtemp())
+            # If ci0 is provided, save it as the initial checkpoint so that
+            # setup() loads it automatically — the same way __main__.py
+            # resumes from an existing checkpoint file.
+            if ci0 is not None:
+                torch.save(ci0, tmpdir / "data.pth")
+
             context = _PyscfRuntimeContext(
                 model_instance=model,
-                folder_path=pathlib.Path(tmpdir),
-                device=self.device,
-                dtype=self.dtype,
-                random_seed=self.random_seed,
-                max_absolute_step=self.max_absolute_step,
+                folder_path=tmpdir,
+                parent_path=rt.parent_path,
+                random_seed=rt.random_seed,
+                checkpoint_interval=rt.checkpoint_interval,
+                device=rt.device,
+                dtype=rt.dtype,
+                max_absolute_step=rt.max_absolute_step,
+                max_relative_step=rt.max_relative_step,
             )
 
-            # setup() creates the folder and sets the grad mode; it returns
-            # the initial checkpoint dict (empty when no ci0 is provided).
-            checkpoint_data: dict[str, typing.Any] = context.setup()
-
-            # If ci0 is supplied, warm-start the network from the saved state.
-            if ci0 is not None:
-                checkpoint_data["network"] = ci0
+            # ------------------------------------------------------------------
+            # 7. Setup context — loads checkpoint / ci0, configures RNG
+            # ------------------------------------------------------------------
+            checkpoint_data = context.setup()
 
             # ------------------------------------------------------------------
-            # 6. Instantiate the algorithm
+            # 8. Instantiate algorithm — mirroring __main__.py
             # ------------------------------------------------------------------
             run = dacite.from_dict(
-                data_class=action_dict[self.action_name],
-                data=self.action_params,
+                data_class=action_dict[self.config.action.name],
+                data=omegaconf.OmegaConf.to_container(
+                    self.config.action.params, resolve=True
+                ),  # type: ignore[arg-type]
                 config=dacite.Config(cast=DACITE_CAST),
             )
 
             # ------------------------------------------------------------------
-            # 7. Run the algorithm; catch the completion signal
+            # 9. Run algorithm — mirroring __main__.py
+            #    The base save() writes the checkpoint to the temp dir and
+            #    then calls sys.exit(0) when max_absolute_step is reached.
+            #    We catch the resulting SystemExit so kernel() can return.
             # ------------------------------------------------------------------
-            final_data: dict[str, typing.Any] = {}
             try:
                 run.main(
                     context=context,
                     runtime_config=runtime_config,
                     checkpoint_data=checkpoint_data,
                 )
-                # If main() returns normally (no max_absolute_step set, or the
-                # algorithm exits for another reason), use whatever data is
-                # available.
-            except _AlgorithmComplete as exc:
-                final_data = exc.data
+            except SystemExit:
+                pass
 
             # ------------------------------------------------------------------
-            # 8. Compute variational energy from the trained network
+            # 10. Read the checkpoint that was saved by save()
             # ------------------------------------------------------------------
-            network = context._network
-            if network is None:
-                raise RuntimeError(
-                    "The algorithm did not create a network. "
-                    "Ensure that the selected action calls context.create_network()."
+            data_path = tmpdir / "data.pth"
+            if data_path.exists():
+                data: dict[str, typing.Any] = torch.load(
+                    data_path, map_location="cpu", weights_only=True
                 )
+            else:
+                data = {}
+
+            # ------------------------------------------------------------------
+            # 11. Extract energy and build ci output
+            # ------------------------------------------------------------------
+            energy = _extract_energy(data, self.config.action.name)
+            # Strip the random engine state — it is stale by the next call.
+            ci: dict[str, typing.Any] = {k: v for k, v in data.items() if k != "random"}
+
+        finally:
+            if tmpdir is not None:
+                shutil.rmtree(tmpdir, ignore_errors=True)
+            torch.set_grad_enabled(prev_grad_enabled)
 
-            # Load the final network weights (the algorithm may have updated
-            # them after context._network was first set).
-            final_state_dict: dict[str, torch.Tensor] = final_data.get("network", {})
-            if final_state_dict:
-                network.load_state_dict(final_state_dict)
+        return energy, ci
 
-            with torch.no_grad():
-                configs, psi, _, _ = network.generate_unique(self.sampling_count)
-                h_psi = model.apply_within(configs, psi, configs)
-                energy_tensor = (psi.conj() @ h_psi) / (psi.conj() @ psi)
-                energy = energy_tensor.real.item()
 
-        # Restore the caller's gradient state.
-        torch.set_grad_enabled(prev_grad_enabled)
+def _extract_energy(data: dict[str, typing.Any], action_name: str) -> float:
+    """
+    Extract the final ground-state energy from algorithm checkpoint data.
 
-        return energy, final_state_dict
+    Parameters
+    ----------
+    data : dict
+        Checkpoint data loaded from the temporary directory.
+    action_name : str
+        Name of the algorithm that produced the checkpoint.
+
+    Returns
+    -------
+    float
+        Ground-state energy, or ``0.0`` if it cannot be determined from the
+        checkpoint (e.g. the algorithm does not store energy in its checkpoint).
+    """
+    if action_name in ("haar", "imag"):
+        haar_data = data.get("haar")
+        if haar_data is None:
+            return 0.0
+        haar_global_step: int = haar_data.get("global", 0)
+        # In haar.py, excited states are stored with key = global_step *before*
+        # the increment, so the last entry key is (global_step - 1).
+        last_key = haar_global_step - 1
+        excited: dict[int, list[typing.Any]] = haar_data.get("excited", {})
+        results = excited.get(last_key)
+        if results:
+            # results is list[(energy_tensor, configs, psi)]; index 0 is ground state.
+            energy_val = results[0][0]
+            if hasattr(energy_val, "item"):
+                return float(energy_val.item())
+    return 0.0

From ed766632aa3176eadf5f89225dc7d1c55be865ed Mon Sep 17 00:00:00 2001
From: Hao Zhang <hzhangxyz@outlook.com>
Date: Fri, 6 Mar 2026 07:42:19 +0800
Subject: [PATCH 4/6] feat(plugins): enhance PySCF solver to return and accept
 CI arrays

---
 qmp/plugins/pyscf.py | 190 +++++++++++++++++++++++++++++++++++++------
 1 file changed, 165 insertions(+), 25 deletions(-)

diff --git a/qmp/plugins/pyscf.py b/qmp/plugins/pyscf.py
index abd0a09..1310f88 100644
--- a/qmp/plugins/pyscf.py
+++ b/qmp/plugins/pyscf.py
@@ -62,6 +62,26 @@
 # ---------------------------------------------------------------------------
 
 
+class _CIArray(numpy.ndarray):
+    """
+    A numpy array subclass that stores a QMP checkpoint dictionary.
+
+    This allows the FCI solver to return a standard-looking CI wavefunction
+    (the array) while still carrying the full state (network, optimizer, etc.)
+    needed for a high-performance warm start in the next call to :meth:`kernel`.
+    """
+
+    def __new__(cls, input_array: numpy.ndarray, checkpoint: dict[str, typing.Any]) -> "_CIArray":
+        obj = numpy.asanyarray(input_array).view(cls)
+        obj.checkpoint = checkpoint
+        return obj
+
+    def __array_finalize__(self, obj: typing.Any) -> None:
+        if obj is None:
+            return
+        self.checkpoint: dict[str, typing.Any] = getattr(obj, "checkpoint", {})
+
+
 class _PyscfRuntimeContext(RuntimeContext):
     """
     A :class:`~qmp.utility.context.RuntimeContext` subclass for programmatic
@@ -97,6 +117,118 @@ def create_model(self, model_config: omegaconf.DictConfig) -> ModelProto:
         return self._model_instance
 
 
+def _to_pyscf_ci(
+    data: dict[str, typing.Any], action_name: str, norb: int, nelec: int | tuple[int, int]
+) -> numpy.ndarray | None:
+    """
+    Convert QMP sampled configurations and amplitudes to a PySCF CI vector.
+    """
+    try:
+        from pyscf.fci import cistring
+    except ImportError:
+        return None
+
+    if action_name not in ("haar", "imag"):
+        return None
+
+    haar_data = data.get("haar")
+    if not haar_data:
+        return None
+    pool = haar_data.get("pool")
+    if not pool:
+        return None
+
+    configs, psi = pool
+    if configs is None or psi is None:
+        return None
+
+    if isinstance(nelec, (int, numpy.integer)):
+        nalpha = (int(nelec) + 1) // 2
+        nbeta = int(nelec) - nalpha
+    else:
+        nalpha, nbeta = int(nelec[0]), int(nelec[1])
+
+    na = cistring.num_strings(norb, nalpha)
+    nb = cistring.num_strings(norb, nbeta)
+
+    # Limit CI vector size to avoid memory overflow (approx 10^8 elements ~ 1.6GB)
+    if na * nb > 10**8:
+        return None
+
+    configs_np = configs.cpu().numpy()
+    psi_np = psi.cpu().numpy()
+
+    # Unpack bits: QMP order is little-endian bit-packed bytes
+    bits = numpy.unpackbits(configs_np, axis=1, bitorder="little")
+
+    # Extract alpha and beta occupancy (alpha=even, beta=odd bits)
+    alpha_bits = bits[:, 0 : 2 * norb : 2]
+    beta_bits = bits[:, 1 : 2 * norb : 2]
+
+    # Convert to integer bitmasks
+    pw2 = 2 ** numpy.arange(norb, dtype=numpy.uint64)
+    alpha_dets = (alpha_bits.astype(numpy.uint64) * pw2).sum(axis=1)
+    beta_dets = (beta_bits.astype(numpy.uint64) * pw2).sum(axis=1)
+
+    # Map bitmasks to PySCF determinant indices
+    try:
+        addr_a = cistring.str2addr(norb, nalpha, alpha_dets)
+        addr_b = cistring.str2addr(norb, nbeta, beta_dets)
+    except Exception:
+        # Fallback for older PySCF or if vectorized call fails
+        addr_a = numpy.array([cistring.str2addr(norb, nalpha, int(d)) for d in alpha_dets])
+        addr_b = numpy.array([cistring.str2addr(norb, nbeta, int(d)) for d in beta_dets])
+
+    ci_array = numpy.zeros((na, nb), dtype=psi_np.dtype)
+    ci_array[addr_a, addr_b] = psi_np
+
+    return ci_array
+
+
+def _from_pyscf_ci(ci_array: numpy.ndarray, norb: int, nelec: int | tuple[int, int]) -> dict[str, typing.Any]:
+    """
+    Convert a PySCF CI vector back to a QMP checkpoint (sampled pool).
+    """
+    try:
+        from pyscf.fci import cistring
+    except ImportError:
+        return {}
+
+    if isinstance(nelec, (int, numpy.integer)):
+        nalpha = (int(nelec) + 1) // 2
+        nbeta = int(nelec) - nalpha
+    else:
+        nalpha, nbeta = int(nelec[0]), int(nelec[1])
+
+    idx_a, idx_b = numpy.where(ci_array != 0)
+    amplitudes = ci_array[idx_a, idx_b]
+
+    str_a = cistring.addrs2str(norb, nalpha, idx_a)
+    str_b = cistring.addrs2str(norb, nbeta, idx_b)
+
+    n_qubytes = (2 * norb + 7) // 8
+    N = len(str_a)
+    configs = numpy.zeros((N, n_qubytes), dtype=numpy.uint8)
+
+    for j in range(norb):
+        a_bit = (str_a >> j) & 1
+        b_bit = (str_b >> j) & 1
+        pos_a = 2 * j
+        pos_b = 2 * j + 1
+        configs[:, pos_a // 8] |= (a_bit << (pos_a % 8)).astype(numpy.uint8)
+        configs[:, pos_b // 8] |= (b_bit << (pos_b % 8)).astype(numpy.uint8)
+
+    return {
+        "haar": {
+            "global": 0,
+            "local": 0,
+            "lanczos": 0,
+            "pool": (torch.from_numpy(configs), torch.from_numpy(amplitudes)),
+            "excited": {},
+        }
+    }
+
+
 # ---------------------------------------------------------------------------
 # Public API
 # ---------------------------------------------------------------------------
@@ -140,11 +272,11 @@ def kernel(
         eri: numpy.ndarray,
         norb: int,
         nelec: int | tuple[int, int],
-        ci0: dict[str, typing.Any] | None = None,
+        ci0: dict[str, typing.Any] | numpy.ndarray | None = None,
         nuclear_repulsion: float = 0.0,
         ref_energy: float = 0.0,
         **kwargs: typing.Any,
-    ) -> tuple[float, dict[str, typing.Any]]:
+    ) -> tuple[float, numpy.ndarray | dict[str, typing.Any]]:
         """
         Solve the electronic structure problem for the given integrals.
 
@@ -159,7 +291,7 @@ def kernel(
 
             * ``(norb, norb, norb, norb)`` – full 4-index array.
             * ``(nij, nij)`` – 4-fold symmetric compressed array
-              (``nij = norb*(norb+1)//2``).
+              (nij = norb*(norb+1)//2).
         norb : int
             Number of spatial orbitals.
         nelec : int or tuple[int, int]
@@ -167,12 +299,10 @@ def kernel(
             as evenly as possible between alpha and beta spin (rounding alpha
             up).  If a 2-tuple ``(nalpha, nbeta)``, each spin count is used
             directly.
-        ci0 : dict, optional
-            Checkpoint data from a previous :meth:`kernel` call (the ``ci``
-            value returned by that call).  When provided it is written to the
-            temporary checkpoint directory so that the algorithm resumes from
-            the saved state, analogous to passing an initial CI vector in a
-            conventional FCI solver.
+        ci0 : dict or numpy.ndarray, optional
+            Initial guess or checkpoint data. If a ``dict``, it is treated as
+            the full QMP checkpoint. If a ``numpy.ndarray``, it is treated as
+            the PySCF CI vector and converted to a QMP pool.
         nuclear_repulsion : float
             Constant nuclear repulsion energy added to the Hamiltonian.
         ref_energy : float
@@ -183,21 +313,18 @@ def kernel(
         Returns
         -------
         energy : float
-            Ground-state energy after ``common.max_absolute_step`` global
-            algorithm steps, extracted from the saved checkpoint.
-        ci : dict
-            Checkpoint data (without random engine state) that can be passed
-            back as ``ci0`` in a subsequent call to continue optimisation.
-            For the ``haar`` algorithm this contains, among other things,
-            ``ci["haar"]["pool"]`` — a ``(configs, psi)`` tuple representing
-            the sampled basis and its Lanczos-optimised amplitudes.
+            Ground-state energy after global algorithm steps.
+        ci : numpy.ndarray or dict
+            CI wavefunction as a numpy array. The array is an instance of
+            ``_CIArray`` which carries the full QMP checkpoint for warm-start
+            compatibility.
         """
         # ------------------------------------------------------------------
         # 1. Resolve electron counts
         # ------------------------------------------------------------------
-        if isinstance(nelec, int):
-            nalpha = (nelec + 1) // 2
-            nbeta = nelec - nalpha
+        if isinstance(nelec, (int, numpy.integer)):
+            nalpha = (int(nelec) + 1) // 2
+            nbeta = int(nelec) - nalpha
         else:
             nalpha, nbeta = int(nelec[0]), int(nelec[1])
 
@@ -260,11 +387,15 @@ def kernel(
         tmpdir: pathlib.Path | None = None
         try:
             tmpdir = pathlib.Path(tempfile.mkdtemp())
-            # If ci0 is provided, save it as the initial checkpoint so that
-            # setup() loads it automatically — the same way __main__.py
-            # resumes from an existing checkpoint file.
+            # Handle ci0: could be full checkpoint dict or amplitudes array
             if ci0 is not None:
-                torch.save(ci0, tmpdir / "data.pth")
+                if isinstance(ci0, _CIArray) and ci0.checkpoint:
+                    checkpoint_to_save = ci0.checkpoint
+                elif isinstance(ci0, numpy.ndarray):
+                    checkpoint_to_save = _from_pyscf_ci(ci0, norb, (nalpha, nbeta))
+                else:
+                    checkpoint_to_save = ci0
+                torch.save(checkpoint_to_save, tmpdir / "data.pth")
 
             context = _PyscfRuntimeContext(
                 model_instance=model,
@@ -325,7 +456,16 @@ def kernel(
             # ------------------------------------------------------------------
             energy = _extract_energy(data, self.config.action.name)
             # Strip the random engine state — it is stale by the next call.
-            ci: dict[str, typing.Any] = {k: v for k, v in data.items() if k != "random"}
+            full_checkpoint: dict[str, typing.Any] = {
+                k: v for k, v in data.items() if k != "random"
+            }
+
+            # Try to convert to PySCF CI vector format
+            ci_array = _to_pyscf_ci(data, self.config.action.name, norb, (nalpha, nbeta))
+            if ci_array is not None:
+                ci: numpy.ndarray | dict[str, typing.Any] = _CIArray(ci_array, full_checkpoint)
+            else:
+                ci = full_checkpoint
 
         finally:
             if tmpdir is not None:

From 509ce9f847c3eb3f3bd2b116903e7bc974d05506 Mon Sep 17 00:00:00 2001
From: Hao Zhang <hzhangxyz@outlook.com>
Date: Fri, 6 Mar 2026 07:57:13 +0800
Subject: [PATCH 5/6] Add pyscf as an optional dependence.

---
 pyproject.toml |  3 +++
 uv.lock        | 24 +++++++++++++++++++++++-
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 909d771..a449bcd 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -38,6 +38,9 @@ dev = [
   "ruff>=0.14.10",
   "types-pyyaml>=6.0.12",
 ]
+pyscf = [
+    "pyscf>=2.12.1",
+]
 
 [project.urls]
 Homepage = "https://github.com/USTC-KnowledgeComputingLab/qmp-kit"
diff --git a/uv.lock b/uv.lock
index db74242..32e3f8e 100644
--- a/uv.lock
+++ b/uv.lock
@@ -1173,6 +1173,24 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/10/5e/1aa9a93198c6b64513c9d7752de7422c06402de6600a8767da1524f9570b/pyparsing-3.2.5-py3-none-any.whl", hash = "sha256:e38a4f02064cf41fe6593d328d0512495ad1f3d8a91c4f73fc401b3079a59a5e", size = 113890, upload-time = "2025-09-21T04:11:04.117Z" },
 ]
 
+[[package]]
+name = "pyscf"
+version = "2.12.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "h5py" },
+    { name = "numpy" },
+    { name = "scipy" },
+    { name = "setuptools" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/ce/c9/555f8a27e2fbc55d73e5034dcf3d1b58db98f127f770ce6038c3eb90dd27/pyscf-2.12.1.tar.gz", hash = "sha256:cae3b026a928ce866965242056a833a17e46b89035d2e3abbf5429a158da4d48", size = 10472353, upload-time = "2026-02-04T01:03:28.034Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/aa/cc/76888e9667b5f3b512dd43d83cc07818339f2a1452708417d5ceacec0eea/pyscf-2.12.1-py3-none-macosx_10_9_x86_64.whl", hash = "sha256:283b94e48699a91c391d6c1de7cf5517d3573e5080b9ec56b69abce4880c1a71", size = 35952312, upload-time = "2026-02-04T00:39:44Z" },
+    { url = "https://files.pythonhosted.org/packages/03/5b/8857c05ef602276b094e7be8695766da1cc68b48dad7128fe1aa882591eb/pyscf-2.12.1-py3-none-macosx_11_0_arm64.whl", hash = "sha256:18f150403afd56b671ec3c337f18f8ff309ec9a7c71fb4bb7a27ebb534e65fc8", size = 35293657, upload-time = "2026-02-04T00:38:02.095Z" },
+    { url = "https://files.pythonhosted.org/packages/b3/9f/310b7e3753f1f9405459b270bbcbc09032a16bb1d2745421b3c5e4b7972a/pyscf-2.12.1-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:e780605991a91d3386b670b1d9fd5e300e6953d02b2958ed1f995904536ada73", size = 44773502, upload-time = "2026-02-04T03:24:39.922Z" },
+    { url = "https://files.pythonhosted.org/packages/8b/12/48c9b68186f8dc12e4220821ae4d3079844d657141b68f6ed7c426c134aa/pyscf-2.12.1-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:025a98c21a9014f1efae21a11f0f80ce14b9a2232cf2d3d934d83691b97b999f", size = 51580734, upload-time = "2026-02-04T01:03:00.43Z" },
+]
+
 [[package]]
 name = "python-dateutil"
 version = "2.9.0.post0"
@@ -1262,6 +1280,9 @@ dev = [
     { name = "ruff" },
     { name = "types-pyyaml" },
 ]
+pyscf = [
+    { name = "pyscf" },
+]
 
 [package.metadata]
 requires-dist = [
@@ -1273,13 +1294,14 @@ requires-dist = [
     { name = "openfermion", specifier = ">=1.7.1" },
     { name = "platformdirs", specifier = ">=4.5.1" },
     { name = "pybind11", specifier = ">=3.0.1" },
+    { name = "pyscf", marker = "extra == 'pyscf'", specifier = ">=2.12.1" },
     { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.14.10" },
     { name = "scipy", specifier = ">=1.16.3" },
     { name = "tensorboard", specifier = ">=2.20.0" },
     { name = "torch", specifier = ">=2.9.1" },
     { name = "types-pyyaml", marker = "extra == 'dev'", specifier = ">=6.0.12" },
 ]
-provides-extras = ["dev"]
+provides-extras = ["dev", "pyscf"]
 
 [[package]]
 name = "requests"

From 0165a56265acc5de2487c07aaaa13dea9d1b6d6f Mon Sep 17 00:00:00 2001
From: Hao Zhang <hzhangxyz@outlook.com>
Date: Fri, 6 Mar 2026 08:01:39 +0800
Subject: [PATCH 6/6] Ruff format

---
 qmp/models/pyscf.py  |  8 ++------
 qmp/plugins/pyscf.py | 12 +++---------
 2 files changed, 5 insertions(+), 15 deletions(-)

diff --git a/qmp/models/pyscf.py b/qmp/models/pyscf.py
index 96368fd..7fe76d8 100644
--- a/qmp/models/pyscf.py
+++ b/qmp/models/pyscf.py
@@ -71,9 +71,7 @@ def _restore_eri(eri: numpy.ndarray, norb: int) -> numpy.ndarray:
     if eri.ndim == 2:
         nij = norb * (norb + 1) // 2
         if eri.shape != (nij, nij):
-            raise ValueError(
-                f"For 2D ERI, expected shape ({nij}, {nij}) for norb={norb}, got {eri.shape}"
-            )
+            raise ValueError(f"For 2D ERI, expected shape ({nij}, {nij}) for norb={norb}, got {eri.shape}")
         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()))
@@ -147,9 +145,7 @@ def __init__(self, args: ModelConfig) -> None:
         interaction_operator: openfermion.InteractionOperator = openfermion.InteractionOperator(
             args.nuclear_repulsion, energy_1_b.numpy(), energy_2_b.numpy()
         )  # type: ignore[no-untyped-call]
-        fermion_operator: openfermion.FermionOperator = openfermion.get_fermion_operator(
-            interaction_operator
-        )  # type: ignore[no-untyped-call]
+        fermion_operator: openfermion.FermionOperator = openfermion.get_fermion_operator(interaction_operator)  # type: ignore[no-untyped-call]
         openfermion_hamiltonian_dict = {
             k: complex(v)
             for k, v in openfermion.normal_ordered(fermion_operator).terms.items()  # type: ignore[no-untyped-call]
diff --git a/qmp/plugins/pyscf.py b/qmp/plugins/pyscf.py
index 1310f88..fe1726b 100644
--- a/qmp/plugins/pyscf.py
+++ b/qmp/plugins/pyscf.py
@@ -419,9 +419,7 @@ def kernel(
             # ------------------------------------------------------------------
             run = dacite.from_dict(
                 data_class=action_dict[self.config.action.name],
-                data=omegaconf.OmegaConf.to_container(
-                    self.config.action.params, resolve=True
-                ),  # type: ignore[arg-type]
+                data=omegaconf.OmegaConf.to_container(self.config.action.params, resolve=True),  # type: ignore[arg-type]
                 config=dacite.Config(cast=DACITE_CAST),
             )
 
@@ -445,9 +443,7 @@ def kernel(
             # ------------------------------------------------------------------
             data_path = tmpdir / "data.pth"
             if data_path.exists():
-                data: dict[str, typing.Any] = torch.load(
-                    data_path, map_location="cpu", weights_only=True
-                )
+                data: dict[str, typing.Any] = torch.load(data_path, map_location="cpu", weights_only=True)
             else:
                 data = {}
 
@@ -456,9 +452,7 @@ def kernel(
             # ------------------------------------------------------------------
             energy = _extract_energy(data, self.config.action.name)
             # Strip the random engine state — it is stale by the next call.
-            full_checkpoint: dict[str, typing.Any] = {
-                k: v for k, v in data.items() if k != "random"
-            }
+            full_checkpoint: dict[str, typing.Any] = {k: v for k, v in data.items() if k != "random"}
 
             # Try to convert to PySCF CI vector format
             ci_array = _to_pyscf_ci(data, self.config.action.name, norb, (nalpha, nbeta))